[mlir] DialectConversion: support erasing blocks

[ftynse]

Summary

PatternRewriter has support for erasing a Block from its parent region, but
this feature has not been implemented for ConversionPatternRewriter that needs
to keep track of and be able to undo block actions. Introduce support for
undoing block erasure in the ConversionPatternRewriter by marking all the ops
it contains for erasure and by detaching the block from its parent region. The
detached block is stored in the action description and is not actually deleted
until the rewrites are applied.

Depends On D80134

[nicolasvasilache]

The part that is modified makes sense to me but I am not qualified to judge how this may break in non-updated places in subtle ways.

[ftynse]

I am not qualified to judge how this may break in non-updated places in subtle ways.

It should not break any existing code because the entry function to all modifications was llvm_unreachable before this change, so it must not have been called.

[rriddle]

mlir/lib/Transforms/DialectConversion.cpp

726

nit: Reverse the if condition here.

[ftynse]

195d8571b9e296e7894d51cf59b7aeed42eaa2d2

[bondhugula]

It would be good to mention about eraseOp and eraseBlock at https://mlir.llvm.org/docs/DialectConversion/. I think the need to use these from within rewriter patterns isn't documented anywhere outside of the code.

mlir/lib/Transforms/DialectConversion.cpp

OpReplacement(Operation *op, ValueRange newValues)
         : op(op), newValues(newValues.begin(), newValues.end()) {}
 
     Operation *op;
     SmallVector<Value, 2> newValues;
   };
 
   /// The kind of the block action performed during the rewrite.  Actions can be
   /// undone if the conversion fails.
-  enum class BlockActionKind { Create, Move, Split, TypeConversion };
+  enum class BlockActionKind { Create, Erase, Move, Split, TypeConversion };
 
   /// Original position of the given block in its parent region.  We cannot use
   /// a region iterator because it could have been invalidated by other region
   /// operations since the position was stored.
   struct BlockPosition {
     Region *region;
     Region::iterator::difference_type position;
   };
 
   /// The storage class for an undoable block action (one of BlockActionKind),
   /// contains the information necessary to undo this action.
   struct BlockAction {
     static BlockAction getCreate(Block *block) {
       return {BlockActionKind::Create, block, {}};
     }
+    static BlockAction getErase(Block *block, BlockPosition originalPos) {
+      return {BlockActionKind::Erase, block, {originalPos}};
+    }
     static BlockAction getMove(Block *block, BlockPosition originalPos) {
       return {BlockActionKind::Move, block, {originalPos}};
     }
     static BlockAction getSplit(Block *block, Block *originalBlock) {
       BlockAction action{BlockActionKind::Split, block, {}};
       action.originalBlock = originalBlock;
       return action;
     }
     static BlockAction getTypeConversion(Block *block) {
       return BlockAction{BlockActionKind::TypeConversion, block, {}};
     }
 
     // The action kind.
     BlockActionKind kind;
 
     // A pointer to the block that was created by the action.
     Block *block;
 
     union {
-      // In use if kind == BlockActionKind::Move and contains a pointer to the
-      // region that originally contained the block as well as the position of
-      // the block in that region.
+      // In use if kind == BlockActionKind::Move or BlockActionKind::Erase, and
+      // contains a pointer to the region that originally contained the block as
+      // well as the position of the block in that region.
       BlockPosition originalPosition;
       // In use if kind == BlockActionKind::Split and contains a pointer to the
       // block that was split into two parts.
       Block *originalBlock;
     };
   };
 
   ConversionPatternRewriterImpl(PatternRewriter &rewriter,
                                 TypeConverter *converter)
       : argConverter(converter, rewriter) {}
 
   /// Return the current state of the rewriter.
   RewriterState getCurrentState();
 
   /// Reset the state of the rewriter to a previously saved point.
   void resetState(RewriterState state);
 
