[Mlir-commits] [mlir] [mlir][Transforms][NFC] Remove InlineBlockRewrite (PR #83262)

[Matthias Springer]

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

When inlining a block during a dialect conversion, an `InlineBlockRewrite` object is stored in the dialect conversion state. This commit removes `InlineBlockRewrite`. Instead, a combination of `EraseBlockRewrite` and multiple `MoveOperationRewrite`/`ReplaceBlockArgRewrite` is used.

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

Note: This change is also towards making `RewriterBase::inlineBlockBefore` non-virtual. (We still need dialect conversion support for `replaceAllUsesWith` to do that.)


>From 7e074504bc9e29acdd96c4ba5023cea5d7dc34c2 Mon Sep 17 00:00:00 2001
From: Matthias Springer <springerm at google.com>
Date: Wed, 28 Feb 2024 13:52:22 +0000
Subject: [PATCH] [mlir][Transforms][NFC] Remove InlineBlockRewrite

---
 .../Transforms/Utils/DialectConversion.cpp    | 63 +++----------------
 1 file changed, 10 insertions(+), 53 deletions(-)

diff --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp b/mlir/lib/Transforms/Utils/DialectConversion.cpp
index 26899301eb742e..b81495a95c80ed 100644
--- a/mlir/lib/Transforms/Utils/DialectConversion.cpp
+++ b/mlir/lib/Transforms/Utils/DialectConversion.cpp
@@ -190,7 +190,6 @@ class IRRewrite {
     // Block rewrites
     CreateBlock,
     EraseBlock,
-    InlineBlock,
     MoveBlock,
     BlockTypeConversion,
     ReplaceBlockArg,
@@ -330,47 +329,6 @@ class EraseBlockRewrite : public BlockRewrite {
   Block *insertBeforeBlock;
 };
 
-/// Inlining of a block. This rewrite is immediately reflected in the IR.
-/// Note: This rewrite represents only the inlining of the operations. The
-/// erasure of the inlined block is a separate rewrite.
-class InlineBlockRewrite : public BlockRewrite {
-public:
-  InlineBlockRewrite(ConversionPatternRewriterImpl &rewriterImpl, Block *block,
-                     Block *sourceBlock, Block::iterator before)
-      : BlockRewrite(Kind::InlineBlock, rewriterImpl, block),
-        sourceBlock(sourceBlock),
-        firstInlinedInst(sourceBlock->empty() ? nullptr
-                                              : &sourceBlock->front()),
-        lastInlinedInst(sourceBlock->empty() ? nullptr : &sourceBlock->back()) {
-  }
-
-  static bool classof(const IRRewrite *rewrite) {
-    return rewrite->getKind() == Kind::InlineBlock;
-  }
-
-  void rollback() override {
-    // Put the operations from the destination block (owned by the rewrite)
-    // back into the source block.
-    if (firstInlinedInst) {
-      assert(lastInlinedInst && "expected operation");
-      sourceBlock->getOperations().splice(sourceBlock->begin(),
-                                          block->getOperations(),
-                                          Block::iterator(firstInlinedInst),
-                                          ++Block::iterator(lastInlinedInst));
-    }
-  }
-
-private:
-  // The block that originally contained the operations.
-  Block *sourceBlock;
-
-  // The first inlined operation.
-  Operation *firstInlinedInst;
-
-  // The last inlined operation.
-  Operation *lastInlinedInst;
-};
-
 /// Moving of a block. This rewrite is immediately reflected in the IR.
 class MoveBlockRewrite : public BlockRewrite {
 public:
@@ -858,10 +816,6 @@ struct ConversionPatternRewriterImpl : public RewriterBase::Listener {
   void notifyBlockInserted(Block *block, Region *previous,
                            Region::iterator previousIt) override;
 
-  /// Notifies that a block is being inlined into another block.
-  void notifyBlockBeingInlined(Block *block, Block *srcBlock,
-                               Block::iterator before);
-
   /// Notifies that a pattern match failed for the given reason.
   void
   notifyMatchFailure(Location loc,
@@ -1494,11 +1448,6 @@ void ConversionPatternRewriterImpl::notifyBlockInserted(
   appendRewrite<MoveBlockRewrite>(block, previous, prevBlock);
 }
 
-void ConversionPatternRewriterImpl::notifyBlockBeingInlined(
-    Block *block, Block *srcBlock, Block::iterator before) {
-  appendRewrite<InlineBlockRewrite>(block, srcBlock, before);
-}
-
 void ConversionPatternRewriterImpl::notifyMatchFailure(
     Location loc, function_ref<void(Diagnostic &)> reasonCallback) {
   LLVM_DEBUG({
@@ -1649,10 +1598,18 @@ void ConversionPatternRewriter::inlineBlockBefore(Block *source, Block *dest,
          "expected 'source' to have no predecessors");
 #endif // NDEBUG
 
-  impl->notifyBlockBeingInlined(dest, source, before);
+  // Replace all uses of block arguments.
+  // TODO: Support `replaceAllUsesWith` in the dialect conversion. Then this
+  // function no longer has to be overridden and can be turned into a
+  // non-virtual function.
   for (auto it : llvm::zip(source->getArguments(), argValues))
     replaceUsesOfBlockArgument(std::get<0>(it), std::get<1>(it));
-  dest->getOperations().splice(before, source->getOperations());
+
+  // Move op by op.
+  while (!source->empty())
+    moveOpBefore(&source->front(), dest, before);
+
+  // Erase the source block.
   eraseBlock(source);
 }