From ddc9892999a42f9fd069f9786b82cb63312eba63 Mon Sep 17 00:00:00 2001 From: mlevesquedion Date: Fri, 29 Mar 2024 13:38:41 -0700 Subject: [PATCH] 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 --- flang/test/Transforms/stack-arrays.fir | 9 +-- .../Utils/GreedyPatternRewriteDriver.cpp | 43 ++++++++++---- ...edy-pattern-rewrite-driver-bottom-up.mlir} | 0 ...reedy-pattern-rewrite-driver-top-down.mlir | 58 +++++++++++++++++++ 4 files changed, 92 insertions(+), 18 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/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 +}