From d54a007d86485a4a4914668fe0eaa1b49c93d6c1 Mon Sep 17 00:00:00 2001 From: Xin Liu Date: Wed, 24 Sep 2025 22:13:51 -0700 Subject: [PATCH 1/6] [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. --- mlir/include/mlir/Analysis/DataFlow/LivenessAnalysis.h | 2 +- mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp | 2 +- mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp | 9 ++++++++- 3 files changed, 10 insertions(+), 3 deletions(-) diff --git a/mlir/include/mlir/Analysis/DataFlow/LivenessAnalysis.h b/mlir/include/mlir/Analysis/DataFlow/LivenessAnalysis.h index cf1fd6e2d48ca..dffa781543243 100644 --- a/mlir/include/mlir/Analysis/DataFlow/LivenessAnalysis.h +++ b/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. diff --git a/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp b/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp index 65df355216f74..9eebc161914c3 100644 --- a/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp +++ b/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,...) diff --git a/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp b/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp index 0d2e2ed85549d..f03290660326a 100644 --- a/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp +++ b/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 argOpOperands = operandsToOpOperands(argOperands); Region *region = callable.getCallableRegion(); + bool isPublicFunc = isa(callableOp) + && cast(callableOp).isPublic(); if (!region || region->empty() || - !getSolverConfig().isInterprocedural()) { + !getSolverConfig().isInterprocedural() || isPublicFunc) { visitExternalCallImpl(call, operandLattices, resultLattices); return success(); } From ed25cb7371626a1497dcb5b8f3bb6a901a243562 Mon Sep 17 00:00:00 2001 From: Xin Liu Date: Wed, 24 Sep 2025 22:26:46 -0700 Subject: [PATCH 2/6] Fix formatter error. --- mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp b/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp index f03290660326a..a509aaaa1e1f1 100644 --- a/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp +++ b/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp @@ -514,8 +514,8 @@ AbstractSparseBackwardDataFlowAnalysis::visitOperation(Operation *op) { MutableArrayRef argOpOperands = operandsToOpOperands(argOperands); Region *region = callable.getCallableRegion(); - bool isPublicFunc = isa(callableOp) - && cast(callableOp).isPublic(); + bool isPublicFunc = isa(callableOp) && + cast(callableOp).isPublic(); if (!region || region->empty() || !getSolverConfig().isInterprocedural() || isPublicFunc) { visitExternalCallImpl(call, operandLattices, resultLattices); From 153e6e630dc0c29016bc5b6539cb6d14b1f99b05 Mon Sep 17 00:00:00 2001 From: Xin Liu Date: Thu, 25 Sep 2025 20:04:42 -0700 Subject: [PATCH 3/6] Update the comment and also add lit test. --- .../mlir/Analysis/DataFlow/LivenessAnalysis.h | 2 +- mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp | 2 +- mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp | 15 ++++++++------- mlir/test/Transforms/remove-dead-values.mlir | 17 +++++++++++++++++ 4 files changed, 27 insertions(+), 9 deletions(-) diff --git a/mlir/include/mlir/Analysis/DataFlow/LivenessAnalysis.h b/mlir/include/mlir/Analysis/DataFlow/LivenessAnalysis.h index dffa781543243..cf1fd6e2d48ca 100644 --- a/mlir/include/mlir/Analysis/DataFlow/LivenessAnalysis.h +++ b/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. diff --git a/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp b/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp index 9eebc161914c3..65df355216f74 100644 --- a/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp +++ b/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,...) diff --git a/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp b/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp index a509aaaa1e1f1..0d87bfe11177a 100644 --- a/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp +++ b/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 argOpOperands = operandsToOpOperands(argOperands); Region *region = callable.getCallableRegion(); - bool isPublicFunc = isa(callableOp) && - cast(callableOp).isPublic(); - if (!region || region->empty() || - !getSolverConfig().isInterprocedural() || isPublicFunc) { + auto isPublicFunction = [&]() { + auto funcOp = dyn_cast(callableOp); + return funcOp && funcOp.isPublic(); + }; + if (!getSolverConfig().isInterprocedural() || !region || + region->empty() || isPublicFunction()) { visitExternalCallImpl(call, operandLattices, resultLattices); return success(); } diff --git a/mlir/test/Transforms/remove-dead-values.mlir b/mlir/test/Transforms/remove-dead-values.mlir index fa2c145bd3701..f4ae5118b966d 100644 --- a/mlir/test/Transforms/remove-dead-values.mlir +++ b/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 + } } // ----- From 9e35848693cec9c2649e38670b4c497dc5e7a430 Mon Sep 17 00:00:00 2001 From: Xin Liu Date: Fri, 26 Sep 2025 09:34:37 -0700 Subject: [PATCH 4/6] Update the comment. --- mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp b/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp index 0d87bfe11177a..ee47ebfeacefb 100644 --- a/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp +++ b/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp @@ -507,8 +507,10 @@ 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, the signature is immutable. - // We need to be conservative and consider all arguments Live. + // If callable is a public function, treat it as external. + // This is because a public function has potential callers we can't + // visit, and thus we need to be conservative and consider all + // arguments live. OperandRange argOperands = call.getArgOperands(); MutableArrayRef argOpOperands = operandsToOpOperands(argOperands); From c166b1f43c8f16500709d0ccf7f9dd79bf355077 Mon Sep 17 00:00:00 2001 From: Xin Liu Date: Fri, 26 Sep 2025 09:44:12 -0700 Subject: [PATCH 5/6] Simplify lit test. use more descriptive name. --- mlir/test/Transforms/remove-dead-values.mlir | 13 ++++--------- 1 file changed, 4 insertions(+), 9 deletions(-) diff --git a/mlir/test/Transforms/remove-dead-values.mlir b/mlir/test/Transforms/remove-dead-values.mlir index f4ae5118b966d..b2269f6a6dfcb 100644 --- a/mlir/test/Transforms/remove-dead-values.mlir +++ b/mlir/test/Transforms/remove-dead-values.mlir @@ -569,21 +569,16 @@ 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> + // the function signature is immutable because it is public. + func.func public @public_fn_with_unused_argument(%unused: i32) -> () { 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) -> () + // CHECK: call @public_fn_with_unused_argument(%[[UNUSED]]) : (i32) -> () func.func @main2(%arg0: i32) -> () { %zero = arith.constant 0 : i32 - call @immutable_fn_return_void_with_unused_argument(%arg0, %zero) : (i32, i32) -> () + call @public_fn_with_unused_argument(%zero) : (i32) -> () return } } From 58302899522c464d8bdd10800f4f57b566f48b9f Mon Sep 17 00:00:00 2001 From: Xin Liu Date: Fri, 26 Sep 2025 12:55:27 -0700 Subject: [PATCH 6/6] Make the lit test minimal. --- mlir/test/Transforms/remove-dead-values.mlir | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlir/test/Transforms/remove-dead-values.mlir b/mlir/test/Transforms/remove-dead-values.mlir index b2269f6a6dfcb..8d50179c863a3 100644 --- a/mlir/test/Transforms/remove-dead-values.mlir +++ b/mlir/test/Transforms/remove-dead-values.mlir @@ -576,7 +576,7 @@ module @return_void_with_unused_argument { // CHECK-LABEL: func.func @main2 // CHECK: %[[UNUSED:.*]] = arith.constant 0 : i32 // CHECK: call @public_fn_with_unused_argument(%[[UNUSED]]) : (i32) -> () - func.func @main2(%arg0: i32) -> () { + func.func @main2() -> () { %zero = arith.constant 0 : i32 call @public_fn_with_unused_argument(%zero) : (i32) -> () return