Add operands to worklist when only used by deleted op (#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.

Fixes #86765

Author: mlevesquedion

flang/test/Transforms/stack-arrays.fir

@@ -127,9 +127,7 @@ func.func @placement1() {
   return
 }
 // CHECK:      func.func @placement1() {
-// CHECK-NEXT:   %[[ONE:.*]] = arith.constant 1 : index
-// CHECK-NEXT:   %[[TWO:.*]] = arith.constant 2 : index
-// CHECK-NEXT:   %[[ARG:.*]] = arith.addi %[[ONE]], %[[TWO]] : index
+// CHECK-NEXT:   %[[ARG:.*]] = arith.constant 3 : index
 // CHECK-NEXT:   %[[MEM:.*]] = fir.alloca !fir.array<?xi32>, %[[ARG]]
 // CHECK-NEXT:   return
 // CHECK-NEXT: }
@@ -204,13 +202,12 @@ func.func @placement4(%arg0 : i1) {
 // CHECK:      func.func @placement4(%arg0: i1) {
 // CHECK-NEXT:   %[[C1:.*]] = arith.constant 1 : index
 // CHECK-NEXT:   %[[C1_I32:.*]] = fir.convert %[[C1]] : (index) -> i32
-// CHECK-NEXT:   %[[C2:.*]] = arith.constant 2 : index
 // CHECK-NEXT:   %[[C10:.*]] = arith.constant 10 : index
 // CHECK-NEXT:   cf.br ^bb1
 // CHECK-NEXT: ^bb1:
-// CHECK-NEXT:   %[[SUM:.*]] = arith.addi %[[C1]], %[[C2]] : index
+// CHECK-NEXT:   %[[C3:.*]] = arith.constant 3 : index
 // CHECK-NEXT:   %[[SP:.*]] = fir.call @llvm.stacksave.p0() : () -> !fir.ref<i8>
-// CHECK-NEXT:   %[[MEM:.*]] = fir.alloca !fir.array<?xi32>, %[[SUM]]
+// CHECK-NEXT:   %[[MEM:.*]] = fir.alloca !fir.array<?xi32>, %[[C3]]
 // CHECK-NEXT:   fir.call @llvm.stackrestore.p0(%[[SP]]) : (!fir.ref<i8>) -> ()
 // CHECK-NEXT:   cf.cond_br %arg0, ^bb1, ^bb2
 // CHECK-NEXT: ^bb2:

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,17 +688,36 @@ 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()))
+void GreedyPatternRewriteDriver::addOperandsToWorklist(Operation *op) {
+  for (Value operand : op->getOperands()) {
+    // If this operand currently has at most 2 users, add its defining op to the
+    // worklist. Indeed, after the op is deleted, then the operand will have at
+    // most 1 user left. If it has 0 users left, it can be deleted too,
+    // and if it has 1 user left, there may be further canonicalization
+    // opportunities.
+    if (!operand)
       continue;
-    if (auto *defOp = operand.getDefiningOp())
-      addToWorklist(defOp);
+
+    auto *defOp = operand.getDefiningOp();
+    if (!defOp)
+      continue;
+
+    Operation *otherUser = nullptr;
+    bool hasMoreThanTwoUses = false;
+    for (auto user : operand.getUsers()) {
+      if (user == op || user == otherUser)
+        continue;
+      if (!otherUser) {
+        otherUser = user;
+        continue;
+      }
+      hasMoreThanTwoUses = true;
+      break;
+    }
+    if (hasMoreThanTwoUses)
+      continue;
+
+    addToWorklist(defOp);
   }
 }
 
@@ -722,7 +741,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)

mlir/test/IR/greedy-pattern-rewriter-driver.mlir renamed to mlir/test/IR/greedy-pattern-rewrite-driver-bottom-up.mlir

@@ -0,0 +1,58 @@
+// RUN: mlir-opt %s -test-patterns="max-iterations=1 top-down=true" \
+// RUN:     --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
+}