[mlir] ensureRegionTerminator: take OpBuilder

[ftynse]

Summary

The SingleBlockImplicitTerminator op trait provides a function
ensureRegionTerminator that injects an appropriate terminator into the block
if necessary, which is used during operation constructing and parsing.
Currently, this function directly modifies the IR using low-level APIs on
Operation and Block. If this function is called from a conversion pattern,
these manipulations are not reflected in the ConversionPatternRewriter and thus
cannot be undone or, worse, lead to tricky memory errors and malformed IR.
Change ensureRegionTerminator to take an instance of OpBuilder instead of
Builder, and use it to construct the block and the terminator when required.
Maintain overloads taking an instance of Builder and creating a simple
OpBuilder to use in parsers, which don't have an OpBuilder and cannot
interact with the dialect conversion mechanism. This change was one of the
reasons to make <OpTy>::build accept an OpBuilder.

Depends On D80137

[nicolasvasilache]

After seeing the stack of CL I am wondering whether RewritePatterns should only operate on Wrapper classes around Region, Block and Operation that would inherit privately.

Thoughts?

[nicolasvasilache]

mlir/include/mlir/IR/OpDefinition.h

1123

would it be better to make it explicit by having a small wrapper class and force users to cast?
Making the non-safe use case hard here and everywhere else would have likely saved many hours of engineering,

[ftynse]

In D80138#2042064, @nicolasvasilache wrote:

After seeing the stack of CL I am wondering whether RewritePatterns should only operate on Wrapper classes around Region, Block and Operation that would inherit privately.

Thoughts?

I thought about this, but it looks infeasible without reimplementing all of the core classes. E.g., if we pass in an OperationWrapper, but we still want specific ops, we then add a possibility to dyn_cast<SomeOp>, from which an Operaiton * can be trivially extracted. Same for the remaining classes. What we seem to be needing is some const-semantics for IR classes, but we decided against that for good reasons. These reasons would still apply to wrappers that essentially achieves const methods.

[ftynse]

mlir/include/mlir/IR/OpDefinition.h

1123

There are two kinds of users for this: ::build and ::parse. The former has OpBuilder readily available and the latter does not (and neither it should). Unless somebody explicitly upcasts the OpBuilder inside ::build for some reason, this should be fine.

[rriddle]

mlir/lib/IR/Operation.cpp

1119

Where is this version called from? Can we remove it?

[ftynse]

mlir/lib/IR/Operation.cpp

1119

From SingleBlockImplicitTerminator::ensureTerminator(Region &, Builder &, Location) overload, which in turn is called by various parsers. Since parsers don't have an OpBuilder, we need to construct it somewhere. If we want to remove this, we need to either give parsers an OpBuilder, or require them to construct one before calling ensureTerminator.

mlir/include/mlir/IR/OpDefinition.h

 // These functions are out-of-line utilities, which avoids them being template
 // instantiated/duplicated.
 namespace impl {
 /// Insert an operation, generated by `buildTerminatorOp`, at the end of the
 /// region's only block if it does not have a terminator already. If the region
 /// is empty, insert a new block first. `buildTerminatorOp` should return the
 /// terminator operation to insert.
 void ensureRegionTerminator(
-    Region &region, Location loc,
-    function_ref<Operation *(OpBuilder &)> buildTerminatorOp);
-/// Templated version that fills the generates the provided operation type.
-template <typename OpTy>
-void ensureRegionTerminator(Region &region, Builder &builder, Location loc) {
-  ensureRegionTerminator(region, loc, [&](OpBuilder &b) {
-    OperationState state(loc, OpTy::getOperationName());
-    OpTy::build(b, state);
-    return Operation::create(state);
-  });
-}
+    Region &region, OpBuilder &builder, Location loc,
+    function_ref<Operation *(OpBuilder &, Location)> buildTerminatorOp);
+void ensureRegionTerminator(
+    Region &region, Builder &builder, Location loc,
+    function_ref<Operation *(OpBuilder &, Location)> buildTerminatorOp);
+
 } // namespace impl
 
 /// This is the concrete base class that holds the operation pointer and has
 /// non-generic methods that only depend on State (to avoid having them
 /// instantiated on template types that don't affect them.
 ///
 /// This also has the fallback implementations of customization hooks for when
 /// they aren't customized.