+  /// Erase any blocks that were unlinked from their regions and stored in block
+  /// actions.
+  void eraseDanglingBlocks();
+
   /// Undo the block actions (motions, splits) one by one in reverse order until
   /// "numActionsToKeep" actions remains.
   void undoBlockActions(unsigned numActionsToKeep = 0);
 
   /// Cleanup and destroy any generated rewrite operations. This method is
   /// invoked when the conversion process fails.
   void discardRewrites();
 
   /// Apply all requested operation rewrites. This method is invoked when the
   /// conversion process succeeds.
   void applyRewrites();
 
   /// Convert the signature of the given block.
   LogicalResult convertBlockSignature(Block *block);
 
   /// Apply a signature conversion on the given region.
   Block *
   applySignatureConversion(Region *region,
                            TypeConverter::SignatureConversion &conversion);
 
   /// PatternRewriter hook for replacing the results of an operation.
   void replaceOp(Operation *op, ValueRange newValues);
 
+  /// Notifies that a block is about to be erased.
+  void notifyBlockIsBeingErased(Block *block);
+
   /// Notifies that a block was created.
   void notifyCreatedBlock(Block *block);
 
   /// Notifies that a block was split.
   void notifySplitBlock(Block *block, Block *continuation);
 
   /// Notifies that the blocks of a region are about to be moved.
   void notifyRegionIsBeingInlinedBefore(Region &region, Region &parent,
while (createdOps.size() != state.numCreatedOps) {
     createdOps.pop_back();
   }
 
   // Pop all of the recorded ignored operations that are no longer valid.
   while (ignoredOps.size() != state.numIgnoredOperations)
     ignoredOps.pop_back();
 }
 
