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

via flang-commits flang-commits at lists.llvm.org
Fri Mar 29 13:38:44 PDT 2024


Author: mlevesquedion
Date: 2024-03-29T21:38:41+01:00
New Revision: ddc9892999a42f9fd069f9786b82cb63312eba63

URL: https://github.com/llvm/llvm-project/commit/ddc9892999a42f9fd069f9786b82cb63312eba63
DIFF: https://github.com/llvm/llvm-project/commit/ddc9892999a42f9fd069f9786b82cb63312eba63.diff

LOG: 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

Added: 
    mlir/test/IR/greedy-pattern-rewrite-driver-bottom-up.mlir
    mlir/test/IR/greedy-pattern-rewrite-driver-top-down.mlir

Modified: 
    flang/test/Transforms/stack-arrays.fir
    mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

Removed: 
    mlir/test/IR/greedy-pattern-rewriter-driver.mlir


################################################################################
diff  --git a/flang/test/Transforms/stack-arrays.fir b/flang/test/Transforms/stack-arrays.fir
index f4fe737e88d785..a2ffe555091ebc 100644
--- a/flang/test/Transforms/stack-arrays.fir
+++ b/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:

diff  --git a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
index 6cb5635e68c922..bbecbdb8566935 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,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)

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..a362d6f99b9478
--- /dev/null
+++ b/mlir/test/IR/greedy-pattern-rewrite-driver-top-down.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
+}


        


More information about the flang-commits mailing list