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

[navyxliu]

Public functions have usage outside of the current module, and thus their
signature is immutable.
Fixes liveness analysis to consider all arguments as live instead.

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[llvmbot]

@llvm/pr-subscribers-mlir

Author: xin liu (navyxliu)

Changes

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.

---

Full diff: https://github.com/llvm/llvm-project/pull/160648.diff

2 Files Affected:

• (modified) mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp (+10-2)
• (modified) mlir/test/Transforms/remove-dead-values.mlir (+17)

diff --git a/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp b/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp
index 0d2e2ed85549d..0d87bfe11177a 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,19 @@ 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, 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();
-      if (!region || region->empty() ||
-          !getSolverConfig().isInterprocedural()) {
+      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();
       }
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
+  }
 }
 
 // -----

[navyxliu]

Hi, Reviewers,

I also learn the race condition issue you are resolving. AbstractSparseBackwardDataFlowAnalysis::visitOperation() could use a minor change to address that issue too. Currently, it calls 'call.resolveCallableInTable(&symbolTable)' too early. I didn't include that change intentionally. We should track it separately.

The other thing is that current change is "intrusive". Like @ftynse mentioned, we should implement this special logic for LivenessAnalysis. One idea is that we define a virtual function visitCallOperation like AbstractSparseForwardDataFlowAnalysis does.

[joker-eph]

I'm wondering now if this is the right fix though.
Making arguments as dead could still allow simplification at call-sites (for example: replace with a constant or an "undef" value), while not touching the signature?

[navyxliu]

I'm wondering now if this is the right fix though. Making arguments as dead could still allow simplification at call-sites (for example: replace with a constant or an "undef" value), while not touching the signature?

yes. that's what I wrote in PR #160242.

With a non-live argument from callee, we could replace it with a dummy placeholder like arith.constant 0 : i32. It would be impossible if we mark all arguments of a public function live.

The problem is whether it is always possible to find cheaper replacement? What about an aggregated type?

if we fail, we need to propagate live arguments in RemoveDeadValues just like PR #160242

[joker-eph]

The problem is whether it is always possible to find cheaper replacement? What about an aggregated type?

It may not always be possible, but it's always safe to do nothing.

if we fail, we need to propagate live arguments in RemoveDeadValues just like PR #160242

Right seems like it may be better handled there: that is "not-live" != "removable" here.

[navyxliu]

The problem is whether it is always possible to find cheaper replacement? What about an aggregated type?

It may not always be possible, but it's always safe to do nothing.

if we fail, we need to propagate live arguments in RemoveDeadValues just like PR #160242

Right seems like it may be better handled there: that is "not-live" != "removable" here.

So I drop this and continue working on #160242?

Allow me to summarize the idea: We intercept callOp whose target is public. We iterate its parameters:

If liveness analysis is intra-procedural or the parameter is NonLive but we manage to find an alternative: do nothing.
Otherwise, we mark the corresponding argument in caller Live. We propagate liveness backward.

[navyxliu]

hi, @joker-eph
We come up yet another idea. How about we clone the public function foo() and mark the clone private?

public void foo(%unused: i32) {
return; 
}
bar() {
%x = ...
foo(%x)
}
=> clone change its uses:
public void foo(%unused: i32) {
return; 
}
private void foo_(%unused: i32) {
return; 
}
bar() {
%x = ...
foo_(%x)
}

then everything is in order.

[navyxliu]

hi, @joker-eph We come up yet another idea. How about we clone the public function foo() and mark the clone private?

public void foo(%unused: i32) {
return; 
}
bar() {
%x = ...
foo(%x)
}
=> clone change its uses:
public void foo(%unused: i32) {
return; 
}
private void foo_(%unused: i32) {
return; 
}
bar() {
%x = ...
foo_(%x)
}

then everything is in order.

I tried this idea. it has one limitation.
we can clone a function, but we can't clone liveness results from the old function to the new.
The dataflow framework doesn't support incremental solver.

If we want to go this route, we need to pre-process call Ops and clone them prior to 'RunLivenessAnalysis'.

[joker-eph]

That's a possibility, but feels like there should be some cost-model somehow maybe? It's not clear that this is always profitable?