public:
   }
 };
 
 /// This class provides APIs and verifiers for ops with regions having a single
 /// block that must terminate with `TerminatorOpType`.
 template <typename TerminatorOpType> struct SingleBlockImplicitTerminator {
   template <typename ConcreteType>
   class Impl : public TraitBase<ConcreteType, Impl> {
+  private:
+    /// Builds a terminator operation without relying on OpBuilder APIs to avoid
+    /// cyclic header inclusion.
+    static Operation *buildTerminator(OpBuilder &builder, Location loc) {
+      OperationState state(loc, TerminatorOpType::getOperationName());
+      TerminatorOpType::build(builder, state);
+      return Operation::create(state);
+    }
+
   public:
     static LogicalResult verifyTrait(Operation *op) {
       for (unsigned i = 0, e = op->getNumRegions(); i < e; ++i) {
         Region &region = op->getRegion(i);
 
         // Empty regions are fine.
         if (region.empty())
           continue;
static LogicalResult verifyTrait(Operation *op) {
                   "'"
                << TerminatorOpType::getOperationName() << '\'';
       }
 
       return success();
     }
 
     /// Ensure that the given region has the terminator required by this trait.
+    /// If OpBuilder is provided, use it to build the terminator and notify the
+    /// OpBuilder litsteners accoridngly. If only a Builder is provided, locally
+    /// construct an OpBuilder with no listeners; this should only be used if no
+    /// OpBuilder is available at the call site, e.g., in the parser.
     static void ensureTerminator(Region &region, Builder &builder,
                                  Location loc) {
-      ::mlir::impl::template ensureRegionTerminator<TerminatorOpType>(
-          region, builder, loc);
+      ::mlir::impl::ensureRegionTerminator(region, builder, loc,
+                                           buildTerminator);
+    }
+    static void ensureTerminator(Region &region, OpBuilder &builder,
+                                 Location loc) {
+      ::mlir::impl::ensureRegionTerminator(region, builder, loc,
+                                           buildTerminator);
     }
 
     Block *getBody(unsigned idx = 0) {
       Region &region = this->getOperation()->getRegion(idx);
       assert(!region.empty() && "unexpected empty region");
       return &region.front();
     }
   };

mlir/lib/IR/Operation.cpp

 // Misc. utils
 //===----------------------------------------------------------------------===//
 
 /// Insert an operation, generated by `buildTerminatorOp`, at the end of the
 /// region's only block if it does not have a terminator already. If the region
 /// is empty, insert a new block first. `buildTerminatorOp` should return the
 /// terminator operation to insert.
 void impl::ensureRegionTerminator(
-    Region &region, Location loc,
-    function_ref<Operation *(OpBuilder &)> buildTerminatorOp) {
+    Region &region, OpBuilder &builder, Location loc,
+    function_ref<Operation *(OpBuilder &, Location)> buildTerminatorOp) {
+  OpBuilder::InsertionGuard guard(builder);
   if (region.empty())
-    region.push_back(new Block);
+    builder.createBlock(&region);
 
   Block &block = region.back();
   if (!block.empty() && block.back().isKnownTerminator())
     return;
 
-  OpBuilder builder(loc.getContext());
-  block.push_back(buildTerminatorOp(builder));
+  builder.setInsertionPointToEnd(&block);
+  builder.insert(buildTerminatorOp(builder, loc));
+}
+
+/// Create a simple OpBuilder and forward to the OpBuilder version of this
+/// function.
+void impl::ensureRegionTerminator(
+    Region &region, Builder &builder, Location loc,
+    function_ref<Operation *(OpBuilder &, Location)> buildTerminatorOp) {
+  OpBuilder opBuilder(builder.getContext());
+  ensureRegionTerminator(region, opBuilder, loc, buildTerminatorOp);
 }
 
 //===----------------------------------------------------------------------===//
 // UseIterator
 //===----------------------------------------------------------------------===//
 
 Operation::UseIterator::UseIterator(Operation *op, bool end)
     : op(op), res(end ? op->result_end() : op->result_begin()) {