[Mlir-commits] [mlir] [mlir] Enable remove-dead-values to delete unused private function (PR #161471)

[lonely eagle]

https://github.com/linuxlonelyeagle updated https://github.com/llvm/llvm-project/pull/161471

>From 8f974c9770eefec75d97fb91db865f9eb36577a9 Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Mon, 6 Oct 2025 15:33:08 +0000
Subject: [PATCH 1/2] add deleteDeadFunction function.

---
 mlir/lib/Transforms/RemoveDeadValues.cpp     | 21 +++++++++
 mlir/test/Transforms/remove-dead-values.mlir | 48 +++++++++++++++++++-
 2 files changed, 68 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Transforms/RemoveDeadValues.cpp b/mlir/lib/Transforms/RemoveDeadValues.cpp
index e0c65b0e09774..3a425a96e1efc 100644
--- a/mlir/lib/Transforms/RemoveDeadValues.cpp
+++ b/mlir/lib/Transforms/RemoveDeadValues.cpp
@@ -235,6 +235,25 @@ static void dropUsesAndEraseResults(Operation *op, BitVector toErase) {
   op->erase();
 }
 
+// Remove the dead functions from moduleOp.
+static void deleteDeadFunction(Operation *module) {
+  bool walkContinue = true;
+  while (walkContinue) {
+    walkContinue = false;
+    module->walk([&](FunctionOpInterface funcOp) {
+      if (funcOp.isPublic() || funcOp.isExternal())
+        return;
+
+      SymbolTable::UseRange uses = *funcOp.getSymbolUses(module);
+      auto callSites = funcOp.getFunctionBody().getOps<CallOpInterface>();
+      if (uses.empty())
+        funcOp.erase();
+      if (uses.empty() && !callSites.empty())
+        walkContinue = true;
+    });
+  }
+}
+
 /// Convert a list of `Operand`s to a list of `OpOperand`s.
 static SmallVector<OpOperand *> operandsToOpOperands(OperandRange operands) {
   OpOperand *values = operands.getBase();
@@ -881,6 +900,8 @@ void RemoveDeadValues::runOnOperation() {
   // end of this pass.
   RDVFinalCleanupList finalCleanupList;
 
+  // Remove the dead function in advance.
+  deleteDeadFunction(module);
   module->walk([&](Operation *op) {
     if (auto funcOp = dyn_cast<FunctionOpInterface>(op)) {
       processFuncOp(funcOp, module, la, deadVals, finalCleanupList);
diff --git a/mlir/test/Transforms/remove-dead-values.mlir b/mlir/test/Transforms/remove-dead-values.mlir
index 56449469dc29f..1efaf2ecc11ce 100644
--- a/mlir/test/Transforms/remove-dead-values.mlir
+++ b/mlir/test/Transforms/remove-dead-values.mlir
@@ -455,7 +455,7 @@ module @llvm_unreachable {
   func.func private @fn_with_llvm_unreachable(%arg0: tensor<4x4xf32>) -> tensor<4x4xi1> {
     llvm.unreachable
   }
-  func.func private @main(%arg0: tensor<4x4xf32>) {
+  func.func @main(%arg0: tensor<4x4xf32>) {
     %0 = call @fn_with_llvm_unreachable(%arg0) : (tensor<4x4xf32>) -> tensor<4x4xi1>
     llvm.return
   }
@@ -649,3 +649,49 @@ func.func @callee(%arg0: index, %arg1: index, %arg2: index) -> index {
   %res = call @mutl_parameter(%arg0, %arg1, %arg2) : (index, index, index) -> (index)
   return %res : index
 }
+
+// -----
+
+// Test the elimination of dead functions.
+
+// CHECK-NOT: func private @single_private_func
+func.func private @single_private_func(%arg0: i64) -> (i64) {
+    %c0_i64 = arith.constant 0 : i64
+    %2 = arith.cmpi eq, %arg0, %c0_i64 : i64
+    cf.cond_br %2, ^bb1, ^bb2
+  ^bb1:  // pred: ^bb0
+    %c1_i64 = arith.constant 1 : i64
+    return %c1_i64 : i64
+  ^bb2:  // pred: ^bb0
+    %c3_i64 = arith.constant 3 : i64
+    return %c3_i64 : i64
+}
+
+// -----
+
+// Test the elimination of dead functions.
+
+// CHECK-NOT: @single_parameter
+func.func private @single_parameter(%arg0: index) {
+  return
+}
+
+// CHECK-NOT: @mutl_parameter
+func.func private @mutl_parameter(%arg0: index, %arg1: index, %arg2: index) -> index {
+  return %arg1 : index
+}
+
+// CHECK-NOT: @eliminate_parameter
+func.func private @eliminate_parameter(%arg0: index, %arg1: index) {
+  call @single_parameter(%arg0) : (index) -> ()
+  return
+}
+
+// CHECK-NOT: @callee
+func.func private @callee(%arg0: index, %arg1: index, %arg2: index) -> index {
+  // CHECK-NOT: call @eliminate_parameter
+  call @eliminate_parameter(%arg0, %arg1) : (index, index) -> ()
+  // CHECK-NOT: call @mutl_parameter 
+  %res = call @mutl_parameter(%arg0, %arg1, %arg2) : (index, index, index) -> (index)
+  return %res : index
+}

>From f31761b829de63a2d8fe724bdaea1f92dd5dbcac Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Tue, 7 Oct 2025 05:59:56 +0000
Subject: [PATCH 2/2] update code.

---
 mlir/lib/Transforms/RemoveDeadValues.cpp | 23 ++++++++++++++---------
 1 file changed, 14 insertions(+), 9 deletions(-)

diff --git a/mlir/lib/Transforms/RemoveDeadValues.cpp b/mlir/lib/Transforms/RemoveDeadValues.cpp
index 3a425a96e1efc..00701f6675bdd 100644
--- a/mlir/lib/Transforms/RemoveDeadValues.cpp
+++ b/mlir/lib/Transforms/RemoveDeadValues.cpp
@@ -237,20 +237,25 @@ static void dropUsesAndEraseResults(Operation *op, BitVector toErase) {
 
 // Remove the dead functions from moduleOp.
 static void deleteDeadFunction(Operation *module) {
-  bool walkContinue = true;
-  while (walkContinue) {
-    walkContinue = false;
-    module->walk([&](FunctionOpInterface funcOp) {
+  auto functions = module->getRegion(0).getOps<FunctionOpInterface>();
+  llvm::DenseSet<FunctionOpInterface> tasks(functions.begin(), functions.end());
+  while (!tasks.empty()) {
+    llvm::DenseSet<FunctionOpInterface> nextTasks;
+    for (FunctionOpInterface funcOp : tasks) {
       if (funcOp.isPublic() || funcOp.isExternal())
         return;
-
       SymbolTable::UseRange uses = *funcOp.getSymbolUses(module);
       auto callSites = funcOp.getFunctionBody().getOps<CallOpInterface>();
-      if (uses.empty())
+      if (uses.empty() && !callSites.empty()) {
+        for (CallOpInterface callOp : callSites) {
+          nextTasks.insert(cast<FunctionOpInterface>(callOp.resolveCallable()));
+        }
+      }
+
+      if (uses.empty() && !nextTasks.contains(funcOp))
         funcOp.erase();
-      if (uses.empty() && !callSites.empty())
-        walkContinue = true;
-    });
+    }
+    tasks = nextTasks;
   }
 }