[Mlir-commits] [mlir] 0edf624 - [mlir][RemoveDeadValues] Replace dead function argument uses with poison (#208881)

[llvmlistbot at llvm.org]

Author: Jinpeng Wang
Date: 2026-10-12T13:11:10+09:00
New Revision: 0edf6243d58840a207d491b3c03d5d7599106bd2

URL: https://github.com/llvm/llvm-project/commit/0edf6243d58840a207d491b3c03d5d7599106bd2
DIFF: https://github.com/llvm/llvm-project/commit/0edf6243d58840a207d491b3c03d5d7599106bd2.diff

LOG: [mlir][RemoveDeadValues] Replace dead function argument uses with poison (#208881)

Dropping all uses of a dead function argument can leave null operands in
surviving side-effecting operations in unreachable code.
Subsequent canonicalization may then crash while folding the invalid
operation.

This PR replaces the remaining uses with `ub.poison` before erasing the
argument and add a regression test.

Fixes #206920, fixes #203226

Added: 
    

Modified: 
    mlir/lib/Transforms/RemoveDeadValues.cpp
    mlir/test/Transforms/remove-dead-values.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Transforms/RemoveDeadValues.cpp b/mlir/lib/Transforms/RemoveDeadValues.cpp
index 0b9321a265221..7a778b049b879 100644
--- a/mlir/lib/Transforms/RemoveDeadValues.cpp
+++ b/mlir/lib/Transforms/RemoveDeadValues.cpp
@@ -652,9 +652,20 @@ static void cleanUpDeadVals(MLIRContext *ctx, RDVFinalCleanupList &list) {
       llvm::interleaveComma(f.nonLiveRets.set_bits(), os);
       os << "]";
     });
-    // Drop all uses of the dead arguments.
-    for (auto deadIdx : f.nonLiveArgs.set_bits())
-      f.funcOp.getArgument(deadIdx).dropAllUses();
+    // Replace remaining uses of the dead arguments with poison values. Simply
+    // dropping the uses would leave null operands behind in ops that survive
+    // the pass (e.g. a side-effecting op in an unreachable function, whose
+    // values are initialized as dead by the liveness analysis).
+    for (auto deadIdx : f.nonLiveArgs.set_bits()) {
+      BlockArgument arg = f.funcOp.getArgument(deadIdx);
+      // Avoid creating an unused poison value if there are no uses to replace.
+      if (arg.use_empty())
+        continue;
+      rewriter.setInsertionPointToStart(arg.getOwner());
+      Value poison =
+          ub::PoisonOp::create(rewriter, arg.getLoc(), arg.getType());
+      rewriter.replaceAllUsesWith(arg, poison);
+    }
     // Some functions may not allow erasing arguments or results. These calls
     // return failure in such cases without modifying the function, so it's okay
     // to proceed.

diff  --git a/mlir/test/Transforms/remove-dead-values.mlir b/mlir/test/Transforms/remove-dead-values.mlir
index 863a0d450b7af..a38d6dc5ca263 100644
--- a/mlir/test/Transforms/remove-dead-values.mlir
+++ b/mlir/test/Transforms/remove-dead-values.mlir
@@ -898,6 +898,55 @@ module @func_with_non_call_users {
 
 // -----
 
+// Verify that the upper bound of a for loop inside an unreachable
+// function is correctly replaced with a poison value.
+// CHECK-LABEL: module @unreachable_func_with_for_loops
+// CHECK-NEXT:    func.func private @main_func() {
+// CHECK-NEXT:      %[[ARG3:.*]] = ub.poison : index
+// CHECK-NEXT:      %[[C0:.*]] = ub.poison : index
+// CHECK-NEXT:      %[[C1:.*]] = ub.poison : index
+// CHECK-NEXT:      %[[C10:.*]] = ub.poison : index
+// CHECK-NEXT:      scf.for %{{.*}} = %[[C0]] to %[[C10]] step %[[C1]] {
+// CHECK-NEXT:        gpu.barrier
+// CHECK-NEXT:      }
+// CHECK-NEXT:      scf.for %{{.*}} = %[[C0]] to %[[ARG3]] step %[[C1]] {
+// CHECK-NEXT:        gpu.barrier
+// CHECK-NEXT:      }
+// CHECK-NEXT:      return
+
+// CHECK-CANONICALIZE-LABEL: module @unreachable_func_with_for_loops
+// CHECK-CANONICALIZE-NEXT:    func.func private @main_func() {
+// CHECK-CANONICALIZE-NEXT:      %[[ARG3:.*]] = ub.poison : index
+// CHECK-CANONICALIZE-NEXT:      %[[C0:.*]] = ub.poison : index
+// CHECK-CANONICALIZE-NEXT:      %[[C1:.*]] = ub.poison : index
+// CHECK-CANONICALIZE-NEXT:      %[[C10:.*]] = ub.poison : index
+// CHECK-CANONICALIZE-NEXT:      scf.for %{{.*}} = %[[C0]] to %[[C10]] step %[[C1]] {
+// CHECK-CANONICALIZE-NEXT:        gpu.barrier
+// CHECK-CANONICALIZE-NEXT:      }
+// CHECK-CANONICALIZE-NEXT:      scf.for %{{.*}} = %[[C0]] to %[[ARG3]] step %[[C1]] {
+// CHECK-CANONICALIZE-NEXT:        gpu.barrier
+// CHECK-CANONICALIZE-NEXT:      }
+// CHECK-CANONICALIZE-NEXT:      return
+module @unreachable_func_with_for_loops {
+  func.func private @main_func(%arg0: index, %arg1: memref<?xf32>, %arg2: index, %arg3: index) {
+    %cst = arith.constant 0.000000e+00 : f32
+    %c0 = arith.constant 0 : index
+    %c1 = arith.constant 1 : index
+    %c10 = arith.constant 10 : index
+    scf.for %arg4 = %c0 to %c10 step %c1 {
+      %0 = memref.load %arg1[%arg4] : memref<?xf32>
+      gpu.barrier
+    }
+    scf.for %arg4 = %c0 to %arg3 step %c1 {
+      %0 = memref.load %arg1[%arg4] : memref<?xf32>
+      gpu.barrier
+    }
+    return
+  }
+}
+
+// -----
+
 // A call op may have results that are produced by the call op itself instead of
 // being forwarded from the callee (`%status` below). Such results are not part
 // of the 1:1 relationship between call results and callee results.