[mlir][Linalg] NFC - Modernize transformation APIs.

[ftynse]

Nice

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

 /// reshapes.
 bool skipUnitDimReshape(const OpResult &producer, OpOperand &consumer);
 
 //===----------------------------------------------------------------------===//
 // Transformations exposed as function calls.
 //===----------------------------------------------------------------------===//
 using LinalgLoops = SmallVector<Operation *, 4>;
 
-/// [DEPRECATED] Populates patterns for vectorization of all ConvN-D ops.
+/// [DEPRECATED] Populate patterns for vectorization of all ConvN-D ops.
 void populateConvVectorizationPatterns(
     MLIRContext *context, SmallVectorImpl<RewritePatternSet> &patterns,
     ArrayRef<int64_t> tileSizes);
 
-/// Populates patterns for vectorizing low-D convolution ops. This is a step in
+/// Populate patterns for vectorizing low-D convolution ops. This is a step in
 /// progressive lowering for convolution ops, it assume high-D convolution ops
 /// were decomposed previously.
 void populateConvolutionVectorizationPatterns(RewritePatternSet &patterns,
                                               PatternBenefit benefit = 1);
 
 /// Populate patterns that convert `ElementwiseMappable` ops to linalg
 /// parallel loops.
 void populateElementwiseToLinalgConversionPatterns(RewritePatternSet &patterns);

 /// collapsing (introducing) unit-extent dimensions.
 void populateFoldUnitDimsReshapeOpsByLinearizationPatterns(
     RewritePatternSet &patterns);
 
 /// Patterns to convert from one named op to another. These can be seen as
 /// canonicalizations of named ops into another named op.
 void populateLinalgNamedOpConversionPatterns(RewritePatternSet &patterns);
 