+void ConversionPatternRewriterImpl::eraseDanglingBlocks() {
+  for (auto &action : blockActions) {
+    if (action.kind != BlockActionKind::Erase)
+      continue;
+    delete action.block;
+  }
+}
+
 void ConversionPatternRewriterImpl::undoBlockActions(
     unsigned numActionsToKeep) {
   for (auto &action :
        llvm::reverse(llvm::drop_begin(blockActions, numActionsToKeep))) {
     switch (action.kind) {
     // Delete the created block.
     case BlockActionKind::Create: {
       // Unlink all of the operations within this block, they will be deleted
       // separately.
       auto &blockOps = action.block->getOperations();
       while (!blockOps.empty())
         blockOps.remove(blockOps.begin());
       action.block->dropAllDefinedValueUses();
       action.block->erase();
       break;
     }
+    // Put the block (owned by action) back into its original position.
+    case BlockActionKind::Erase: {
+      auto &blockList = action.originalPosition.region->getBlocks();
+      blockList.insert(
+          std::next(blockList.begin(), action.originalPosition.position),
+          action.block);
+      break;
+    }
     // Move the block back to its original position.
     case BlockActionKind::Move: {
       Region *originalRegion = action.originalPosition.region;
       originalRegion->getBlocks().splice(
           std::next(originalRegion->begin(), action.originalPosition.position),
           action.block->getParent()->getBlocks(), action.block);
       break;
     }
void ConversionPatternRewriterImpl::applyRewrites() {
 
   // In a second pass, erase all of the replaced operations in reverse. This
   // allows processing nested operations before their parent region is
   // destroyed.
   for (auto &repl : llvm::reverse(replacements))
     repl.op->erase();
 
   argConverter.applyRewrites(mapping);
+
+  // Now that the ops have been erased, also erase dangling blocks.
+  eraseDanglingBlocks();
 }
 
 LogicalResult
 ConversionPatternRewriterImpl::convertBlockSignature(Block *block) {
   // Check to see if this block should not be converted:
   // * There is no type converter.
   // * The block has already been converted.
   // * This is an entry block, these are converted explicitly via patterns.
void ConversionPatternRewriterImpl::replaceOp(Operation *op,
   // Record the requested operation replacement.
   replacements.emplace_back(op, newValues);
 
   /// Mark this operation as recursively ignored so that we don't need to
   /// convert any nested operations.
   markNestedOpsIgnored(op);
 }
 
+void ConversionPatternRewriterImpl::notifyBlockIsBeingErased(Block *block) {
+  Region *region = block->getParent();
+  auto position = std::distance(region->begin(), Region::iterator(block));
+  blockActions.push_back(BlockAction::getErase(block, {region, position}));
+}
+
 void ConversionPatternRewriterImpl::notifyCreatedBlock(Block *block) {
   blockActions.push_back(BlockAction::getCreate(block));
 }
 
 void ConversionPatternRewriterImpl::notifySplitBlock(Block *block,
                                                      Block *continuation) {
   blockActions.push_back(BlockAction::getSplit(continuation, block));
 }
LLVM_DEBUG({
     impl->logger.startLine()
         << "** Erase   : '" << op->getName() << "'(" << op << ")\n";
   });
   SmallVector<Value, 1> nullRepls(op->getNumResults(), nullptr);
   impl->replaceOp(op, nullRepls);
 }
 
 void ConversionPatternRewriter::eraseBlock(Block *block) {
-  llvm_unreachable("erasing blocks for dialect conversion not implemented");
+  impl->notifyBlockIsBeingErased(block);
+
+  // Mark all ops for erasure.
+  for (Operation &op : *block)
+    eraseOp(&op);
+
+  // Unlink the block from its parent region. The block is kept in the block
+  // action and will be actually destroyed when rewrites are applied. This
+  // allows us to keep the operations in the block live and undo the removal by
+  // re-inserting the block.
+  block->getParent()->getBlocks().remove(block);
 }
 
 /// Apply a signature conversion to the entry block of the given region.
 Block *ConversionPatternRewriter::applySignatureConversion(
     Region *region, TypeConverter::SignatureConversion &conversion) {
   return impl->applySignatureConversion(region, conversion);
 }
 
#endif
   // If the pattern moved or created any blocks, try to legalize their types.
   // This ensures that the types of the block arguments are legal for the region
   // they were moved into.
   for (unsigned i = curState.numBlockActions,
                 e = rewriterImpl.blockActions.size();
        i != e; ++i) {
     auto &action = rewriterImpl.blockActions[i];
     if (action.kind ==
-        ConversionPatternRewriterImpl::BlockActionKind::TypeConversion)
+            ConversionPatternRewriterImpl::BlockActionKind::TypeConversion ||
+        action.kind == ConversionPatternRewriterImpl::BlockActionKind::Erase)
       continue;
 
     // Convert the block signature.
     if (failed(rewriterImpl.convertBlockSignature(action.block))) {
       LLVM_DEBUG(logFailure(rewriterImpl.logger,
                             "failed to convert types of moved block"));
       return cleanupFailure();
     }

mlir/test/Transforms/test-legalizer.mlir

^bb0(%arg0: i32):
     "test.return"(%arg0) : (i32) -> ()
   }) : () -> ()
   // expected-remark@+1 {{op 'std.return' is not legalizable}}
   return
 }
 
 // -----
 
+// The op in this function is rewritten to itself (and thus remains illegal) by
+// a pattern that removes its second block after adding an operation into it.
+// Check that we can undo block removal succesfully.
+// CHECK-LABEL: @undo_block_erase
+func @undo_block_erase() {
+  // CHECK: test.undo_block_erase
+  "test.undo_block_erase"() ({
+    // expected-remark@-1 {{not legalizable}}
+    // CHECK: "unregistered.return"()[^[[BB:.*]]]
+    "unregistered.return"()[^bb1] : () -> ()
+    // expected-remark@-1 {{not legalizable}}
+  // CHECK: ^[[BB]]
+  ^bb1:
+    // CHECK: unregistered.return
+    "unregistered.return"() : () -> ()
+    // expected-remark@-1 {{not legalizable}}
+  }) : () -> ()
+}
+
+// -----
+
 // The op in this function is attempted to be rewritten to another illegal op
 // with an attached region containing an invalid terminator. The terminator is
 // created before the parent op. The deletion should not crash when deleting
 // created ops in the inverse order, i.e. deleting the parent op and then the
 // child op.
 // CHECK-LABEL: @undo_child_created_before_parent
 func @undo_child_created_before_parent() {
   // expected-remark@+1 {{is not legalizable}}
   "test.illegal_op_with_region_anchor"() : () -> ()
   // expected-remark@+1 {{op 'std.return' is not legalizable}}
   return
 }

