Do not assume Blocks belong to Functions

    Fix Block::splitBlock and Block::eraseFromFunction that erronously assume
    blocks belong to functions.  They now belong to regions.  When splitting, new
    blocks should be created in the same region as the existing block.  When
    erasing a block, it should be removed from the region rather than from the
    function body that transitively contains the region.

    Also rename Block::eraseFromFunction to Block::erase for consistency with other
    IR containers.

--

PiperOrigin-RevId: 250278272

diff --git a/mlir/include/mlir/IR/Block.h b/mlir/include/mlir/IR/Block.h
index 646e546..93d4aff 100644 (file)
--- a/mlir/include/mlir/IR/Block.h
+++ b/mlir/include/mlir/IR/Block.h
@@ -111,8 +111,8 @@ public:
   /// the specified block.
   void insertBefore(Block *block);
 
-  /// Unlink this Block from its Function and delete it.
-  void eraseFromFunction();
+  /// Unlink this Block from its parent region and delete it.
+  void erase();
 
   //===--------------------------------------------------------------------===//
   // Block argument management

diff --git a/mlir/lib/IR/Block.cpp b/mlir/lib/IR/Block.cpp
index 2ed8fd0..cf85cc8 100644 (file)
--- a/mlir/lib/IR/Block.cpp
+++ b/mlir/lib/IR/Block.cpp
@@ -70,9 +70,9 @@ void Block::insertBefore(Block *block) {
 }
 
 /// Unlink this Block from its Function and delete it.
-void Block::eraseFromFunction() {
-  assert(getFunction() && "Block has no parent");
-  getFunction()->getBlocks().erase(this);
+void Block::erase() {
+  assert(getParent() && "Block has no parent");
+  getParent()->getBlocks().erase(this);
 }
 
 /// Returns 'op' if 'op' lies in this block, or otherwise finds the
@@ -258,7 +258,7 @@ Block *Block::splitBlock(iterator splitBefore) {
   // Start by creating a new basic block, and insert it immediate after this
   // one in the containing function.
   auto newBB = new Block();
-  getFunction()->getBlocks().insert(++Function::iterator(this), newBB);
+  getParent()->getBlocks().insert(std::next(Region::iterator(this)), newBB);
 
   // Move all of the operations from the split point to the end of the function
   // into the new block.

diff --git a/mlir/lib/Transforms/LowerAffine.cpp b/mlir/lib/Transforms/LowerAffine.cpp
index 364905a..4dcc82f 100644 (file)
--- a/mlir/lib/Transforms/LowerAffine.cpp
+++ b/mlir/lib/Transforms/LowerAffine.cpp
@@ -556,7 +556,7 @@ static LogicalResult lowerAffineIf(AffineIfOp ifOp) {
   condBlock = builder.getInsertionBlock();
   if (condBlock->empty()) {
     condBlock->replaceAllUsesWith(thenBlock);
-    condBlock->eraseFromFunction();
+    condBlock->erase();
   } else {
     builder.create<BranchOp>(loc, thenBlock);
   }