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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-core

Author: Prathamesh Tagore (meshtag)

<details>
<summary>Changes</summary>

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.

Fixes https://github.com/llvm/llvm-project/issues/179944

---
Full diff: https://github.com/llvm/llvm-project/pull/181013.diff


2 Files Affected:

- (modified) mlir/lib/Transforms/RemoveDeadValues.cpp (+33) 
- (added) mlir/test/Transforms/remove-dead-values-invalid.mlir (+25) 


``````````diff
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..a85c52777ccf7
--- /dev/null
+++ b/mlir/test/Transforms/remove-dead-values-invalid.mlir
@@ -0,0 +1,25 @@
+// 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
+    // 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
+}

``````````

</details>


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