[Mlir-commits] [mlir] [mlir][Transforms] --remove-dead-values: keep values used in public function calls (PR #83249)

[Artem Tyurin]

https://github.com/agentcooper created https://github.com/llvm/llvm-project/pull/83249

Fixes #82788

>From a11b37b03a7be329412c92cf906217a1058378c0 Mon Sep 17 00:00:00 2001
From: Artem Tyurin <artem.tyurin at gmail.com>
Date: Wed, 28 Feb 2024 12:18:41 +0100
Subject: [PATCH] [mlir][Transforms] --remove-dead-values: keep values used in
 public function calls

Fixes #82788
---
 mlir/lib/Transforms/RemoveDeadValues.cpp     | 16 +++++++++++-
 mlir/test/Transforms/remove-dead-values.mlir | 27 ++++++++++++++++++++
 2 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Transforms/RemoveDeadValues.cpp b/mlir/lib/Transforms/RemoveDeadValues.cpp
index 055256903a1522..4fad27e5adabda 100644
--- a/mlir/lib/Transforms/RemoveDeadValues.cpp
+++ b/mlir/lib/Transforms/RemoveDeadValues.cpp
@@ -169,7 +169,10 @@ static SmallVector<OpOperand *> operandsToOpOperands(OperandRange operands) {
 /// Here, cleaning means:
 ///   (1) Dropping all its uses, AND
 ///   (2) Erasing it
-/// iff it has no memory effects and none of its results are live.
+/// iff:
+///   (1) It has no memory effects
+///   (2) None of its results are live
+///   (3) It is not used as an argument in a public function call
 ///
 /// It is assumed that `op` is simple. Here, a simple op is one which isn't a
 /// symbol op, a symbol-user op, a region branch op, a branch op, a region
@@ -178,6 +181,17 @@ static void cleanSimpleOp(Operation *op, RunLivenessAnalysis &la) {
   if (!isMemoryEffectFree(op) || hasLive(op->getResults(), la))
     return;
 
+  for (auto &use : op->getUses()) {
+    if (auto callOp = dyn_cast<CallOpInterface>(use.getOwner())) {
+      if (auto funcOp =
+              dyn_cast<FunctionOpInterface>(callOp.resolveCallable())) {
+        if (funcOp.isPublic()) {
+          return;
+        }
+      }
+    }
+  }
+
   op->dropAllUses();
   op->erase();
 }
diff --git a/mlir/test/Transforms/remove-dead-values.mlir b/mlir/test/Transforms/remove-dead-values.mlir
index 69426fdb620832..de83e80acc3ca6 100644
--- a/mlir/test/Transforms/remove-dead-values.mlir
+++ b/mlir/test/Transforms/remove-dead-values.mlir
@@ -357,3 +357,30 @@ func.func @kernel(%arg0: memref<18xf32>) {
 // CHECK: gpu.launch blocks
 // CHECK: memref.store
 // CHECK-NEXT: gpu.terminator
+
+// -----
+
+// Even though %arg0 is not used in func0,
+// we can't remove the parameter, as the function is public,
+// so `%true = arith.constant true` needs to be kept in func1 as well.
+//
+module {
+// CHECK:       func.func @func0(%arg0: i1, %arg1: i1) -> (i1, i1) {
+// CHECK-NEXT:    %true = arith.constant true
+// CHECK-NEXT:    return %arg1, %true : i1, i1
+// CHECK-NEXT:  }
+  func.func @func0(%arg0: i1, %arg1: i1) -> (i1, i1) {
+    %true = arith.constant true
+    return %arg1, %true : i1, i1
+  }
+// CHECK:       func.func @func1(%arg0: i1) -> (i1, i1) {
+// CHECK-NEXT:    %true = arith.constant true
+// CHECK-NEXT:    %0:2 = call @func0(%true, %arg0) : (i1, i1) -> (i1, i1)
+// CHECK-NEXT:    return %0#0, %0#1 : i1, i1
+// CHECK-NEXT:  }
+  func.func @func1(%arg0: i1) -> (i1, i1) {
+    %true = arith.constant true
+    %0:2 = call @func0(%true, %arg0) : (i1, i1) -> (i1, i1)
+    return %0#0, %0#1 : i1, i1
+  }
+}