Add operands to worklist when only used by deleted op

[mlevesquedion]

I believe the existing check to determine if an operand should be added is incorrect: operand.use_empty() || operand.hasOneUse(). This is because these checks do not take into account the fact that the op is being deleted. It hasn't been deleted yet, so operand.use_empty() cannot be true, and operand.hasOneUse() may be true if the op being deleted is the only user of the operand and it only uses it once, but it will fail if the operand is used more than once (e.g. something like add %0, %0).

Instead, check if the op being deleted is the only user of the operand. If so, add the operand to the worklist.

There is a TODO comment above the check and I have tried to preserve its spirit. If the operand will have a single use after the op is deleted, the operand may need to be added to the worklist because there may be further canonicalization opportunities for it.

Fixes #86765

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-core

Author: mlevesquedion (mlevesquedion)

Changes

I believe the existing check to determine if an operand should be added is incorrect: operand.use_empty() || operand.hasOneUse(). This is because these checks do not take into account the fact that the op is being deleted. It hasn't been deleted yet, so operand.use_empty() cannot be true, and operand.hasOneUse() may be true if the op being deleted is the only user of the operand and it only uses it once, but it will fail if the operand is used more than once (e.g. something like add %0, %0).

Instead, check if the op being deleted is the only user of the operand. If so, add the operand to the worklist.

There is a TODO comment above the check and I have tried to preserve its spirit. If the operand will have a single use after the op is deleted, the operand may need to be added to the worklist because there may be further canonicalization opportunities for it.

Fixes #86765

---

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

3 Files Affected:

• (modified) mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp (+13-11)
• (renamed) mlir/test/IR/greedy-pattern-rewrite-driver-bottom-up.mlir ()
• (added) mlir/test/IR/greedy-pattern-rewrite-driver-top-down.mlir (+60)

diff --git a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
index 6cb5635e68c922..2f2a15a0a28e51 100644
--- a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
+++ b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
@@ -377,7 +377,7 @@ class GreedyPatternRewriteDriver : public PatternRewriter,
   /// be re-added to the worklist. This function should be called when an
   /// operation is modified or removed, as it may trigger further
   /// simplifications.
-  void addOperandsToWorklist(ValueRange operands);
+  void addOperandsToWorklist(Operation* op);
 
   /// Notify the driver that the given block was inserted.
   void notifyBlockInserted(Block *block, Region *previous,
@@ -688,15 +688,17 @@ void GreedyPatternRewriteDriver::notifyOperationModified(Operation *op) {
   addToWorklist(op);
 }
 
-void GreedyPatternRewriteDriver::addOperandsToWorklist(ValueRange operands) {
-  for (Value operand : operands) {
-    // If the use count of this operand is now < 2, we re-add the defining
-    // operation to the worklist.
-    // TODO: This is based on the fact that zero use operations
-    // may be deleted, and that single use values often have more
-    // canonicalization opportunities.
-    if (!operand || (!operand.use_empty() && !operand.hasOneUse()))
-      continue;
+void GreedyPatternRewriteDriver::addOperandsToWorklist(Operation* op) {
+  for (Value operand : op->getOperands()) {
+    // If this operand was only used by the op under consideration, we re-add
+    // the operation that defined it to the worklist. Indeed, if the op is about
+    // to be deleted and it was the sole user of the operand, the operand may
+    // also be deleted.
+    // TODO: if the operand has a single use besides the op under consideration,
+    // there may be further canonicalization opportunities, so it should be
+    // added to the worklist.
+    if (!operand) continue;
+    if (!llvm::all_of(operand.getUsers(), [&op](auto u) { return u == op; })) continue;
     if (auto *defOp = operand.getDefiningOp())
       addToWorklist(defOp);
   }
@@ -722,7 +724,7 @@ void GreedyPatternRewriteDriver::notifyOperationErased(Operation *op) {
   if (config.listener)
     config.listener->notifyOperationErased(op);
 
-  addOperandsToWorklist(op->getOperands());
+  addOperandsToWorklist(op);
   worklist.remove(op);
 
   if (config.strictMode != GreedyRewriteStrictness::AnyOp)
diff --git a/mlir/test/IR/greedy-pattern-rewriter-driver.mlir b/mlir/test/IR/greedy-pattern-rewrite-driver-bottom-up.mlir
similarity index 100%
rename from mlir/test/IR/greedy-pattern-rewriter-driver.mlir
rename to mlir/test/IR/greedy-pattern-rewrite-driver-bottom-up.mlir
diff --git a/mlir/test/IR/greedy-pattern-rewrite-driver-top-down.mlir b/mlir/test/IR/greedy-pattern-rewrite-driver-top-down.mlir
new file mode 100644
index 00000000000000..5a0d58c62845bf
--- /dev/null
+++ b/mlir/test/IR/greedy-pattern-rewrite-driver-top-down.mlir
@@ -0,0 +1,60 @@
+// RUN: mlir-opt %s -test-patterns="max-iterations=1 top-down=true" \
+// RUN:     -allow-unregistered-dialect --split-input-file | FileCheck %s
+
+// Tests for https://github.com/llvm/llvm-project/issues/86765. Ensure
+// that operands of a dead op are added to the worklist even if the same value
+// appears multiple times as an operand.
+
+// -----
+
+// 2 uses of the same operand
+
+// CHECK:       func.func @f(%arg0: i1) {
+// CHECK-NEXT:    return
+// CHECK-NEXT:  }
+func.func @f(%arg0: i1) {
+  %0 = arith.constant 0 : i32
+  %if = scf.if %arg0 -> (i32) {
+    scf.yield %0 : i32
+  } else {
+    scf.yield %0 : i32
+  }
+  %dead_leaf = arith.addi %if, %if : i32
+  return
+}
+
+// -----
+
+// 3 uses of the same operand
+
+// CHECK:       func.func @f() {
+// CHECK-NEXT:    return
+// CHECK-NEXT:  }
+func.func @f() {
+  %0 = arith.constant 0 : i1
+  %if = scf.if %0 -> (i1) {
+    scf.yield %0 : i1
+  } else {
+    scf.yield %0 : i1
+  }
+  %dead_leaf = arith.select %if, %if, %if : i1
+  return
+}
+
+// -----
+
+// 2 uses of the same operand, op has 3 operands
+
+// CHECK:       func.func @f() {
+// CHECK-NEXT:    return
+// CHECK-NEXT:  }
+func.func @f() {
+  %0 = arith.constant 0 : i1
+  %if = scf.if %0 -> (i1) {
+    scf.yield %0 : i1
+  } else {
+    scf.yield %0 : i1
+  }
+  %dead_leaf = arith.select %0, %if, %if : i1
+  return
+}

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[zero9178]

LGTM. Adding some other reviewers just in case for visibility

[mlevesquedion]

Looks like this caused a few test cases to fail under flang because some canonicalizations are happening that used to not happen (some additions): https://buildkite.com/llvm-project/github-pull-requests/builds/51398#018e8759-932d-4e11-8f10-052eb29b7346

I fixed the test cases.

[mlevesquedion]

Thanks for the reviews, everyone! :) I don't have commit access so feel free to merge if you can, unless we want @jpienaar to review also.

[jpienaar]

Thanks for the reviews, everyone! :) I don't have commit access so feel free to merge if you can, unless we want @jpienaar to review also.

All good, all 3 the reviewers know this section well (and I was on vacation :)).