[mlir][IR] Adjust insertion block when splitting blocks / moving ops #150819

matthias-springer · 2025-07-27T08:38:25Z

Similar to #146955, adjust the insertion block when splitting a block or moving an operation. This is the ensure that the insertion point is not left in an invalid state.

llvmbot · 2025-07-27T08:38:56Z

@llvm/pr-subscribers-mlir-gpu
@llvm/pr-subscribers-mlir-core

@llvm/pr-subscribers-mlir

Author: Matthias Springer (matthias-springer)

Changes

Similar to #146955, adjust the insertion block when splitting a block or moving an operation. This is the ensure that the insertion point is not left in an invalid state.

Full diff: https://github.com/llvm/llvm-project/pull/150819.diff

2 Files Affected:

(modified) mlir/include/mlir/IR/PatternMatch.h (+15)
(modified) mlir/lib/IR/PatternMatch.cpp (+17-3)

diff --git a/mlir/include/mlir/IR/PatternMatch.h b/mlir/include/mlir/IR/PatternMatch.h
index b5a93a0c5a898..fa87d6987c52c 100644
--- a/mlir/include/mlir/IR/PatternMatch.h
+++ b/mlir/include/mlir/IR/PatternMatch.h
@@ -576,24 +576,39 @@ class RewriterBase : public OpBuilder {
 
   /// Split the operations starting at "before" (inclusive) out of the given
   /// block into a new block, and return it.
+  ///
+  /// If the current insertion point is before the split point, the insertion
+  /// point is adjusted to the new block.
   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
   /// function.
+  ///
+  /// If the insertion point is before the moved operation, the insertion block
+  /// is adjusted to the block of `existingOp`.
   void moveOpBefore(Operation *op, Operation *existingOp);
 
   /// Unlink this operation from its current block and insert it right before
   /// `iterator` in the specified block.
+  ///
+  /// If the insertion point is before the moved operation, the insertion block
+  /// is adjusted to the specified block.
   void moveOpBefore(Operation *op, Block *block, Block::iterator iterator);
 
   /// Unlink this operation from its current block and insert it right after
   /// `existingOp` which may be in the same or another block in the same
   /// function.
+  ///
+  /// If the insertion point is before the moved operation, the insertion block
+  /// is adjusted to the block of `existingOp`.
   void moveOpAfter(Operation *op, Operation *existingOp);
 
   /// Unlink this operation from its current block and insert it right after
   /// `iterator` in the specified block.
+  ///
+  /// If the insertion point is before the moved operation, the insertion block
+  /// is adjusted to the specified block.
   void moveOpAfter(Operation *op, Block *block, Block::iterator iterator);
 
   /// Unlink this block and insert it right before `existingBlock`.
diff --git a/mlir/lib/IR/PatternMatch.cpp b/mlir/lib/IR/PatternMatch.cpp
index 9332f55bd9393..016569718275e 100644
--- a/mlir/lib/IR/PatternMatch.cpp
+++ b/mlir/lib/IR/PatternMatch.cpp
@@ -9,6 +9,7 @@
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/Iterators.h"
 #include "mlir/IR/RegionKindInterface.h"
+#include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/SmallPtrSet.h"
 
 using namespace mlir;
@@ -348,14 +349,21 @@ void RewriterBase::mergeBlocks(Block *source, Block *dest,
 /// Split the operations starting at "before" (inclusive) out of the given
 /// block into a new block, and return it.
 Block *RewriterBase::splitBlock(Block *block, Block::iterator before) {
+  Block *newBlock;
+  auto adjustInsertionPoint = llvm::make_scope_exit([&]() {
+    // If the current insertion point is before the split point, adjust the
+    // insertion point to the new block.
+    if (getInsertionPoint() == before)
+      setInsertionPoint(newBlock, before);
+  });
+
   // Fast path: If no listener is attached, split the block directly.
   if (!listener)
-    return block->splitBlock(before);
+    return newBlock = block->splitBlock(before);
 
   // `createBlock` sets the insertion point at the beginning of the new block.
   InsertionGuard g(*this);
-  Block *newBlock =
-      createBlock(block->getParent(), std::next(block->getIterator()));
+  newBlock = createBlock(block->getParent(), std::next(block->getIterator()));
 
   // If `before` points to end of the block, no ops should be moved.
   if (before == block->end())
@@ -413,6 +421,12 @@ void RewriterBase::moveOpBefore(Operation *op, Block *block,
   Block *currentBlock = op->getBlock();
   Block::iterator nextIterator = std::next(op->getIterator());
   op->moveBefore(block, iterator);
+
+  // If the current insertion point is before the moved operation, we may have
+  // to adjust the insertion block.
+  if (getInsertionPoint() == op->getIterator())
+    setInsertionPoint(block, op->getIterator());
+
   if (listener)
     listener->notifyOperationInserted(
         op, /*previous=*/InsertPoint(currentBlock, nextIterator));

Copilot

Pull Request Overview

This PR improves the robustness of MLIR's rewriter infrastructure by ensuring that insertion points remain valid when splitting blocks or moving operations. The changes prevent the insertion point from being left in an invalid state during these transformations.

Key changes:

Add insertion point adjustment logic to splitBlock() and moveOpBefore() methods
Introduce a new Block::isBeforeInBlock() method for iterator comparison
Replace duplicate iterator comparison logic with calls to the new centralized method

Reviewed Changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
mlir/lib/IR/PatternMatch.cpp	Implements insertion point adjustment in `splitBlock()` and `moveOpBefore()` methods
mlir/include/mlir/IR/PatternMatch.h	Updates documentation for move/split methods to describe insertion point behavior
mlir/lib/IR/Block.cpp	Implements new `isBeforeInBlock()` method for iterator comparison
mlir/include/mlir/IR/Block.h	Declares the new `isBeforeInBlock()` method
mlir/lib/IR/Dominance.cpp	Replaces local iterator comparison logic with calls to `Block::isBeforeInBlock()`
mlir/lib/Dialect/GPU/Transforms/AllReduceLowering.cpp	Adds explicit insertion point adjustment after block splitting

Copilot · 2025-07-27T09:58:33Z

mlir/lib/IR/PatternMatch.cpp

  if (!listener)
-    return block->splitBlock(before);
+    return newBlock = block->splitBlock(before);



[nitpick] The assignment within the return statement makes the code harder to read. Consider separating the assignment from the return statement for clarity.

Copilot · 2025-07-27T09:58:33Z

mlir/lib/IR/PatternMatch.cpp

+      // of the new block.
+      setInsertionPointToEnd(newBlock);
+    } else if (moveIpToNewBlock) {
+      setInsertionPoint(newBlock, getInsertionPoint());


Using getInsertionPoint() as the second argument may reference an iterator from the original block, which could be invalid after the split. The iterator should be adjusted to reference the corresponding position in the new block.

zero9178 · 2025-07-27T10:14:18Z

mlir/include/mlir/IR/Block.h

+  /// Return if the iterator `a` is before `b`. Both iterators must point into
+  /// this block.
+  bool isBeforeInBlock(iterator a, iterator b);


Suggested change

/// Return if the iterator `a` is before `b`. Both iterators must point into

/// this block.

bool isBeforeInBlock(iterator a, iterator b);

/// Return if the iterator 'a' is before 'b'. Both iterators must point into

/// this block.

bool isBeforeInBlock(iterator a, iterator b);

zero9178 · 2025-07-27T10:20:21Z

mlir/include/mlir/IR/PatternMatch.h

@@ -576,24 +576,39 @@ class RewriterBase : public OpBuilder {

  /// Split the operations starting at "before" (inclusive) out of the given
  /// block into a new block, and return it.
+  ///
+  /// If the current insertion point is before the split point, the insertion
+  /// point is adjusted to the new block.


I think this could be made clearer in that the insertion point is set to the exact same location in the new block.
Maybe something like:

Suggested change

/// point is adjusted to the new block.

///

/// The insertion point is updated to insert before the same operation as prior to the split.

/// If the insertion point was at the end of 'block', the new insertion point is at the end of the returned block.

zero9178 · 2025-07-27T10:24:56Z

mlir/include/mlir/IR/PatternMatch.h

+  ///
+  /// If the insertion point is before the moved operation, the insertion block
+  /// is adjusted to the block of `existingOp`.
  void moveOpBefore(Operation *op, Operation *existingOp);

  /// Unlink this operation from its current block and insert it right before
  /// `iterator` in the specified block.
+  ///
+  /// If the insertion point is before the moved operation, the insertion block
+  /// is adjusted to the specified block.
  void moveOpBefore(Operation *op, Block *block, Block::iterator iterator);

  /// Unlink this operation from its current block and insert it right after
  /// `existingOp` which may be in the same or another block in the same
  /// function.
+  ///
+  /// If the insertion point is before the moved operation, the insertion block
+  /// is adjusted to the block of `existingOp`.
  void moveOpAfter(Operation *op, Operation *existingOp);

  /// Unlink this operation from its current block and insert it right after
  /// `iterator` in the specified block.
+  ///
+  /// If the insertion point is before the moved operation, the insertion block
+  /// is adjusted to the specified block.


I am not fully confident this is the best behaviour for these operations but I fear this depends on ones iternal model of the insertion point.

In my mind, erasing an operation and moving an operation from the POV of the current insertion point are not much different: In both cases the erased/moved op disappears from the current insertion point and arguably the insertion point shouldn't care what happens to it after.

My expected behaviour would then be similar to eraseOp, which is that the insertion point remains exactly where it was, relatively speaking, and doesn't magically "travel" with the operation.

I agree, it is a bit odd that the insertion point "jumps", potentially even into a different block.

zero9178 · 2025-07-27T10:25:55Z

mlir/lib/IR/PatternMatch.cpp

+  // If the current insertion point is at or after the split point, adjust the
+  // insertion point to the new block.
+  bool moveIpToNewBlock = getBlock() == block &&
+                          !block->isBeforeInBlock(getInsertionPoint(), before);
+  auto adjustInsertionPoint = llvm::make_scope_exit([&]() {
+    if (getInsertionPoint() == block->end()) {
+      // If the insertion point is at the end of the block, move it to the end
+      // of the new block.
+      setInsertionPointToEnd(newBlock);
+    } else if (moveIpToNewBlock) {
+      setInsertionPoint(newBlock, getInsertionPoint());
+    }
+  });


Is there an easy way we can test the logic here?

joker-eph · 2025-07-27T20:14:31Z

mlir/lib/IR/Block.cpp

@@ -68,6 +68,16 @@ void Block::erase() {
  getParent()->getBlocks().erase(this);
 }

+bool Block::isBeforeInBlock(iterator a, iterator b) {


Can you add an assertion that they are in the same block here?

matthias-springer requested review from jpienaar, joker-eph and zero9178 July 27, 2025 08:38

llvmbot added mlir:core MLIR Core Infrastructure mlir labels Jul 27, 2025

matthias-springer force-pushed the users/matthias-springer/ip_split_move branch from 7adbc4c to 98fdd7c Compare July 27, 2025 09:17

llvmbot added the mlir:gpu label Jul 27, 2025

matthias-springer marked this pull request as draft July 27, 2025 09:30

matthias-springer force-pushed the users/matthias-springer/ip_split_move branch from 98fdd7c to 36065f5 Compare July 27, 2025 09:53

matthias-springer requested a review from Copilot July 27, 2025 09:58

Copilot AI reviewed Jul 27, 2025

View reviewed changes

matthias-springer marked this pull request as ready for review July 27, 2025 09:58

[mlir][IR] Move insertion point when splitting blocks / moving ops

b7bd595

matthias-springer force-pushed the users/matthias-springer/ip_split_move branch from 36065f5 to b7bd595 Compare July 27, 2025 09:59

zero9178 reviewed Jul 27, 2025

View reviewed changes

joker-eph reviewed Jul 27, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[mlir][IR] Adjust insertion block when splitting blocks / moving ops #150819

[mlir][IR] Adjust insertion block when splitting blocks / moving ops #150819

Uh oh!

matthias-springer commented Jul 27, 2025

Uh oh!

llvmbot commented Jul 27, 2025 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI Jul 27, 2025

Uh oh!

Copilot AI Jul 27, 2025

Uh oh!

zero9178 Jul 27, 2025

Uh oh!

zero9178 Jul 27, 2025

Uh oh!

zero9178 Jul 27, 2025

Uh oh!

matthias-springer Jul 27, 2025

Uh oh!

zero9178 Jul 27, 2025

Uh oh!

joker-eph Jul 27, 2025

Uh oh!

Uh oh!

-  /// point is adjusted to the new block.
+  ///
+  /// The insertion point is updated to insert before the same operation as prior to the split.
+  /// If the insertion point was at the end of 'block', the new insertion point is at the end of the returned block.

[mlir][IR] Adjust insertion block when splitting blocks / moving ops #150819

Are you sure you want to change the base?

[mlir][IR] Adjust insertion block when splitting blocks / moving ops #150819

Uh oh!

Conversation

matthias-springer commented Jul 27, 2025

Uh oh!

llvmbot commented Jul 27, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

Uh oh!

Copilot AI Jul 27, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jul 27, 2025

Choose a reason for hiding this comment

Uh oh!

zero9178 Jul 27, 2025

Choose a reason for hiding this comment

Uh oh!

zero9178 Jul 27, 2025

Choose a reason for hiding this comment

Uh oh!

zero9178 Jul 27, 2025

Choose a reason for hiding this comment

Uh oh!

matthias-springer Jul 27, 2025

Choose a reason for hiding this comment

Uh oh!

zero9178 Jul 27, 2025

Choose a reason for hiding this comment

Uh oh!

joker-eph Jul 27, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

llvmbot commented Jul 27, 2025 •

edited

Loading