[Mlir-commits] [mlir] Add operands to worklist when only used by deleted op (PR #86990)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Mar 28 13:29:03 PDT 2024
https://github.com/mlevesquedion updated https://github.com/llvm/llvm-project/pull/86990
>From 27500e1e09bcebed3413d0114b829d2acff45643 Mon Sep 17 00:00:00 2001
From: Michael Levesque-Dion <mlevesquedion at google.com>
Date: Thu, 28 Mar 2024 11:53:07 -0700
Subject: [PATCH 1/2] Add operands to worklist when only used by deleted op
Fixes #86765
---
.../Utils/GreedyPatternRewriteDriver.cpp | 24 ++++----
...edy-pattern-rewrite-driver-bottom-up.mlir} | 0
...reedy-pattern-rewrite-driver-top-down.mlir | 60 +++++++++++++++++++
3 files changed, 73 insertions(+), 11 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/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
index 6cb5635e68c922..2f2a15a0a28e51 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,15 +688,17 @@ 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()))
- continue;
+void GreedyPatternRewriteDriver::addOperandsToWorklist(Operation* op) {
+ for (Value operand : op->getOperands()) {
+ // If this operand was only used by the op under consideration, we re-add
+ // the operation that defined it to the worklist. Indeed, if the op is about
+ // to be deleted and it was the sole user of the operand, the operand may
+ // also be deleted.
+ // TODO: if the operand has a single use besides the op under consideration,
+ // there may be further canonicalization opportunities, so it should be
+ // added to the worklist.
+ if (!operand) continue;
+ if (!llvm::all_of(operand.getUsers(), [&op](auto u) { return u == op; })) continue;
if (auto *defOp = operand.getDefiningOp())
addToWorklist(defOp);
}
@@ -722,7 +724,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..5a0d58c62845bf
--- /dev/null
+++ b/mlir/test/IR/greedy-pattern-rewrite-driver-top-down.mlir
@@ -0,0 +1,60 @@
+// RUN: mlir-opt %s -test-patterns="max-iterations=1 top-down=true" \
+// RUN: -allow-unregistered-dialect --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
+}
>From e2fd08df237d418e6c5b431b6cb0f05380957a7a Mon Sep 17 00:00:00 2001
From: Michael Levesque-Dion <mlevesquedion at google.com>
Date: Thu, 28 Mar 2024 13:28:31 -0700
Subject: [PATCH 2/2] Fix formatting
---
.../Transforms/Utils/GreedyPatternRewriteDriver.cpp | 10 ++++++----
1 file changed, 6 insertions(+), 4 deletions(-)
diff --git a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
index 2f2a15a0a28e51..e7a94fab82af43 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(Operation* op);
+ void addOperandsToWorklist(Operation *op);
/// Notify the driver that the given block was inserted.
void notifyBlockInserted(Block *block, Region *previous,
@@ -688,7 +688,7 @@ void GreedyPatternRewriteDriver::notifyOperationModified(Operation *op) {
addToWorklist(op);
}
-void GreedyPatternRewriteDriver::addOperandsToWorklist(Operation* op) {
+void GreedyPatternRewriteDriver::addOperandsToWorklist(Operation *op) {
for (Value operand : op->getOperands()) {
// If this operand was only used by the op under consideration, we re-add
// the operation that defined it to the worklist. Indeed, if the op is about
@@ -697,8 +697,10 @@ void GreedyPatternRewriteDriver::addOperandsToWorklist(Operation* op) {
// TODO: if the operand has a single use besides the op under consideration,
// there may be further canonicalization opportunities, so it should be
// added to the worklist.
- if (!operand) continue;
- if (!llvm::all_of(operand.getUsers(), [&op](auto u) { return u == op; })) continue;
+ if (!operand)
+ continue;
+ if (!llvm::all_of(operand.getUsers(), [&op](auto u) { return u == op; }))
+ continue;
if (auto *defOp = operand.getDefiningOp())
addToWorklist(defOp);
}
More information about the Mlir-commits
mailing list