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

[Matthias Springer via flang-commits]

================
@@ -688,17 +688,32 @@ 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;
+    if (!llvm::all_of(operand.getUsers(), [&](Operation *user) {
----------------
matthias-springer wrote:

nit: I find the `!llvm::all_of` a bit difficult to read. How about spelling out the loop? It could also be a bit more efficient because you can break out of the loop early.

```c++
Operation *otherUser = nullptr;
bool hasMoreThanTwoUsers = false;
for (Operation *user : operand.getUsers()) {
  if (user == op || user == otherUser)
    continue;
  if (!otherUser) {
    otherUser = user;
    continue;
  }
  hasMoreThanTwoUsers = true;
  break;
}
if (hasMoreThanTwoUsers)
  continue;
```


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