[mlir][Linalg] NFC - Purge OpBuilder uses in favor of RewriterBase in places unrelated to op definitions

[nicolasvasilache]

Summary

RewriterBase is the proper builder to use so one can listen to IR modifications (i.e. not just creation).

[nicolasvasilache]

Drop spurious text changes.

[nicolasvasilache]

More spurious text change reversal.

[springerm]

What's the reason for this? I've been "cleaning up" some code in the past by using OpBuilder instead of PatternRewriter/RewriterBase/... when possible. The reason for this was better compatibility (both builders and rewriters can be used) and code documentation (when the function decl takes an OpBuilder, the caller can be sure that the function does only create new IR and never deletes/modifies existing IR).

If using RewriterBase is actually a best practice, we should document it somewhere, otherwise, somebody may come along with the same idea as me in a few months and undo your change.

[mravishankar]

I dont see a strong reason to do this. My thumb rule has been whenever there is need to erase operations use RewriterBase. If this is just building IR, OpBuilder usage is more stable IMO and more widely applicable.

[nicolasvasilache]

In D137922#3924882, @springerm wrote:

What's the reason for this? I've been "cleaning up" some code in the past by using OpBuilder instead of PatternRewriter/RewriterBase/... when possible. The reason for this was better compatibility (both builders and rewriters can be used) and code documentation (when the function decl takes an OpBuilder, the caller can be sure that the function does only create new IR and never deletes/modifies existing IR).

If using RewriterBase is actually a best practice, we should document it somewhere, otherwise, somebody may come along with the same idea as me in a few months and undo your change.

In D137922#3925266, @mravishankar wrote:

I dont see a strong reason to do this. My thumb rule has been whenever there is need to erase operations use RewriterBase. If this is just building IR, OpBuilder usage is more stable IMO and more widely applicable.

There are multiple reasons IMO.

https://reviews.llvm.org/D137690 significantly increases the capabilities of the vectorizer by adding masking, as part of that the following code is introduced:

Operation *maskOpTerminator = &maskOp.getMaskRegion().front().back();
for (auto &en : llvm::enumerate(opToMask->getResults()))
  en.value().replaceAllUsesExcept(maskOp.getResult(en.index()),
                                  maskOpTerminator);

This is an enormous footgun and it needs to go through a RewriterBase updateRootInPlace which will permeate through the whole file all the way to the top.

Using a RewriterBase also encourages proper use of the debugging infra and avoids gross workarounds such as IRRewriter(b).notifyMatchFailure().

Lastly, from a consistency and cognitive overhead perspective, using a single thing is simpler: if you write a transform use RewriterBase, if you write a builder use OpBuilder.
Otherwise the rule of thumb is: use a RewriterBase *in the whole call stack to any leaf that modifies IR* (either erase, inplace update or update of uses) otherwise it may be ok to use OpBuilder but you may have to modify the whole call stack if in the future IR modifications occur (otherwise very subtle bugs may appear).

[springerm]

In D137922#3963080, @nicolasvasilache wrote:

https://reviews.llvm.org/D137690 significantly increases the capabilities of the vectorizer by adding masking, as part of that the following code is introduced:

Operation *maskOpTerminator = &maskOp.getMaskRegion().front().back();
for (auto &en : llvm::enumerate(opToMask->getResults()))
  en.value().replaceAllUsesExcept(maskOp.getResult(en.index()),
                                  maskOpTerminator);

This is an enormous footgun and it needs to go through a RewriterBase updateRootInPlace which will permeate through the whole file all the way to the top.

Such code should be flagged during review (as you did). I've seen people using Value::replaceAllUses despite having a RewriterBase readily available in the same function, so whether the function takes an OpBuilder or a RewriterBase may not make a difference here.

Lastly, from a consistency and cognitive overhead perspective, using a single thing is simpler: if you write a transform use RewriterBase, if you write a builder use OpBuilder.

Some helper functions in Transforms/Vectorization.cpp are builders, e.g., buildVectorWrite.

[nicolasvasilache]

In D137922#3963091, @springerm wrote:

In D137922#3963080, @nicolasvasilache wrote:

https://reviews.llvm.org/D137690 significantly increases the capabilities of the vectorizer by adding masking, as part of that the following code is introduced:

Operation *maskOpTerminator = &maskOp.getMaskRegion().front().back();
for (auto &en : llvm::enumerate(opToMask->getResults()))
  en.value().replaceAllUsesExcept(maskOp.getResult(en.index()),
                                  maskOpTerminator);

This is an enormous footgun and it needs to go through a RewriterBase updateRootInPlace which will permeate through the whole file all the way to the top.

Such code should be flagged during review (as you did). I've seen people using Value::replaceAllUses despite having a RewriterBase readily available in the same function, so whether the function takes an OpBuilder or a RewriterBase may not make a difference here.

Really this should be type safe and not best effort hope reviewers catch it .. we should make these private, expose to a very limited set of friends and force everything else through RewriterBase..
The current situation is just terrible.

Lastly, from a consistency and cognitive overhead perspective, using a single thing is simpler: if you write a transform use RewriterBase, if you write a builder use OpBuilder.

Some helper functions in Transforms/Vectorization.cpp are builders, e.g., buildVectorWrite.

Good point, some of these need to move indeed.

[mravishankar]

In D137922#3963080, @nicolasvasilache wrote:

In D137922#3924882, @springerm wrote:

What's the reason for this? I've been "cleaning up" some code in the past by using OpBuilder instead of PatternRewriter/RewriterBase/... when possible. The reason for this was better compatibility (both builders and rewriters can be used) and code documentation (when the function decl takes an OpBuilder, the caller can be sure that the function does only create new IR and never deletes/modifies existing IR).

If using RewriterBase is actually a best practice, we should document it somewhere, otherwise, somebody may come along with the same idea as me in a few months and undo your change.

In D137922#3925266, @mravishankar wrote:

