[MLIR] Add getBody() method to SingleImplicitBlockTerminator op trait.

[pifon2a]

Summary

Many ops with this trait have getBody() and getBodyBuilder() methods defined in extraClassDeclaration in tablegen. getBody() implementation is the same accross all these ops, but getBodyBuilder() can return builders with varying insertion points set. In this PR, getBody() is moved into SingleImplicitBlockTerminator struct and getBodyBuilder() is replaced with OpBuilder::atBlock(End|Terminator)(op.getBody);.

[mehdi_amini]

LGTM, thanks!

mlir/include/mlir/IR/Builders.h

214

laster -> last?

[rriddle]

mlir/include/mlir/IR/Builders.h

212

This seems weird. I would expect you to just do block->getTerminator()

mlir/include/mlir/IR/OpDefinition.h

1098

This is not applicable to all users of this trait, e.g. IfOp. SingleBlockImplicitTerminator does not imply a single region.

[pifon2a]

Addressed River's comment.

[pifon2a]

More comments.

[herhut]

Nice, thanks!

[rriddle]

Looks great, thanks for the cleanup! LGTM after resolving the comments about assert messages

mlir/include/mlir/IR/Builders.h

213

nit: Please update all of the asserts to start with lowercase and end with no period.

[pifon2a]

Addressed the comments.

flang/include/flang/Optimizer/Dialect/FIROps.td