-/// Populates the given list with patterns to bufferize linalg ops.
+/// Populate the given list with patterns to bufferize linalg ops.
 void populateLinalgBufferizePatterns(
     bufferization::BufferizeTypeConverter &converter,
     RewritePatternSet &patterns);
 
 /// Create linalg op on buffers given the original tensor-based operation and
 /// the buffers for the outputs.
 LinalgOp createLinalgOpOnBuffers(ConversionPatternRewriter &rewriter,
                                  LinalgOp linalgOp, ValueRange inputs,
struct LinalgElementwiseFusionOptions {
   ControlElementwiseOpsFusionFn controlFoldingReshapesFn = skipUnitDimReshape;
 
   LinalgElementwiseFusionOptions &
   setControlFoldingReshapes(ControlElementwiseOpsFusionFn fun) {
     controlFoldingReshapesFn = std::move(fun);
     return *this;
   }
 
-  /// Function that allows the caller to control when to stop fusion. Once a
+  /// Function to allow the caller to control when to stop fusion. Once a
   /// producer is deemed fusable with the consumer (structurally), this callback
   /// can be used to abort the fusion based on non-structural constraints. This
   /// is the hook for cost models to control the amount of fusion done.
   ControlElementwiseOpsFusionFn controlElementwiseOpsFusionFn =
       [](const OpResult & /*producer */, OpOperand & /*consumer */) {
         return true;
       };
 
   LinalgElementwiseFusionOptions &
   setControlElementwiseOpsFusionFn(ControlElementwiseOpsFusionFn fun) {
     controlElementwiseOpsFusionFn = std::move(fun);
     return *this;
   }
 };
 
 /// Patterns for fusing linalg operation on tensors.
 void populateElementwiseOpsFusionPatterns(
     RewritePatternSet &patterns,
     LinalgElementwiseFusionOptions options = LinalgElementwiseFusionOptions());
 
 /// Patterns to push reshape op towards the end of the graph in order to expose
 /// more fusion opportunities.
 void populatePushReshapeOpsPatterns(RewritePatternSet &patterns);
 
-/// Performs standalone tiling of a single LinalgOp by `tileSizes`.
+/// Perform standalone tiling of a single LinalgOp by `tileSizes`.
 /// and permute the loop nest according to `interchangeVector`
 /// The permutation is expressed as a list of integers that specify
 /// the new ordering of the loop nest. The length of `interchangeVector`
 /// must be equal to the length of `tileSizes`.
 /// An empty vector is interpreted as the identity permutation and the
 /// transformation returns early.
 ///
-/// Returns a struct containing the tiled loops in the specified order
+/// Return a struct containing the tiled loops in the specified order
 /// and the cloned op if successful, llvm::None otherwise.
 ///
 /// E.g. the permutation `(i,j,k) -> (j,k,i)` is expressed by
 /// `interchangeVector = [1,2,0]`. All values in `interchangeVector` must be
 /// integers, in the range 0..`tileSizes.size()` without duplications
 /// (i.e. `[1,1,2]` is an invalid permutation).
 struct TiledLinalgOp {
   LinalgOp op;
struct TiledAndFusedLinalgOps {
   /// The fused loop generated.
   SmallVector<Operation *, 4> fusedLoops;
 };
 FailureOr<TiledAndFusedLinalgOps>
 tileAndFuseLinalgOps(OpBuilder &builder, ArrayRef<LinalgOp> ops,
                      const LinalgDependenceGraph &dependenceGraph,
                      const LinalgTilingOptions &tilingOptions);
 
-/// Interchanges the `iterator_types` and `iterator_maps` dimensions and adapts
+/// Interchange the `iterator_types` and `iterator_maps` dimensions and adapts
 /// the index accesses of `op`. This is an in-place transformation controlled by
 /// `interchangeVector`. An empty vector is interpreted as the identity
 /// permutation and the transformation returns early.
 ///
 /// E.g. the permutation `(i,j,k) -> (j,k,i)` is expressed with
 /// `interchangeVector = [1,2,0]`. All values in `interchangeVector` must be
 /// integers, in the range 0..`op.rank` without duplications
 /// (i.e. `[1,1,2]` is an invalid permutation).
-void interchangeGenericOp(PatternRewriter &rewriter, GenericOp genericOp,
+FailureOr<GenericOp> interchangeGenericOp(RewriterBase &rewriter,
+                                          GenericOp genericOp,
                                           ArrayRef<unsigned> interchangeVector);
 
-/// Creates a GenericOp from the given named operation `namedOp`. Assumes
-/// `namedOp` is not a GenericOp and has a region builder.
-GenericOp generalizeNamedOp(PatternRewriter &rewriter, LinalgOp namedOp);
+/// Create a GenericOp from the given named operation `namedOp` and replace
+/// namedOp.
+/// Return failure if `namedOp` is a GenericOp or misses a region builder.
+FailureOr<GenericOp> generalizeNamedOp(RewriterBase &rewriter,
+                                       LinalgOp namedOp);
 
 /// Callback function type used to perform the allocation for the promoted
 /// `subView`. In `boundingSubViewsize` a best attempt is made to find the
 /// smallest constant value for the size of the buffer needed for each
 /// dimension. If that is not possible, contains the dynamic size of the
 /// subview. The call back should return the buffer to use.
 using AllocBufferCallbackFn = std::function<Optional<Value>(
     OpBuilder &b, memref::SubViewOp subView,
struct LinalgPromotionOptions {
   LinalgPromotionOptions &setCopyInOutFns(CopyCallbackFn const &copyIn,
                                           CopyCallbackFn const &copyOut) {
     copyInFn = copyIn;
     copyOutFn = copyOut;
     return *this;
   }
 };
 
-/// Creates a new buffer using the `allocationFn` provided. The size of this
+/// Create a new buffer using the `allocationFn` provided. The size of this
 /// buffer is the smallest constant bounding size along each dimension that can
 /// be computed for the size of the result of `subView`. Returns the allocated
 /// buffer as `fullLocalView` and the view that matches the size of the result
 /// of subview operation as `partialLocalView`.
 struct PromotionInfo {
   Value fullLocalView;
   Value partialLocalView;
 };
 FailureOr<PromotionInfo>
 promoteSubviewAsNewBuffer(OpBuilder &b, Location loc, memref::SubViewOp subView,
                           const AllocBufferCallbackFn &allocationFn,
                           DataLayout &layout);
 
-/// Promotes the `subViews` into a new buffer allocated at the insertion point
+/// Promote the `subViews` into a new buffer allocated at the insertion point
 /// `b`. Promotion occurs in 3 steps:
 ///   1. Create a new buffer for a full tile (i.e. not clipped at the boundary).
 ///   2. Take a full view on the buffer.
 ///   3. Take a partial slice of the full view in step 2. and copy into it.
 /// Infers statically sized buffers from subViews unless `dynamicBuffers` is
 /// true.
 ///
-/// Returns the modified linalg op (the modification happens in place) as well
+/// Return the modified linalg op (the modification happens in place) as well
 /// as all the copy ops created.
 FailureOr<LinalgOp> promoteSubViews(OpBuilder &b, LinalgOp op,
                                     const LinalgPromotionOptions &options);
 
 /// Emit a suitable vector form for a Linalg op with fully static shape.
-LogicalResult vectorizeLinalgOp(OpBuilder &builder, Operation *op,
-                                SmallVectorImpl<Value> &newResults);
+LogicalResult vectorize(RewriterBase &builder, LinalgOp linalgOp);
 
-/// Emits a loop nest of `scf.for` with the proper body for `linalgOp`.
+/// Emit a loop nest of `scf.for` with the proper body for `linalgOp`.
 FailureOr<LinalgLoops> linalgOpToLoops(PatternRewriter &rewriter,
                                        LinalgOp linalgOp);
 
-/// Emits a loop nest of `scf.parallel` with the proper body for `linalgOp`.
+/// Emit a loop nest of `scf.parallel` with the proper body for `linalgOp`.
 FailureOr<LinalgLoops> linalgOpToParallelLoops(PatternRewriter &rewriter,
                                                LinalgOp linalgOp);
 
-/// Emits a loop nest of `affine.for` with the proper body for `linalgOp`.
+/// Emit a loop nest of `affine.for` with the proper body for `linalgOp`.
 FailureOr<LinalgLoops> linalgOpToAffineLoops(PatternRewriter &rewriter,
                                              LinalgOp linalgOp);
 
 //===----------------------------------------------------------------------===//
 // Preconditions that ensure the corresponding transformation succeeds and can
 // be applied as a rewrite pattern.
 //===----------------------------------------------------------------------===//
-/// Emits a `generic` operation with the `indexing_maps` and `iterator_types`
-/// permutated according to `permutation`.
-LogicalResult
-interchangeGenericOpPrecondition(GenericOp genericOp,
-                                 ArrayRef<unsigned> interchangeVector);
-
-/// Generalize named operations to generic operations.
-LogicalResult generalizeNamedOpPrecondition(Operation *op);
-
-/// Promote std.subviews feeding linalg operations.
+/// Promote memref.subviews feeding linalg-on-buffers operations.
 LogicalResult promoteSubviewsPrecondition(Operation *op,
                                           LinalgPromotionOptions options);
 
-/// Return success if the operation can be vectorized.
-LogicalResult vectorizeLinalgOpPrecondition(Operation *op);
-
-/// Return success if `op` can be vectorized assuming it is static. This allows
-/// checking if an op will be vectorizable once all the dimensions are folded to
-/// static values.
-/// It is the same as `vectorizeLinalgOpPrecondition` for static shapes.
-LogicalResult vectorizeStaticLinalgOpPrecondition(LinalgOp op);
-
 //===----------------------------------------------------------------------===//
 // Transformations exposed as rewrite patterns.
 //===----------------------------------------------------------------------===//
 // Marker used as attribute name in generated Linalg rewriting transformations.
 struct LinalgTransforms {
   static const StringLiteral kLinalgTransformMarker;
 };
 

 };
 
 /// Canonicalization patterns relevant to apply after tiling patterns. These are
 /// applied automatically by the tiling pass but need to be applied manually
 /// when tiling is called programmatically.
 RewritePatternSet getLinalgTilingCanonicalizationPatterns(MLIRContext *ctx);
 void populateLinalgTilingCanonicalizationPatterns(RewritePatternSet &patterns);
 
-/// Base pattern that applied the tiling transformation specified by `options`.
+/// Base pattern that applies the tiling transformation specified by `options`.
 /// Abort and return failure in 2 cases:
 ///   1. if the tiling specification is invalid and tiling fails to occur.
 ///   2. if tiling occurs but `options.paddingValueComputationFunction` is set
 ///      and some operand shape cannot be bounded statically.
 struct LinalgBaseTilingPattern : public RewritePattern {
   // Entry point to match any LinalgOp OpInterface.
   LinalgBaseTilingPattern(
       MLIRContext *context, LinalgTilingOptions options,

 private:
   /// LinalgTransformMarker handles special attribute manipulations.
   LinalgTransformationFilter filter;
   /// Tile sizes and interchange used to tile the root operation.
   LinalgTilingAndFusionOptions options;
 };
 
 ///
-/// Linalg generic interchage pattern.
+/// Linalg generic interchange pattern.
 ///
-/// Apply the `interchange` transformation as a pattern.
+/// Apply the `interchange` transformation on a RewriterBase.
 /// `filter` controls LinalgTransformMarker matching and update when specified.
 /// See `interchange` for more details.
 struct GenericOpInterchangePattern : public OpRewritePattern<GenericOp> {
   using OpRewritePattern<GenericOp>::OpRewritePattern;
   GenericOpInterchangePattern(
       MLIRContext *context, ArrayRef<unsigned> interchangeVector,
       LinalgTransformationFilter filter = LinalgTransformationFilter(),
       PatternBenefit benefit = 1);
LinalgPromotionPattern(
       LinalgTransformationFilter filter = LinalgTransformationFilter(),
       PatternBenefit benefit = 1)
       : LinalgBasePromotionPattern(opName, context, options, filter, benefit) {}
 };
 
 ///
 /// Linalg vectorization patterns.
 ///
-/// Apply the `vectorizeLinalgOp` transformation as a pattern.
-/// `filter` controls LinalgTransformMarker matching and update when specified.
-/// See `vectorizeLinalgOp` for more details.
-
 /// Empty for now, used for SFINAE purposes only.
 struct LinalgVectorizationOptions {};
 
+/// `filter` controls LinalgTransformMarker matching and update when specified.
+/// See `vectorizeLinalgOp` for more details.
 struct LinalgBaseVectorizationPattern : public RewritePattern {
   /// MatchAnyOpTag-based constructor with a mandatory `filter`.
   LinalgBaseVectorizationPattern(MLIRContext *context,
                                  LinalgTransformationFilter filter,
                                  PatternBenefit benefit = 1);
   /// Name-based constructor with an optional `filter`.
   LinalgBaseVectorizationPattern(
       StringRef opName, MLIRContext *context,

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

 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/Debug.h"
 
 #define DEBUG_TYPE "linalg-generalization"
 
 using namespace mlir;
 using namespace mlir::linalg;
 
-LogicalResult mlir::linalg::generalizeNamedOpPrecondition(Operation *op) {
+static LogicalResult generalizeNamedOpPrecondition(Operation *op) {
   LinalgOp namedOp = dyn_cast<LinalgOp>(op);
   // Check if the operation is a LinalgOp but not a GenericOp.
   if (!namedOp || isa<GenericOp>(op))
     return failure();
   // Check if the operation has a region builder.
   if (!namedOp.getRegionBuilder())
     return failure();
   return success();
 }
 
-GenericOp mlir::linalg::generalizeNamedOp(PatternRewriter &rewriter,
+FailureOr<GenericOp> mlir::linalg::generalizeNamedOp(RewriterBase &rewriter,
                                                      LinalgOp namedOp) {
+  if (failed(generalizeNamedOpPrecondition(namedOp)))
+    return rewriter.notifyMatchFailure(namedOp, "preconditions not met");
+
   SmallVector<Value> inputOperands = namedOp.getInputOperands();
   SmallVector<Value> outputOperands = namedOp.getOutputOperands();
   SmallVector<AffineMap> indexingMaps = namedOp.getIndexingMaps();
   SmallVector<StringRef> iterators = llvm::to_vector<4>(
       namedOp.iterator_types().getAsValueRange<StringAttr>());
   SmallVector<RankedTensorType> resultTypes = namedOp.getOutputTensorTypes();
   SmallVector<Type> types(resultTypes.begin(), resultTypes.end());
 
   // All named ops have a region attached that can be inlined.
   assert(namedOp->getNumRegions() == 1 &&
          "expect named op to have one region attached");
   GenericOp genericOp =
       rewriter.create<GenericOp>(namedOp.getLoc(), types, inputOperands,
                                  outputOperands, indexingMaps, iterators);
   rewriter.inlineRegionBefore(namedOp->getRegion(0), genericOp.region(),
                               genericOp.region().begin());
+  rewriter.replaceOp(namedOp, genericOp->getResults());
   return genericOp;
 }
 
 namespace {
 
 struct LinalgGeneralizationPass
     : public LinalgGeneralizationBase<LinalgGeneralizationPass> {
   void runOnFunction() override;

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

 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
 #include "mlir/Dialect/Vector/VectorOps.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Support/LLVM.h"
+#include "llvm/ADT/ScopeExit.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 #include <type_traits>
 
 #define DEBUG_TYPE "linalg-interchange"
 
 using namespace mlir;
 using namespace mlir::linalg;
 
-LogicalResult mlir::linalg::interchangeGenericOpPrecondition(
-    GenericOp genericOp, ArrayRef<unsigned> interchangeVector) {
+static LogicalResult
+interchangeGenericOpPrecondition(GenericOp genericOp,
+                                 ArrayRef<unsigned> interchangeVector) {
   // Interchange vector must be non-empty and match the number of loops.
   if (interchangeVector.empty() ||
       genericOp.getNumLoops() != interchangeVector.size())
     return failure();
   // Permutation map must be invertible.
   if (!inversePermutation(AffineMap::getPermutationMap(interchangeVector,
                                                        genericOp.getContext())))
     return failure();
   return success();
 }
 
-void mlir::linalg::interchangeGenericOp(PatternRewriter &rewriter,
-                                        GenericOp genericOp,
-                                        ArrayRef<unsigned> interchangeVector) {
-  // 1. Compute the inverse permutation map.
+FailureOr<GenericOp>
+mlir::linalg::interchangeGenericOp(RewriterBase &rewriter, GenericOp genericOp,
+  if (failed(interchangeGenericOpPrecondition(genericOp, interchangeVector)))
+    return rewriter.notifyMatchFailure(genericOp, "preconditions not met");
+
+  // 1. Compute the inverse permutation map, it must be non-null since the
+  // preconditions are satisfied.
   MLIRContext *context = genericOp.getContext();
   AffineMap permutationMap = inversePermutation(
       AffineMap::getPermutationMap(interchangeVector, context));
-  assert(permutationMap && "expected permutation to be invertible");
-  assert(interchangeVector.size() == genericOp.getNumLoops() &&
-         "expected interchange vector to have entry for every loop");
+  assert(permutationMap && "unexpected null map");
+
+  // Start a guarded inplace update.
+  rewriter.startRootUpdate(genericOp);
+  auto guard =
+      llvm::make_scope_exit([&]() { rewriter.finalizeRootUpdate(genericOp); });
 
   // 2. Compute the interchanged indexing maps.
-  SmallVector<Attribute, 4> newIndexingMaps;
+  SmallVector<AffineMap> newIndexingMaps;
   for (OpOperand *opOperand : genericOp.getInputAndOutputOperands()) {
     AffineMap m = genericOp.getTiedIndexingMap(opOperand);
     if (!permutationMap.isEmpty())
       m = m.compose(permutationMap);
-    newIndexingMaps.push_back(AffineMapAttr::get(m));
+    newIndexingMaps.push_back(m);
   }
   genericOp->setAttr(getIndexingMapsAttrName(),
-                     ArrayAttr::get(context, newIndexingMaps));
+                     rewriter.getAffineMapArrayAttr(newIndexingMaps));
 
   // 3. Compute the interchanged iterator types.
   ArrayRef<Attribute> itTypes = genericOp.iterator_types().getValue();
-  SmallVector<Attribute, 4> itTypesVector;
+  SmallVector<Attribute> itTypesVector;
   llvm::append_range(itTypesVector, itTypes);
   SmallVector<int64_t> permutation(interchangeVector.begin(),
                                    interchangeVector.end());
   applyPermutationToVector(itTypesVector, permutation);
   genericOp->setAttr(getIteratorTypesAttrName(),
                      ArrayAttr::get(context, itTypesVector));
 
   // 4. Transform the index operations by applying the permutation map.
   if (genericOp.hasIndexSemantics()) {
     OpBuilder::InsertionGuard guard(rewriter);
     for (IndexOp indexOp :
          llvm::make_early_inc_range(genericOp.getBody()->getOps<IndexOp>())) {
       rewriter.setInsertionPoint(indexOp);
       SmallVector<Value> allIndices;
       allIndices.reserve(genericOp.getNumLoops());
       llvm::transform(llvm::seq<uint64_t>(0, genericOp.getNumLoops()),
                       std::back_inserter(allIndices), [&](uint64_t dim) {
                         return rewriter.create<IndexOp>(indexOp->getLoc(), dim);
                       });
       rewriter.replaceOpWithNewOp<AffineApplyOp>(
           indexOp, permutationMap.getSubMap(indexOp.dim()), allIndices);
     }
   }
+
+  return genericOp;
 }

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

LogicalResult matchAndRewrite(DepthwiseConv2DNhwcHwcmQOp op,
     return matchAndReplaceDepthwiseConv(operation, input, kernel, iZp, kZp,
                                         init, stride, dilation, rewriter);
   }
 };
 
 struct LinalgNamedOpConversionPass
     : public LinalgNamedOpConversionBase<LinalgNamedOpConversionPass> {
   LinalgNamedOpConversionPass() = default;
-  LinalgNamedOpConversionPass(const LinalgNamedOpConversionPass &) {}
+  LinalgNamedOpConversionPass(const LinalgNamedOpConversionPass &) = default;
 
   void runOnOperation() override {
     Operation *op = getOperation();
     RewritePatternSet patterns(op->getContext());
     populateLinalgNamedOpConversionPatterns(patterns);
     if (failed(applyPatternsAndFoldGreedily(op, std::move(patterns))))
       return signalPassFailure();
   }

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

mlir::linalg::GenericOpInterchangePattern::GenericOpInterchangePattern(
     LinalgTransformationFilter filter, PatternBenefit benefit)
     : OpRewritePattern(context, benefit), filter(std::move(filter)),
       interchangeVector(interchangeVector.begin(), interchangeVector.end()) {}
 
 LogicalResult mlir::linalg::GenericOpInterchangePattern::matchAndRewrite(
     GenericOp genericOp, PatternRewriter &rewriter) const {
   if (failed(filter.checkAndNotify(rewriter, genericOp)))
     return failure();
-  if (failed(interchangeGenericOpPrecondition(genericOp, interchangeVector)))
-    return failure();
 
-  // TODO: figure out how this interplays with named ops. In particular this
-  // should break the named op property.
-  rewriter.updateRootInPlace(genericOp, [&]() {
+  FailureOr<GenericOp> transformedOp =
       interchangeGenericOp(rewriter, genericOp, interchangeVector);
-    // New filter if specified.
-    filter.replaceLinalgTransformationFilter(rewriter, genericOp);
-  });
+  if (failed(transformedOp))
+    return failure();
+
   return success();
 }
 
 /// Linalg generalization pattern.
 mlir::linalg::LinalgGeneralizationPattern::LinalgGeneralizationPattern(
     MLIRContext *context, LinalgTransformationFilter filter,
     PatternBenefit benefit)
     : RewritePattern(MatchAnyOpTypeTag(), benefit, context),
       filter(std::move(filter)) {}
 
 mlir::linalg::LinalgGeneralizationPattern::LinalgGeneralizationPattern(
     StringRef opName, MLIRContext *context, LinalgTransformationFilter filter,
     PatternBenefit benefit)
     : RewritePattern(opName, benefit, context, {}), filter(std::move(filter)) {}
 
 LogicalResult mlir::linalg::LinalgGeneralizationPattern::matchAndRewrite(
     Operation *op, PatternRewriter &rewriter) const {
   if (failed(filter.checkAndNotify(rewriter, op)))
     return failure();
-  if (failed(generalizeNamedOpPrecondition(op)))
+  FailureOr<GenericOp> genericOp = generalizeNamedOp(rewriter, op);
+  if (failed(genericOp))
     return failure();
-
-  GenericOp genericOp = generalizeNamedOp(rewriter, op);
-  rewriter.replaceOp(op, genericOp.getResults());
-  filter.replaceLinalgTransformationFilter(rewriter, genericOp);
+  filter.replaceLinalgTransformationFilter(rewriter, *genericOp);
   return success();
 }
 
 mlir::linalg::LinalgBasePromotionPattern::LinalgBasePromotionPattern(
     MLIRContext *context, LinalgTransformationFilter filter,
     LinalgPromotionOptions options, PatternBenefit benefit)
     : RewritePattern(MatchAnyOpTypeTag(), benefit, context),
       filter(std::move(filter)), options(std::move(options)) {}

 
 mlir::linalg::LinalgBaseVectorizationPattern::LinalgBaseVectorizationPattern(
     StringRef opName, MLIRContext *context, LinalgTransformationFilter filter,
     PatternBenefit benefit)
     : RewritePattern(opName, benefit, context, {}), filter(std::move(filter)) {}
 
 LogicalResult mlir::linalg::LinalgBaseVectorizationPattern::matchAndRewrite(
     Operation *op, PatternRewriter &rewriter) const {
+  // TODO: Interface-based rewrite.
   LinalgOp linalgOp = dyn_cast<LinalgOp>(op);
   if (!linalgOp)
     return failure();
-  if (failed(filter.checkAndNotify(rewriter, linalgOp)))
-    return failure();
-  SmallVector<Value> newResults;
-  if (failed(vectorizeLinalgOp(rewriter, op, newResults)))
+  if (failed(filter.checkAndNotify(rewriter, op)))
     return failure();
-  if (!newResults.empty())
-    rewriter.replaceOp(op, newResults);
-  else
-    rewriter.eraseOp(op);
-  return success();
+  return vectorize(rewriter, linalgOp);
 }
 
 LogicalResult mlir::linalg::applyStagedPatterns(
     Operation *op, ArrayRef<FrozenRewritePatternSet> stage1Patterns,
     const FrozenRewritePatternSet &stage2Patterns,
     function_ref<LogicalResult(Operation *)> stage3Lambda) {
   unsigned iteration = 0;
   (void)iteration;
LogicalResult mlir::linalg::applyStagedPatterns(
   }
   return success();
 }
 
 static SmallVector<StringRef> getNParallelLoopsAttrs(unsigned nParallelLoops) {
   return SmallVector<StringRef>(nParallelLoops, getParallelIteratorTypeName());
 }
 
-/// Rewrite a PadTensorOp into a sequence of InitTensorOp, FillOp (to initialize
-/// with pad_val) and GenericOp (to copy contents).
+/// Rewrite a PadTensorOp into a sequence of InitTensorOp, FillOp (to
+/// initialize with pad_val) and GenericOp (to copy contents).
 LogicalResult PadTensorOpTransformationPattern::matchAndRewrite(
     linalg::PadTensorOp padOp, PatternRewriter &rewriter) const {
 
   auto inputShapedType = padOp.source().getType().cast<ShapedType>();
   auto resultShapedType = padOp.result().getType().cast<ShapedType>();
 
   // Bail on non-static shapes.
   if (!inputShapedType.hasStaticShape())

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

for (OpOperand *opOperand : op.getOutputOperands()) {
     if (!reduceOp || !getKindForOp(reduceOp)) {
       LDBG("reduction precondition failed: reduction detection failed");
       return failure();
     }
   }
   return success();
 }
 
-LogicalResult
-mlir::linalg::vectorizeStaticLinalgOpPrecondition(linalg::LinalgOp op) {
+static LogicalResult vectorizeStaticLinalgOpPrecondition(linalg::LinalgOp op) {
   if (isElementwise(op))
     return success();
   // TODO: isaConvolutionOpInterface that can also infer from generic features.
   // But we will still need stride/dilation attributes that will be annoying to
   // reverse-engineer...
   if (isa<ConvolutionOpInterface>(op.getOperation()))
     return success();
   // TODO: the common vector shape is equal to the static loop sizes only when
   // all indexing maps are projected permutations. For convs and stencils the
   // logic will need to evolve.
   if (!allIndexingsAreProjectedPermutation(op)) {
     LDBG("precondition failed: not projected permutations");
     return failure();
   }
   if (failed(reductionPreconditions(op))) {
     LDBG("precondition failed: reduction preconditions");
     return failure();
   }
   return success();
 }
 
-LogicalResult mlir::linalg::vectorizeLinalgOpPrecondition(Operation *op) {
-  auto linalgOp = cast<linalg::LinalgOp>(op);
+static LogicalResult vectorizeLinalgOpPrecondition(LinalgOp linalgOp) {
   // All types must be static shape to go to vector.
   if (linalgOp.hasDynamicShape()) {
     LDBG("precondition failed: dynamic shape");
     return failure();
   }
   return vectorizeStaticLinalgOpPrecondition(linalgOp);
 }
 
-LogicalResult
-mlir::linalg::vectorizeLinalgOp(OpBuilder &b, Operation *op,
-                                SmallVectorImpl<Value> &newResults) {
-  if (failed(vectorizeLinalgOpPrecondition(op)))
+LogicalResult mlir::linalg::vectorize(RewriterBase &rewriter,
+                                      LinalgOp linalgOp) {
+  if (failed(vectorizeLinalgOpPrecondition(linalgOp)))
     return failure();
 
-  auto linalgOp = cast<LinalgOp>(op);
-
-  // TODO: isaConvolutionOpInterface that can also infer from generic features.
-  // But we will still need stride/dilation attributes that will be annoying to
-  // reverse-engineer...
-  if (auto convOp = dyn_cast<ConvolutionOpInterface>(op)) {
-    FailureOr<Operation *> resultOrFail = vectorizeConvolution(b, convOp);
-    if (failed(resultOrFail))
+  SmallVector<Value> results;
+  // TODO: isaConvolutionOpInterface that can also infer from generic
+  // features. Will require stride/dilation attributes inference.
+  if (auto convOp = dyn_cast<ConvolutionOpInterface>(linalgOp.getOperation())) {
+    LDBG("Vectorize as a conv: " << linalgOp);
+    FailureOr<Operation *> convOr = vectorizeConvolution(rewriter, convOp);
+    if (failed(convOr))
+      return failure();
+    llvm::append_range(results, (*convOr)->getResults());
+  } else {
+    LDBG("Vectorize generic by broadcasting to a common shape: " << linalgOp);
+    if (failed(vectorizeAsLinalgGeneric(rewriter, linalgOp, results)))
       return failure();
-    Operation *newOp = *resultOrFail;
-    llvm::append_range(newResults, newOp->getResults());
-    return success();
   }
 
-  LDBG(""
-       << "Vectorize linalg op as a generic by broadcasting to "
-          "maximal common shape: "
-       << *op);
-  return vectorizeAsLinalgGeneric(b, linalgOp, newResults);
+  if (!results.empty())
+    rewriter.replaceOp(linalgOp, results);
+  else
+    rewriter.eraseOp(linalgOp);
+
+  return success();
 }
 
 //----------------------------------------------------------------------------//
 // Misc. vectorization patterns.
 //----------------------------------------------------------------------------//
 
 /// 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.
+/// 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,
                                            ArrayRef<OpFoldResult> ofrs) {
   SmallVector<Value> result;
   llvm::for_each(ofrs, [&](auto o) {
     if (auto val = o.template dyn_cast<Value>()) {
       result.push_back(val);
     } else {
       result.push_back(builder.create<arith::ConstantIndexOp>(
           loc, getIntFromAttr(o.template get<Attribute>())));
     }
   });
   return result;
 }
 
 /// Rewrite a PadTensorOp into a sequence of InitTensorOp, FillOp and
 /// InsertSliceOp. For now, only constant padding values are supported.
 /// If there is enough static type information, TransferReadOps and
 /// TransferWriteOps may be generated instead of InsertSliceOps.
 struct GenericPadTensorOpVectorizationPattern
     : public GeneralizePadTensorOpPattern {
   GenericPadTensorOpVectorizationPattern(MLIRContext *context,
                                          PatternBenefit benefit = 1)
       : GeneralizePadTensorOpPattern(context, tryVectorizeCopy, benefit) {}
-  /// Vectorize the copying of a PadTensorOp's source. This is possible if each
-  /// dimension size is statically know in the source type or the result type
-  /// (or both).
+  /// Vectorize the copying of a PadTensorOp's source. This is possible if
+  /// each dimension size is statically know in the source type or the result
+  /// type (or both).
   static LogicalResult tryVectorizeCopy(PatternRewriter &rewriter,
                                         PadTensorOp padOp, Value dest) {
     auto sourceType = padOp.getSourceType();
     auto resultType = padOp.getResultType();
 
     // Copy cannot be vectorized if pad value is non-constant and source shape
     // is dynamic. In case of a dynamic source shape, padding must be appended
     // by TransferReadOp, but TransferReadOp supports only constant padding.
     auto padValue = padOp.getConstantPaddingValue();
     if (!padValue) {
       if (!sourceType.hasStaticShape())
         return failure();
       // Create dummy padding value.
       auto elemType = sourceType.getElementType();
       padValue = rewriter.create<arith::ConstantOp>(
           padOp.getLoc(), elemType, rewriter.getZeroAttr(elemType));
     }
 
     SmallVector<int64_t> vecShape;
     SmallVector<bool> readInBounds;
     SmallVector<bool> writeInBounds;
     for (unsigned i = 0; i < sourceType.getRank(); ++i) {
       if (!sourceType.isDynamicDim(i)) {
         vecShape.push_back(sourceType.getDimSize(i));
-        // Source shape is statically known: Neither read nor write are out-of-
-        // bounds.
+        // Source shape is statically known: Neither read nor write are
+        // out-of- bounds.
         readInBounds.push_back(true);
         writeInBounds.push_back(true);
       } else if (!resultType.isDynamicDim(i)) {
-        // Source shape is not statically known, but result shape is. Vectorize
-        // with size of result shape. This may be larger than the source size.
+        // Source shape is not statically known, but result shape is.
+        // Vectorize with size of result shape. This may be larger than the
+        // source size.
         vecShape.push_back(resultType.getDimSize(i));
         // Read may be out-of-bounds because the result size could be larger
         // than the source size.
         readInBounds.push_back(false);
         // Write is out-of-bounds if low padding > 0.
         writeInBounds.push_back(
             getConstantIntValue(padOp.getMixedLowPad()[i]) ==
             static_cast<int64_t>(0));
       } else {
         // Neither source nor result dim of padOp is static. Cannot vectorize
         // the copy.
         return failure();
       }
     }
     auto vecType = VectorType::get(vecShape, sourceType.getElementType());
 
     // Generate TransferReadOp.
     SmallVector<Value> readIndices(
         vecType.getRank(),
         rewriter.create<arith::ConstantIndexOp>(padOp.getLoc(), 0));
     auto read = rewriter.create<vector::TransferReadOp>(
         padOp.getLoc(), vecType, padOp.source(), readIndices, padValue,
         ArrayRef<bool>{readInBounds});
 
-    // If `dest` is a FillOp and the TransferWriteOp would overwrite the entire
-    // tensor, write directly to the FillOp's operand.
+    // If `dest` is a FillOp and the TransferWriteOp would overwrite the
+    // entire tensor, write directly to the FillOp's operand.
     if (llvm::equal(vecShape, resultType.getShape()) &&
         llvm::all_of(writeInBounds, [](bool b) { return b; }))
       if (auto fill = dest.getDefiningOp<FillOp>())
         dest = fill.output();
 
     // Generate TransferWriteOp.
     auto writeIndices =
         ofrToIndexValues(rewriter, padOp.getLoc(), padOp.getMixedLowPad());
     rewriter.replaceOpWithNewOp<vector::TransferWriteOp>(
         padOp, read, dest, writeIndices, ArrayRef<bool>{writeInBounds});
 
     return success();
   }
 };
 
-/// Base pattern for rewriting PadTensorOps whose result is consumed by a given
-/// operation type OpTy.
+/// Base pattern for rewriting PadTensorOps whose result is consumed by a
+/// given operation type OpTy.
 template <typename OpTy>
 struct VectorizePadTensorOpUserPattern : public OpRewritePattern<PadTensorOp> {
   using OpRewritePattern<PadTensorOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(PadTensorOp padOp,
                                 PatternRewriter &rewriter) const final {
     bool changed = false;
     // Insert users in vector, because some users may be replaced/removed.
rewriter.updateRootInPlace(xferOp, [&]() {
       xferOp.paddingMutable().assign(padValue);
     });
 
     return success();
   }
 };
 
 /// Rewrite use of PadTensorOp result in TransferWriteOp.
-/// This pattern rewrites TransferWriteOps that write to a padded tensor value,
-/// where the same amount of padding is immediately removed again after the
-/// write. In such cases, the TransferWriteOp can write to the non-padded tensor
-/// value and apply out-of-bounds masking. E.g.:
+/// This pattern rewrites TransferWriteOps that write to a padded tensor
+/// value, where the same amount of padding is immediately removed again after
+/// the write. In such cases, the TransferWriteOp can write to the non-padded
+/// tensor value and apply out-of-bounds masking. E.g.:
 /// ```
 /// %0 = tensor.extract_slice ...[...] [%s0, %s1] [1, 1]
 ///     : tensor<...> to tensor<?x?xf32>
 /// %1 = linalg.pad_tensor %0 ... : tensor<?x?xf32> to tensor<17x5xf32>
 /// %2 = vector.transfer_write %vec, %1[...]
 ///     : vector<17x5xf32>, tensor<17x5xf32>
 /// %r = tensor.extract_slice %2[0, 0] [%s0, %s1] [1, 1]
 ///     : tensor<17x5xf32> to tensor<?x?xf32>
 /// ```
 /// is rewritten to:
 /// ```
 /// %0 = tensor.extract_slice ...[...] [%s0, %s1] [1, 1]
 ///     : tensor<...> to tensor<?x?xf32>
-/// %r = vector.transfer_write %vec, %0[...] : vector<17x5xf32>, tensor<?x?xf32>
+/// %r = vector.transfer_write %vec, %0[...] : vector<17x5xf32>,
+/// tensor<?x?xf32>
 /// ```
 /// Note: It is important that the ExtractSliceOp %r resizes the result of the
-/// TransferWriteOp to the same size as the input of the TensorPadOp (or an even
-/// smaller size). Otherwise, %r's new (dynamic) dimensions would differ from
-/// %r's old dimensions.
+/// TransferWriteOp to the same size as the input of the TensorPadOp (or an
+/// even smaller size). Otherwise, %r's new (dynamic) dimensions would differ
+/// from %r's old dimensions.
 ///
 /// This rewrite is possible if:
 /// - Low padding is static 0.
 /// - `xferOp` has exactly one use, which is an ExtractSliceOp. This
-///   ExtractSliceOp trims the same amount of padding that was added beforehand.
+///   ExtractSliceOp trims the same amount of padding that was added
+///   beforehand.
 /// - Single, scalar padding value.
 struct PadTensorOpVectorizationWithTransferWritePattern
     : public VectorizePadTensorOpUserPattern<vector::TransferWriteOp> {
   using VectorizePadTensorOpUserPattern<
       vector::TransferWriteOp>::VectorizePadTensorOpUserPattern;
 
   LogicalResult rewriteUser(PatternRewriter &rewriter, PadTensorOp padOp,
                             vector::TransferWriteOp xferOp) const override {
struct PadTensorOpVectorizationWithTransferWritePattern
   /// dimensions, this function tries to infer the (static) tensor size by
   /// looking at the defining op and utilizing op-specific knowledge.
   ///
   /// This is a conservative analysis. In case equal tensor sizes cannot be
   /// proven statically, this analysis returns `false` even though the tensor
   /// sizes may turn out to be equal at runtime.
   bool hasSameTensorSize(Value beforePadding,
                          tensor::ExtractSliceOp afterTrimming) const {
-    // If the input to PadTensorOp is a CastOp, try with with both CastOp result
-    // and CastOp operand.
+    // If the input to PadTensorOp is a CastOp, try with with both CastOp
+    // result and CastOp operand.
     if (auto castOp = beforePadding.getDefiningOp<tensor::CastOp>())
       if (hasSameTensorSize(castOp.source(), afterTrimming))
         return true;
 
     auto t1 = beforePadding.getType().dyn_cast<RankedTensorType>();
     auto t2 = afterTrimming.getType().dyn_cast<RankedTensorType>();
     // Only RankedTensorType supported.
     if (!t1 || !t2)
for (unsigned i = 0; i < t1.getRank(); ++i) {
       if (!t1.isDynamicDim(i) && t1.getDimSize(i) != t2.getDimSize(i))
         return false;
     }
 
     // Nothing more to check if all dimensions are static.
     if (t1.getNumDynamicDims() == 0)
       return true;
 
-    // All dynamic sizes must be the same. The only supported case at the moment
-    // is when `beforePadding` is an ExtractSliceOp (or a cast thereof).
+    // All dynamic sizes must be the same. The only supported case at the
+    // moment is when `beforePadding` is an ExtractSliceOp (or a cast
+    // thereof).
 
     // Apart from CastOp, only ExtractSliceOp is supported.
     auto beforeSlice = beforePadding.getDefiningOp<tensor::ExtractSliceOp>();
     if (!beforeSlice)
       return false;
 
     assert(static_cast<size_t>(t1.getRank()) ==
            beforeSlice.getMixedSizes().size());
if (!llvm::all_of(
               return getConstantIntValue(std::get<0>(it)) == std::get<1>(it);
             }))
       return failure();
 
     // Insert the TransferReadOp and TransferWriteOp at the position of the
     // InsertSliceOp.
     rewriter.setInsertionPoint(insertOp);
 
-    // Generate TransferReadOp: Read entire source tensor and add high padding.
+    // Generate TransferReadOp: Read entire source tensor and add high
+    // padding.
     SmallVector<Value> readIndices(
         vecRank, rewriter.create<arith::ConstantIndexOp>(padOp.getLoc(), 0));
     auto read = rewriter.create<vector::TransferReadOp>(
         padOp.getLoc(), vecType, padOp.source(), readIndices, padValue);
 
     // Generate TransferWriteOp: Write to InsertSliceOp's dest tensor at
     // specified offsets. Write is fully in-bounds because a InsertSliceOp's
     // source must fit into the destination at the specified offsets.
void mlir::linalg::populateConvVectorizationPatterns(
   patterns.push_back(std::move(promotion));
   patterns.push_back(std::move(vectorization));
 }
 
 //----------------------------------------------------------------------------//
 // Forwarding patterns
 //----------------------------------------------------------------------------//
 
-/// Check whether there is any interleaved use of any `values` between `firstOp`
-/// and `secondOp`. Conservatively return `true` if any op or value is in a
-/// different block.
+/// Check whether there is any interleaved use of any `values` between
+/// `firstOp` and `secondOp`. Conservatively return `true` if any op or value
+/// is in a different block.
 static bool mayExistInterleavedUses(Operation *firstOp, Operation *secondOp,
                                     ValueRange values) {
   if (firstOp->getBlock() != secondOp->getBlock() ||
       !firstOp->isBeforeInBlock(secondOp)) {
     LDBG("interleavedUses precondition failed, firstOp: "
          << *firstOp << ", second op: " << *secondOp);
     return true;
   }
for (auto &u : v.getUses()) {
       LDBG(" found interleaved op " << *owner << ", firstOp: " << *firstOp
                                     << ", second op: " << *secondOp);
       return true;
     }
   }
   return false;
 }
 
-/// Return the unique subview use of `v` if it is indeed unique, null otherwise.
+/// Return the unique subview use of `v` if it is indeed unique, null
+/// otherwise.
 static memref::SubViewOp getSubViewUseIfUnique(Value v) {
   memref::SubViewOp subViewOp;
   for (auto &u : v.getUses()) {
     if (auto newSubViewOp = dyn_cast<memref::SubViewOp>(u.getOwner())) {
       if (subViewOp)
         return memref::SubViewOp();
       subViewOp = newSubViewOp;
     }
if (auto newCopyOp = dyn_cast<CopyOp>(u.getOwner())) {
       copyOp = newCopyOp;
       break;
     }
   }
   if (!copyOp)
     return failure();
   LDBG("with copy " << *copyOp);
 
-  // Find the fill into `viewOrAlloc` without interleaved uses before the copy.
+  // Find the fill into `viewOrAlloc` without interleaved uses before the
+  // copy.
   FillOp maybeFillOp;
   for (auto &u : viewOrAlloc.getUses()) {
     if (auto newFillOp = dyn_cast<FillOp>(u.getOwner())) {
       assert(newFillOp.output().getType().isa<MemRefType>());
       if (newFillOp.output() != viewOrAlloc)
         continue;
       LDBG("fill candidate " << *newFillOp);
       if (mayExistInterleavedUses(newFillOp, copyOp, {viewOrAlloc, subView}))
struct Conv1DNwcGenerator : public StructuredGenerator<LinalgOp> {
 
   /// Generate a vector implementation for:
   /// ```
   ///   Op def: (     n,     w,     c,    kw,    f  )
   ///    Iters: ({Par(), Par(), Par(), Red(), Red()})
   ///   Layout: {{n, strideW * w + dilationW * kw, c}, {kw, c, f}, {n, w, f}}
   /// ```
   /// kw is always unrolled.
-  /// TODO: w (resp. kw) is unrolled when the strideW ( resp. dilationW) is > 1.
+  /// TODO: w (resp. kw) is unrolled when the strideW ( resp. dilationW) is
+  /// > 1.
   FailureOr<Operation *> conv() {
     if (!valid)
       return failure();
 
     int nSize = lhsShapedType.getShape()[0];
     int wSize = resShapedType.getShape()[1];
     int cSize = lhsShapedType.getShape()[2];
     int kwSize = rhsShapedType.getShape()[0];
     int fSize = rhsShapedType.getShape()[2];
 
     vector::TransferWriteOp write;
     Value zero = builder.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
+    // 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)
          // Perform the proper inclusive -> exclusive -> inclusive.
          ((wSize - 1) * strideW + 1) + ((kwSize - 1) * dilationW + 1) - 1,
          cSize},
         lhsEltType);
     VectorType rhsType = VectorType::get({kwSize, cSize, fSize}, rhsEltType);
     VectorType resType = VectorType::get({nSize, wSize, fSize}, resEltType);
 
-    // Read lhs slice of size {w * strideW + kw * dilationW, c, f} @ [0, 0, 0].
+    // 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});
     // Read rhs slice of size {kw, c, f} @ [0, 0, 0].
     Value rhs = builder.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});
struct Conv1DNwcGenerator : public StructuredGenerator<LinalgOp> {
 
   /// Generate a vector implementation for:
   /// ```
   ///   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 always unrolled.
-  /// TODO: w (resp. kw) is unrolled when the strideW ( resp. dilationW) is > 1.
+  /// TODO: w (resp. kw) is unrolled when the strideW ( resp. dilationW) is
+  /// > 1.
   FailureOr<Operation *> dilatedConv() {
     if (!valid)
       return failure();
 
     int nSize = lhsShapedType.getShape()[0];
     int wSize = resShapedType.getShape()[1];
     int cSize = lhsShapedType.getShape()[2];
     int kwSize = rhsShapedType.getShape()[0];
 
     vector::TransferWriteOp write;
     Value zero = builder.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
+    // 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].
+    // 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});
     // Read rhs slice of size {kw, c} @ [0, 0].
     Value rhs = builder.create<vector::TransferReadOp>(loc, rhsType, rhsShaped,
                                                        ValueRange{zero, zero});
     // Read res slice of size {n, w, c} @ [0, 0, 0].
     Value res = builder.create<vector::TransferReadOp>(
         loc, resType, resShaped, ValueRange{zero, zero, zero});