I dont see a strong reason to do this. My thumb rule has been whenever there is need to erase operations use RewriterBase. If this is just building IR, OpBuilder usage is more stable IMO and more widely applicable.

There are multiple reasons IMO.

https://reviews.llvm.org/D137690 significantly increases the capabilities of the vectorizer by adding masking, as part of that the following code is introduced:

Operation *maskOpTerminator = &maskOp.getMaskRegion().front().back();
for (auto &en : llvm::enumerate(opToMask->getResults()))
  en.value().replaceAllUsesExcept(maskOp.getResult(en.index()),
                                  maskOpTerminator);

This is an enormous footgun and it needs to go through a RewriterBase updateRootInPlace which will permeate through the whole file all the way to the top.

Using a RewriterBase also encourages proper use of the debugging infra and avoids gross workarounds such as IRRewriter(b).notifyMatchFailure().

+1 to using notifyMatchFailure. That should be probably even moved out of RewriterBase . Its a good debugging tool irrespective of this change.

Lastly, from a consistency and cognitive overhead perspective, using a single thing is simpler: if you write a transform use RewriterBase, if you write a builder use OpBuilder.
Otherwise the rule of thumb is: use a RewriterBase *in the whole call stack to any leaf that modifies IR* (either erase, inplace update or update of uses) otherwise it may be ok to use OpBuilder but you may have to modify the whole call stack if in the future IR modifications occur (otherwise very subtle bugs may appear).

I was writing an explanation as to why this shouldnt be changed, but actually realized those explanations dont hold water. This looks good to me.

[dcaballe]

Thanks for the feedback! Please note that some of us come from LLVM where builders and UD chain manipulation is the daily basis :). Not sure I totally agree with the level of usage restriction suggested for some of these methods but I'm happy to try to do the same with a rewrite. Note, though, that this code is not doing a replacement but it's building a new mask operation and moving another existing operation into the mask operation region. Can we do that with updateRootInPlace?

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

1426

s/b/rewriter? (+readability). I think it's important to state that it's a rewriter and not a builder in the code.

[nicolasvasilache]

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

1426

+1, thanks!

[nicolasvasilache]

b -> rewriter

[nicolasvasilache]

Address commetns.

mlir/include/mlir/Dialect/Utils/StructuredOpsUtils.h

 #include "mlir/IR/Location.h"
 #include "mlir/Support/LLVM.h"
 
 // Pull in all enum type definitions and utility function declarations.
 #include "mlir/Dialect/Utils/DialectUtilsEnums.h.inc"
 
 namespace mlir {
 
-class OpBuilder;
+class RewriterBase;
 
 /// Tests whether the given maps describe a row major matmul. The test is
 /// permutation-invariant. Note that this only checks the affine maps from an
 /// operation, so does not perform any checks on the math being performed within
 /// the reduction.
 bool isRowMajorMatmul(ArrayAttr indexingMaps);
 
 /// Tests whether the given maps describe a column major matmul. The test is
public:
   };
   struct Par : public IteratorType {
     Par() : IteratorType(IteratorTypeT::parallel) {}
   };
   struct Red : public IteratorType {
     Red() : IteratorType(IteratorTypeT::reduction) {}
   };
 
-  StructuredGenerator(OpBuilder &builder, StructuredOpInterface op)
-      : builder(builder), ctx(op.getContext()), loc(op.getLoc()),
+  StructuredGenerator(RewriterBase &b, StructuredOpInterface op)
+      : b(b), ctx(op.getContext()), loc(op.getLoc()),
         iterators(op.getIteratorTypesArray()), maps(op.getIndexingMapsArray()),
         op(op) {}
 
   bool iters(ArrayRef<IteratorType> its) {
     if (its.size() != iterators.size())
       return false;
     for (int i = 0, e = its.size(); i != e; ++i) {
       if (!its[i].isOfType(iterators[i]))
         return false;
     }
     return true;
   }
 
   bool layout(MapList l) {
     auto infer = [](MapList m) { return AffineMap::inferFromExprList(m); };
     return maps == infer(l);
   }
 
 protected:
-  OpBuilder &builder;
+  RewriterBase &b;
   MLIRContext *ctx;
   Location loc;
   SmallVector<IteratorTypeT> iterators;
   SmallVector<AffineMap, 4> maps;
   Operation *op;
 };
 
 } // namespace mlir
 
 #endif // MLIR_DIALECT_UTILS_STRUCTUREDOPSUTILS_H

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

 using namespace mlir::linalg;
 
 #define DEBUG_TYPE "linalg-vectorization"
 
 #define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ")
 #define LDBG(X) LLVM_DEBUG(DBGS() << X)
 
 /// Try to vectorize `convOp` as a convolution.
