[Mlir-commits] [mlir] Add operands to worklist when only used by deleted op (PR #86990)

[llvmlistbot at llvm.org]

https://github.com/mlevesquedion created https://github.com/llvm/llvm-project/pull/86990

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

>From 27500e1e09bcebed3413d0114b829d2acff45643 Mon Sep 17 00:00:00 2001
From: Michael Levesque-Dion <mlevesquedion at google.com>
Date: Thu, 28 Mar 2024 11:53:07 -0700
Subject: [PATCH] Add operands to worklist when only used by deleted op

Fixes #86765
---
 .../Utils/GreedyPatternRewriteDriver.cpp      | 24 ++++----
 ...edy-pattern-rewrite-driver-bottom-up.mlir} |  0
 ...reedy-pattern-rewrite-driver-top-down.mlir | 60 +++++++++++++++++++
 3 files changed, 73 insertions(+), 11 deletions(-)
 rename mlir/test/IR/{greedy-pattern-rewriter-driver.mlir => greedy-pattern-rewrite-driver-bottom-up.mlir} (100%)
 create mode 100644 mlir/test/IR/greedy-pattern-rewrite-driver-top-down.mlir

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
+}