[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 83e1fa92f3d6742ce755da8f0dc3cd4c4aed5d1c 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 operation results with
 live-uses by poison before erasure

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..761c14bb43927 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 return 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
+}