[Mlir-commits] [mlir] [MLIR] Fix LivenessAnalysis/RemoveDeadValues handling of dead function arguments (PR #160755)

[Mehdi Amini]

https://github.com/joker-eph created https://github.com/llvm/llvm-project/pull/160755

In #153973 I added the correctly handling of block arguments, unfortunately this was gated on operation that also have results. This wasn't intentional and this excluded operations like function from being correctly processed.

>From 0dafe0018e85d7c4b2fe3fc1f628a7822d09eb92 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Thu, 25 Sep 2025 11:37:59 -0700
Subject: [PATCH] [MLIR] Fix LivenessAnalysis/RemoveDeadValues handling of dead
 function arguments

In #153973 I added the correctly handling of block arguments, unfortunately
this was gated on operation that also have results. This wasn't intentional
and this excluded operations like function from being correctly processed.
---
 .../Analysis/DataFlow/LivenessAnalysis.cpp    |  2 --
 mlir/lib/Transforms/RemoveDeadValues.cpp      |  3 +-
 mlir/test/Transforms/remove-dead-values.mlir  | 34 +++++++++++++++++++
 3 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp b/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp
index 65df355216f74..fdb97d5963299 100644
--- a/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp
+++ b/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp
@@ -299,8 +299,6 @@ RunLivenessAnalysis::RunLivenessAnalysis(Operation *op) {
   // The framework doesn't visit operations in dead blocks, so we need to
   // explicitly mark them as dead.
   op->walk([&](Operation *op) {
-    if (op->getNumResults() == 0)
-      return;
     for (auto result : llvm::enumerate(op->getResults())) {
       if (getLiveness(result.value()))
         continue;
diff --git a/mlir/lib/Transforms/RemoveDeadValues.cpp b/mlir/lib/Transforms/RemoveDeadValues.cpp
index 0e84b6dd17f29..59983365acfb8 100644
--- a/mlir/lib/Transforms/RemoveDeadValues.cpp
+++ b/mlir/lib/Transforms/RemoveDeadValues.cpp
@@ -287,7 +287,8 @@ static void processSimpleOp(Operation *op, RunLivenessAnalysis &la,
 static void processFuncOp(FunctionOpInterface funcOp, Operation *module,
                           RunLivenessAnalysis &la, DenseSet<Value> &nonLiveSet,
                           RDVFinalCleanupList &cl) {
-  LDBG() << "Processing function op: " << funcOp.getOperation()->getName();
+  LDBG() << "Processing function op: "
+         << OpWithFlags(funcOp, OpPrintingFlags().skipRegions());
   if (funcOp.isPublic() || funcOp.isExternal()) {
     LDBG() << "Function is public or external, skipping: "
            << funcOp.getOperation()->getName();
diff --git a/mlir/test/Transforms/remove-dead-values.mlir b/mlir/test/Transforms/remove-dead-values.mlir
index fa2c145bd3701..33afd54d8113d 100644
--- a/mlir/test/Transforms/remove-dead-values.mlir
+++ b/mlir/test/Transforms/remove-dead-values.mlir
@@ -615,3 +615,37 @@ module @last_block_not_exit {
   // CHECK-LABEL: @call_private_but_not_use
   // CHECK: call @terminated_with_condbr(%false, %true) : (i1, i1)
 }
+
+// -----
+
+// Test the elimination of function arguments.
+
+// CHECK-LABEL: func private @single_parameter
+// CHECK-SAME: () {
+func.func private @single_parameter(%arg0: index) {
+  return
+}
+
+// CHECK-LABEL: func.func private @mutl_parameter(
+// CHECK-SAME: %[[ARG0:.*]]: index)
+// CHECK: return %[[ARG0]]
+func.func private @mutl_parameter(%arg0: index, %arg1: index, %arg2: index) -> index {
+  return %arg1 : index
+}
+
+// CHECK-LABEL: func private @eliminate_parameter
+// CHECK-SAME: () {
+func.func private @eliminate_parameter(%arg0: index, %arg1: index) {
+  call @single_parameter(%arg0) : (index) -> ()
+  return
+}
+
+// CHECK-LABEL: func @callee
+// CHECK-SAME: (%[[ARG0:.*]]: index, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index)
+func.func @callee(%arg0: index, %arg1: index, %arg2: index) -> index {
+// CHECK: call @eliminate_parameter() : () -> ()
+  call @eliminate_parameter(%arg0, %arg1) : (index, index) -> ()
+// CHECK: call @mutl_parameter(%[[ARG1]]) : (index) -> index
+  %res = call @mutl_parameter(%arg0, %arg1, %arg2) : (index, index, index) -> (index)
+  return %res : index
+}
\ No newline at end of file