[Mlir-commits] [mlir] [mlir][IR] Send missing notifications when inlining a block (PR #79593)
Matthias Springer
llvmlistbot at llvm.org
Fri Jan 26 05:15:07 PST 2024
https://github.com/matthias-springer created https://github.com/llvm/llvm-project/pull/79593
When a block is inlined into another block, the nested operations are moved into another block and the `notifyOperationInserted` callback should be triggered. This commit adds the missing notifications for:
* `RewriterBase::inlineBlockBefore`
* `RewriterBase::mergeBlocks`
Depends on #79579. Review only the top commit.
>From bd20abf8e5ba6c429bd5a08ca175070d9f588d77 Mon Sep 17 00:00:00 2001
From: Matthias Springer <springerm at google.com>
Date: Fri, 26 Jan 2024 11:52:44 +0000
Subject: [PATCH 1/2] [mlir][IR] Add `RewriterBase::moveBlockBefore` and fix
bug in `moveOpBefore`
This commit adds a new method to the rewriter API: `moveBlockBefore`. This op is utilized by `inlineRegionBefore` and covered by dialect conversion test cases.
Also fixes a bug in `moveOpBefore`, where the previous op location was not passed correctly. Adds a test case to `test-strict-pattern-driver.mlir`.
---
mlir/include/mlir/IR/Block.h | 4 +++
mlir/include/mlir/IR/PatternMatch.h | 7 ++++
mlir/lib/IR/Block.cpp | 9 +++--
mlir/lib/IR/PatternMatch.cpp | 34 +++++++++++-------
.../test-strict-pattern-driver.mlir | 20 +++++++++++
mlir/test/lib/Dialect/Test/TestPatterns.cpp | 35 +++++++++++++++++--
6 files changed, 92 insertions(+), 17 deletions(-)
diff --git a/mlir/include/mlir/IR/Block.h b/mlir/include/mlir/IR/Block.h
index 4139dcaeea81bb9..c14e9aad8f6d1e2 100644
--- a/mlir/include/mlir/IR/Block.h
+++ b/mlir/include/mlir/IR/Block.h
@@ -67,6 +67,10 @@ class Block : public IRObjectWithUseList<BlockOperand>,
/// specific block.
void moveBefore(Block *block);
+ /// Unlink this block from its current region and insert it right before the
+ /// block that the given iterator points to in the region region.
+ void moveBefore(Region *region, llvm::iplist<Block>::iterator iterator);
+
/// Unlink this Block from its parent region and delete it.
void erase();
diff --git a/mlir/include/mlir/IR/PatternMatch.h b/mlir/include/mlir/IR/PatternMatch.h
index 8eb129206b95ef6..72cbc85e9081d7e 100644
--- a/mlir/include/mlir/IR/PatternMatch.h
+++ b/mlir/include/mlir/IR/PatternMatch.h
@@ -614,6 +614,13 @@ class RewriterBase : public OpBuilder {
virtual void moveOpAfter(Operation *op, Block *block,
Block::iterator iterator);
+ /// Unlink this block and insert it right before `existingBlock`.
+ void moveBlockBefore(Block *block, Block *anotherBlock);
+
+ /// Unlink this block and insert it right before the location that the given
+ /// iterator points to in the given region.
+ void moveBlockBefore(Block *block, Region *region, Region::iterator iterator);
+
/// This method is used to notify the rewriter that an in-place operation
/// modification is about to happen. A call to this function *must* be
/// followed by a call to either `finalizeOpModification` or
diff --git a/mlir/lib/IR/Block.cpp b/mlir/lib/IR/Block.cpp
index 82ea303cf0171f3..65099f8ff15a6f7 100644
--- a/mlir/lib/IR/Block.cpp
+++ b/mlir/lib/IR/Block.cpp
@@ -52,8 +52,13 @@ void Block::insertAfter(Block *block) {
/// specific block.
void Block::moveBefore(Block *block) {
assert(block->getParent() && "cannot insert before a block without a parent");
- block->getParent()->getBlocks().splice(
- block->getIterator(), getParent()->getBlocks(), getIterator());
+ moveBefore(block->getParent(), block->getIterator());
+}
+
+/// Unlink this block from its current region and insert it right before the
+/// block that the given iterator points to in the region region.
+void Block::moveBefore(Region *region, llvm::iplist<Block>::iterator iterator) {
+ region->getBlocks().splice(iterator, getParent()->getBlocks(), getIterator());
}
/// Unlink this Block from its parent Region and delete it.
diff --git a/mlir/lib/IR/PatternMatch.cpp b/mlir/lib/IR/PatternMatch.cpp
index 817bbb363e0d585..ee285f094cdd1d3 100644
--- a/mlir/lib/IR/PatternMatch.cpp
+++ b/mlir/lib/IR/PatternMatch.cpp
@@ -350,11 +350,8 @@ void RewriterBase::inlineRegionBefore(Region ®ion, Region &parent,
}
// Move blocks from the beginning of the region one-by-one.
- while (!region.empty()) {
- Block *block = ®ion.front();
- parent.getBlocks().splice(before, region.getBlocks(), block->getIterator());
- listener->notifyBlockInserted(block, ®ion, region.begin());
- }
+ while (!region.empty())
+ moveBlockBefore(®ion.front(), &parent, before);
}
void RewriterBase::inlineRegionBefore(Region ®ion, Block *before) {
inlineRegionBefore(region, *before->getParent(), before->getIterator());
@@ -378,6 +375,21 @@ void RewriterBase::cloneRegionBefore(Region ®ion, Block *before) {
cloneRegionBefore(region, *before->getParent(), before->getIterator());
}
+void RewriterBase::moveBlockBefore(Block *block, Block *anotherBlock) {
+ moveBlockBefore(block, anotherBlock->getParent(),
+ anotherBlock->getIterator());
+}
+
+void RewriterBase::moveBlockBefore(Block *block, Region *region,
+ Region::iterator iterator) {
+ Region *currentRegion = block->getParent();
+ Region::iterator nextIterator = std::next(block->getIterator());
+ block->moveBefore(region, iterator);
+ if (listener)
+ listener->notifyBlockInserted(block, /*previous=*/currentRegion,
+ /*previousIt=*/nextIterator);
+}
+
void RewriterBase::moveOpBefore(Operation *op, Operation *existingOp) {
moveOpBefore(op, existingOp->getBlock(), existingOp->getIterator());
}
@@ -385,11 +397,11 @@ void RewriterBase::moveOpBefore(Operation *op, Operation *existingOp) {
void RewriterBase::moveOpBefore(Operation *op, Block *block,
Block::iterator iterator) {
Block *currentBlock = op->getBlock();
- Block::iterator currentIterator = op->getIterator();
+ Block::iterator nextIterator = std::next(op->getIterator());
op->moveBefore(block, iterator);
if (listener)
listener->notifyOperationInserted(
- op, /*previous=*/InsertPoint(currentBlock, currentIterator));
+ op, /*previous=*/InsertPoint(currentBlock, nextIterator));
}
void RewriterBase::moveOpAfter(Operation *op, Operation *existingOp) {
@@ -398,10 +410,6 @@ void RewriterBase::moveOpAfter(Operation *op, Operation *existingOp) {
void RewriterBase::moveOpAfter(Operation *op, Block *block,
Block::iterator iterator) {
- Block *currentBlock = op->getBlock();
- Block::iterator currentIterator = op->getIterator();
- op->moveAfter(block, iterator);
- if (listener)
- listener->notifyOperationInserted(
- op, /*previous=*/InsertPoint(currentBlock, currentIterator));
+ assert(iterator != block->end() && "cannot move after end of block");
+ moveOpBefore(op, block, std::next(iterator));
}
diff --git a/mlir/test/Transforms/test-strict-pattern-driver.mlir b/mlir/test/Transforms/test-strict-pattern-driver.mlir
index a5ab8f97c74ce33..2ebc66b4f26d0e8 100644
--- a/mlir/test/Transforms/test-strict-pattern-driver.mlir
+++ b/mlir/test/Transforms/test-strict-pattern-driver.mlir
@@ -24,6 +24,7 @@ func.func @test_erase() {
// -----
+// CHECK-EN: notifyOperationInserted: test.insert_same_op, was unlinked
// CHECK-EN-LABEL: func @test_insert_same_op
// CHECK-EN-SAME: {pattern_driver_all_erased = false, pattern_driver_changed = true}
// CHECK-EN: "test.insert_same_op"() {skip = true}
@@ -35,6 +36,7 @@ func.func @test_insert_same_op() {
// -----
+// CHECK-EN: notifyOperationInserted: test.new_op, was unlinked
// CHECK-EN-LABEL: func @test_replace_with_new_op
// CHECK-EN-SAME: {pattern_driver_all_erased = true, pattern_driver_changed = true}
// CHECK-EN: %[[n:.*]] = "test.new_op"
@@ -49,6 +51,9 @@ func.func @test_replace_with_new_op() {
// -----
+// CHECK-EN: notifyOperationInserted: test.erase_op, was unlinked
+// CHECK-EN: notifyOperationRemoved: test.replace_with_new_op
+// CHECK-EN: notifyOperationRemoved: test.erase_op
// CHECK-EN-LABEL: func @test_replace_with_erase_op
// CHECK-EN-SAME: {pattern_driver_all_erased = true, pattern_driver_changed = true}
// CHECK-EN-NOT: "test.replace_with_new_op"
@@ -229,3 +234,18 @@ func.func @test_remove_diamond(%c: i1) {
}) : () -> ()
return
}
+
+// -----
+
+// CHECK-AN: notifyOperationInserted: test.move_before_parent_op, previous = test.dummy_terminator
+// CHECK-AN-LABEL: func @test_move_op_before(
+// CHECK-AN: test.move_before_parent_op
+// CHECK-AN: test.op_with_region
+// CHECK-AN: test.dummy_terminator
+func.func @test_move_op_before() {
+ "test.op_with_region"() ({
+ "test.move_before_parent_op"() : () -> ()
+ "test.dummy_terminator"() : () ->()
+ }) : () -> ()
+ return
+}
diff --git a/mlir/test/lib/Dialect/Test/TestPatterns.cpp b/mlir/test/lib/Dialect/Test/TestPatterns.cpp
index 89b9d1ce78a52b6..2165b388d559c5e 100644
--- a/mlir/test/lib/Dialect/Test/TestPatterns.cpp
+++ b/mlir/test/lib/Dialect/Test/TestPatterns.cpp
@@ -198,6 +198,21 @@ struct HoistEligibleOps : public OpRewritePattern<test::OneRegionOp> {
}
};
+/// This pattern moves "test.move_before_parent_op" before the parent op.
+struct MoveBeforeParentOp : public RewritePattern {
+ MoveBeforeParentOp(MLIRContext *context)
+ : RewritePattern("test.move_before_parent_op", /*benefit=*/1, context) {}
+
+ LogicalResult matchAndRewrite(Operation *op,
+ PatternRewriter &rewriter) const override {
+ // Do not hoist past functions.
+ if (isa<FunctionOpInterface>(op->getParentOp()))
+ return failure();
+ rewriter.moveOpBefore(op, op->getParentOp());
+ return success();
+ }
+};
+
struct TestPatternDriver
: public PassWrapper<TestPatternDriver, OperationPass<>> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestPatternDriver)
@@ -238,6 +253,20 @@ struct TestPatternDriver
};
struct DumpNotifications : public RewriterBase::Listener {
+ void notifyOperationInserted(Operation *op,
+ OpBuilder::InsertPoint previous) override {
+ llvm::outs() << "notifyOperationInserted: " << op->getName();
+ if (!previous.isSet()) {
+ llvm::outs() << ", was unlinked\n";
+ } else {
+ if (previous.getPoint() == previous.getBlock()->end()) {
+ llvm::outs() << ", was last in block\n";
+ } else {
+ llvm::outs() << ", previous = " << previous.getPoint()->getName()
+ << "\n";
+ }
+ }
+ }
void notifyOperationRemoved(Operation *op) override {
llvm::outs() << "notifyOperationRemoved: " << op->getName() << "\n";
}
@@ -267,14 +296,16 @@ struct TestStrictPatternDriver
ReplaceWithNewOp,
EraseOp,
ChangeBlockOp,
- ImplicitChangeOp
+ ImplicitChangeOp,
+ MoveBeforeParentOp
// clang-format on
>(ctx);
SmallVector<Operation *> ops;
getOperation()->walk([&](Operation *op) {
StringRef opName = op->getName().getStringRef();
if (opName == "test.insert_same_op" || opName == "test.change_block_op" ||
- opName == "test.replace_with_new_op" || opName == "test.erase_op") {
+ opName == "test.replace_with_new_op" || opName == "test.erase_op" ||
+ opName == "test.move_before_parent_op") {
ops.push_back(op);
}
});
>From b682d0fd6b79e61720edbb40dcf872cf90f6c4d2 Mon Sep 17 00:00:00 2001
From: Matthias Springer <springerm at google.com>
Date: Fri, 26 Jan 2024 13:12:55 +0000
Subject: [PATCH 2/2] [mlir][IR] Send missing notifications when inlining a
block
When a block is inlined into another block, the operations are moved into another block and the `notifyOperationInserted` callback should be triggered. This commit adds the missing notifications for:
* `RewriterBase::inlineBlockBefore`
* `RewriterBase::mergeBlocks`
---
mlir/lib/IR/PatternMatch.cpp | 11 +++++++-
.../Dialect/Affine/simplify-structures.mlir | 2 --
.../test-strict-pattern-driver.mlir | 20 ++++++++++++++
mlir/test/lib/Dialect/Test/TestPatterns.cpp | 26 +++++++++++++++++--
4 files changed, 54 insertions(+), 5 deletions(-)
diff --git a/mlir/lib/IR/PatternMatch.cpp b/mlir/lib/IR/PatternMatch.cpp
index ee285f094cdd1d3..0bf79f5817b5af4 100644
--- a/mlir/lib/IR/PatternMatch.cpp
+++ b/mlir/lib/IR/PatternMatch.cpp
@@ -317,7 +317,16 @@ void RewriterBase::inlineBlockBefore(Block *source, Block *dest,
// Move operations from the source block to the dest block and erase the
// source block.
- dest->getOperations().splice(before, source->getOperations());
+ if (!listener) {
+ // Fast path: If no listener is attached, move all operations at once.
+ dest->getOperations().splice(before, source->getOperations());
+ } else {
+ while (!source->empty())
+ moveOpBefore(&source->front(), dest, before);
+ }
+
+ // Erase the source block.
+ assert(source->empty() && "expected 'source' to be empty");
eraseBlock(source);
}
diff --git a/mlir/test/Dialect/Affine/simplify-structures.mlir b/mlir/test/Dialect/Affine/simplify-structures.mlir
index 2c693ea1551c013..92d3d86bc93068f 100644
--- a/mlir/test/Dialect/Affine/simplify-structures.mlir
+++ b/mlir/test/Dialect/Affine/simplify-structures.mlir
@@ -411,8 +411,6 @@ func.func @test_trivially_false_returning_two_results(%arg0: index) -> (index, i
// CHECK: %[[c13:.*]] = arith.constant 13 : index
%c7 = arith.constant 7 : index
%c13 = arith.constant 13 : index
- // CHECK: %[[c2:.*]] = arith.constant 2 : index
- // CHECK: %[[c3:.*]] = arith.constant 3 : index
%res:2 = affine.if affine_set<(d0, d1) : (5 >= 0, -2 >= 0)> (%c7, %c13) -> (index, index) {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
diff --git a/mlir/test/Transforms/test-strict-pattern-driver.mlir b/mlir/test/Transforms/test-strict-pattern-driver.mlir
index 2ebc66b4f26d0e8..4785795121653b3 100644
--- a/mlir/test/Transforms/test-strict-pattern-driver.mlir
+++ b/mlir/test/Transforms/test-strict-pattern-driver.mlir
@@ -249,3 +249,23 @@ func.func @test_move_op_before() {
}) : () -> ()
return
}
+
+// -----
+
+// CHECK-AN: notifyOperationInserted: test.op_1, previous = test.op_2
+// CHECK-AN: notifyOperationInserted: test.op_2, previous = test.op_3
+// CHECK-AN: notifyOperationInserted: test.op_3, was last in block
+// CHECK-AN-LABEL: func @test_inline_block_before(
+// CHECK-AN: test.op_1
+// CHECK-AN: test.op_2
+// CHECK-AN: test.op_3
+// CHECK-AN: test.inline_blocks_into_parent
+// CHECK-AN: return
+func.func @test_inline_block_before() {
+ "test.inline_blocks_into_parent"() ({
+ "test.op_1"() : () -> ()
+ "test.op_2"() : () -> ()
+ "test.op_3"() : () -> ()
+ }) : () -> ()
+ return
+}
diff --git a/mlir/test/lib/Dialect/Test/TestPatterns.cpp b/mlir/test/lib/Dialect/Test/TestPatterns.cpp
index 2165b388d559c5e..950f28e28bd7306 100644
--- a/mlir/test/lib/Dialect/Test/TestPatterns.cpp
+++ b/mlir/test/lib/Dialect/Test/TestPatterns.cpp
@@ -213,6 +213,26 @@ struct MoveBeforeParentOp : public RewritePattern {
}
};
+/// This pattern inlines blocks that are nested in
+/// "test.inline_blocks_into_parent" into the parent block.
+struct InlineBlocksIntoParent : public RewritePattern {
+ InlineBlocksIntoParent(MLIRContext *context)
+ : RewritePattern("test.inline_blocks_into_parent", /*benefit=*/1,
+ context) {}
+
+ LogicalResult matchAndRewrite(Operation *op,
+ PatternRewriter &rewriter) const override {
+ bool changed = false;
+ for (Region &r : op->getRegions()) {
+ while (!r.empty()) {
+ rewriter.inlineBlockBefore(&r.front(), op);
+ changed = true;
+ }
+ }
+ return success(changed);
+ }
+};
+
struct TestPatternDriver
: public PassWrapper<TestPatternDriver, OperationPass<>> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestPatternDriver)
@@ -297,7 +317,8 @@ struct TestStrictPatternDriver
EraseOp,
ChangeBlockOp,
ImplicitChangeOp,
- MoveBeforeParentOp
+ MoveBeforeParentOp,
+ InlineBlocksIntoParent
// clang-format on
>(ctx);
SmallVector<Operation *> ops;
@@ -305,7 +326,8 @@ struct TestStrictPatternDriver
StringRef opName = op->getName().getStringRef();
if (opName == "test.insert_same_op" || opName == "test.change_block_op" ||
opName == "test.replace_with_new_op" || opName == "test.erase_op" ||
- opName == "test.move_before_parent_op") {
+ opName == "test.move_before_parent_op" ||
+ opName == "test.inline_blocks_into_parent") {
ops.push_back(op);
}
});
More information about the Mlir-commits
mailing list