Update the comment and also add lit test.

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 an argument of or is returned by a public function OR
+///   (2) 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 an argument of or is returned by a public function OR
+///   (2) 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

@@ -507,17 +507,18 @@ AbstractSparseBackwardDataFlowAnalysis::visitOperation(Operation *op) {
       // 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.
+      // If callable is a public function, the signature is immutable.
+      // We need to be conservative and consider all arguments Live.
       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() || isPublicFunc) {
+      auto isPublicFunction = [&]() {
+        auto funcOp = dyn_cast<FunctionOpInterface>(callableOp);
+        return funcOp && funcOp.isPublic();
+      };
+      if (!getSolverConfig().isInterprocedural() || !region ||
+          region->empty() || isPublicFunction()) {
         visitExternalCallImpl(call, operandLattices, resultLattices);
         return success();
       }

mlir/test/Transforms/remove-dead-values.mlir

@@ -569,6 +569,23 @@ module @return_void_with_unused_argument {
     call @fn_return_void_with_unused_argument(%arg0, %unused) : (i32, memref<4xi32>) -> ()
     return %unused : memref<4xi32>
   }
+  // the function is immutable because it is public.
+  func.func public @immutable_fn_return_void_with_unused_argument(%arg0: i32, %unused: i32) -> () {
+    %sum = arith.addi %arg0, %arg0 : i32
+    %c0 = arith.constant 0 : index
+    %buf = memref.alloc() : memref<1xi32>
+    memref.store %sum, %buf[%c0] : memref<1xi32>
+    return
+  }
+  // CHECK-LABEL: func.func @main2
+  // CHECK-SAME: (%[[ARG0_MAIN:.*]]: i32)
+  // CHECK: %[[UNUSED:.*]] = arith.constant 0 : i32
+  // CHECK: call @immutable_fn_return_void_with_unused_argument(%[[ARG0_MAIN]], %[[UNUSED]]) : (i32, i32) -> ()
+  func.func @main2(%arg0: i32) -> () {
+    %zero = arith.constant 0 : i32
+    call @immutable_fn_return_void_with_unused_argument(%arg0, %zero) : (i32, i32) -> ()
+    return
+  }
 }
 
 // -----