mlir/test/lib/Dialect/Test/TestPatterns.cpp

auto illegalOp =
         rewriter.create<ILLegalOpF>(op->getLoc(), rewriter.getF32Type());
     rewriter.replaceUsesOfBlockArgument(op->getRegion(0).front().getArgument(0),
                                         illegalOp);
     rewriter.updateRootInPlace(op, [] {});
     return success();
   }
 };
 
+/// A rewrite pattern that tests the undo mechanism when erasing a block.
+struct TestUndoBlockErase : public ConversionPattern {
+  TestUndoBlockErase(MLIRContext *ctx)
+      : ConversionPattern("test.undo_block_erase", /*benefit=*/1, ctx) {}
+
+  LogicalResult
+  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const final {
+    Block *secondBlock = &*std::next(op->getRegion(0).begin());
+    rewriter.setInsertionPointToStart(secondBlock);
+    rewriter.create<ILLegalOpF>(op->getLoc(), rewriter.getF32Type());
+    rewriter.eraseBlock(secondBlock);
+    rewriter.updateRootInPlace(op, [] {});
+    return success();
+  }
+};
+
 //===----------------------------------------------------------------------===//
 // Type-Conversion Rewrite Testing
 
 /// This patterns erases a region operation that has had a type conversion.
 struct TestDropOpSignatureConversion : public ConversionPattern {
   TestDropOpSignatureConversion(MLIRContext *ctx, TypeConverter &converter)
       : ConversionPattern("test.drop_region_op", 1, ctx), converter(converter) {
   }
struct TestLegalizePatternDriver
   enum class ConversionMode { Analysis, Full, Partial };
 
   TestLegalizePatternDriver(ConversionMode mode) : mode(mode) {}
 
   void runOnOperation() override {
     TestTypeConverter converter;
     mlir::OwningRewritePatternList patterns;
     populateWithGenerated(&getContext(), &patterns);
-    patterns.insert<TestRegionRewriteBlockMovement, TestRegionRewriteUndo,
-                    TestCreateBlock, TestCreateIllegalBlock,
-                    TestUndoBlockArgReplace, TestPassthroughInvalidOp,
-                    TestSplitReturnType, TestChangeProducerTypeI32ToF32,
-                    TestChangeProducerTypeF32ToF64,
-                    TestChangeProducerTypeF32ToInvalid, TestUpdateConsumerType,
-                    TestNonRootReplacement, TestBoundedRecursiveRewrite,
-                    TestNestedOpCreationUndoRewrite>(&getContext());
+    patterns.insert<
+        TestRegionRewriteBlockMovement, TestRegionRewriteUndo, TestCreateBlock,
+        TestCreateIllegalBlock, TestUndoBlockArgReplace, TestUndoBlockErase,
+        TestPassthroughInvalidOp, TestSplitReturnType,
+        TestChangeProducerTypeI32ToF32, TestChangeProducerTypeF32ToF64,
     patterns.insert<TestDropOpSignatureConversion>(&getContext(), converter);
     mlir::populateFuncOpTypeConversionPattern(patterns, &getContext(),
                                               converter);
     mlir::populateCallOpTypeConversionPattern(patterns, &getContext(),
                                               converter);
 
     // Define the conversion target used for the test.
     ConversionTarget target(getContext());