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, %[[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 -// CHECK-NEXT: %[[MEM:.*]] = fir.alloca !fir.array, %[[SUM]] +// CHECK-NEXT: %[[MEM:.*]] = fir.alloca !fir.array, %[[C3]] // CHECK-NEXT: fir.call @llvm.stackrestore.p0(%[[SP]]) : (!fir.ref) -> () // 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 +}