[Mlir-commits] [mlir] [mlir][Transforms] Fix crash in `-remove-dead-values` on private functions (PR #169269)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-core

Author: Matthias Springer (matthias-springer)

<details>
<summary>Changes</summary>

This commit fixes a crash in the `-remove-dead-values` pass. This pass removes dead block arguments from functions. However, it did not remove side-effecting operations, even if they are dead. This lead to invalid IR: a block argument was removed, but some of its uses survived.

With this commit, a "simple" operation is erased when it has a dead operand. If the operation were not dead, the liveness analysis would not have marked one of its operands as "dead".


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


2 Files Affected:

- (modified) mlir/lib/Transforms/RemoveDeadValues.cpp (+38) 
- (modified) mlir/test/Transforms/remove-dead-values.mlir (+11) 


``````````diff
diff --git a/mlir/lib/Transforms/RemoveDeadValues.cpp b/mlir/lib/Transforms/RemoveDeadValues.cpp
index 989c614ef6617..9d4d24c39c116 100644
--- a/mlir/lib/Transforms/RemoveDeadValues.cpp
+++ b/mlir/lib/Transforms/RemoveDeadValues.cpp
@@ -141,6 +141,33 @@ static bool hasLive(ValueRange values, const DenseSet<Value> &nonLiveSet,
   return false;
 }
 
+/// Return true iff at least one value in `values` is dead, given the liveness
+/// information in `la`.
+static bool hasDead(ValueRange values, const DenseSet<Value> &nonLiveSet,
+                    RunLivenessAnalysis &la) {
+  for (Value value : values) {
+    if (nonLiveSet.contains(value)) {
+      LDBG() << "Value " << value << " is already marked non-live (dead)";
+      return true;
+    }
+
+    const Liveness *liveness = la.getLiveness(value);
+    if (!liveness) {
+      LDBG() << "Value " << value
+             << " has no liveness info, conservatively considered live";
+      continue;
+    }
+    if (liveness->isLive) {
+      LDBG() << "Value " << value << " is live according to liveness analysis";
+      continue;
+    } else {
+      LDBG() << "Value " << value << " is dead according to liveness analysis";
+      return true;
+    }
+  }
+  return false;
+}
+
 /// Return a BitVector of size `values.size()` where its i-th bit is 1 iff the
 /// i-th value in `values` is live, given the liveness information in `la`.
 static BitVector markLives(ValueRange values, const DenseSet<Value> &nonLiveSet,
@@ -260,6 +287,17 @@ static SmallVector<OpOperand *> operandsToOpOperands(OperandRange operands) {
 static void processSimpleOp(Operation *op, RunLivenessAnalysis &la,
                             DenseSet<Value> &nonLiveSet,
                             RDVFinalCleanupList &cl) {
+  if (hasDead(op->getOperands(), nonLiveSet, la)) {
+    LDBG() << "Simple op has dead operands, so the op must be dead: "
+           << OpWithFlags(op, OpPrintingFlags().skipRegions());
+    assert(!hasLive(op->getResults(), nonLiveSet, la) &&
+           "expected the op to have no live results");
+    cl.operations.push_back(op);
+    collectNonLiveValues(nonLiveSet, op->getResults(),
+                         BitVector(op->getNumResults(), true));
+    return;
+  }
+
   if (!isMemoryEffectFree(op) || hasLive(op->getResults(), nonLiveSet, la)) {
     LDBG() << "Simple op is not memory effect free or has live results, "
               "preserving it: "
diff --git a/mlir/test/Transforms/remove-dead-values.mlir b/mlir/test/Transforms/remove-dead-values.mlir
index 4bae85dcf4f7d..af157fc8bc5b0 100644
--- a/mlir/test/Transforms/remove-dead-values.mlir
+++ b/mlir/test/Transforms/remove-dead-values.mlir
@@ -118,6 +118,17 @@ func.func @main(%arg0 : i32) {
 
 // -----
 
+// CHECK-LABEL: func.func private @clean_func_op_remove_side_effecting_op() {
+// CHECK-NEXT:    return
+// CHECK-NEXT:  }
+func.func private @clean_func_op_remove_side_effecting_op(%arg0: i32) -> (i32) {
+  // vector.print has a side effect but the op is dead.
+  vector.print %arg0 : i32
+  return %arg0 : i32
+}
+
+// -----
+
 // %arg0 is not live because it is never used. %arg1 is not live because its
 // user `arith.addi` doesn't have any uses and the value that it is forwarded to
 // (%non_live_0) also doesn't have any uses.

``````````

</details>


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