let extraClassDeclaration = [{
     static constexpr const char *stepAttrName() { return "step"; }
 
     /// Is this an unordered loop?
     bool isUnordered() { return getAttr(unorderedAttrName()).isa<UnitAttr>(); }
 
     /// Does loop set (and return) the final value of the control variable?
     bool hasLastValue() { return lastVal().size(); }
 
-    /// Get the body of the loop
-    mlir::Block *getBody() { return &region().front(); }
-
     /// Get the block argument corresponding to the loop control value (PHI)
     mlir::Value getInductionVar() { return getBody()->getArgument(0); }
 
-    /// Get a builder to insert operations into the LoopOp
-    mlir::OpBuilder getBodyBuilder() {
-      return mlir::OpBuilder(getBody(), std::prev(getBody()->end()));
-    }
-
     void setLowerBound(mlir::Value bound) {
       getOperation()->setOperand(0, bound);
     }
 
     void setUpperBound(mlir::Value bound) {
       getOperation()->setOperand(1, bound);
     }
 

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

OpBuilder<"Builder *builder, OperationState &result, "
               "int64_t step = 1">
   ];
 
   let extraClassDeclaration = [{
     static StringRef getStepAttrName() { return "step"; }
     static StringRef getLowerBoundAttrName() { return "lower_bound"; }
     static StringRef getUpperBoundAttrName() { return "upper_bound"; }
 
-    Block *getBody() { return &region().front(); }
     Value getInductionVar() { return getBody()->getArgument(0); }
-    OpBuilder getBodyBuilder() {
-      return OpBuilder(getBody(), std::prev(getBody()->end()));
-    }
 
     // TODO: provide iterators for the lower and upper bound operands
     // if the current access via getLowerBound(), getUpperBound() is too slow.
 
     /// Returns operands for the lower bound map.
     operand_range getLowerBoundOperands();
 
     /// Returns operands for the upper bound map.

mlir/include/mlir/Dialect/LoopOps/LoopOps.td

def ForOp : Loop_Op<"for",
   let skipDefaultBuilders = 1;
   let builders = [
     OpBuilder<"Builder *builder, OperationState &result, "
               "Value lowerBound, Value upperBound, Value step, "
               "ValueRange iterArgs = llvm::None">
   ];
 
   let extraClassDeclaration = [{
-    Block *getBody() { return &region().front(); }
     Value getInductionVar() { return getBody()->getArgument(0); }
-    OpBuilder getBodyBuilder() {
-      return OpBuilder(getBody(), std::prev(getBody()->end()));
-    }
     Block::BlockArgListType getRegionIterArgs() {
       return getBody()->getArguments().drop_front();
     }
     Operation::operand_range getIterOperands() {
       return getOperands().drop_front(getNumControlOperands());
     }
 
     void setLowerBound(Value bound) { getOperation()->setOperand(0, bound); }
let builders = [
     OpBuilder<"Builder *builder, OperationState &result, "
               "Value cond, bool withElseRegion">,
     OpBuilder<"Builder *builder, OperationState &result, "
               "TypeRange resultTypes, Value cond, bool withElseRegion">
   ];
 
   let extraClassDeclaration = [{
     OpBuilder getThenBodyBuilder() {
-      assert(!thenRegion().empty() && "Unexpected empty 'then' region.");
-      Block &body = thenRegion().front();
-      return OpBuilder(&body,
-                       results().empty() ? std::prev(body.end()) : body.end());
+      Block* body = getBody(0);
+      return results().empty() ? OpBuilder::atBlockTerminator(body)
+                               : OpBuilder::atBlockEnd(body);
     }
     OpBuilder getElseBodyBuilder() {
-      assert(!elseRegion().empty() && "Unexpected empty 'else' region.");
-      Block &body = elseRegion().front();
-      return OpBuilder(&body,
-                       results().empty() ? std::prev(body.end()) : body.end());
+      Block* body = getBody(1);
+      return results().empty() ? OpBuilder::atBlockTerminator(body)
+                               : OpBuilder::atBlockEnd(body);
     }
   }];
 }
 
 def ParallelOp : Loop_Op<"parallel",
     [AttrSizedOperandSegments,
      DeclareOpInterfaceMethods<LoopLikeOpInterface>,
      RecursiveSideEffects,
def ParallelOp : Loop_Op<"parallel",
   let skipDefaultBuilders = 1;
   let builders = [
     OpBuilder<"Builder *builder, OperationState &result, "
               "ValueRange lowerBounds, ValueRange upperBounds, "
               "ValueRange steps, ValueRange initVals = {}">,
   ];
 
   let extraClassDeclaration = [{
-    Block *getBody() { return &region().front(); }
     ValueRange getInductionVars() {
       return getBody()->getArguments();
     }
     unsigned getNumLoops() { return step().size(); }
     unsigned getNumReductions() { return initVals().size(); }
   }];
 }
 

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

def GenericAtomicRMWOp : Std_Op<"generic_atomic_rmw", [
 
   let skipDefaultBuilders = 1;
   let builders = [
     OpBuilder<"Builder *builder, OperationState &result, "
               "Value memref, ValueRange ivs">
   ];
 
   let extraClassDeclaration = [{
-    OpBuilder getBodyBuilder() {
-      assert(!body().empty() && "Unexpected empty 'body' region.");
-      Block &block = body().front();
-      return OpBuilder(&block, block.end());
-    }
     // The value stored in memref[ivs].
     Value getCurrentValue() {
       return body().front().getArgument(0);
     }
     MemRefType getMemRefType() {
       return memref().getType().cast<MemRefType>();
     }
   }];

mlir/include/mlir/IR/Builders.h

public:
   }
 
   OpBuilder(Block *block, Block::iterator insertPoint)
       : OpBuilder(block->getParent()) {
     setInsertionPoint(block, insertPoint);
   }
 
   /// Create a builder and set the insertion point to before the first operation
-  /// in the block but still inside th block.
+  /// in the block but still inside the block.
   static OpBuilder atBlockBegin(Block *block) {
     return OpBuilder(block, block->begin());
   }
 
   /// Create a builder and set the insertion point to after the last operation
   /// in the block but still inside the block.
   static OpBuilder atBlockEnd(Block *block) {
     return OpBuilder(block, block->end());
   }
 
+  /// Create a builder and set the insertion point to before the block
+  /// terminator.
+  static OpBuilder atBlockTerminator(Block *block) {
+    auto *terminator = block->getTerminator();
+    assert(terminator != nullptr && "the block has no terminator");
+    return OpBuilder(block, terminator->getIterator());
+  }
+
   /// This class represents a saved insertion point.
   class InsertPoint {
   public:
     /// Creates a new insertion point which doesn't point to anything.
     InsertPoint() = default;
 
     /// Creates a new insertion point at the given location.
     InsertPoint(Block *insertBlock, Block::iterator insertPt)
public:
 
   /// Create an operation of specific op type at the current insertion point.
   template <typename OpTy, typename... Args>
   OpTy create(Location location, Args &&... args) {
     OperationState state(location, OpTy::getOperationName());
     OpTy::build(this, state, std::forward<Args>(args)...);
     auto *op = createOperation(state);
     auto result = dyn_cast<OpTy>(op);
-    assert(result && "Builder didn't return the right type");
+    assert(result && "builder didn't return the right type");
     return result;
   }
 
   /// Create an operation of specific op type at the current insertion point,
   /// and immediately try to fold it. This functions populates 'results' with
   /// the results after folding the operation.
   template <typename OpTy, typename... Args>
   void createOrFold(SmallVectorImpl<Value> &results, Location location,

mlir/include/mlir/IR/OpDefinition.h

public:
     }
 
     /// Ensure that the given region has the terminator required by this trait.
     static void ensureTerminator(Region &region, Builder &builder,
                                  Location loc) {
       ::mlir::impl::template ensureRegionTerminator<TerminatorOpType>(
           region, builder, loc);
     }
+
+    Block *getBody(unsigned idx = 0) {
+      Region &region = this->getOperation()->getRegion(idx);
+      assert(!region.empty() && "unexpected empty region");
+      return &region.front();
+    }
   };
 };
 
 /// This class provides a verifier for ops that are expecting a specific parent.
 template <typename ParentOpType> struct HasParent {
   template <typename ConcreteType>
   class Impl : public TraitBase<ConcreteType, Impl> {
   public:

mlir/lib/Transforms/Utils/LoopUtils.cpp

static AffineForOp generateShiftedLoop(
 
   auto loopChunk = b.create<AffineForOp>(srcForOp.getLoc(), lbOperands, lbMap,
                                          ubOperands, ubMap, srcForOp.getStep());
   auto loopChunkIV = loopChunk.getInductionVar();
   auto srcIV = srcForOp.getInductionVar();
 
   BlockAndValueMapping operandMap;
 
-  OpBuilder bodyBuilder = loopChunk.getBodyBuilder();
+  auto bodyBuilder = OpBuilder::atBlockTerminator(loopChunk.getBody());
   for (auto it = opGroupQueue.begin() + offset, e = opGroupQueue.end(); it != e;
        ++it) {
     uint64_t shift = it->first;
     auto ops = it->second;
     // All 'same shift' operations get added with their operands being
     // remapped to results of cloned operations, and their IV used remapped.
     // Generate the remapping if the shift is not zero: remappedIV = newIV -
     // shift.
LogicalResult mlir::loopUnrollByFactor(AffineForOp forOp,
   }
 
   // Scale the step of loop being unrolled by unroll factor.
   int64_t step = forOp.getStep();
   forOp.setStep(step * unrollFactor);
 
   // Builder to insert unrolled bodies just before the terminator of the body of
   // 'forOp'.
-  OpBuilder builder = forOp.getBodyBuilder();
+  auto builder = OpBuilder::atBlockTerminator(forOp.getBody());
 
   // Keep a pointer to the last non-terminator operation in the original block
   // so that we know what to clone (since we are doing this in-place).
   Block::iterator srcBlockEnd = std::prev(forOp.getBody()->end(), 2);
 
   // Unroll the contents of 'forOp' (append unrollFactor - 1 additional copies).
   auto forOpIV = forOp.getInductionVar();
   for (unsigned i = 1; i < unrollFactor; i++) {
stripmineSink(AffineForOp forOp, uint64_t factor,
   SmallVector<Value, 4> ubOperands(forOp.getUpperBoundOperands());
   augmentMapAndBounds(b, forOp.getInductionVar(), &ubMap, &ubOperands,
                       /*offset=*/scaledStep);
 
   auto iv = forOp.getInductionVar();
   SmallVector<AffineForOp, 8> innerLoops;
   for (auto t : targets) {
     // Insert newForOp before the terminator of `t`.
-    OpBuilder b = t.getBodyBuilder();
+    auto b = OpBuilder::atBlockTerminator(t.getBody());
     auto newForOp = b.create<AffineForOp>(t.getLoc(), lbOperands, lbMap,
                                           ubOperands, ubMap, originalStep);
     auto begin = t.getBody()->begin();
     // Skip terminator and `newForOp` which is just before the terminator.
     auto nOps = t.getBody()->getOperations().size() - 2;
     newForOp.getBody()->getOperations().splice(
         newForOp.getBody()->getOperations().begin(),
         t.getBody()->getOperations(), begin, std::next(begin, nOps));
static Loops stripmineSink(loop::ForOp forOp, Value factor,
 
   Loops innerLoops;
   for (auto t : targets) {
     // Save information for splicing ops out of t when done
     auto begin = t.getBody()->begin();
     auto nOps = t.getBody()->getOperations().size();
 
     // Insert newForOp before the terminator of `t`.
-    OpBuilder b(t.getBodyBuilder());
+    auto b = OpBuilder::atBlockTerminator((t.getBody()));
     Value stepped = b.create<AddIOp>(t.getLoc(), iv, forOp.step());
     Value less = b.create<CmpIOp>(t.getLoc(), CmpIPredicate::slt,
                                   forOp.upperBound(), stepped);
     Value ub =
         b.create<SelectOp>(t.getLoc(), less, forOp.upperBound(), stepped);
 
     // Splice [begin, begin + nOps - 1) into `newForOp` and replace uses.
     auto newForOp = b.create<loop::ForOp>(t.getLoc(), iv, ub, originalStep);
generatePointWiseCopy(Location loc, Value memref, Value fastMemRef,
   AffineForOp copyNestRoot;
   SmallVector<AffineApplyOp, 4> mayBeDeadApplys;
   for (unsigned d = 0; d < rank; ++d) {
     auto forOp = createCanonicalizedAffineForOp(b, loc, lbOperands, lbMaps[d],
                                                 ubOperands, ubMaps[d]);
     if (d == 0)
       copyNestRoot = forOp;
 
-    b = forOp.getBodyBuilder();
+    b = OpBuilder::atBlockTerminator(forOp.getBody());
 
     auto fastBufOffsetMap =
         AffineMap::get(lbOperands.size(), 0, fastBufOffsets[d]);
     auto offset = b.create<AffineApplyOp>(loc, fastBufOffsetMap, lbOperands);
 
     // Construct the subscript for the fast memref being copied into/from:
     // x - offset_x.
     fastBufExprs.push_back(b.getAffineDimExpr(2 * d + 1) -
for (auto loop : inputNest) {
     cst.setDimSymbolSeparation(/*newSymbolCount=*/0);
 
     AffineValueMap lbVmap, ubVmap;
     cst.getIneqAsAffineValueMap(/*pos=*/0, lbPos, lbVmap, b.getContext());
     cst.getIneqAsAffineValueMap(/*pos=*/0, ubPos, ubVmap, b.getContext());
     AffineForOp fullTileLoop = createCanonicalizedAffineForOp(
         b, loop.getLoc(), lbVmap.getOperands(), lbVmap.getAffineMap(),
         ubVmap.getOperands(), ubVmap.getAffineMap());
-    b = fullTileLoop.getBodyBuilder();
+    b = OpBuilder::atBlockTerminator(fullTileLoop.getBody());
     fullTileLoops.push_back(fullTileLoop);
   }
 
   // Add the body for the full tile loop nest.
   BlockAndValueMapping operandMap;
   for (auto loopEn : llvm::enumerate(inputNest))
     operandMap.map(loopEn.value().getInductionVar(),
                    fullTileLoops[loopEn.index()].getInductionVar());
-  b = fullTileLoops.back().getBodyBuilder();
+  b = OpBuilder::atBlockTerminator(fullTileLoops.back().getBody());
   for (auto &op : inputNest.back().getBody()->without_terminator())
     b.clone(op, operandMap);
   return success();
 }
 
 LogicalResult
 mlir::separateFullTiles(MutableArrayRef<AffineForOp> inputNest,
                         SmallVectorImpl<AffineForOp> *fullTileNest) {