-static FailureOr<Operation *> vectorizeConvolution(OpBuilder &b,
+static FailureOr<Operation *> vectorizeConvolution(RewriterBase &b,
                                                    LinalgOp convOp);
 
 /// Return the unique instance of OpType in `block` if it is indeed unique.
 /// Return null if none or more than 1 instances exist.
 template <typename OpType>
 static OpType getSingleOpOfType(Block &block) {
   OpType res;
   block.walk([&](OpType op) {
if (!matchReduction(linalgOp.getRegionOutputArgs(), outputPos, combinerOps) ||
     return nullptr;
 
   // Return the combiner operation.
   return combinerOps[0];
 }
 
 /// Broadcast `value` to a vector of `shape` if possible. Return value
 /// otherwise.
-static Value broadcastIfNeeded(OpBuilder &b, Value value,
+static Value broadcastIfNeeded(RewriterBase &b, Value value,
                                ArrayRef<int64_t> shape) {
   // If no shape to broadcast to, just return `value`.
   if (shape.empty())
     return value;
   VectorType targetVectorType =
       VectorType::get(shape, getElementTypeOrSelf(value));
   if (vector::isBroadcastableTo(value.getType(), targetVectorType) !=
       vector::BroadcastableToResult::Success)
     return value;
   Location loc = b.getInsertionPoint()->getLoc();
   return b.createOrFold<vector::BroadcastOp>(loc, targetVectorType, value);
 }
 
 /// Create MultiDimReductionOp to compute the reduction for `reductionOp`. This
 /// assumes that `reductionOp` has two operands and one of them is the reduction
 /// initial value.
-static Operation *buildMultiDimReduce(OpBuilder &b, Operation *reduceOp,
+static Operation *buildMultiDimReduce(RewriterBase &b, Operation *reduceOp,
                                       Value valueToReduce, Value acc,
                                       const SmallVector<bool> &reductionMask) {
   auto maybeKind = getCombinerOpKind(reduceOp);
   assert(maybeKind && "Failed precondition: could not get reduction kind");
   return b.create<vector::MultiDimReductionOp>(
       reduceOp->getLoc(), valueToReduce, acc, reductionMask, *maybeKind);
 }
 
 static SmallVector<bool> getReductionMask(LinalgOp linalgOp) {
   return llvm::to_vector(
       llvm::map_range(linalgOp.getIteratorTypesArray(), isReductionIterator));
 }
 
 /// Build a vector.transfer_write of `value` into `outputOperand` at indices set
 /// to all `0`; where `outputOperand` is an output operand of the LinalgOp
 /// currently being vectorized. If `dest` has null rank, build an memref.store.
 /// Return the produced value or null if no value is produced.
-static Value buildVectorWrite(OpBuilder &b, Value value,
+static Value buildVectorWrite(RewriterBase &b, Value value,
                               OpOperand *outputOperand) {
   Operation *write;
   Location loc = value.getLoc();
   auto linalgOp = cast<LinalgOp>(outputOperand->getOwner());
   ArrayRef<int64_t> shape = linalgOp.getShape(outputOperand);
   auto vectorType = VectorType::get(
       shape, getElementTypeOrSelf(outputOperand->get().getType()));
   if (vectorType.getRank() > 0) {
static Value buildVectorWrite(RewriterBase &b, Value value,
   }
   LDBG("vectorized op: " << *write);
   if (!write->getResults().empty())
     return write->getResult(0);
   return Value();
 }
 
 // Custom vectorization precondition function type. This is intented to be used
-// with CustomVectorizationHook. Returns success if the correpsonding custom
+// with CustomVectorizationHook. Returns success if the corresponding custom
 // hook can vectorize the op.
 using CustomVectorizationPrecondition =
     std::function<LogicalResult(Operation *)>;
 
 // Custom vectorization function type. Produce a vector form of Operation*
 // assuming all its vectorized operands are already in the BlockAndValueMapping.
 // Return nullptr if the Operation cannot be vectorized.
 using CustomVectorizationHook = std::function<VectorizationResult(
     Operation *, const BlockAndValueMapping &)>;
 
 /// Helper function to vectorize the terminator of a `linalgOp`. New result
 /// vector values are appended to `newResults`. Return
 /// VectorizationStatus::NoReplace to signal the vectorization algorithm that it
 /// should not try to map produced operations and instead return the results
-/// using the `newResults` vector making them available to the
-/// vectorization algorithm for RAUW. This function is meant to be used as a
+/// using the `newResults` vector making them available to the vectorization
+/// algorithm for RAUW. This function is meant to be used as a
 /// CustomVectorizationHook.
 static VectorizationResult
-vectorizeLinalgYield(OpBuilder &b, Operation *op,
+vectorizeLinalgYield(RewriterBase &b, Operation *op,
                      const BlockAndValueMapping &bvm, LinalgOp linalgOp,
                      SmallVectorImpl<Value> &newResults) {
   auto yieldOp = dyn_cast<linalg::YieldOp>(op);
   if (!yieldOp)
     return VectorizationResult{VectorizationStatus::Failure, nullptr};
   for (const auto &outputs : llvm::enumerate(yieldOp.getValues())) {
     // TODO: Scan for an opportunity for reuse.
     // TODO: use a map.
     Value vectorValue = bvm.lookup(outputs.value());
     Value newResult = buildVectorWrite(
         b, vectorValue, linalgOp.getDpsInitOperand(outputs.index()));
     if (newResult)
       newResults.push_back(newResult);
   }
   return VectorizationResult{VectorizationStatus::NoReplace, nullptr};
 }
 
 /// Helper function to vectorize the index operations of a `linalgOp`. Return
 /// VectorizationStatus::NewOp to signal the vectorization algorithm that it
 /// should map the produced operations. This function is meant to be used as a
 /// CustomVectorizationHook.
-static VectorizationResult vectorizeLinalgIndex(OpBuilder &b, Operation *op,
+static VectorizationResult vectorizeLinalgIndex(RewriterBase &b, Operation *op,
                                                 LinalgOp linalgOp) {
   IndexOp indexOp = dyn_cast<linalg::IndexOp>(op);
   if (!indexOp)
     return VectorizationResult{VectorizationStatus::Failure, nullptr};
   auto loc = indexOp.getLoc();
   // Compute the static loop sizes of the index op.
   auto targetShape = linalgOp.computeStaticLoopSizes();
   // Compute a one-dimensional index vector for the index op dimension.
static LogicalResult tensorExtractVectorizationPrecondition(Operation *op) {
   return success();
 }
 
 /// Helper function to vectorize the tensor.extract operations. Returns
 /// VectorizationStatus::NewOp to signal the vectorization algorithm that it
 /// should map the produced operations. This function is meant to be used as a
 /// CustomVectorizationHook.
 static VectorizationResult
-vectorizeTensorExtract(OpBuilder &b, Operation *op, LinalgOp linalgOp,
+vectorizeTensorExtract(RewriterBase &b, Operation *op, LinalgOp linalgOp,
                        const BlockAndValueMapping &bvm) {
   tensor::ExtractOp extractOp = dyn_cast<tensor::ExtractOp>(op);
   if (!extractOp)
     return VectorizationResult{VectorizationStatus::Failure, nullptr};
   auto loc = extractOp.getLoc();
 
   // Currently only supports extraction with an 1-D index. Checked in the
   // tensorExtractVectorizationPrecondition.
auto gatherOp = b.create<vector::GatherOp>(
       loc, resultType, extractOp.getTensor(), gatherIndices, indexVec,
       maskConstantOp, passThruConstantOp);
 
   return VectorizationResult{VectorizationStatus::NewOp, gatherOp};
 }
 
 /// Emit reduction operations if the shapes of the value to reduce is different
 /// that the result shape.
-static Operation *reduceIfNeeded(OpBuilder &b, LinalgOp linalgOp, Operation *op,
-                                 Value reduceValue, Value initialValue,
+static Operation *reduceIfNeeded(RewriterBase &b, LinalgOp linalgOp,
+                                 Operation *op, Value reduceValue,
+                                 Value initialValue,
                                  const BlockAndValueMapping &bvm) {
   Value reduceVec = bvm.lookup(reduceValue);
   Value outputVec = bvm.lookup(initialValue);
   auto reduceType = reduceVec.getType().dyn_cast<VectorType>();
   auto outputType = outputVec.getType().dyn_cast<VectorType>();
   // Reduce only if needed as the value may already have been reduce for
   // contraction vectorization.
   if (!reduceType ||

 ///   3. Fail on any remaining non `ElementwiseMappable` op. It is the purpose
 ///   of the `customVectorizationHooks` to cover such cases.
 ///   4. Clone `op` in vector form to a vector of shape prescribed by the first
 ///   operand of maximal rank. Other operands have smaller rank and are
 ///   broadcast accordingly. It is assumed this broadcast is always legal,
 ///   otherwise, it means one of the `customVectorizationHooks` is incorrect.
 ///
 /// This function assumes all operands of `op` have been vectorized and are in
-/// the `bvm` mapping. As a consequence, this function is meant to be called on
+/// the `bvm` mapping. As a consequence, this function is meant to be called  on
 /// a topologically-sorted list of ops.
 /// This function does not update `bvm` but returns a VectorizationStatus that
 /// instructs the caller what `bvm` update needs to occur.
 static VectorizationResult
-vectorizeOneOp(OpBuilder &b, LinalgOp linalgOp, Operation *op,
+vectorizeOneOp(RewriterBase &b, LinalgOp linalgOp, Operation *op,
                const BlockAndValueMapping &bvm,
                ArrayRef<CustomVectorizationHook> customVectorizationHooks) {
   LDBG("vectorize op " << *op);
 
   // 1. Try to apply any CustomVectorizationHook.
   if (!customVectorizationHooks.empty()) {
     for (auto &customFunc : customVectorizationHooks) {
       VectorizationResult result = customFunc(op, bvm);

 /// iteration space. This eager broadcasting is introduced in the
 /// permutation_map of the vector.transfer_read operations. The eager
 /// broadcasting makes it trivial to detrmine where broadcast, transposes and
 /// reductions should occur, without any bookkeeping. The tradeoff is that, in
 /// the absence of good canonicalizations, the amount of work increases.
 /// This is not deemed a problem as we expect canonicalizations and foldings to
 /// aggressively clean up the useless work.
 static LogicalResult
-vectorizeAsLinalgGeneric(OpBuilder &b, LinalgOp linalgOp,
+vectorizeAsLinalgGeneric(RewriterBase &b, LinalgOp linalgOp,
                          SmallVectorImpl<Value> &newResults) {
   Block *block = linalgOp.getBlock();
 
   // 2. Values defined above the region can only be broadcast for now. Make them
   // map to themselves.
   BlockAndValueMapping bvm;
   SetVector<Value> valuesSet;
   mlir::getUsedValuesDefinedAbove(linalgOp->getRegion(0), valuesSet);

 /// Helper function that retrieves the value of an IntegerAttr.
 static int64_t getIntFromAttr(Attribute attr) {
   return attr.cast<IntegerAttr>().getInt();
 }
 
 /// Given an ArrayRef of OpFoldResults, return a vector of Values.
 /// IntegerAttrs are converted to ConstantIndexOps. Other attribute types are
 /// not supported.
-static SmallVector<Value> ofrToIndexValues(OpBuilder &builder, Location loc,
+static SmallVector<Value> ofrToIndexValues(RewriterBase &b, Location loc,
                                            ArrayRef<OpFoldResult> ofrs) {
   SmallVector<Value> result;
   for (auto o : ofrs) {
     if (auto val = o.template dyn_cast<Value>()) {
       result.push_back(val);
     } else {
-      result.push_back(builder.create<arith::ConstantIndexOp>(
+      result.push_back(b.create<arith::ConstantIndexOp>(
           loc, getIntFromAttr(o.template get<Attribute>())));
     }
   }
   return result;
 }
 
 /// Rewrite a tensor::PadOp into a sequence of EmptyOp, FillOp and
 /// InsertSliceOp. For now, only constant padding values are supported.

 /// ```
 ///   Op def: (     n,     w,     c,    kw )
 ///    Iters: ({Par(), Par(), Par(), Red()})
 ///   Layout: {{n, strideW * w + dilationW * kw, c}, {kw, c}, {n, w, c}}
 /// ```
 /// kw is unrolled, w is unrolled iff dilationW > 1.
 struct Conv1DGenerator
     : public StructuredGenerator<LinalgOp, utils::IteratorType> {
-  Conv1DGenerator(OpBuilder &builder, LinalgOp linalgOp, int strideW,
+  Conv1DGenerator(RewriterBase &b, LinalgOp linalgOp, int strideW,
                   int dilationW)
-      : StructuredGenerator<LinalgOp, utils::IteratorType>(builder, linalgOp),
+      : StructuredGenerator<LinalgOp, utils::IteratorType>(b, linalgOp),
         strideW(strideW), dilationW(dilationW) {
     // Determine whether `linalgOp` can be generated with this generator
     if (linalgOp.getNumDpsInputs() != 2 || linalgOp.getNumDpsInits() != 1)
       return;
     lhsShaped = linalgOp.getDpsInputOperand(0)->get();
     rhsShaped = linalgOp.getDpsInputOperand(1)->get();
     resShaped = linalgOp.getDpsInitOperand(0)->get();
     lhsShapedType = lhsShaped.getType().dyn_cast<ShapedType>();
case Conv1DOpOrder::Ncw:
                   ((wSize - 1) * strideW + 1) + ((kwSize - 1) * dilationW + 1) -
                       1};
       rhsShape = {fSize, cSize, kwSize};
       resShape = {nSize, fSize, wSize};
       break;
     }
 
     vector::TransferWriteOp write;
-    Value zero = builder.create<arith::ConstantIndexOp>(loc, 0);
+    Value zero = b.create<arith::ConstantIndexOp>(loc, 0);
 
     // w is unrolled (i.e. wSizeStep == 1) iff strideW > 1.
     // When strideW == 1, we can batch the contiguous loads and avoid
     // unrolling
     int64_t wSizeStep = strideW == 1 ? wSize : 1;
 
     Type lhsEltType = lhsShapedType.getElementType();
     Type rhsEltType = rhsShapedType.getElementType();
     Type resEltType = resShapedType.getElementType();
     auto lhsType = VectorType::get(lhsShape, lhsEltType);
     auto rhsType = VectorType::get(rhsShape, rhsEltType);
     auto resType = VectorType::get(resShape, resEltType);
     // Read lhs slice of size {w * strideW + kw * dilationW, c, f} @ [0, 0,
     // 0].
-    Value lhs = builder.create<vector::TransferReadOp>(
-        loc, lhsType, lhsShaped, ValueRange{zero, zero, zero});
+    Value lhs = b.create<vector::TransferReadOp>(loc, lhsType, lhsShaped,
+                                                 ValueRange{zero, zero, zero});
     // Read rhs slice of size {kw, c, f} @ [0, 0, 0].
-    Value rhs = builder.create<vector::TransferReadOp>(
-        loc, rhsType, rhsShaped, ValueRange{zero, zero, zero});
+    Value rhs = b.create<vector::TransferReadOp>(loc, rhsType, rhsShaped,
+                                                 ValueRange{zero, zero, zero});
     // Read res slice of size {n, w, f} @ [0, 0, 0].
-    Value res = builder.create<vector::TransferReadOp>(
-        loc, resType, resShaped, ValueRange{zero, zero, zero});
+    Value res = b.create<vector::TransferReadOp>(loc, resType, resShaped,
+                                                 ValueRange{zero, zero, zero});
 
     // The base vectorization case is input: {n,w,c}, weight: {kw,c,f}, output:
     // {n,w,f}. To reuse the base pattern vectorization case, we do pre
     // transpose on input, weight, and output.
     switch (conv1DOpOrder) {
     case Conv1DOpOrder::Nwc:
       // Base case, so no transposes necessary.
       break;
     case Conv1DOpOrder::Ncw: {
       // To match base vectorization case, we pre-transpose current case.
       // ncw -> nwc
       static constexpr std::array<int64_t, 3> permLhs = {0, 2, 1};
-      lhs = builder.create<vector::TransposeOp>(loc, lhs, permLhs);
+      lhs = b.create<vector::TransposeOp>(loc, lhs, permLhs);
       // fcw -> wcf
       static constexpr std::array<int64_t, 3> permRhs = {2, 1, 0};
-      rhs = builder.create<vector::TransposeOp>(loc, rhs, permRhs);
+      rhs = b.create<vector::TransposeOp>(loc, rhs, permRhs);
       // nfw -> nwf
       static constexpr std::array<int64_t, 3> permRes = {0, 2, 1};
-      res = builder.create<vector::TransposeOp>(loc, res, permRes);
+      res = b.create<vector::TransposeOp>(loc, res, permRes);
       break;
     }
     }
 
     //===------------------------------------------------------------------===//
     // Begin vector-only rewrite part
     //===------------------------------------------------------------------===//
     // Unroll along kw and read slices of lhs and rhs.
     SmallVector<Value> lhsVals, rhsVals, resVals;
     // Extract lhs slice of size {n, wSizeStep, c} @ [0, sw * w + dw * kw, 0].
     for (int64_t kw = 0; kw < kwSize; ++kw) {
       for (int64_t w = 0; w < wSize; w += wSizeStep) {
-        lhsVals.push_back(builder.create<vector::ExtractStridedSliceOp>(
+        lhsVals.push_back(b.create<vector::ExtractStridedSliceOp>(
             loc, lhs,
             /*offsets=*/ArrayRef<int64_t>{0, w * strideW + kw * dilationW, 0},
             /*sizes=*/ArrayRef<int64_t>{nSize, wSizeStep, cSize},
             /*strides=*/ArrayRef<int64_t>{1, 1, 1}));
       }
     }
     // Extract rhs slice of size {c, f} @ [kw].
     for (int64_t kw = 0; kw < kwSize; ++kw) {
-      rhsVals.push_back(builder.create<vector::ExtractOp>(
+      rhsVals.push_back(b.create<vector::ExtractOp>(
           loc, rhs, /*offsets=*/ArrayRef<int64_t>{kw}));
     }
     // Extract res slice: {n, wSizeStep, f} @ [0, w, 0].
     for (int64_t w = 0; w < wSize; w += wSizeStep) {
-      resVals.push_back(builder.create<vector::ExtractStridedSliceOp>(
+      resVals.push_back(b.create<vector::ExtractStridedSliceOp>(
           loc, res,
           /*offsets=*/ArrayRef<int64_t>{0, w, 0},
           /*sizes=*/ArrayRef<int64_t>{nSize, wSizeStep, fSize},
           /*strides=*/ArrayRef<int64_t>{1, 1, 1}));
     }
 
     auto linearIndex = [&](int64_t kw, int64_t w) {
       return kw * (wSize / wSizeStep) + w;
     };
 
     // Compute contraction: O{n, w, f} += I{n, sw * w + dw * kw, c} * F{c, f}
     for (int64_t kw = 0; kw < kwSize; ++kw) {
       for (int64_t w = 0; w < wSize; w += wSizeStep) {
         resVals[w] = conv1dSliceAsContraction(
-            builder, loc, lhsVals[linearIndex(kw, w)], rhsVals[kw], resVals[w]);
+            b, loc, lhsVals[linearIndex(kw, w)], rhsVals[kw], resVals[w]);
       }
     }
 
     // Write back res slice: {n, wSizeStep, f} @ [0, w, 0].
     // This does not depend on kw.
     for (int64_t w = 0; w < wSize; w += wSizeStep) {
-      res = builder.create<vector::InsertStridedSliceOp>(
+      res = b.create<vector::InsertStridedSliceOp>(
           loc, resVals[w], res,
           /*offsets=*/ArrayRef<int64_t>{0, w, 0},
           /*strides=*/ArrayRef<int64_t>{1, 1, 1});
     }
     //===------------------------------------------------------------------===//
     // End vector-only rewrite part
     //===------------------------------------------------------------------===//
 
     // The base vectorization case is output: {n,w,f}
     // To reuse the result from base pattern vectorization case, we post
     // transpose the base case result.
     switch (conv1DOpOrder) {
     case Conv1DOpOrder::Nwc:
       // Base case, so no transposes necessary.
       break;
     case Conv1DOpOrder::Ncw: {
       // nwf -> nfw
       static constexpr std::array<int64_t, 3> perm = {0, 2, 1};
-      res = builder.create<vector::TransposeOp>(loc, res, perm);
+      res = b.create<vector::TransposeOp>(loc, res, perm);
       break;
     }
     }
 
     // Write back res slice of size {n, w, f} @ [0, 0, 0].
-    return builder
+    return b
         .create<vector::TransferWriteOp>(loc, res, resShaped,
                                          ValueRange{zero, zero, zero})
         .getOperation();
   }
 
   // Create a contraction: lhs{n, w, c} * rhs{c, f} -> res{n, w, f}
-  Value conv1dSliceAsContraction(OpBuilder &b, Location loc, Value lhs,
+  Value conv1dSliceAsContraction(RewriterBase &b, Location loc, Value lhs,
                                  Value rhs, Value res) {
     vector::IteratorType par = vector::IteratorType::parallel;
     vector::IteratorType red = vector::IteratorType::reduction;
     AffineExpr n, w, f, c;
     bindDims(ctx, n, w, f, c);
-    return builder.create<vector::ContractionOp>(
+    return b.create<vector::ContractionOp>(
         loc, lhs, rhs, res,
         /*indexingMaps=*/MapList{{n, w, c}, {c, f}, {n, w, f}},
         /*iteratorTypes=*/ArrayRef<vector::IteratorType>{par, par, par, red});
   }
 
   /// Generate a vector implementation for:
   /// ```
   ///   Op def: (     n,     w,     c,    kw)
FailureOr<Operation *> depthwiseConv() {
 
     int64_t nSize, wSize, cSize, kwSize;
     // kernel{kw, c}
     bindShapeDims(rhsShapedType, kwSize, cSize);
     // out{n, w, c}
     bindShapeDims(resShapedType, nSize, wSize);
 
     vector::TransferWriteOp write;
-    Value zero = builder.create<arith::ConstantIndexOp>(loc, 0);
+    Value zero = b.create<arith::ConstantIndexOp>(loc, 0);
 
     // w is unrolled (i.e. wSizeStep == 1) iff strideW > 1.
     // When strideW == 1, we can batch the contiguous loads and avoid
     // unrolling
     int64_t wSizeStep = strideW == 1 ? wSize : 1;
 
     Type lhsEltType = lhsShapedType.getElementType();
     Type rhsEltType = rhsShapedType.getElementType();
     Type resEltType = resShapedType.getElementType();
     VectorType lhsType = VectorType::get(
         {nSize,
          // iw = ow * sw + kw *  dw - 1
          //   (i.e. 16 convolved with 3 (@stride 1 dilation 1) -> 14)
          ((wSize - 1) * strideW + 1) + ((kwSize - 1) * dilationW + 1) - 1,
          cSize},
         lhsEltType);
     VectorType rhsType = VectorType::get({kwSize, cSize}, rhsEltType);
     VectorType resType = VectorType::get({nSize, wSize, cSize}, resEltType);
 
     // Read lhs slice of size {n, w * strideW + kw * dilationW, c} @ [0, 0,
     // 0].
-    Value lhs = builder.create<vector::TransferReadOp>(
-        loc, lhsType, lhsShaped, ValueRange{zero, zero, zero});
+    Value lhs = b.create<vector::TransferReadOp>(loc, lhsType, lhsShaped,
+                                                 ValueRange{zero, zero, zero});
     // Read rhs slice of size {kw, c} @ [0, 0].
-    Value rhs = builder.create<vector::TransferReadOp>(loc, rhsType, rhsShaped,
-                                                       ValueRange{zero, zero});
+    Value rhs = b.create<vector::TransferReadOp>(loc, rhsType, rhsShaped,
     // Read res slice of size {n, w, c} @ [0, 0, 0].
-    Value res = builder.create<vector::TransferReadOp>(
-        loc, resType, resShaped, ValueRange{zero, zero, zero});
+    Value res = b.create<vector::TransferReadOp>(loc, resType, resShaped,
+                                                 ValueRange{zero, zero, zero});
 
     //===------------------------------------------------------------------===//
     // Begin vector-only rewrite part
     //===------------------------------------------------------------------===//
     // Unroll along kw and read slices of lhs and rhs.
     SmallVector<Value> lhsVals, rhsVals, resVals;
     // Extract lhs slice of size {n, wSizeStep, c}
     //   @ [0, sw * w + dw * kw, 0].
     for (int64_t kw = 0; kw < kwSize; ++kw) {
       for (int64_t w = 0; w < wSize; w += wSizeStep) {
-        lhsVals.push_back(builder.create<vector::ExtractStridedSliceOp>(
+        lhsVals.push_back(b.create<vector::ExtractStridedSliceOp>(
             loc, lhs,
             /*offsets=*/ArrayRef<int64_t>{0, w * strideW + kw * dilationW, 0},
             /*sizes=*/ArrayRef<int64_t>{nSize, wSizeStep, cSize},
             /*strides=*/ArrayRef<int64_t>{1, 1, 1}));
       }
     }
     // Extract rhs slice of size {c} @ [kw].
     for (int64_t kw = 0; kw < kwSize; ++kw) {
-      rhsVals.push_back(builder.create<vector::ExtractOp>(
+      rhsVals.push_back(b.create<vector::ExtractOp>(
           loc, rhs, /*offsets=*/ArrayRef<int64_t>{kw}));
     }
     // Extract res slice: {n, wSizeStep, c} @ [0, w, 0].
     for (int64_t w = 0; w < wSize; w += wSizeStep) {
-      resVals.push_back(builder.create<vector::ExtractStridedSliceOp>(
+      resVals.push_back(b.create<vector::ExtractStridedSliceOp>(
           loc, res,
           /*offsets=*/ArrayRef<int64_t>{0, w, 0},
           /*sizes=*/ArrayRef<int64_t>{nSize, wSizeStep, cSize},
           /*strides=*/ArrayRef<int64_t>{1, 1, 1}));
     }
 
     auto linearIndex = [&](int64_t kw, int64_t w) {
       return kw * (wSize / wSizeStep) + w;
     };
 
     // Compute contraction: O{n, w, c} += I{n, sw * w + dw * kw, c} * F{c}
     for (int64_t kw = 0; kw < kwSize; ++kw) {
       for (int64_t w = 0; w < wSize; w += wSizeStep) {
         resVals[w] = depthwiseConv1dSliceAsMulAcc(
-            builder, loc, lhsVals[linearIndex(kw, w)], rhsVals[kw], resVals[w]);
+            b, loc, lhsVals[linearIndex(kw, w)], rhsVals[kw], resVals[w]);
       }
     }
 
     // Its possible we failed to create the Fma
     for (auto v : resVals) {
       if (!v)
         return failure();
     }
 
     // Write back res slice: {n, wSizeStep, c} @ [0, w, 0].
     // This does not depend on kw.
     for (int64_t w = 0; w < wSize; w += wSizeStep) {
-      res = builder.create<vector::InsertStridedSliceOp>(
+      res = b.create<vector::InsertStridedSliceOp>(
           loc, resVals[w], res,
           /*offsets=*/ArrayRef<int64_t>{0, w, 0},
           /*strides=*/ArrayRef<int64_t>{1, 1, 1});
     }
     //===------------------------------------------------------------------===//
     // End vector-only rewrite part
     //===------------------------------------------------------------------===//
 
     // Write back res slice of size {n, w, c} @ [0, 0, 0].
-    return builder
+    return b
         .create<vector::TransferWriteOp>(loc, res, resShaped,
                                          ValueRange{zero, zero, zero})
         .getOperation();
   }
 
   // Take a value of element type T and widen to the destination type.
-  Value promote(OpBuilder &b, Location loc, Value val, Type ty) {
+  Value promote(RewriterBase &b, Location loc, Value val, Type ty) {
     if (val.getType() == ty)
       return val;
 
     const int64_t srcWidth =
         getElementTypeOrSelf(val.getType()).getIntOrFloatBitWidth();
     const int64_t destWidth = getElementTypeOrSelf(ty).getIntOrFloatBitWidth();
 
     if (getElementTypeOrSelf(ty).isa<FloatType>() && srcWidth < destWidth)
-      return builder.create<arith::ExtFOp>(loc, ty, val);
+      return b.create<arith::ExtFOp>(loc, ty, val);
 
     if (getElementTypeOrSelf(ty).isa<IntegerType>() && srcWidth < destWidth)
-      return builder.create<arith::ExtSIOp>(loc, ty, val);
+      return b.create<arith::ExtSIOp>(loc, ty, val);
 
     return nullptr;
   }
 
   /// Lower lhs{n, w, c} * rhs{c} -> res{n, w, c} to MulAcc
-  Value depthwiseConv1dSliceAsMulAcc(OpBuilder &b, Location loc, Value lhs,
+  Value depthwiseConv1dSliceAsMulAcc(RewriterBase &b, Location loc, Value lhs,
                                      Value rhs, Value res) {
     auto rhsTy = rhs.getType().cast<ShapedType>();
     auto resTy = res.getType().cast<ShapedType>();
 
     // TODO(suderman): Change this to use a vector.ima intrinsic.
     lhs = promote(b, loc, lhs, resTy);
 
-    rhs = builder.create<vector::BroadcastOp>(
+    rhs = b.create<vector::BroadcastOp>(
         loc, resTy.clone(rhsTy.getElementType()), rhs);
     rhs = promote(b, loc, rhs, resTy);
 
     if (!lhs || !rhs)
       return nullptr;
 
     if (resTy.getElementType().isa<FloatType>())
       return b.create<vector::FMAOp>(loc, lhs, rhs, res);
private:
   int strideW, dilationW;
   Value lhsShaped, rhsShaped, resShaped;
   ShapedType lhsShapedType, rhsShapedType, resShapedType;
 };
 } // namespace
 
 /// Helper function to vectorize a LinalgOp with convolution semantics.
 // TODO: extend the generic vectorization to support windows and drop this.
-static FailureOr<Operation *> vectorizeConvolution(OpBuilder &b, LinalgOp op) {
+static FailureOr<Operation *> vectorizeConvolution(RewriterBase &b,
+                                                   LinalgOp op) {
   // The ConvolutionOpInterface gives us guarantees of existence for
   // strides/dilations. However, we do not need to rely on those, we can simply
   // use them if present, otherwise use the default and let the generic conv.
   // matcher in the ConvGenerator succeed or fail.
   auto strides = op->getAttrOfType<DenseIntElementsAttr>("strides");
   auto dilations = op->getAttrOfType<DenseIntElementsAttr>("dilations");
   auto stride = strides ? *strides.getValues<uint64_t>().begin() : 1;
   auto dilation = dilations ? *dilations.getValues<uint64_t>().begin() : 1;

mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp

 }
 
 namespace {
 
 /// Generate a vector implementation for matmat, matvec and tmatvec.
 /// This unrolls outer-products along the reduction dimension.
 struct UnrolledOuterProductGenerator
     : public StructuredGenerator<vector::ContractionOp, vector::IteratorType> {
-  UnrolledOuterProductGenerator(OpBuilder &builder, vector::ContractionOp op)
-      : StructuredGenerator<vector::ContractionOp, vector::IteratorType>(
-            builder, op),
+  UnrolledOuterProductGenerator(RewriterBase &b, vector::ContractionOp op)
+      : StructuredGenerator<vector::ContractionOp, vector::IteratorType>(b, op),
         kind(op.getKind()), lhs(op.getLhs()), rhs(op.getRhs()),
         res(op.getAcc()), lhsType(op.getLhsType()) {}
 
   Value t(Value v) {
     static constexpr std::array<int64_t, 2> perm = {1, 0};
-    return builder.create<vector::TransposeOp>(loc, v, perm);
+    return b.create<vector::TransposeOp>(loc, v, perm);
   }
 
   Value promote(Value v, Type dstElementType) {
     Type elementType = v.getType();
     auto vecType = elementType.dyn_cast<VectorType>();
     if (vecType)
       elementType = vecType.getElementType();
     if (elementType == dstElementType)
       return v;
     Type promotedType = dstElementType;
     if (vecType)
       promotedType = VectorType::get(vecType.getShape(), promotedType);
     if (dstElementType.isa<FloatType>())
-      return builder.create<arith::ExtFOp>(loc, promotedType, v);
-    return builder.create<arith::ExtSIOp>(loc, promotedType, v);
+      return b.create<arith::ExtFOp>(loc, promotedType, v);
+    return b.create<arith::ExtSIOp>(loc, promotedType, v);
   }
 
   Value outerProd(Value lhs, Value rhs, Value res, int reductionSize) {
     assert(reductionSize > 0);
     Type resElementType = res.getType().cast<VectorType>().getElementType();
     for (int64_t k = 0; k < reductionSize; ++k) {
-      Value a = builder.create<vector::ExtractOp>(loc, lhs, k);
-      Value b = builder.create<vector::ExtractOp>(loc, rhs, k);
-      a = promote(a, resElementType);
-      b = promote(b, resElementType);
-      res = builder.create<vector::OuterProductOp>(loc, res.getType(), a, b,
-                                                   res, kind);
+      Value extractA = b.create<vector::ExtractOp>(loc, lhs, k);
+      Value extractB = b.create<vector::ExtractOp>(loc, rhs, k);
+      extractA = promote(extractA, resElementType);
+      extractB = promote(extractB, resElementType);
+      res = b.create<vector::OuterProductOp>(loc, res.getType(), extractA,
+                                             extractB, res, kind);
     }
     return res;
   }
 
   /// Two outer parallel, one inner reduction (matmat flavor).
   FailureOr<Value> matmat() {
     if (!iters({Par(), Par(), Red()}))
       return failure();
     // Set up the parallel/reduction structure in the right form.
     AffineExpr m, n, k;
-    bindDims(builder.getContext(), m, n, k);
+    bindDims(b.getContext(), m, n, k);
     // Classical row-major matmul:  Just permute the lhs.
     if (layout({{m, k}, {k, n}, {m, n}}))
       return outerProd(t(lhs), rhs, res, lhsType.getDimSize(1));
     // TODO: may be better to fail and use some vector<k> -> scalar reduction.
     if (layout({{m, k}, {n, k}, {m, n}})) {
       Value tlhs = t(lhs);
       return outerProd(tlhs, t(rhs), res, lhsType.getDimSize(1));
     }
FailureOr<Value> matmat() {
     return failure();
   }
 
   /// One outer parallel, one inner reduction (matvec flavor)
   FailureOr<Value> matvec() {
     if (!iters({Par(), Red()}))
       return failure();
     AffineExpr m, k;
-    bindDims(builder.getContext(), m, k);
+    bindDims(b.getContext(), m, k);
 
     // Case mat-vec: transpose.
     if (layout({{m, k}, {k}, {m}}))
       return outerProd(t(lhs), rhs, res, lhsType.getDimSize(1));
     // Case mat-trans-vec: ready to go.
     if (layout({{k, m}, {k}, {m}}))
       return outerProd(lhs, rhs, res, lhsType.getDimSize(0));
     // Case vec-mat: swap and transpose.
     if (layout({{k}, {m, k}, {m}}))
       return outerProd(t(rhs), lhs, res, lhsType.getDimSize(0));
     // Case vec-mat-trans: swap and ready to go.
     if (layout({{k}, {k, m}, {m}}))
       return outerProd(rhs, lhs, res, lhsType.getDimSize(0));
     return failure();
   }
 
   //
   // One outer reduction, one inner parallel (tmatvec flavor)
   //
   FailureOr<Value> tmatvec() {
     if (!iters({Red(), Par()}))
       return failure();
     AffineExpr k, m;
-    bindDims(builder.getContext(), k, m);
+    bindDims(b.getContext(), k, m);
 
     // Case mat-vec: transpose.
     if (layout({{m, k}, {k}, {m}}))
       return outerProd(t(lhs), rhs, res, lhsType.getDimSize(1));
     // Case mat-trans-vec: ready to go.
     if (layout({{k, m}, {k}, {m}}))
       return outerProd(lhs, rhs, res, lhsType.getDimSize(0));
     // Case vec-mat: swap and transpose.