[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 d07c58cce82474425d4c06359b52ccb57ada0a71 Mon Sep 17 00:00:00 2001
From: Prathamesh Tagore <prathameshtagore at gmail.com>
Date: Wed, 11 Feb 2026 13:33:00 +0100
Subject: [PATCH] [mlir][remove-dead-values] Fix crash due to dangling operands

The commit fixes a crash in the RemoveDeadValues pass that occurs when an
operation has a side effect and we erase some of its operands. In this case,
we end up with an operation that has some null operands, which are considered
dangling. The fix involves checking for dangling operands after the cleanup
phase and skipping canonicalization if any are found, since canonicalization
patterns expect a well-formed IR.
Note: We are not skipping any pass-specific modifications since
canonicalization happens separately after all dead values are removed by the
pass.
---
 mlir/lib/Transforms/RemoveDeadValues.cpp      | 33 +++++++++++++++++++
 .../remove-dead-values-invalid.mlir           | 27 +++++++++++++++
 2 files changed, 60 insertions(+)
 create mode 100644 mlir/test/Transforms/remove-dead-values-invalid.mlir

diff --git a/mlir/lib/Transforms/RemoveDeadValues.cpp b/mlir/lib/Transforms/RemoveDeadValues.cpp
index 66f369e8a5f65..b80318140edd9 100644
--- a/mlir/lib/Transforms/RemoveDeadValues.cpp
+++ b/mlir/lib/Transforms/RemoveDeadValues.cpp
@@ -733,6 +733,27 @@ static void cleanUpDeadVals(MLIRContext *ctx, RDVFinalCleanupList &list) {
   LDBG() << "Finished cleanup of dead values";
 }
 
+/// Return true if the operation has any null operands.
+static bool hasNullOperands(Operation *operation) {
+  MutableArrayRef<OpOperand> operands = operation->getOpOperands();
+  for (OpOperand &operand : operands)
+    if (!operand.get())
+      return true;
+  return false;
+}
+
+/// Return true if any operation in the module has a dangling operand.
+/// We consider an operand dangling if it is null. This can happen when an
+/// operand is erased without erasing the operation that consumes it.
+static bool hasDanglingOperands(Operation *operation) {
+  WalkResult walk = operation->walk([&](Operation *op) {
+    if (hasNullOperands(op))
+      return WalkResult::interrupt();
+    return WalkResult::advance();
+  });
+  return walk.wasInterrupted();
+}
+
 struct RemoveDeadValues : public impl::RemoveDeadValuesBase<RemoveDeadValues> {
   void runOnOperation() override;
 };
@@ -774,6 +795,18 @@ void RemoveDeadValues::runOnOperation() {
   if (!canonicalize)
     return;
 
+  // If there are dangling operands, it means that some operations were not
+  // fully erased. This can happen when an operation has a side effect and we
+  // don't erase it, but we still erase some of its operands. In this case, we
+  // end up with an operation that has some null operands, which are considered
+  // dangling. We can't run canonicalization patterns in this state because they
+  // expect a well-formed IR, hence we skip canonicalization.
+  // Note: We are not skipping any pass-specific modifications since
+  // canonicalization happens separately after all dead values are removed by
+  // the pass.
+  if (hasDanglingOperands(module))
+    return;
+
   // Canonicalize all region branch ops.
   SmallVector<Operation *> opsToCanonicalize;
   module->walk([&](RegionBranchOpInterface regionBranchOp) {
diff --git a/mlir/test/Transforms/remove-dead-values-invalid.mlir b/mlir/test/Transforms/remove-dead-values-invalid.mlir
new file mode 100644
index 0000000000000..bd25a51709087
--- /dev/null
+++ b/mlir/test/Transforms/remove-dead-values-invalid.mlir
@@ -0,0 +1,27 @@
+// RUN: mlir-opt %s --remove-dead-values -verify-diagnostics
+
+func.func @main() {
+  %0 = vector.step : vector<1xindex>
+  %1 = scf.while (%arg0 = %0) : (vector<1xindex>) -> vector<1xindex> {
+    %cond = arith.constant true
+    scf.condition(%cond) %arg0 : vector<1xindex>
+  } do {
+  ^bb0(%arg0: vector<1xindex>):
+    scf.yield %arg0 : vector<1xindex>
+  }
+  %2 = scf.while (%arg0 = %1) : (vector<1xindex>) -> vector<1xindex> {
+    %cond = arith.constant true
+    // The pass should not crash during canonicalization because of
+    // the null operand here.
+    // expected-error @+1 {{null operand found}}
+    scf.condition(%cond) %arg0 : vector<1xindex>
+  } do {
+  ^bb0(%arg0: vector<1xindex>):
+    scf.yield %arg0 : vector<1xindex>
+  }
+  smt.solver(%2) : (vector<1xindex>) -> () {
+  ^bb0(%arg0: vector<1xindex>):
+    smt.yield
+  }
+  return
+}