[MLIR][LivenessAnalysis] Treat a public function as an external

This change treats a public function as an external function in the
inter-procedural liveness analysis. This prohibits NonLive arguments
from propagating to caller. RemoveDeadValues deletes unused values based
on the results of liveness Analysis. On the other side, it can not
change the signature of a public function.

Author: navyxliu

mlir/include/mlir/Analysis/DataFlow/LivenessAnalysis.h

@@ -36,7 +36,7 @@ namespace mlir::dataflow {
 ///
 /// A value is considered "live" iff it:
 ///   (1) has memory effects OR
-///   (2) is returned by a public function OR
+///   (2) is an argument of or is returned by a public function OR
 ///   (3) is used to compute a value of type (1) or (2).
 /// It is also to be noted that a value could be of multiple types (1/2/3) at
 /// the same time.

mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp

@@ -52,7 +52,7 @@ ChangeResult Liveness::meet(const AbstractSparseLattice &other) {
 ///
 /// A value is considered "live" iff it:
 ///   (1) has memory effects OR
-///   (2) is returned by a public function OR
+///   (2) is an argument of or is returned by a public function OR
 ///   (3) is used to compute a value of type (1) or (2) OR
 ///   (4) is returned by a return-like op whose parent isn't a callable
 ///       nor a RegionBranchOpInterface (e.g.: linalg.yield, gpu.yield,...)

mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp

@@ -17,6 +17,7 @@
 #include "mlir/IR/ValueRange.h"
 #include "mlir/Interfaces/CallInterfaces.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
+#include "mlir/Interfaces/FunctionInterfaces.h"
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/DebugLog.h"
@@ -505,12 +506,18 @@ AbstractSparseBackwardDataFlowAnalysis::visitOperation(Operation *op) {
       // If the call invokes an external function (or a function treated as
       // external due to config), defer to the corresponding extension hook.
       // By default, it just does `visitCallOperand` for all operands.
+      //
+      // If callable is a public function, treat it as an external function.
+      // Transforms like RemoveDeadValues cannot change the arguments or returns
+      // of it.
       OperandRange argOperands = call.getArgOperands();
       MutableArrayRef<OpOperand> argOpOperands =
           operandsToOpOperands(argOperands);
       Region *region = callable.getCallableRegion();
+      bool isPublicFunc = isa<FunctionOpInterface>(callableOp)
+                       && cast<FunctionOpInterface>(callableOp).isPublic();
       if (!region || region->empty() ||
-          !getSolverConfig().isInterprocedural()) {
+          !getSolverConfig().isInterprocedural() || isPublicFunc) {
         visitExternalCallImpl(call, operandLattices, resultLattices);
         return success();
       }