[mlir][EDSC] Make use of InsertGuard

[nicolasvasilache]

Summary

This revision cleans up a layer of complexity in ScopedContext and uses InsertGuard instead of previously manual bookkeeping.
The method getBuilder is renamed to getBuilderRef and spurious copies of OpBuilder are tracked.

This results in some canonicalizations not happening anymore in the Linalg matmul to vector test. This test is retired because relying on DRRs for this has been shaky at best. The solution will be better support to write fused passes in C++ with more idiomatic pattern composition and application.

[rriddle]

mlir/lib/EDSC/Builders.cpp

86

You should be able to just do res.block->addArguments(argTypes)

[nicolasvasilache]

ADdress comment.

[nicolasvasilache]

Update.

mlir/include/mlir/Dialect/Linalg/EDSC/FoldedIntrinsics.h

 namespace edsc {
 namespace intrinsics {
 
 template <typename Op>
 struct FoldedValueBuilder {
   // Builder-based
   template <typename... Args>
   FoldedValueBuilder(OperationFolder *folder, Args... args) {
-    value = folder ? folder->create<Op>(ScopedContext::getBuilder(),
+    value = folder ? folder->create<Op>(ScopedContext::getBuilderRef(),
                                         ScopedContext::getLocation(), args...)
-                   : ScopedContext::getBuilder().create<Op>(
+                   : ScopedContext::getBuilderRef().create<Op>(
                          ScopedContext::getLocation(), args...);
   }
 
   operator Value() { return value; }
   Value value;
 };
 
 using folded_std_constant_index = FoldedValueBuilder<ConstantIndexOp>;

mlir/include/mlir/Dialect/Linalg/Transforms/CMakeLists.txt

 set(LLVM_TARGET_DEFINITIONS LinalgTransformPatterns.td)
-mlir_tablegen(TestLinalgMatmulToVectorPatterns.h.inc -gen-rewriters)
 mlir_tablegen(LinalgTransformPatterns.h.inc -gen-rewriters)
 add_public_tablegen_target(MLIRLinalgTransformPatternsIncGen)
 
 # Including Linalg in TableGen requires to depends on generated files
 add_dependencies(MLIRLinalgTransformPatternsIncGen LinalgOdsGen)
 

mlir/include/mlir/EDSC/Builders.h

 
 /// Helper class to transparently handle builder insertion points by RAII.
 /// As its name indicates, a ScopedContext is means to be used locally in a
 /// scoped fashion. This abstracts away all the boilerplate related to
 /// checking proper usage of captures, NestedBuilders as well as handling the
 /// setting and restoring of insertion points.
 class ScopedContext {
 public:
-  ScopedContext(OpBuilder &builder, Location location);
+  ScopedContext(OpBuilder &b, Location location);
 
   /// Sets the insertion point of the builder to 'newInsertPt' for the duration
   /// of the scope. The existing insertion point of the builder is restored on
   /// destruction.
-  ScopedContext(OpBuilder &builder, OpBuilder::InsertPoint newInsertPt,
+  ScopedContext(OpBuilder &b, OpBuilder::InsertPoint newInsertPt,
                 Location location);
   ~ScopedContext();
 
   static MLIRContext *getContext();
-  static OpBuilder &getBuilder();
+  static OpBuilder &getBuilderRef();
   static Location getLocation();
 
 private:
   /// Only NestedBuilder (which is used to create an operation with a body)
   /// may access private members in order to implement scoping.
   friend class NestedBuilder;
 
   ScopedContext() = delete;
   ScopedContext(const ScopedContext &) = delete;
   ScopedContext &operator=(const ScopedContext &) = delete;
 
   static ScopedContext *&getCurrentScopedContext();
 
   /// Top level OpBuilder.
   OpBuilder &builder;
-  /// The previous insertion point of the builder.
-  Optional<OpBuilder::InsertPoint> prevBuilderInsertPoint;
+  /// Guard to the previous insertion point.
+  OpBuilder::InsertionGuard guard;
   /// Current location.
   Location location;
   /// Parent context we return into.
   ScopedContext *enclosingScopedContext;
-  /// Defensively keeps track of the current NestedBuilder to ensure proper
-  /// scoping usage.
-  NestedBuilder *nestedBuilder;
 };
 
 template <typename Op>
 struct ValueBuilder {
   template <typename... Args>
   ValueBuilder(Args... args) {
-    value = ScopedContext::getBuilder()
+    value = ScopedContext::getBuilderRef()
                 .create<Op>(ScopedContext::getLocation(), args...)
                 .getResult();
   }
   operator Value() { return value; }
   Value value;
 };
 
 template <typename Op>
 struct OperationBuilder {
   template <typename... Args>
   OperationBuilder(Args... args) {
-    op = ScopedContext::getBuilder().create<Op>(ScopedContext::getLocation(),
+    op = ScopedContext::getBuilderRef().create<Op>(ScopedContext::getLocation(),
                                                    args...);
   }
   operator Op() { return op; }
   operator Operation *() { return op.getOperation(); }
   Op op;
 };
 
 /// A NestedBuilder is a scoping abstraction to create an idiomatic syntax
 /// embedded in C++ that serves the purpose of building nested MLIR.
protected:
 
   /// Enter an mlir::Block and setup a ScopedContext to insert operations at
   /// the end of it. Since we cannot use c++ language-level scoping to implement
   /// scoping itself, we use enter/exit pairs of operations.
   /// As a consequence we must allocate a new OpBuilder + ScopedContext and
   /// let the escape.
   void enter(mlir::Block *block) {
     bodyScope = new ScopedContext(
-        ScopedContext::getBuilder(),
+        ScopedContext::getBuilderRef(),
         OpBuilder::InsertPoint(block, std::prev(block->end())),
         ScopedContext::getLocation());
     if (!block->empty()) {
       auto &termOp = block->back();
       if (termOp.isKnownTerminator())
-        ScopedContext::getBuilder().setInsertionPoint(&termOp);
+        ScopedContext::getBuilderRef().setInsertionPoint(&termOp);
     }
-    bodyScope->nestedBuilder = this;
   }
 
   /// Exit the current mlir::Block by explicitly deleting the dynamically
   /// allocated OpBuilder and ScopedContext.
   void exit() {
-    // Reclaim now to exit the scope.
-    bodyScope->nestedBuilder = nullptr;
     delete bodyScope;
     bodyScope = nullptr;
   }
 
   /// Custom destructor does nothing because we already destroyed bodyScope
   /// manually in `exit`. Insert an assertion to defensively guard against
   /// improper usage of scoping.
   ~NestedBuilder() {

mlir/lib/Conversion/LinalgToLLVM/LinalgToLLVM.cpp

public:
   void setStride(unsigned i, Value v) { d.setStride(rewriter(), loc(), i, v); }
   void setConstantStride(unsigned i, int64_t v) {
     d.setConstantStride(rewriter(), loc(), i, v);
   }
 
   operator Value() { return d; }
 
 private:
-  OpBuilder &rewriter() { return ScopedContext::getBuilder(); }
+  OpBuilder &rewriter() { return ScopedContext::getBuilderRef(); }
   Location loc() { return ScopedContext::getLocation(); }
 
   MemRefDescriptor d;
 };
 
 // RangeOp creates a new range descriptor.
 class RangeOpConversion : public ConvertToLLVMPattern {
 public:

mlir/lib/Conversion/VectorToLoops/ConvertVectorToLoops.cpp

void NDTransferOpHelper<ConcreteOp>::emitInBounds(
   for (auto it : llvm::zip(majorIvs, majorOffsets, memrefBounds.getUbs())) {
     Value iv = std::get<0>(it), off = std::get<1>(it), ub = std::get<2>(it);
     using namespace mlir::edsc::op;
     majorIvsPlusOffsets.push_back(iv + off);
     Value inBounds2 = majorIvsPlusOffsets.back() < ub;
     inBounds = inBounds && inBounds2;
   }
 
-  auto ifOp = ScopedContext::getBuilder().create<loop::IfOp>(
+  auto ifOp = ScopedContext::getBuilderRef().create<loop::IfOp>(
       ScopedContext::getLocation(), TypeRange{}, inBounds,
       /*withElseRegion=*/std::is_same<ConcreteOp, TransferReadOp>());
   BlockBuilder(&ifOp.thenRegion().front(),
                Append())([&] { thenBlockBuilder(majorIvsPlusOffsets); });
   if (std::is_same<ConcreteOp, TransferReadOp>())
     BlockBuilder(&ifOp.elseRegion().front(),
                  Append())([&] { elseBlockBuilder(majorIvsPlusOffsets); });
 }

mlir/lib/Dialect/Affine/EDSC/Builders.cpp

static Optional<Value> emitStaticFor(ArrayRef<Value> lbs, ArrayRef<Value> ubs,
   auto *ubDef = ubs.front().getDefiningOp();
   if (!lbDef || !ubDef)
     return Optional<Value>();
 
   auto lbConst = dyn_cast<ConstantIndexOp>(lbDef);
   auto ubConst = dyn_cast<ConstantIndexOp>(ubDef);
   if (!lbConst || !ubConst)
     return Optional<Value>();
-  return ScopedContext::getBuilder()
+  return ScopedContext::getBuilderRef()
       .create<AffineForOp>(ScopedContext::getLocation(), lbConst.getValue(),
                            ubConst.getValue(), step)
       .getInductionVar();
 }
 
 LoopBuilder mlir::edsc::makeAffineLoopBuilder(Value *iv, ArrayRef<Value> lbs,
                                               ArrayRef<Value> ubs,
                                               int64_t step) {
   mlir::edsc::LoopBuilder result;
-  if (auto staticForIv = emitStaticFor(lbs, ubs, step)) {
+  if (auto staticForIv = emitStaticFor(lbs, ubs, step))
     *iv = staticForIv.getValue();
-  } else {
-    auto b = ScopedContext::getBuilder();
-    *iv =
-        Value(b.create<AffineForOp>(ScopedContext::getLocation(), lbs,
-                                    b.getMultiDimIdentityMap(lbs.size()), ubs,
-                                    b.getMultiDimIdentityMap(ubs.size()), step)
-                  .getInductionVar());
-  }
+  else
+    *iv = ScopedContext::getBuilderRef()
+              .create<AffineForOp>(
+                  ScopedContext::getLocation(), lbs,
+                  ScopedContext::getBuilderRef().getMultiDimIdentityMap(
+                      lbs.size()),
+                  ubs,
+                  ScopedContext::getBuilderRef().getMultiDimIdentityMap(
+                      ubs.size()),
+                  step)
+              .getInductionVar();
 
   auto *body = getForInductionVarOwner(*iv).getBody();
   result.enter(body);
   return result;
 }
 
 mlir::edsc::AffineLoopNestBuilder::AffineLoopNestBuilder(Value *iv,
                                                          ArrayRef<Value> lbs,
static Value createBinaryIndexHandle(
   if (v0)
     operands.push_back(v0);
   if (v1)
     operands.push_back(v1);
   auto map = AffineMap::get(numDims, numSymbols, affCombiner(d0, d1));
 
   // TODO: createOrFold when available.
   Operation *op =
-      makeComposedAffineApply(ScopedContext::getBuilder(),
+      makeComposedAffineApply(ScopedContext::getBuilderRef(),
                               ScopedContext::getLocation(), map, operands)
           .getOperation();
   assert(op->getNumResults() == 1 && "Expected single result AffineApply");
   return op->getResult(0);
 }
 
 template <typename IOp, typename FOp>
 static Value createBinaryHandle(
static Value createIComparisonExpr(CmpIPredicate predicate, Value lhs,
   auto lhsType = lhs.getType();
   auto rhsType = rhs.getType();
   (void)lhsType;
   (void)rhsType;
   assert(lhsType == rhsType && "cannot mix types in operators");
   assert((lhsType.isa<IndexType>() || lhsType.isSignlessInteger()) &&
          "only integer comparisons are supported");
 
-  return ScopedContext::getBuilder().create<CmpIOp>(
+  return ScopedContext::getBuilderRef().create<CmpIOp>(
       ScopedContext::getLocation(), predicate, lhs, rhs);
 }
 
 static Value createFComparisonExpr(CmpFPredicate predicate, Value lhs,
                                    Value rhs) {
   auto lhsType = lhs.getType();
   auto rhsType = rhs.getType();
   (void)lhsType;
   (void)rhsType;
   assert(lhsType == rhsType && "cannot mix types in operators");
   assert(lhsType.isa<FloatType>() && "only float comparisons are supported");
 
-  return ScopedContext::getBuilder().create<CmpFOp>(
+  return ScopedContext::getBuilderRef().create<CmpFOp>(
       ScopedContext::getLocation(), predicate, lhs, rhs);
 }
 
 // All floating point comparison are ordered through EDSL
 Value mlir::edsc::op::eq(Value lhs, Value rhs) {
   auto type = lhs.getType();
   return type.isa<FloatType>()
              ? createFComparisonExpr(CmpFPredicate::OEQ, lhs, rhs)

mlir/lib/Dialect/Linalg/EDSC/Builders.cpp

Operation *mlir::edsc::makeGenericLinalgOp(
     ArrayRef<IteratorType> iteratorTypes, ArrayRef<StructuredIndexed> inputs,
     ArrayRef<StructuredIndexed> outputs,
     function_ref<void(ArrayRef<BlockArgument>)> regionBuilder,
     ArrayRef<Value> otherValues, ArrayRef<Attribute> otherAttributes) {
   for (unsigned i = 0, e = outputs.size(); i + 1 < e; ++i)
     assert(!(outputs[i].getType().isa<RankedTensorType>() &&
              outputs[i + 1].getType().isa<MemRefType>()) &&
            "output tensors must be passed after output buffers");
-  auto &builder = edsc::ScopedContext::getBuilder();
+  auto &builder = edsc::ScopedContext::getBuilderRef();
   auto *ctx = builder.getContext();
   unsigned nInputs = inputs.size();
   unsigned nOutputs = outputs.size();
 
   SmallVector<SmallVector<AffineExpr, 4>, 4> exprsList;
   exprsList.reserve(nInputs + nOutputs);
   for (auto structuredIndexed : inputs)
     exprsList.emplace_back(structuredIndexed.getExprs().begin(),
Operation *mlir::edsc::makeGenericLinalgOp(
   SmallVector<Type, 4> types;
   std::copy_if(outputs.begin(), outputs.end(), std::back_inserter(types),
                [](StructuredIndexed s) { return !s.hasValue(); });
 
   auto iteratorStrTypes =
       llvm::to_vector<8>(llvm::map_range(iteratorTypes, toString));
   // clang-format off
   auto *op =
-      edsc::ScopedContext::getBuilder()
+      edsc::ScopedContext::getBuilderRef()
           .create<linalg::GenericOp>(
               edsc::ScopedContext::getLocation(),
               types,
               values,
               IntegerAttr::get(IntegerType::get(64, ctx), nInputs),
               IntegerAttr::get(IntegerType::get(64, ctx), nOutputs),
               builder.getAffineMapArrayAttr(maps),
               builder.getStrArrayAttr(iteratorStrTypes),

mlir/lib/Dialect/Linalg/Transforms/LinalgToLoops.cpp

for (auto e : map.getResults()) {
     canonicalizeMapAndOperands(&exprMap, &operands);
     res.push_back(affine_apply(exprMap, operands));
   }
   return res;
 }
 
 static SmallVector<Value, 4> permuteIvs(ArrayRef<Value> ivs,
                                         Optional<AffineMap> permutation) {
-  return permutation ? applyMapToValues(ScopedContext::getBuilder(),
+  return permutation ? applyMapToValues(ScopedContext::getBuilderRef(),
                                         ScopedContext::getLocation(),
                                         permutation.getValue(), ivs)
                      : SmallVector<Value, 4>(ivs.begin(), ivs.end());
 }
 
 // Creates a number of ranges equal to the number of results in `map`.
 // The returned ranges correspond to the loop ranges, in the proper order, for
 // which new loops will be created.
SmallVector<Value, 4> emitLoopRanges(OpBuilder &b, Location loc, AffineMap map,
   return res;
 }
 
 template <typename OpType>
 static void
 inlineRegionAndEmitStdStore(OpType op, ArrayRef<Value> indexedValues,
                             ArrayRef<SmallVector<Value, 8>> indexing,
                             ArrayRef<Value> outputBuffers) {
-  auto &b = ScopedContext::getBuilder();
+  auto &b = ScopedContext::getBuilderRef();
   auto &block = op.region().front();
   BlockAndValueMapping map;
   map.map(block.getArguments(), indexedValues);
   for (auto &op : block.without_terminator()) {
     assert(op.getNumRegions() == 0 && "expected a non-nested region");
     auto *newOp = b.clone(op, map);
     map.map(op.getResults(), newOp->getResults());
   }

 // SingleInputPoolingOp `op`.
 struct InputAndOutputIndices {
   SmallVector<Value, 8> inputs;
   SmallVector<Value, 8> outputs;
 };
 template <typename SingleInputPoolingOp>
 static InputAndOutputIndices getInputAndOutputIndices(ArrayRef<Value> allIvs,
                                                       SingleInputPoolingOp op) {
-  auto &b = ScopedContext::getBuilder();
+  auto &b = ScopedContext::getBuilderRef();
   auto loc = ScopedContext::getLocation();
   auto mapsRange = op.indexing_maps().template getAsRange<AffineMapAttr>();
   auto maps = llvm::to_vector<8>(
       llvm::map_range(mapsRange, [](AffineMapAttr a) { return a.getValue(); }));
   return InputAndOutputIndices{
       makeCanonicalAffineApplies(b, loc, maps[0], allIvs),
       makeCanonicalAffineApplies(b, loc, maps[2], allIvs)};
 }
 
 namespace {
 
 // Generic loop emitter, to be specialized on an op-per op basis.
 // TODO: Hook up to named ops interface and, later, retire when all named ops
 // are auto-generated.
 template <typename IndexedValueType, typename LinalgOpType>
 class LinalgScopedEmitter {
 public:
   static void emitScalarImplementation(ArrayRef<Value> allIvs,
                                        LinalgOpType linalgOp) {
     assert(linalgOp.hasBufferSemantics() &&
            "expected linalg op with buffer semantics");
     llvm_unreachable("NYI");
-    linalgOp.emitScalarImplementation()(ScopedContext::getBuilder(),
+    linalgOp.emitScalarImplementation()(ScopedContext::getBuilderRef(),
                                         ScopedContext::getLocation(), allIvs);
   }
 };
 
 template <typename IndexedValueType>
 class LinalgScopedEmitter<IndexedValueType, CopyOp> {
 public:
   static void emitScalarImplementation(ArrayRef<Value> allIvs, CopyOp copyOp) {
for (auto iter : llvm::enumerate(imIdx)) {
       auto maxMap = AffineMap::get(/*dimCount=*/1, 0,
                                    {getAffineDimExpr(/*position=*/0, context),
                                     getAffineConstantExpr(0, context)},
                                    context);
       clampedImIdx.push_back(
           affine_max(dim.getType(), maxMap, ValueRange{dim}));
     }
 
-    auto b = ScopedContext::getBuilder();
+    auto &b = ScopedContext::getBuilderRef();
     Type type = convOp.input().getType().cast<MemRefType>().getElementType();
     Value zero = std_constant(type, b.getZeroAttr(type));
     Value readInput = im(clampedImIdx);
     return conds.empty() ? readInput
                          : (Value)std_select(conds.back(), zero, readInput);
   }
 
   static void emitScalarImplementation(ArrayRef<Value> allIvs, ConvOp convOp) {
     assert(convOp.hasBufferSemantics() &&
            "expected linalg op with buffer semantics");
-    auto b = ScopedContext::getBuilder();
+    auto &b = ScopedContext::getBuilderRef();
     auto loc = ScopedContext::getLocation();
     auto mapsRange = convOp.indexing_maps().getAsRange<AffineMapAttr>();
     auto maps = llvm::to_vector<8>(llvm::map_range(
         mapsRange, [](AffineMapAttr a) { return a.getValue(); }));
     SmallVector<Value, 8> fIdx(
         makeCanonicalAffineApplies(b, loc, maps[0], allIvs));
     SmallVector<Value, 8> imIdx(
         makeCanonicalAffineApplies(b, loc, maps[1], allIvs));

 // ```
 template <typename IndexedValueType>
 class LinalgScopedEmitter<IndexedValueType, GenericOp> {
 public:
   static void emitScalarImplementation(ArrayRef<Value> allIvs,
                                        GenericOp genericOp) {
     assert(genericOp.hasBufferSemantics() &&
            "expected linalg op with buffer semantics");
-    auto b = ScopedContext::getBuilder();
+    auto &b = ScopedContext::getBuilderRef();
     auto loc = ScopedContext::getLocation();
     unsigned nInputs = genericOp.getNumInputs();
     unsigned nOutputs = genericOp.getNumOutputs();
     SmallVector<Value, 4> indexedValues(nInputs + nOutputs);
 
     // 1.a. Emit std_load from input views.
     for (unsigned i = 0; i < nInputs; ++i) {
       auto indexing = makeCanonicalAffineApplies(

 // ```
 template <typename IndexedValueType>
 class LinalgScopedEmitter<IndexedValueType, IndexedGenericOp> {
 public:
   static void emitScalarImplementation(ArrayRef<Value> allIvs,
                                        IndexedGenericOp indexedGenericOp) {
     assert(indexedGenericOp.hasBufferSemantics() &&
            "expected linalg op with buffer semantics");
-    auto b = ScopedContext::getBuilder();
+    auto &b = ScopedContext::getBuilderRef();
     auto loc = ScopedContext::getLocation();
     unsigned nInputs = indexedGenericOp.getNumInputs();
     unsigned nOutputs = indexedGenericOp.getNumOutputs();
     unsigned nLoops = allIvs.size();
     SmallVector<Value, 4> indexedValues(nLoops + nInputs + nOutputs);
 
     for (unsigned i = 0; i < nLoops; ++i) {
       indexedValues[i] = allIvs[i];
LinalgOpToLoopsImpl<LoopTy, ConcreteOpTy>::doit(Operation *op,
   if (invertedMap.isEmpty()) {
     LinalgScopedEmitter<IndexedValueTy, ConcreteOpTy>::emitScalarImplementation(
         {}, linalgOp);
     return LinalgLoops();
   }
 
   SmallVector<Value, 4> allIvs(nLoops);
   auto loopRanges =
-      emitLoopRanges(scope.getBuilder(), scope.getLocation(), invertedMap,
+      emitLoopRanges(scope.getBuilderRef(), scope.getLocation(), invertedMap,
                      getViewSizes(rewriter, linalgOp));
   assert(loopRanges.size() == allIvs.size());
   Impl::doit(linalgOp, loopRanges, allIvs);
   // Number of loop ops might be different from the number of ivs since some
   // loops like affine.parallel and loop.parallel have multiple ivs.
   llvm::SetVector<Operation *> loopSet;
   for (Value iv : allIvs) {
     if (!iv)

mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp

Optional<TiledLinalgOp> static tileLinalgOpImpl(OpBuilder &b, LinalgOp op,
   SmallVector<Value, 4> ivs(loopRanges.size());
   // Convert SubViewOp::Range to linalg_range.
   SmallVector<Value, 4> linalgRanges;
   for (auto &range : loopRanges) {
     linalgRanges.push_back(
         linalg_range(range.offset, range.size, range.stride));
   }
   GenericLoopNestRangeBuilder<LoopTy>(ivs, linalgRanges)([&] {
-    auto b = ScopedContext::getBuilder();
+    auto &b = ScopedContext::getBuilderRef();
     auto loc = ScopedContext::getLocation();
     SmallVector<Value, 4> ivValues(ivs.begin(), ivs.end());
 
     // If we have to apply a permutation to the tiled loop nest, we have to
     // reorder the induction variables This permutation is the right one
     // assuming that loopRanges have previously been permuted by
     // (i,j,k)->(k,i,j) So this permutation should be the inversePermutation of
     // that one: (d0,d1,d2)->(d2,d0,d1)

mlir/lib/EDSC/Builders.cpp

 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 
 #include "llvm/ADT/Optional.h"
 
 using namespace mlir;
 using namespace mlir::edsc;
 
-mlir::edsc::ScopedContext::ScopedContext(OpBuilder &builder, Location location)
-    : builder(builder), location(location),
-      enclosingScopedContext(ScopedContext::getCurrentScopedContext()),
-      nestedBuilder(nullptr) {
+mlir::edsc::ScopedContext::ScopedContext(OpBuilder &b, Location location)
+    : builder(b), guard(builder), location(location),
+      enclosingScopedContext(ScopedContext::getCurrentScopedContext()) {
   getCurrentScopedContext() = this;
 }
 
 /// Sets the insertion point of the builder to 'newInsertPt' for the duration
 /// of the scope. The existing insertion point of the builder is restored on
 /// destruction.
-mlir::edsc::ScopedContext::ScopedContext(OpBuilder &builder,
+mlir::edsc::ScopedContext::ScopedContext(OpBuilder &b,
                                          OpBuilder::InsertPoint newInsertPt,
                                          Location location)
-    : builder(builder), prevBuilderInsertPoint(builder.saveInsertionPoint()),
-      location(location),
-      enclosingScopedContext(ScopedContext::getCurrentScopedContext()),
-      nestedBuilder(nullptr) {
+    : builder(b), guard(builder), location(location),
+      enclosingScopedContext(ScopedContext::getCurrentScopedContext()) {
   getCurrentScopedContext() = this;
   builder.restoreInsertionPoint(newInsertPt);
 }
 
 mlir::edsc::ScopedContext::~ScopedContext() {
-  assert(!nestedBuilder &&
-         "Active NestedBuilder must have been exited at this point!");
-  if (prevBuilderInsertPoint)
-    builder.restoreInsertionPoint(*prevBuilderInsertPoint);
   getCurrentScopedContext() = enclosingScopedContext;
 }
 
 ScopedContext *&mlir::edsc::ScopedContext::getCurrentScopedContext() {
   thread_local ScopedContext *context = nullptr;
   return context;
 }
 
-OpBuilder &mlir::edsc::ScopedContext::getBuilder() {
+OpBuilder &mlir::edsc::ScopedContext::getBuilderRef() {
   assert(ScopedContext::getCurrentScopedContext() &&
          "Unexpected Null ScopedContext");
   return ScopedContext::getCurrentScopedContext()->builder;
 }
 
 Location mlir::edsc::ScopedContext::getLocation() {
   assert(ScopedContext::getCurrentScopedContext() &&
          "Unexpected Null ScopedContext");
   return ScopedContext::getCurrentScopedContext()->location;
 }
 
 MLIRContext *mlir::edsc::ScopedContext::getContext() {
-  return getBuilder().getContext();
+  return getBuilderRef().getContext();
 }
 
 BlockHandle mlir::edsc::BlockHandle::create(ArrayRef<Type> argTypes) {
-  auto &currentB = ScopedContext::getBuilder();
+  auto &currentB = ScopedContext::getBuilderRef();
   auto *ib = currentB.getInsertionBlock();
   auto ip = currentB.getInsertionPoint();
   BlockHandle res;
-  res.block = ScopedContext::getBuilder().createBlock(ib->getParent());
+  res.block = ScopedContext::getBuilderRef().createBlock(ib->getParent());
   // createBlock sets the insertion point inside the block.
   // We do not want this behavior when using declarative builders with nesting.
   currentB.setInsertionPoint(ib, ip);
   for (auto t : argTypes) {
     res.block->addArgument(t);
   }
   return res;
 }
 
 BlockHandle mlir::edsc::BlockHandle::createInRegion(Region &region,
                                                     ArrayRef<Type> argTypes) {
-  auto &currentB = ScopedContext::getBuilder();
   BlockHandle res;
   region.push_back(new Block);
   res.block = &region.back();
   // createBlock sets the insertion point inside the block.
   // We do not want this behavior when using declarative builders with nesting.
-  OpBuilder::InsertionGuard g(currentB);
-  currentB.setInsertionPoint(res.block, res.block->begin());
-  for (auto t : argTypes) {
+  OpBuilder::InsertionGuard g(ScopedContext::getBuilderRef());
+  ScopedContext::getBuilderRef().setInsertionPoint(res.block,
+                                                   res.block->begin());
+  for (auto t : argTypes)
     res.block->addArgument(t);
-  }
   return res;
 }
 
 void mlir::edsc::LoopBuilder::operator()(function_ref<void(void)> fun) {
   // Call to `exit` must be explicit and asymmetric (cannot happen in the
   // destructor) because of ordering wrt comma operator.
   /// The particular use case concerns nested blocks:
   ///

mlir/test/Dialect/Linalg/matmul-to-vector.mlir

-// RUN: mlir-opt %s -linalg-matmul-to-vector | FileCheck %s
-
-func @matmul_perm(%A: memref<1584x1584xf32, offset: 0, strides: [1584, 1]>,
-                  %B: memref<1584x1584xf32, offset: 0, strides: [1584, 1]>,
-                  %C: memref<1584x1584xf32, offset: 0, strides: [1584, 1]>) {
-  linalg.matmul(%A, %B, %C) {__internal_linalg_transform__ = "__with_perm__"} :
-    memref<1584x1584xf32, offset: 0, strides: [1584, 1]>,
-    memref<1584x1584xf32, offset: 0, strides: [1584, 1]>,
-    memref<1584x1584xf32, offset: 0, strides: [1584, 1]>
-  return
-}
-
-// CHECK-LABEL:func @matmul_perm
-//      CHECK:        vector.contract
-// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]
-// CHECK-SAME: : vector<8x16xf32>, vector<16x12xf32> into vector<8x12xf32>

mlir/test/lib/DeclarativeTransforms/CMakeLists.txt

 set(LLVM_TARGET_DEFINITIONS TestLinalgTransformPatterns.td)
 mlir_tablegen(TestLinalgTransformPatterns.h.inc -gen-rewriters)
 add_public_tablegen_target(MLIRTestLinalgTransformPatternsIncGen)
 # Including Linalg in TableGen requires to depends on generated files
 add_dependencies(MLIRTestLinalgTransformPatternsIncGen LinalgOdsGen)
 
 set(LLVM_TARGET_DEFINITIONS TestVectorTransformPatterns.td)
 mlir_tablegen(TestVectorTransformPatterns.h.inc -gen-rewriters)
 add_public_tablegen_target(MLIRTestVectorTransformPatternsIncGen)
-
-set(LLVM_TARGET_DEFINITIONS TestLinalgMatmulToVectorPatterns.td)
-mlir_tablegen(TestLinalgMatmulToVectorPatterns.h.inc -gen-rewriters)
-add_public_tablegen_target(MLIRTestLinalgMatmulToVectorPatternsIncGen)
-# Including Linalg in TableGen requires to depends on generated files
-add_dependencies(MLIRTestLinalgTransformPatternsIncGen LinalgOdsGen)

mlir/test/lib/DeclarativeTransforms/TestLinalgMatmulToVectorPatterns.td

-//===- TestLinalgMatmulToVectorPatterns.td - Test patterns -*- tablegen -*-===//
-//
-// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
-// See https://llvm.org/LICENSE.txt for license information.
-// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-//
-//===----------------------------------------------------------------------===//
-//
-// This is the pattern definition file for declarative Linalg transformations
-// tests.
-//
-//===----------------------------------------------------------------------===//
-
-#ifndef TEST_LINALG_MATMUL_TO_VECTOR_PATTERNS
-#define TEST_LINALG_MATMUL_TO_VECTOR_PATTERNS
-
-include "mlir/Dialect/Linalg/Transforms/LinalgTransformPatterns.td"
-include "mlir/Dialect/Vector/VectorTransformPatterns.td"
-
-//===----------------------------------------------------------------------===//
-// Linalg tiling and permutation patterns.
-//===----------------------------------------------------------------------===//
-def : Pat<(MatmulOp:$op $_, $_, $_),
-          (TileLinalgOp<[768, 264, 768], "L2__with_perm__", [1, 2, 0]>),
-         [(Constraint<HasLinalgTransformMarker<"__with_perm__">>)]>;
-def : Pat<(MatmulOp:$op $_, $_, $_),
-          (TileLinalgOp<[8, 12, 16], "L1__with_perm__", [1, 0, 2]>),
-         [(Constraint<HasLinalgTransformMarker<"L2__with_perm__">>)]>;
-def : Pat<(MatmulOp:$op $_, $_, $_),
-          (PromoteSubviewsLinalgOp),
-         [(Constraint<HasOperandsOfType<"SubViewOp">>),
-          (Constraint<HasLinalgTransformMarker<"L1__with_perm__">>)]>;
-
-//===----------------------------------------------------------------------===//
-// Linalg to vector contraction patterns.
-//===----------------------------------------------------------------------===//
-def : Pattern<(MatmulOp:$op $_, $_, $_),
-              [(VectorizeLinalgOp)],
-              [(Constraint<And<[
-                HasLinalgTransformMarker<"L1__with_perm__">,
-                PreconditionVectorizeLinalgOp]>>)]>;
-
-#endif // TEST_LINALG_MATMUL_TO_VECTOR_PATTERNS

mlir/test/lib/Transforms/CMakeLists.txt

 add_llvm_library(MLIRTestTransforms
   TestAllReduceLowering.cpp
   TestBufferPlacement.cpp
   TestCallGraph.cpp
   TestConstantFold.cpp
   TestConvertGPUKernelToCubin.cpp
   TestDominance.cpp
   TestLoopFusion.cpp
   TestGpuMemoryPromotion.cpp
   TestGpuParallelLoopMapping.cpp
   TestInlining.cpp
-  TestLinalgMatmulToVector.cpp
   TestLinalgTransforms.cpp
   TestLiveness.cpp
   TestLoopMapping.cpp
   TestLoopParametricTiling.cpp
   TestOpaqueLoc.cpp
   TestMemRefBoundCheck.cpp
   TestMemRefDependenceCheck.cpp
   TestMemRefStrideCalculation.cpp
   TestVectorToLoopsConversion.cpp
   TestVectorTransforms.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Transforms
 
   DEPENDS
   MLIRStandardOpsIncGen
-  MLIRTestLinalgMatmulToVectorPatternsIncGen
   MLIRTestLinalgTransformPatternsIncGen
   MLIRTestVectorTransformPatternsIncGen
 )
 
 include_directories(${CMAKE_CURRENT_SOURCE_DIR}/../Dialect/Test)
 include_directories(${CMAKE_CURRENT_BINARY_DIR}/../Dialect/Test)
 include_directories(${CMAKE_CURRENT_SOURCE_DIR}/../DeclarativeTransforms)
 include_directories(${CMAKE_CURRENT_BINARY_DIR}/../DeclarativeTransforms)

mlir/test/lib/Transforms/TestLinalgMatmulToVector.cpp

-//===- TestLinalgMatmulToVector.cpp - Test VectorTransfers lowering -------===//
-//
-// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
-// See https://llvm.org/LICENSE.txt for license information.
-// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-//
-//===----------------------------------------------------------------------===//
-
-#include <type_traits>
-
-#include "mlir/Dialect/Affine/IR/AffineOps.h"
-#include "mlir/Dialect/Linalg/IR/LinalgOps.h"
-#include "mlir/Dialect/Linalg/Transforms/LinalgTransforms.h"
-#include "mlir/Dialect/Vector/VectorOps.h"
-#include "mlir/Dialect/Vector/VectorTransforms.h"
-#include "mlir/IR/PatternMatch.h"
-#include "mlir/IR/StandardTypes.h"
-#include "mlir/Pass/Pass.h"
-
-using namespace mlir;
-using namespace mlir::linalg;
-using namespace mlir::vector;
-
-namespace {
-#include "TestLinalgMatmulToVectorPatterns.h.inc"
-
-struct DeclarativeTransforms
-    : public PassWrapper<DeclarativeTransforms, FunctionPass> {
-  void runOnFunction() override {
-    OwningRewritePatternList patterns;
-    auto *context = &getContext();
-    AffineApplyOp::getCanonicalizationPatterns(patterns, context);
-    AffineMinOp::getCanonicalizationPatterns(patterns, context);
-    AffineMaxOp::getCanonicalizationPatterns(patterns, context);
-    AllocOp::getCanonicalizationPatterns(patterns, context);
-    SubViewOp::getCanonicalizationPatterns(patterns, context);
-    ViewOp::getCanonicalizationPatterns(patterns, context);
-    populateWithGenerated(context, &patterns);
-    applyPatternsAndFoldGreedily(getFunction(), patterns);
-  }
-};
-} // end anonymous namespace
-
-namespace mlir {
-void registerTestLinalgMatmulToVectorPass() {
-  PassRegistration<DeclarativeTransforms> pass(
-      "linalg-matmul-to-vector",
-      "Test declarative transform patterns for matmul 3-D tiling + promotion"
-      " + vectorization");
-}
-} // namespace mlir

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerPrintOpAvailabilityPass();
 void registerSideEffectTestPasses();
 void registerSimpleParametricTilingPass();
 void registerSymbolTestPasses();
 void registerTestAffineDataCopyPass();
 void registerTestAllReduceLoweringPass();
 void registerTestAffineLoopUnswitchingPass();
 void registerTestBufferPlacementPreparationPass();
-void registerTestLinalgMatmulToVectorPass();
 void registerTestLoopPermutationPass();
 void registerTestCallGraphPass();
 void registerTestConstantFold();
 void registerTestConvertGPUKernelToCubinPass();
 void registerTestDominancePass();
 void registerTestFunc();
 void registerTestGpuMemoryPromotionPass();
 void registerTestLinalgTransforms();
void registerTestPasses() {
   registerPatternsTestPass();
   registerPrintOpAvailabilityPass();
   registerSideEffectTestPasses();
   registerSimpleParametricTilingPass();
   registerSymbolTestPasses();
   registerTestAffineDataCopyPass();
   registerTestAllReduceLoweringPass();
   registerTestAffineLoopUnswitchingPass();
-  registerTestLinalgMatmulToVectorPass();
   registerTestLoopPermutationPass();
   registerTestCallGraphPass();
   registerTestConstantFold();
 #if MLIR_CUDA_CONVERSIONS_ENABLED
   registerTestConvertGPUKernelToCubinPass();
 #endif
   registerTestBufferPlacementPreparationPass();
   registerTestDominancePass();