[mlir] Erase or clear blocks through ConversionPatternRewriter when applicable

[ftynse]

Summary

Multiple places in the code base were erasing Blocks or operations in them
using in-place modifications (Block::erase or Block::clear) unknown to
ConversionPatternRewriter. These operations could not be undone if the pattern
failed and could lead to inconsistent in-memory state of the IR with dangling
pointers. Use ConversionPatternRewriter::eraseOp and ::eraseBlock instead.

Depends On D80135

[rriddle]

mlir/examples/toy/Ch5/mlir/LowerToAffineLoops.cpp

75–76

This was because originally ensureTerminator wasn't tracked, so this could now be removed entirely. The create<AffineTerminatorOp> below would also need to be removed.

[ftynse]

mlir/examples/toy/Ch5/mlir/LowerToAffineLoops.cpp

75–76

Thanks, done in b596ecdd574e3ec717f71950370d2f58ae44296f

mlir/examples/toy/Ch5/mlir/LowerToAffineLoops.cpp

static void lowerOpToLoops(Operation *op, ArrayRef<Value> operands,
   // Insert an allocation and deallocation for the result of this operation.
   auto memRefType = convertTensorToMemRef(tensorType);
   auto alloc = insertAllocAndDealloc(memRefType, loc, rewriter);
 
   // Create an empty affine loop for each of the dimensions within the shape.
   SmallVector<Value, 4> loopIvs;
   for (auto dim : tensorType.getShape()) {
     auto loop = rewriter.create<AffineForOp>(loc, /*lb=*/0, dim, /*step=*/1);
-    loop.getBody()->clear();
+    for (Operation &nested : *loop.getBody())
+      rewriter.eraseOp(&nested);
     loopIvs.push_back(loop.getInductionVar());
 
     // Terminate the loop body and update the rewriter insertion point to the
     // beginning of the loop.
     rewriter.setInsertionPointToStart(loop.getBody());
     rewriter.create<AffineTerminatorOp>(loc);
     rewriter.setInsertionPointToStart(loop.getBody());
   }

mlir/examples/toy/Ch6/mlir/LowerToAffineLoops.cpp

static void lowerOpToLoops(Operation *op, ArrayRef<Value> operands,
   // Insert an allocation and deallocation for the result of this operation.
   auto memRefType = convertTensorToMemRef(tensorType);
   auto alloc = insertAllocAndDealloc(memRefType, loc, rewriter);
 
   // Create an empty affine loop for each of the dimensions within the shape.
   SmallVector<Value, 4> loopIvs;
   for (auto dim : tensorType.getShape()) {
     auto loop = rewriter.create<AffineForOp>(loc, /*lb=*/0, dim, /*step=*/1);
-    loop.getBody()->clear();
+    for (Operation &nested : *loop.getBody())
+      rewriter.eraseOp(&nested);
     loopIvs.push_back(loop.getInductionVar());
 
     // Terminate the loop body and update the rewriter insertion point to the
     // beginning of the loop.
     rewriter.setInsertionPointToStart(loop.getBody());
     rewriter.create<AffineTerminatorOp>(loc);
     rewriter.setInsertionPointToStart(loop.getBody());
   }

mlir/examples/toy/Ch6/mlir/LowerToLLVM.cpp

matchAndRewrite(Operation *op, ArrayRef<Value> operands,
     // Create a loop for each of the dimensions within the shape.
     SmallVector<Value, 4> loopIvs;
     for (unsigned i = 0, e = memRefShape.size(); i != e; ++i) {
       auto lowerBound = rewriter.create<ConstantIndexOp>(loc, 0);
       auto upperBound = rewriter.create<ConstantIndexOp>(loc, memRefShape[i]);
       auto step = rewriter.create<ConstantIndexOp>(loc, 1);
       auto loop =
           rewriter.create<scf::ForOp>(loc, lowerBound, upperBound, step);
-      loop.getBody()->clear();
+      for (Operation &nested : *loop.getBody())
+        rewriter.eraseOp(&nested);
       loopIvs.push_back(loop.getInductionVar());
 
       // Terminate the loop body.
       rewriter.setInsertionPointToEnd(loop.getBody());
 
       // Insert a newline after each of the inner dimensions of the shape.
       if (i != e - 1)
         rewriter.create<CallOp>(loc, printfRef, rewriter.getIntegerType(32),

mlir/examples/toy/Ch7/mlir/LowerToAffineLoops.cpp

static void lowerOpToLoops(Operation *op, ArrayRef<Value> operands,
   // Insert an allocation and deallocation for the result of this operation.
   auto memRefType = convertTensorToMemRef(tensorType);
   auto alloc = insertAllocAndDealloc(memRefType, loc, rewriter);
 
   // Create an empty affine loop for each of the dimensions within the shape.
   SmallVector<Value, 4> loopIvs;
   for (auto dim : tensorType.getShape()) {
     auto loop = rewriter.create<AffineForOp>(loc, /*lb=*/0, dim, /*step=*/1);
-    loop.getBody()->clear();
+    for (Operation &nested : *loop.getBody())
+      rewriter.eraseOp(&nested);
     loopIvs.push_back(loop.getInductionVar());
 
     // Terminate the loop body and update the rewriter insertion point to the
     // beginning of the loop.
     rewriter.setInsertionPointToStart(loop.getBody());
     rewriter.create<AffineTerminatorOp>(loc);
     rewriter.setInsertionPointToStart(loop.getBody());
   }

mlir/examples/toy/Ch7/mlir/LowerToLLVM.cpp

matchAndRewrite(Operation *op, ArrayRef<Value> operands,
     // Create a loop for each of the dimensions within the shape.
     SmallVector<Value, 4> loopIvs;
     for (unsigned i = 0, e = memRefShape.size(); i != e; ++i) {
       auto lowerBound = rewriter.create<ConstantIndexOp>(loc, 0);
       auto upperBound = rewriter.create<ConstantIndexOp>(loc, memRefShape[i]);
       auto step = rewriter.create<ConstantIndexOp>(loc, 1);
       auto loop =
           rewriter.create<scf::ForOp>(loc, lowerBound, upperBound, step);
-      loop.getBody()->clear();
+      for (Operation &nested : *loop.getBody())
+        rewriter.eraseOp(&nested);
       loopIvs.push_back(loop.getInductionVar());
 
       // Terminate the loop body.
       rewriter.setInsertionPointToEnd(loop.getBody());
 
       // Insert a newline after each of the inner dimensions of the shape.
       if (i != e - 1)
         rewriter.create<CallOp>(loc, printfRef, rewriter.getIntegerType(32),

mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp

public:
 
   LogicalResult matchAndRewrite(AffineForOp op,
                                 PatternRewriter &rewriter) const override {
     Location loc = op.getLoc();
     Value lowerBound = lowerAffineLowerBound(op, rewriter);
     Value upperBound = lowerAffineUpperBound(op, rewriter);
     Value step = rewriter.create<ConstantIndexOp>(loc, op.getStep());
     auto f = rewriter.create<scf::ForOp>(loc, lowerBound, upperBound, step);
-    f.region().getBlocks().clear();
+    rewriter.eraseBlock(f.getBody());
     rewriter.inlineRegionBefore(op.region(), f.region(), f.region().end());
     rewriter.eraseOp(op);
     return success();
   }
 };
 
 class AffineIfLowering : public OpRewritePattern<AffineIfOp> {
 public:
for (unsigned i = 0, e = integerSet.getNumConstraints(); i < e; ++i) {
           cond ? rewriter.create<AndOp>(loc, cond, cmpVal).getResult() : cmpVal;
     }
     cond = cond ? cond
                 : rewriter.create<ConstantIntOp>(loc, /*value=*/1, /*width=*/1);
 
     bool hasElseRegion = !op.elseRegion().empty();
     auto ifOp = rewriter.create<scf::IfOp>(loc, cond, hasElseRegion);
     rewriter.inlineRegionBefore(op.thenRegion(), &ifOp.thenRegion().back());
-    ifOp.thenRegion().back().erase();
+    rewriter.eraseBlock(&ifOp.thenRegion().back());
     if (hasElseRegion) {
       rewriter.inlineRegionBefore(op.elseRegion(), &ifOp.elseRegion().back());
-      ifOp.elseRegion().back().erase();
+      rewriter.eraseBlock(&ifOp.elseRegion().back());
     }
 
     // Ok, we're done!
     rewriter.eraseOp(op);
     return success();
   }
 };
 

mlir/lib/Conversion/GPUToSPIRV/ConvertGPUToSPIRV.cpp

LogicalResult GPUModuleConversion::matchAndRewrite(
 
   // Move the region from the module op into the SPIR-V module.
   Region &spvModuleRegion = spvModule.body();
   rewriter.inlineRegionBefore(moduleOp.body(), spvModuleRegion,
                               spvModuleRegion.begin());
   // The spv.module build method adds a block with a terminator. Remove that
   // block. The terminator of the module op in the remaining block will be
   // legalized later.
-  spvModuleRegion.back().erase();
+  rewriter.eraseBlock(&spvModuleRegion.back());
   rewriter.eraseOp(moduleOp);
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // GPU return inside kernel functions to SPIR-V return.
 //===----------------------------------------------------------------------===//