[Mlir-commits] [mlir] [mlir][remove-dead-values] Fix crash due to dangling operands (PR #181013)

[Prathamesh Tagore]

https://github.com/meshtag updated https://github.com/llvm/llvm-project/pull/181013

>From 6e92ff12c5ec472e0fe701323e825e3cd955294c Mon Sep 17 00:00:00 2001
From: Prathamesh Tagore <prathameshtagore at gmail.com>
Date: Thu, 12 Feb 2026 17:28:35 +0100
Subject: [PATCH] [mlir][remove-dead-values] Replace appropriate operation
 results with poison

Before erasing the operation, replace all result values with live-uses by
ub.poison values. This is important to maintain IR validity. For example,
if we have an op with one of its results used by another op, erasing the
op without replacing its corresponding result would leave us with a
dangling operand in the user op. By replacing the result with a ub.poison
value, we ensure that the user op still has a valid operand, even though
it's a poison value which will be cleaned up later if it can be cleaned
up. This keeps the IR valid for further simplification and
canonicalization while fixing a related crash in the canonicalizer.
---
 mlir/lib/Transforms/RemoveDeadValues.cpp     | 26 +++++++++++++++
 mlir/test/Transforms/remove-dead-values.mlir | 33 ++++++++++++++++++++
 2 files changed, 59 insertions(+)

diff --git a/mlir/lib/Transforms/RemoveDeadValues.cpp b/mlir/lib/Transforms/RemoveDeadValues.cpp
index 66f369e8a5f65..8b71d62221bd0 100644
--- a/mlir/lib/Transforms/RemoveDeadValues.cpp
+++ b/mlir/lib/Transforms/RemoveDeadValues.cpp
@@ -720,6 +720,32 @@ static void cleanUpDeadVals(MLIRContext *ctx, RDVFinalCleanupList &list) {
       // When erasing a terminator, insert an unreachable op in its place.
       ub::UnreachableOp::create(rewriter, op->getLoc());
     }
+
+    // Before erasing the operation, replace all result values with live-uses by
+    // ub.poison values. This is important to maintain IR validity. For example,
+    // if we have an op with one of its results used by another op, erasing the
+    // op without replacing its corresponding result would leave us with a
+    // dangling operand in the user op. By replacing the result with a ub.poison
+    // value, we ensure that the user op still has a valid operand, even though
+    // it's a poison value which will be cleaned up later if it can be cleaned
+    // up. This keeps the IR valid for further simplification and
+    // canonicalization.
+    auto opResults = op->getResults();
+    for (Value opResult : opResults) {
+      // Early continue for the case where the op result has no uses. No need to
+      // create a poison op here.
+      if (opResult.use_empty())
+        continue;
+
+      rewriter.setInsertionPoint(op);
+      Value poisonedValue = createPoisonedValues(rewriter, opResult).front();
+      auto resultUses = opResult.getUses();
+      for (OpOperand &use : llvm::make_early_inc_range(resultUses)) {
+        rewriter.setInsertionPoint(use.getOwner());
+        use.set(poisonedValue);
+      }
+    }
+
     op->dropAllUses();
     rewriter.eraseOp(op);
   }
diff --git a/mlir/test/Transforms/remove-dead-values.mlir b/mlir/test/Transforms/remove-dead-values.mlir
index ae83eac0c376f..8866a7ff8f1eb 100644
--- a/mlir/test/Transforms/remove-dead-values.mlir
+++ b/mlir/test/Transforms/remove-dead-values.mlir
@@ -796,3 +796,36 @@ func.func @scf_while_dead_iter_args() -> i32 {
   }
   return %result#0 : i32
 }
+
+// -----
+
+// CHECK-LABEL: func.func @replace_dead_operarion_result_with_poison
+func.func @replace_dead_operarion_result_with_poison() {
+  %0 = gpu.dynamic_shared_memory : memref<?xi8, #gpu.address_space<workgroup>>
+  %1 = vector.step : vector<1xindex>
+  %2 = scf.while (%arg0 = %1) : (vector<1xindex>) -> vector<1xindex> {
+    %cond = arith.constant true
+    scf.condition(%cond) %arg0 : vector<1xindex>
+  } do {
+    ^bb0(%arg0: vector<1xindex>):
+    scf.yield %arg0 : vector<1xindex>
+  }
+  %3 = scf.while (%arg0 = %2) : (vector<1xindex>) -> vector<1xindex> {
+    %cond = arith.constant true
+    // Check that the binary value in condition is replaced with poison, and
+    // the condition itself is well-formed IR. This prevents a crash in the
+    // canonicalization phase which happens after the dead value removal phase.
+    // CHECK-CANONICALIZE:      %[[VAL1:.*]] = ub.poison : i1
+    // CHECK-CANONICALIZE-NEXT: %[[VAL2:.*]] = ub.poison : vector<1xindex>
+    // CHECK-CANONICALIZE-NEXT: scf.condition(%[[VAL1]]) %[[VAL2]]
+    scf.condition(%cond) %arg0 : vector<1xindex>
+  } do {
+    ^bb0(%arg0: vector<1xindex>):
+    scf.yield %arg0 : vector<1xindex>
+  }
+  smt.solver(%3) : (vector<1xindex>) -> () {
+    ^bb0(%arg0: vector<1xindex>):
+    smt.yield
+  }
+  return
+}