[Mlir-commits] [mlir] [mlir][Transforms][NFC] Remove `SplitBlockRewrite` (PR #82777)

[Matthias Springer]

https://github.com/matthias-springer created https://github.com/llvm/llvm-project/pull/82777

When splitting a block during a dialect conversion, a `SplitBlockRewrite` object is stored in the dialect conversion state. This commit removes `SplitBlockRewrite`. Instead, a combination of `CreateBlockRewrite` and multiple `MoveOperationRewrite` is used.

This change simplifies the internal state of the dialect conversion and is also needed to properly support listeners.

`RewriteBase::splitBlock` is now no longer virtual. All necessary information for committing/rolling back a split block rewrite can be deduced from `Listener::notifyBlockInserted` and `Listener::notifyOperationInserted` (which is also called when moving an operation).

>From 9b2eb3746277bab4c04d64f384b3545a77fe11e6 Mon Sep 17 00:00:00 2001
From: Matthias Springer <springerm at google.com>
Date: Fri, 23 Feb 2024 15:35:39 +0000
Subject: [PATCH] [mlir][Transforms][NFC] Remove `SplitBlockRewrite`

When splitting a block during a dialect conversion, a `SplitBlockRewrite` object is stored in the dialect conversion state. This commit removes `SplitBlockRewrite`. Instead, a combination of `CreateBlockRewrite` and multiple `MoveOperationRewrite` is used.

This change simplifies the internal state of the dialect conversion and is also needed to properly support listeners.

`RewriteBase::splitBlock` is now no longer virtual. All necessary information for committing/rolling back a split block rewrite can be deduced from `Listener::notifyBlockInserted` and `Listener::notifyOperationInserted` (which is also called when moving an operation).
---
 mlir/include/mlir/IR/PatternMatch.h           |  2 +-
 .../mlir/Transforms/DialectConversion.h       |  3 --
 .../Transforms/Utils/DialectConversion.cpp    | 40 -------------------
 3 files changed, 1 insertion(+), 44 deletions(-)

diff --git a/mlir/include/mlir/IR/PatternMatch.h b/mlir/include/mlir/IR/PatternMatch.h
index 2ce3bc3fc2e783..f8d22cfb22afd0 100644
--- a/mlir/include/mlir/IR/PatternMatch.h
+++ b/mlir/include/mlir/IR/PatternMatch.h
@@ -579,7 +579,7 @@ class RewriterBase : public OpBuilder {
 
   /// Split the operations starting at "before" (inclusive) out of the given
   /// block into a new block, and return it.
-  virtual Block *splitBlock(Block *block, Block::iterator before);
+  Block *splitBlock(Block *block, Block::iterator before);
 
   /// Unlink this operation from its current block and insert it right before
   /// `existingOp` which may be in the same or another block in the same
diff --git a/mlir/include/mlir/Transforms/DialectConversion.h b/mlir/include/mlir/Transforms/DialectConversion.h
index 7e8e67a9d17824..84396529eb7c2e 100644
--- a/mlir/include/mlir/Transforms/DialectConversion.h
+++ b/mlir/include/mlir/Transforms/DialectConversion.h
@@ -741,9 +741,6 @@ class ConversionPatternRewriter final : public PatternRewriter {
   /// implemented for dialect conversion.
   void eraseBlock(Block *block) override;
 
-  /// PatternRewriter hook for splitting a block into two parts.
-  Block *splitBlock(Block *block, Block::iterator before) override;
-
   /// PatternRewriter hook for inlining the ops of a block into another block.
   void inlineBlockBefore(Block *source, Block *dest, Block::iterator before,
                          ValueRange argValues = std::nullopt) override;
diff --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp b/mlir/lib/Transforms/Utils/DialectConversion.cpp
index d015bd52901233..a83a01c44b4ed6 100644
--- a/mlir/lib/Transforms/Utils/DialectConversion.cpp
+++ b/mlir/lib/Transforms/Utils/DialectConversion.cpp
@@ -189,7 +189,6 @@ class IRRewrite {
     EraseBlock,
     InlineBlock,
     MoveBlock,
-    SplitBlock,
     BlockTypeConversion,
     ReplaceBlockArg,
     // Operation rewrites
@@ -397,30 +396,6 @@ class MoveBlockRewrite : public BlockRewrite {
   Block *insertBeforeBlock;
 };
 
-/// Splitting of a block. This rewrite is immediately reflected in the IR.
-class SplitBlockRewrite : public BlockRewrite {
-public:
-  SplitBlockRewrite(ConversionPatternRewriterImpl &rewriterImpl, Block *block,
-                    Block *originalBlock)
-      : BlockRewrite(Kind::SplitBlock, rewriterImpl, block),
-        originalBlock(originalBlock) {}
-
-  static bool classof(const IRRewrite *rewrite) {
-    return rewrite->getKind() == Kind::SplitBlock;
-  }
-
-  void rollback() override {
-    // Merge back the block that was split out.
-    originalBlock->getOperations().splice(originalBlock->end(),
-                                          block->getOperations());
-    eraseBlock(block);
-  }
-
-private:
-  // The original block from which this block was split.
-  Block *originalBlock;
-};
-
 /// This structure contains the information pertaining to an argument that has
 /// been converted.
 struct ConvertedArgInfo {
@@ -881,9 +856,6 @@ struct ConversionPatternRewriterImpl : public RewriterBase::Listener {
   void notifyBlockInserted(Block *block, Region *previous,
                            Region::iterator previousIt) override;
 
-  /// Notifies that a block was split.
-  void notifySplitBlock(Block *block, Block *continuation);
-
   /// Notifies that a block is being inlined into another block.
   void notifyBlockBeingInlined(Block *block, Block *srcBlock,
                                Block::iterator before);
@@ -1526,11 +1498,6 @@ void ConversionPatternRewriterImpl::notifyBlockInserted(
   appendRewrite<MoveBlockRewrite>(block, previous, prevBlock);
 }
 
-void ConversionPatternRewriterImpl::notifySplitBlock(Block *block,
-                                                     Block *continuation) {
-  appendRewrite<SplitBlockRewrite>(continuation, block);
-}
-
 void ConversionPatternRewriterImpl::notifyBlockBeingInlined(
     Block *block, Block *srcBlock, Block::iterator before) {
   appendRewrite<InlineBlockRewrite>(block, srcBlock, before);
@@ -1657,13 +1624,6 @@ ConversionPatternRewriter::getRemappedValues(ValueRange keys,
                            results);
 }
 
-Block *ConversionPatternRewriter::splitBlock(Block *block,
-                                             Block::iterator before) {
-  auto *continuation = block->splitBlock(before);
-  impl->notifySplitBlock(block, continuation);
-  return continuation;
-}
-
 void ConversionPatternRewriter::inlineBlockBefore(Block *source, Block *dest,
                                                   Block::iterator before,
                                                   ValueRange argValues) {