[mlir][Linalg] NFC - Apply cleanups to transforms

[nicolasvasilache]

Summary

Depends on: D144656

[nicolasvasilache]

Rebase

[nicolasvasilache]

Landed as feccbc6bb721d8c64e6aedb4489a36e01a05d9ed

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

 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Interfaces/TilingInterface.h"
 #include "mlir/Support/LogicalResult.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "llvm/ADT/SmallBitVector.h"
 #include "llvm/ADT/SmallSet.h"
 
 namespace mlir {
-namespace bufferization {
-class BufferizeTypeConverter;
-} // namespace bufferization
+namespace linalg {
 
-class FrozenRewritePatternSet;
+class LinalgOp;
 
-namespace linalg {
+//===----------------------------------------------------------------------===//
+// Utils.
+//===----------------------------------------------------------------------===//
 
-struct LinalgElementwiseFusionOptions;
-struct LinalgFusionOptions;
-struct LinalgTilingOptions;
+/// Return vector::CombiningKind for the given op.
+std::optional<vector::CombiningKind> getCombinerOpKind(Operation *combinerOp);
 
 //===----------------------------------------------------------------------===//
-// Transformations exposed as function calls.
+// Structs that configure the behavior of various transformations.
 //===----------------------------------------------------------------------===//
+
+using TileSizeComputationFunction =
+    std::function<SmallVector<Value, 4>(OpBuilder &, Operation *)>;
+
+struct LinalgTilingOptions {
+  /// Computation function that returns the tile sizes for each operation.
+  /// Delayed construction of constant tile sizes should occur to interoperate
+  /// with folding.
+  TileSizeComputationFunction tileSizeComputationFunction = nullptr;
+
+  LinalgTilingOptions &
+  setTileSizeComputationFunction(TileSizeComputationFunction fun) {
+    tileSizeComputationFunction = std::move(fun);
+    return *this;
+  }
+  /// Set the `tileSizeComputationFunction` to return the values `ts`. The
+  /// values must not fold away when tiling. Otherwise, use a more robust
+  /// `tileSizeComputationFunction`.
+  LinalgTilingOptions &setTileSizes(const SmallVector<Value, 4> &ts) {
+    tileSizeComputationFunction = [=](OpBuilder &, Operation *) { return ts; };
+    return *this;
+  }
+  /// Convenience function to set the `tileSizeComputationFunction` to a
+  /// function that computes tile sizes at the point they are needed. Allows
+  /// proper interaction with folding.
+  LinalgTilingOptions &setTileSizes(ArrayRef<int64_t> ts);
+
+  /// Tile all dynamic dimensions by 1. I.e., scalarize those dimensions.
+  /// Note: `scalarizeDynamicDims` and `setTileSizes` cannot be used together.
+  LinalgTilingOptions &scalarizeDynamicDims();
+
+  /// The interchange vector to reorder the tiled loops.
+  SmallVector<unsigned, 4> interchangeVector = {};
+
+  LinalgTilingOptions &setInterchange(ArrayRef<unsigned> interchange) {
+    interchangeVector.assign(interchange.begin(), interchange.end());
+    return *this;
+  }
+
+  /// The type of tile loops to generate.
+  LinalgTilingLoopType loopType = LinalgTilingLoopType::Loops;
+
+  LinalgTilingOptions &setLoopType(LinalgTilingLoopType lt) {
+    loopType = lt;
+    return *this;
+  }
+
+  /// When specified, specifies distribution of generated tile loops to
+  /// processors.
+  std::optional<LinalgLoopDistributionOptions> distribution;
+
+  LinalgTilingOptions &
+  setDistributionOptions(LinalgLoopDistributionOptions distributionOptions) {
+    distribution = std::move(distributionOptions);
+    return *this;
+  }
+
+  /// Specification markers of how to distribute the `linalg.tiled_loop`.
+  SmallVector<StringRef, 2> distributionTypes = {};
+
+  LinalgTilingOptions &setDistributionTypes(ArrayRef<StringRef> types) {
+    distributionTypes.assign(types.begin(), types.end());
+    return *this;
+  }
+
+  /// Peel the specified loops.
+  SmallVector<int64_t> peeledLoops;
+
+  LinalgTilingOptions &setPeeledLoops(ArrayRef<int64_t> loops) {
+    peeledLoops.clear();
+    peeledLoops.append(loops.begin(), loops.end());
+    return *this;
+  }
+};
+
+struct LinalgTilingAndFusionOptions {
+  /// Tile sizes used to tile the root operation.
+  SmallVector<int64_t> tileSizes;
+  LinalgTilingAndFusionOptions &setTileSizes(ArrayRef<int64_t> ts) {
+    tileSizes.assign(ts.begin(), ts.end());
+    return *this;
+  }
+  /// Tile interchange used to permute the tile loops.
+  SmallVector<int64_t> tileInterchange;
+  /// When specified, specifies distribution of generated tile loops to
+  /// processors.
+  std::optional<LinalgLoopDistributionOptions> tileDistribution;
+  LinalgTilingAndFusionOptions &
+  setDistributionOptions(LinalgLoopDistributionOptions distributionOptions) {
+    tileDistribution = std::move(distributionOptions);
+    return *this;
+  }
+};
+
+struct LinalgPaddingOptions {
+  /// A padding value for every operand.
+  SmallVector<Attribute> paddingValues;
+  LinalgPaddingOptions &setPaddingValues(ArrayRef<Attribute> pv) {
+    paddingValues.assign(pv.begin(), pv.end());
+    return *this;
+  }
+  /// A list of iterator dimensions to pad.
+  SmallVector<int64_t> paddingDimensions;
+  LinalgPaddingOptions &setPaddingDimensions(ArrayRef<int64_t> pd) {
+    paddingDimensions.assign(pd.begin(), pd.end());
+    return *this;
+  }
+  /// A flag for every operand to mark the PadOp as nofold which enables
+  /// packing for statically shaped operands.
+  SmallVector<bool> packPaddings;
+  LinalgPaddingOptions &setPackPaddings(ArrayRef<bool> pp) {
+    packPaddings.assign(pp.begin(), pp.end());
+    return *this;
+  }
+  /// A number of loops to hoist the PadOp out for every operand.
+  SmallVector<int64_t> hoistPaddings;
+  LinalgPaddingOptions &setHoistPaddings(ArrayRef<int64_t> hp) {
+    hoistPaddings.assign(hp.begin(), hp.end());
+    return *this;
+  }
+  /// A permutation vector for every operand used to transpose the packed
+  /// PadOp results.
+  SmallVector<SmallVector<int64_t>> transposePaddings;
+  LinalgPaddingOptions &
+  setTransposePaddings(ArrayRef<SmallVector<int64_t>> tp) {
+    transposePaddings.assign(tp.begin(), tp.end());
+    return *this;
+  }
+};
+
+/// 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<std::optional<Value>(
+    OpBuilder &b, memref::SubViewOp subView,
+    ArrayRef<Value> boundingSubViewSize, DataLayout &layout)>;
+
+/// Callback function type used to deallocate the buffers used to hold the
+/// promoted subview.
+using DeallocBufferCallbackFn =
+    std::function<LogicalResult(OpBuilder &b, Value buffer)>;
+
+/// Callback function type used to insert copy from original subview to
+/// subview of the promoted region for the read operands/subview of promoted
+/// region to original subview for the results. The copy has to happen from
+/// `src` to `dst`.
+using CopyCallbackFn =
+    std::function<LogicalResult(OpBuilder &b, Value src, Value dst)>;
+
+struct LinalgPromotionOptions {
+  /// Indices of subViews to promote. If `std::nullopt`, try to promote all
+  /// operands.
+  std::optional<DenseSet<unsigned>> operandsToPromote;
+  LinalgPromotionOptions &setOperandsToPromote(ArrayRef<int64_t> operands) {
+    operandsToPromote = DenseSet<unsigned>();
+    operandsToPromote->insert(operands.begin(), operands.end());
+    return *this;
+  }
+  /// If ith element of `useFullTiles` is true the full view should be used
+  /// for the promoted buffer of the ith operand in `operandsToPromote`.
+  /// Otherwise the partial view will be used. The decision is defaulted to
+  /// `useFullTileBuffersDefault` when `useFullTileBuffers` is None and for
+  /// operands missing from `useFullTileBuffers`.
+  std::optional<llvm::SmallBitVector> useFullTileBuffers;
+  LinalgPromotionOptions &setUseFullTileBuffers(ArrayRef<bool> useFullTiles) {
+    unsigned size = useFullTiles.size();
+    llvm::SmallBitVector tmp(size, false);
+    for (unsigned i = 0; i < size; ++i)
+      tmp[i] = useFullTiles[i];
+    useFullTileBuffers = tmp;
+    return *this;
+  }
+  /// If true all operands unspecified by `useFullTileBuffers` will use the
+  /// full view, otherwise the partial view.
+  bool useFullTileBuffersDefault = false;
+  LinalgPromotionOptions &setUseFullTileBuffersByDefault(bool use) {
+    useFullTileBuffersDefault = use;
+    return *this;
+  }
+  /// Alignment of promoted buffer. If `std::nullopt` do not specify alignment.
+  std::optional<unsigned> alignment;
+  LinalgPromotionOptions &setAlignment(unsigned align) {
+    alignment = align;
+    return *this;
+  }
+  /// Use alloca with the default allocation scheme.
+  bool useAlloca = false;
+  LinalgPromotionOptions &setUseAlloca(bool use) {
+    useAlloca = use;
+    return *this;
+  }
+  /// Callback function to do the allocation of the promoted buffer. If
+  /// std::nullopt, then the default allocation scheme of allocating a
+  /// memref<?xi8> buffer followed by a view operation is used.
+  std::optional<AllocBufferCallbackFn> allocationFn;
+  std::optional<DeallocBufferCallbackFn> deallocationFn;
+  LinalgPromotionOptions &
+  setAllocationDeallocationFns(AllocBufferCallbackFn const &allocFn,
+                               DeallocBufferCallbackFn const &deallocFn) {
+    allocationFn = allocFn;
+    deallocationFn = deallocFn;
+    return *this;
+  }
+  /// Callback function to do the copy of data to and from the promoted
+  /// subview. If std::nullopt then a memref.copy is used.
+  std::optional<CopyCallbackFn> copyInFn;
+  std::optional<CopyCallbackFn> copyOutFn;
+  LinalgPromotionOptions &setCopyInOutFns(CopyCallbackFn const &copyIn,
+                                          CopyCallbackFn const &copyOut) {
+    copyInFn = copyIn;
+    copyOutFn = copyOut;
+    return *this;
+  }
+};
+
+/// Split Reduction options.
+struct SplitReductionOptions {
+  // Ratio used to split the reduction dimension.  If the ratio is <= 1,
+  // nothing will be done.
+  int64_t ratio = 0;
+  // Index where the extra dimension is added to the intermediate tensor
+  // shape.
+  unsigned index = 0;
+  // If the inner dimension after splitting is parallel or reduction.
+  bool innerParallel = false;
+};
+
+/// Function signature to control reduction splitting. This returns
+/// `SplitReductionOptions`.
+// TODO: don't use unsigned unless doing bit manipulation.
+using ControlSplitReductionFn =
+    std::function<SplitReductionOptions(LinalgOp op)>;
+
+//===----------------------------------------------------------------------===//
+// Preconditions that ensure the corresponding transformation succeeds and can
+// be applied as a rewrite pattern.
+//===----------------------------------------------------------------------===//
+
+/// Return true if two `linalg.generic` operations with producer/consumer
+/// relationship through `fusedOperand` can be fused using elementwise op
+/// fusion.
+bool areElementwiseOpsFusable(OpOperand *fusedOperand);
+
+/// Promote memref.subviews feeding linalg-on-buffers operations.
+LogicalResult promoteSubviewsPrecondition(Operation *op,
+                                          LinalgPromotionOptions options);
+
+/// Return success if the operation can be vectorized.
+LogicalResult
+vectorizeLinalgOpPrecondition(LinalgOp linalgOp,
+                              ArrayRef<int64_t> inputVectorSizes = {},
+                              bool vectorizeNDExtract = false);
+
+//===----------------------------------------------------------------------===//
+// Transformations exposed as functional-style API calls.
+//===----------------------------------------------------------------------===//
+
 using LinalgLoops = SmallVector<Operation *, 4>;
 
 /// Materialize a buffer allocation for the given tensor.pad op and lower the
 /// op to linalg.fill/linalg.generic + memref.tensor_store. E.g.:
 ///
 /// %0 = tensor.pad low[%l] high[%h] %t ...
 ///
 /// is lowered to:

 /// memref.tensor_store %value, %alloc
 /// %0 = bufferization.to_tensor %alloc restrict writable
 ///
 /// In case `value` is a tensor.pad result, the corresponding overload is used
 /// internally to produce a better bufferization.
 Value bufferizeToAllocation(RewriterBase &rewriter, Value value,
                             Attribute memorySpace = {});
 
-void populatePadTensorTilingPatterns(RewritePatternSet &patterns,
-                                     const LinalgTilingOptions &options);
-
-/// Populate patterns for splitting a `LinalgOp` with multiple statements within
-/// its payload into multiple `GenericOp` that have a single statement.
-/// The option `removeDeadArgsAndResults` adds patterns to remove dead arguments
-/// and results from the generated decomposed ops. This is default `true` since
-/// the core decomposition patterns relies on these clean up patterns. It is set
-/// to false only for testing purposes.
-void populateDecomposeLinalgOpsPattern(RewritePatternSet &patterns,
-                                       bool removeDeadArgsAndResults = true);
-
-/// Populate patterns that convert non-destination-style ops to destination
-/// style ops.
-void populateConvertToDestinationStylePatterns(RewritePatternSet &patterns);
-
-/// 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);
-
-/// Populate patterns that are only useful in the context of sparse tensors.
-void populateSparseTensorRewriting(RewritePatternSet &patterns);
-
-/// Function type which is used to control when to stop fusion. It is expected
-/// that OpOperand is not modified in the callback. The OpOperand is not marked
-/// as const to allow callers to use non-const methods.
-using ControlFusionFn = std::function<bool(OpOperand *fusedOperand)>;
-
-/// Patterns for fusing linalg operation on tensors.
-
-/// Pattern to fuse `linalg.generic` -> `linalg.generic` operations
-/// when both operations are fusable elementwise operations.
-void populateElementwiseOpsFusionPatterns(
-    RewritePatternSet &patterns,
-    const ControlFusionFn &controlElementwiseOpFusion);
-
-/// Patterns to bubble up or down data layout ops across other operations.
-void populateDataLayoutPropagationPatterns(RewritePatternSet &patterns);
-
-/// Pattern to remove dead operands and results of `linalg.generic` operations.
-/// This is effectively DCE for a linalg op.
-void populateEraseUnusedOperandsAndResultsPatterns(RewritePatternSet &patterns);
-
-/// Patterns to promote inputs to outputs and remove unused inputs of
-/// `linalg.generic` ops.
-void populateEraseUnnecessaryInputsPatterns(RewritePatternSet &patterns);
-
-/// Function type to control generic op dimension collapsing. It is expected
-/// to return an array of `ReassociationIndices` representing dimensions that
-/// should be merged.
-using GetCollapsableDimensionsFn =
-    std::function<SmallVector<ReassociationIndices>(linalg::GenericOp)>;
-
-/// Pattern to collapse dimensions in a linalg.generic op. This will collapse
-/// tensor operands when needed and expand back the result tensors.
-void populateCollapseDimensions(
-    RewritePatternSet &patterns,
-    const GetCollapsableDimensionsFn &controlCollapseDimensions);
-
-/// Patterns to fold an expanding (collapsing) tensor_reshape operation with its
-/// producer (consumer) generic operation by expanding the dimensionality of the
-/// loop in the generic op.
-void populateFoldReshapeOpsByExpansionPatterns(
-    RewritePatternSet &patterns, const ControlFusionFn &controlFoldingReshapes);
-
-/// Patterns to fold an expanding tensor.expand_shape operation with its
-/// producer generic operation by collapsing the dimensions of the generic op.
-void populateFoldReshapeOpsByCollapsingPatterns(
-    RewritePatternSet &patterns, const ControlFusionFn &controlFoldingReshapes);
-
-/// Patterns to constant fold Linalg operations.
-void populateConstantFoldLinalgOperations(RewritePatternSet &patterns,
-                                          const ControlFusionFn &controlFn);
-
-/// Pattern to fuse a `tensor.pad` operation with the producer of its source,
-/// if the producer is a `linalg` operation with all parallel iterator types.
-void populateFuseTensorPadWithProducerLinalgOpPatterns(
-    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);
-
-/// Patterns to fold unit-extent dimensions in operands/results of linalg ops on
-/// tensors via reassociative reshape ops.
-void populateFoldUnitExtentDimsViaReshapesPatterns(RewritePatternSet &patterns);
-
-/// Patterns to fold unit-extent dimensions in operands/results of linalg ops on
-/// tensors via rank-reducing slices.
-void populateFoldUnitExtentDimsViaSlicesPatterns(RewritePatternSet &patterns);
-
-/// A pattern that converts init operands to input operands.
-void populateMoveInitOperandsToInputPattern(RewritePatternSet &patterns);
-
-/// Patterns that are used to inline constant operands into linalg generic ops.
-void populateInlineConstantOperandsPatterns(RewritePatternSet &patterns);
-
-/// Patterns that are used to bubble up extract slice op above linalg op.
-void populateBubbleUpExtractSliceOpPatterns(RewritePatternSet &patterns);
-
-/// Adds patterns that waps tensor.extract_slice(linalg.fill(%cst, %init)) into
-/// linalg.fill(%cst, tensor.extract_slice(%init)).
-void populateSwapExtractSliceWithFillPatterns(RewritePatternSet &patterns);
-
-/// Return true if two `linalg.generic` operations with producer/consumer
-/// relationship through `fusedOperand` can be fused using elementwise op
-/// fusion.
-bool areElementwiseOpsFusable(OpOperand *fusedOperand);
-
 /// Fuse two `linalg.generic` operations that have a producer-consumer
 /// relationship captured through `fusedOperand`. The method expects
 /// that `areElementwiseOpsFusable` returns true for the given `fusedOperand`.
 struct ElementwiseOpFusionResult {
   Operation *fusedOp;
   llvm::DenseMap<Value, Value> replacements;
 };
 FailureOr<ElementwiseOpFusionResult>
 fuseElementwiseOps(RewriterBase &rewriter, OpOperand *fusedOperand);
 
+/// Try to peel and canonicalize loop `op` and return the new result.
+/// Also applies affine_min/max bounds simplification on the fly where relevant.
+// TODO: Add support for scf.parallel and affine.for loops.
+SmallVector<Value> peelLoop(RewriterBase &rewriter, Operation *op);
+
+/// Peel 'loops' and applies affine_min/max bounds simplification on the fly
+/// where relevant.
+void peelLoops(RewriterBase &rewriter, ArrayRef<scf::ForOp> loops);
+
+/// Pad the iterator dimensions `paddingDimensions` of all `opToPad` operands
+/// to a static bounding box. Use `paddingValues` and `packPaddings` to set
+/// padding value and nofold attribute of the created tensor::PadOps,
+/// respectively. Update `paddedOp` to the cloned operation with statically
+/// shaped `paddingDimensions` and return the extracted dynamically shaped
+/// results. If padding fails, return failure.
+FailureOr<SmallVector<Value>>
+rewriteAsPaddedOp(OpBuilder &b, LinalgOp opToPad,
+                  ArrayRef<int64_t> paddingDimensions,
+                  ArrayRef<Attribute> paddingValues,
+                  ArrayRef<bool> packPaddings, LinalgOp &paddedOp);
+
+/// Apply padding to `linalgOp`
+FailureOr<LinalgOp> padLinalgOp(RewriterBase &rewriter, LinalgOp linalgOp,
+                                LinalgPaddingOptions options);
+
 /// Split the given `op` into two parts along the given iteration space
 /// `dimension` at the specified `splitPoint`, and return the two parts.
 /// If the second part is statically known to be empty, do not create it
 /// and return nullptr instead. Error state is signalled by returning
 /// a pair of nullptrs.
 ///
 /// For example, the following op:
 ///

 struct TiledLinalgOp {
   LinalgOp op;
   SmallVector<Operation *, 8> loops;
   SmallVector<Value, 4> tensorResults;
 };
 FailureOr<TiledLinalgOp> tileLinalgOp(RewriterBase &b, LinalgOp op,
                                       const LinalgTilingOptions &options);
 
-/// Try to peel and canonicalize loop `op` and return the new result.
-// TODO: Add support for scf.parallel and affine.for loops.
-SmallVector<Value> peelLoop(RewriterBase &rewriter, Operation *op);
-/// Peel and canonicalize 'loops'.
-void peelLoops(RewriterBase &rewriter, ArrayRef<scf::ForOp> loops);
-
 /// 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).
 ///
 /// Return failure if the permutation is not valid.
 FailureOr<GenericOp> interchangeGenericOp(RewriterBase &rewriter,
                                           GenericOp genericOp,
                                           ArrayRef<unsigned> interchangeVector);
 
 /// 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<std::optional<Value>(
-    OpBuilder &b, memref::SubViewOp subView,
-    ArrayRef<Value> boundingSubViewSize, DataLayout &layout)>;
-
-/// Callback function type used to deallocate the buffers used to hold the
-/// promoted subview.
-using DeallocBufferCallbackFn =
-    std::function<LogicalResult(OpBuilder &b, Value buffer)>;
-
-/// Callback function type used to insert copy from original subview to
-/// subview of the promoted region for the read operands/subview of promoted
-/// region to original subview for the results. The copy has to happen from
-/// `src` to `dst`.
-using CopyCallbackFn =
-    std::function<LogicalResult(OpBuilder &b, Value src, Value dst)>;
-
-struct LinalgPromotionOptions {
-  /// Indices of subViews to promote. If `std::nullopt`, try to promote all
-  /// operands.
-  std::optional<DenseSet<unsigned>> operandsToPromote;
-  LinalgPromotionOptions &setOperandsToPromote(ArrayRef<int64_t> operands) {
-    operandsToPromote = DenseSet<unsigned>();
-    operandsToPromote->insert(operands.begin(), operands.end());
-    return *this;
-  }
-  /// If ith element of `useFullTiles` is true the full view should be used
-  /// for the promoted buffer of the ith operand in `operandsToPromote`.
-  /// Otherwise the partial view will be used. The decision is defaulted to
-  /// `useFullTileBuffersDefault` when `useFullTileBuffers` is None and for
-  /// operands missing from `useFullTileBuffers`.
-  std::optional<llvm::SmallBitVector> useFullTileBuffers;
-  LinalgPromotionOptions &setUseFullTileBuffers(ArrayRef<bool> useFullTiles) {
-    unsigned size = useFullTiles.size();
-    llvm::SmallBitVector tmp(size, false);
-    for (unsigned i = 0; i < size; ++i)
-      tmp[i] = useFullTiles[i];
-    useFullTileBuffers = tmp;
-    return *this;
-  }
-  /// If true all operands unspecified by `useFullTileBuffers` will use the
-  /// full view, otherwise the partial view.
-  bool useFullTileBuffersDefault = false;
-  LinalgPromotionOptions &setUseFullTileBuffersByDefault(bool use) {
-    useFullTileBuffersDefault = use;
-    return *this;
-  }
-  /// Alignment of promoted buffer. If `std::nullopt` do not specify alignment.
-  std::optional<unsigned> alignment;
-  LinalgPromotionOptions &setAlignment(unsigned align) {
-    alignment = align;
-    return *this;
-  }
-  /// Use alloca with the default allocation scheme.
-  bool useAlloca = false;
-  LinalgPromotionOptions &setUseAlloca(bool use) {
-    useAlloca = use;
-    return *this;
-  }
-  /// Callback function to do the allocation of the promoted buffer. If
-  /// std::nullopt, then the default allocation scheme of allocating a
-  /// memref<?xi8> buffer followed by a view operation is used.
-  std::optional<AllocBufferCallbackFn> allocationFn;
-  std::optional<DeallocBufferCallbackFn> deallocationFn;
-  LinalgPromotionOptions &
-  setAllocationDeallocationFns(AllocBufferCallbackFn const &allocFn,
-                               DeallocBufferCallbackFn const &deallocFn) {
-    allocationFn = allocFn;
-    deallocationFn = deallocFn;
-    return *this;
-  }
-  /// Callback function to do the copy of data to and from the promoted
-  /// subview. If std::nullopt then a memref.copy is used.
-  std::optional<CopyCallbackFn> copyInFn;
-  std::optional<CopyCallbackFn> copyOutFn;
-  LinalgPromotionOptions &setCopyInOutFns(CopyCallbackFn const &copyIn,
-                                          CopyCallbackFn const &copyOut) {
-    copyInFn = copyIn;
-    copyOutFn = copyOut;
-    return *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;

 /// Emit a loop nest of `scf.parallel` with the proper body for `linalgOp`.
 FailureOr<LinalgLoops> linalgOpToParallelLoops(PatternRewriter &rewriter,
                                                LinalgOp 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.
-//===----------------------------------------------------------------------===//
-/// Promote memref.subviews feeding linalg-on-buffers operations.
-LogicalResult promoteSubviewsPrecondition(Operation *op,
-                                          LinalgPromotionOptions options);
-
-/// Return success if the operation can be vectorized.
-LogicalResult
-vectorizeLinalgOpPrecondition(LinalgOp linalgOp,
-                              ArrayRef<int64_t> inputVectorSizes = {},
-                              bool vectorizeNDExtract = false);
-
-//===----------------------------------------------------------------------===//
-// Transformations exposed as rewrite patterns.
-//===----------------------------------------------------------------------===//
-
-using TileSizeComputationFunction =
-    std::function<SmallVector<Value, 4>(OpBuilder &, Operation *)>;
-
 /// Creates a number of ranges equal to the number of non-zero in `tileSizes`.
 /// One for each loop of the LinalgOp that is tiled. The `tileSizes` argument
 /// has one entry per surrounding loop. It uses zero as the convention that a
 /// particular loop is not tiled. This convention simplifies implementations
 /// by avoiding affine map manipulations. The returned ranges correspond to
 /// the loop ranges, in the proper order, that are tiled and for which new
 /// loops will be created. Also the function returns a map from loop indices
 /// of the LinalgOp to the corresponding non-empty range indices of newly

 /// }
 ///
 /// TODO: Investigate whether mixing implicit and explicit indices
 /// does not lead to losing information.
 void transformIndexOps(RewriterBase &b, LinalgOp op,
                        SmallVectorImpl<Value> &ivs,
                        const LoopIndexToRangeIndexMap &loopIndexToRangeIndex);
 
-struct LinalgPaddingOptions {
-  /// A padding value for every operand.
-  SmallVector<Attribute> paddingValues;
-  LinalgPaddingOptions &setPaddingValues(ArrayRef<Attribute> pv) {
-    paddingValues.assign(pv.begin(), pv.end());
-    return *this;
-  }
-  /// A list of iterator dimensions to pad.
-  SmallVector<int64_t> paddingDimensions;
-  LinalgPaddingOptions &setPaddingDimensions(ArrayRef<int64_t> pd) {
-    paddingDimensions.assign(pd.begin(), pd.end());
-    return *this;
-  }
-  /// A flag for every operand to mark the PadOp as nofold which enables
-  /// packing for statically shaped operands.
-  SmallVector<bool> packPaddings;
-  LinalgPaddingOptions &setPackPaddings(ArrayRef<bool> pp) {
-    packPaddings.assign(pp.begin(), pp.end());
-    return *this;
-  }
-  /// A number of loops to hoist the PadOp out for every operand.
-  SmallVector<int64_t> hoistPaddings;
-  LinalgPaddingOptions &setHoistPaddings(ArrayRef<int64_t> hp) {
-    hoistPaddings.assign(hp.begin(), hp.end());
-    return *this;
-  }
-  /// A permutation vector for every operand used to transpose the packed
-  /// PadOp results.
-  SmallVector<SmallVector<int64_t>> transposePaddings;
-  LinalgPaddingOptions &
-  setTransposePaddings(ArrayRef<SmallVector<int64_t>> tp) {
-    transposePaddings.assign(tp.begin(), tp.end());
-    return *this;
-  }
-};
-
-struct LinalgTilingAndFusionOptions {
-  /// Tile sizes used to tile the root operation.
-  SmallVector<int64_t> tileSizes;
-  LinalgTilingAndFusionOptions &setTileSizes(ArrayRef<int64_t> ts) {
-    tileSizes.assign(ts.begin(), ts.end());
-    return *this;
-  }
-  /// Tile interchange used to permute the tile loops.
-  SmallVector<int64_t> tileInterchange;
-  /// When specified, specifies distribution of generated tile loops to
-  /// processors.
-  std::optional<LinalgLoopDistributionOptions> tileDistribution;
-  LinalgTilingAndFusionOptions &
-  setDistributionOptions(LinalgLoopDistributionOptions distributionOptions) {
-    tileDistribution = std::move(distributionOptions);
-    return *this;
-  }
+/// Apply transformation to split the single linalg op reduction into a
+/// parallel and reduction dimension. Then create a new linalg.generic op
+/// doing the rest of the reduction. Return the new linalg op with an extra
+/// parallel dimension or failure if the transformation didn't happen.
+///
+/// Example:
+/// ```
+///  %r = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>,
+///                                        affine_map<(d0) -> ()>],
+///       iterator_types = ["reduction"]}
+///  ins(%in : tensor<32xf32>)
+///  outs(%out : tensor<f32>) {
+///  ^bb0(%arg1: f32, %arg2: f32):
+///    %y = arith.addf %arg1, %arg2 : f32
+///    linalg.yield %y : f32
+///  } -> tensor<f32>
+/// ```
+/// To:
+/// ```
+///  %cst = arith.constant 0.000000e+00 : f32
+///  %0 = tensor.expand_shape %in [[0, 1]] : tensor<32xf32> into
+///  tensor<4x8xf32> %1 = tensor.empty [4] : tensor<4xf32> %2 = linalg.fill
+///  ins(%cst : f32) outs(%1 : tensor<4xf32>) -> tensor<4xf32> %3 =
+///  linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
+///                                        affine_map<(d0, d1) -> (d0)>],
+///    iterator_types = ["parallel", "reduction"]}
+///    ins(%0 : tensor<4x8xf32>) outs(%2 : tensor<4xf32>) {
+///    ^bb0(%arg3: f32, %arg5: f32):
+///    %5 = arith.addf %arg3, %arg4 : f32
+///    linalg.yield %5 : f32
+///  } -> tensor<4xf32>
+/// %r = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>,
+///                                       affine_map<(d0) -> ()>],
+///   iterator_types = ["reduction"]}
+///   ins(%3 : tensor<4xf32>) outs(%out : tensor<f32>) {
+///   ^bb0(%arg3: f32, %arg4: f32):
+///   %5 = arith.addf %arg3, %arg4 : f32
+///   linalg.yield %5 : f32
+/// } -> tensor<f32>
+/// ```
+struct SplitReductionResult {
+  Operation *initOrAlloc;
+  FillOp fillOp;
+  LinalgOp splitLinalgOp;
+  LinalgOp resultCombiningLinalgOp;
 };
+FailureOr<SplitReductionResult>
+splitReduction(PatternRewriter &b, LinalgOp op,
+               const ControlSplitReductionFn &controlSplitReductionFn,
+               bool useAlloc = false);
 
-struct LinalgTilingOptions {
-  /// Computation function that returns the tile sizes for each operation.
-  /// Delayed construction of constant tile sizes should occur to interoperate
-  /// with folding.
-  TileSizeComputationFunction tileSizeComputationFunction = nullptr;
-
-  LinalgTilingOptions &
-  setTileSizeComputationFunction(TileSizeComputationFunction fun) {
-    tileSizeComputationFunction = std::move(fun);
-    return *this;
-  }
-  /// Set the `tileSizeComputationFunction` to return the values `ts`. The
-  /// values must not fold away when tiling. Otherwise, use a more robust
-  /// `tileSizeComputationFunction`.
-  LinalgTilingOptions &setTileSizes(const SmallVector<Value, 4> &ts) {
-    tileSizeComputationFunction = [=](OpBuilder &, Operation *) { return ts; };
-    return *this;
-  }
-  /// Convenience function to set the `tileSizeComputationFunction` to a
-  /// function that computes tile sizes at the point they are needed. Allows
-  /// proper interaction with folding.
-  LinalgTilingOptions &setTileSizes(ArrayRef<int64_t> ts);
-
-  /// Tile all dynamic dimensions by 1. I.e., scalarize those dimensions.
-  /// Note: `scalarizeDynamicDims` and `setTileSizes` cannot be used together.
-  LinalgTilingOptions &scalarizeDynamicDims();
-
-  /// The interchange vector to reorder the tiled loops.
-  SmallVector<unsigned, 4> interchangeVector = {};
-
-  LinalgTilingOptions &setInterchange(ArrayRef<unsigned> interchange) {
-    interchangeVector.assign(interchange.begin(), interchange.end());
-    return *this;
-  }
-
-  /// The type of tile loops to generate.
-  LinalgTilingLoopType loopType = LinalgTilingLoopType::Loops;
-
-  LinalgTilingOptions &setLoopType(LinalgTilingLoopType lt) {
-    loopType = lt;
-    return *this;
-  }
+/// Scaling-based implementation of the split reduction transformation.
+/// Instead of introducing an ExpandShapeOp, this rewrites a reduction
+/// dimension `k` into `k * scale + kk`.
+///
+/// Example:
+/// ```
+///  %0 = linalg.matmul ins(%A, %B: tensor<16x256xf32>, tensor<256x32xf32>)
+///    outs(%C: tensor<16x32xf32>) -> tensor<16x32xf32>
+/// ```
+///
+/// Is transformed to:
+///
+/// ```
+///  #map0 = affine_map<(d0, d1, d2, d3) -> (d0, d2 * 4 + d3)>
+///  #map1 = affine_map<(d0, d1, d2, d3) -> (d2 * 4 + d3, d1)>
+///  #map2 = affine_map<(d0, d1, d2, d3) -> (d2, d3)>
+///  #map3 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+///  #map4 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+///  #map5 = affine_map<(d0, d1, d2) -> (d0, d1)>
+///  %0 = tensor.empty [16, 32, 64] : tensor<16x32x64xf32>
+///  %cst = arith.constant 0.000000e+00 : f32
+///  %1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<16x32x64xf32>) ->
+///     tensor<16x32x64xf32>
+///  %2 = tensor.empty [64, 4] : tensor<64x4xi1>
+///
+///  %3 = linalg.generic {indexing_maps = [#map0, #map1, #map2, #map3],
+///    iterator_types = ["parallel", "parallel", "parallel", "reduction"]}
+///    ins(%A, %B, %2 : tensor<16x256xf32>, tensor<256x32xf32>,
+///    tensor<64x4xi1>)
+///   outs(%1 : tensor<16x32x64xf32>) {
+///      ^bb0(%arg3: f32, %arg4: f32, %arg5: i1, %arg6: f32):
+///        %5 = arith.mulf %arg3, %arg4 : f32
+///        %6 = arith.addf %arg6, %5 : f32
+///        linalg.yield %6 : f32
+///  } -> tensor<16x32x64xf32>
+///
+///  %4 = linalg.generic {indexing_maps = [#map4, #map5],
+///    iterator_types = ["parallel", "parallel", "reduction"]}
+//     ins(%3 : tensor<16x32x64xf32>)
+///    outs(%C : tensor<16x32xf32>) {
+///      ^bb0(%arg3: f32, %arg4: f32):
+///        %5 = arith.addf %arg3, %arg4 : f32
+///        linalg.yield %5 : f32
+///  } -> tensor<16x32xf32>
+///
+///  return %4 : tensor<16x32xf32>
+/// ```
+FailureOr<SplitReductionResult>
+splitReductionByScaling(PatternRewriter &b, LinalgOp op,
+                        const ControlSplitReductionFn &controlSplitReductionFn,
+                        bool useAlloc = false);
 
-  /// When specified, specifies distribution of generated tile loops to
-  /// processors.
-  std::optional<LinalgLoopDistributionOptions> distribution;
+/// Collapses dimensions of linalg.generic operation. It also collapses inputs
+/// before the op and expands outputs after the op.
+FailureOr<SmallVector<Value>> collapseGenericOpIterationDims(
+    GenericOp genericOp, ArrayRef<ReassociationIndices> foldedIterationDims,
+    RewriterBase &rewriter);
 
-  LinalgTilingOptions &
-  setDistributionOptions(LinalgLoopDistributionOptions distributionOptions) {
-    distribution = std::move(distributionOptions);
-    return *this;
-  }
+/// Struct to hold the result of a `pack` call.
+struct PackResult {
+  SmallVector<tensor::PackOp> packOps;
+  linalg::LinalgOp packedLinalgOp;
+  SmallVector<tensor::UnPackOp> unPackOps;
+};
+/// Implement packing of a single LinalgOp by `packedSizes`.
+/// There must be one packedSizes entry per `linalgOp` iterator.
+/// Return the packed Linalg op on success, failure otherwise.
+FailureOr<PackResult> pack(RewriterBase &rewriter, linalg::LinalgOp linalgOp,
+                           ArrayRef<OpFoldResult> packedSizes);
 
-  /// Specification markers of how to distribute the `linalg.tiled_loop`.
-  SmallVector<StringRef, 2> distributionTypes = {};
+/// Struct to hold the result of a `packTranspose` call.
+struct PackTransposeResult {
+  tensor::PackOp transposedPackOp;
+  linalg::LinalgOp transposedLinalgOp;
+  tensor::UnPackOp transposedUnPackOp;
+};
+/// Transpose a single PackOp -> LinalgOp -> UnPackOp chain and return the
+/// transposed PackOp -> LinalgOp -> UnPackOp chain after replacements.
+/// Return failure if either:
+///   1. the `packOp` does not have the `linalgOp` as its unique use.
+///   2. the `maybeUnPackOp`, if specified must be a consumer of the result tied
+///      to the unique `packOp` use.
+///   3. `outerPerm` (resp. `innerPerm`) must be valid permutations of
+///      `packOp.getOuterDimsPerm` (resp. `packOp.getInnerDimsPerm`) or empty.
+FailureOr<PackTransposeResult>
+packTranspose(RewriterBase &rewriter, tensor::PackOp packOp,
+              linalg::LinalgOp linalgOp, tensor::UnPackOp maybeUnPackOp,
+              ArrayRef<int64_t> outerPerm, ArrayRef<int64_t> innerPerm);
 
-  LinalgTilingOptions &setDistributionTypes(ArrayRef<StringRef> types) {
-    distributionTypes.assign(types.begin(), types.end());
-    return *this;
-  }
+/// Rewrite tensor.from_elements to linalg.generic.
+FailureOr<Operation *>
+rewriteInDestinationPassingStyle(RewriterBase &rewriter,
+                                 tensor::FromElementsOp fromElementsOp);
 
-  /// Peel the specified loops.
-  SmallVector<int64_t> peeledLoops;
+/// Rewrite tensor.generate to linalg.generic.
+FailureOr<Operation *>
+rewriteInDestinationPassingStyle(RewriterBase &rewriter,
+                                 tensor::GenerateOp generateOp);
 
-  LinalgTilingOptions &setPeeledLoops(ArrayRef<int64_t> loops) {
-    peeledLoops.clear();
-    peeledLoops.append(loops.begin(), loops.end());
-    return *this;
-  }
-};
+/// Rewrite tensor.pad to linalg.generic + tensor.insert_slice.
+FailureOr<Operation *> rewriteInDestinationPassingStyle(RewriterBase &rewriter,
+                                                        tensor::PadOp padOp);
 
-/// 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);
+//===----------------------------------------------------------------------===//
+// Rewrite patterns wrapping transformations.
+// TODO: every single such pattern should be a close to noop wrapper around a
+// functional-stye API call.
+//===----------------------------------------------------------------------===//
 
 ///
 /// Linalg padding pattern.
 ///
 /// Apply the `padding` transformation as a pattern.
 /// See `padding` for more details.
 struct LinalgPaddingPattern : public OpInterfaceRewritePattern<LinalgOp> {
   LinalgPaddingPattern(MLIRContext *context,
                        LinalgPaddingOptions options = LinalgPaddingOptions(),
                        PatternBenefit benefit = 1);
 
-  /// `matchAndRewrite` implementation that returns the significant
-  /// transformed pieces of IR.
-  FailureOr<LinalgOp> returningMatchAndRewrite(LinalgOp op,
-                                               PatternRewriter &rewriter) const;
-
   LogicalResult matchAndRewrite(LinalgOp op,
-                                PatternRewriter &rewriter) const override {
-    return returningMatchAndRewrite(op, rewriter);
-  }
+                                PatternRewriter &rewriter) const override;
 
 private:
   /// Options to control padding and hoisting.
   LinalgPaddingOptions options;
 };
 
 /// Rewrites 2-D convolution ops with size-1 window dimensions into 1-D
 /// convolution ops.

 /// Vectorization pattern for memref::CopyOp.
 struct CopyVectorizationPattern : public OpRewritePattern<memref::CopyOp> {
   using OpRewritePattern<memref::CopyOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(memref::CopyOp copyOp,
                                 PatternRewriter &rewriter) const override;
 };
 
-/// Return vector::CombiningKind for the given op.
-std::optional<vector::CombiningKind> getCombinerOpKind(Operation *combinerOp);
-
-//===----------------------------------------------------------------------===//
-// Transformations exposed as rewrite patterns.
-//===----------------------------------------------------------------------===//
-
-/// Linalg generalization patterns
-
-/// Populates `patterns` with patterns to convert spec-generated named ops to
-/// linalg.generic ops.
-void populateLinalgNamedOpsGeneralizationPatterns(RewritePatternSet &patterns);
-
-/// Linalg decompose convolutions patterns
-
-/// Populates patterns to decompose high-D convolution ops into low-D ones.
-/// This is a step in progressive lowering for convolution ops, afterwards we
-/// can vectorize the low-D convolution ops.
-void populateDecomposeConvolutionPatterns(RewritePatternSet &patterns,
-                                          PatternBenefit benefit = 1);
-
-//===----------------------------------------------------------------------===//
-// Op-specific patterns.
-//===----------------------------------------------------------------------===//
-
 /// tensor::PadOp is not canonicalized away yet, so we provide a
 /// transformation to `linalg.generic`.
 struct PadOpTransformationPattern : public OpRewritePattern<tensor::PadOp> {
   using OpRewritePattern<tensor::PadOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(tensor::PadOp padOp,
                                 PatternRewriter &rewriter) const override;
 };
 
-/// Pad the iterator dimensions `paddingDimensions` of all `opToPad` operands
-/// to a static bounding box. Use `paddingValues` and `packPaddings` to set
-/// padding value and nofold attribute of the created tensor::PadOps,
-/// respectively. Update `paddedOp` to the cloned operation with statically
-/// shaped `paddingDimensions` and return the extracted dynamically shaped
-/// results. If padding fails, return failure.
-FailureOr<SmallVector<Value>>
-rewriteAsPaddedOp(OpBuilder &b, LinalgOp opToPad,
-                  ArrayRef<int64_t> paddingDimensions,
-                  ArrayRef<Attribute> paddingValues,
-                  ArrayRef<bool> packPaddings, LinalgOp &paddedOp);
-
 using OptimizeCopyFn =
     std::function<LogicalResult(PatternRewriter &, tensor::PadOp, Value)>;
 
 /// Rewrite a tensor::PadOp into a sequence of EmptyOp, FillOp and
 /// InsertSliceOp. For now, only constant padding values are supported.
 /// `OptimizeCopyFn` can be used to customize copying step optimization.
 struct GeneralizePadOpPattern : public OpRewritePattern<tensor::PadOp> {
   GeneralizePadOpPattern(MLIRContext *context,

 /// being all 1s.
 struct GeneralizeOuterUnitDimsUnPackOpPattern
     : public OpRewritePattern<tensor::UnPackOp> {
   using OpRewritePattern<tensor::UnPackOp>::OpRewritePattern;
   LogicalResult matchAndRewrite(tensor::UnPackOp unpackOp,
                                 PatternRewriter &rewriter) const override;
 };
 
-/// Populates `patterns` with patterns that vectorize tensor.pad.
-/// These patterns are meant to apply in a complementary fashion. Benefits
-/// are used to encode a certain ordering of pattern application. To avoid
-/// scattering magic constants throughout the code base, the patterns must be
-/// added with this function. `baseBenefit` can be used to offset the benefit
-/// of all tensor::PadOp vectorization patterns by a certain value.
-void populatePadOpVectorizationPatterns(RewritePatternSet &patterns,
-                                        PatternBenefit baseBenefit = 1);
-
-void populateExtractOpVectorizationPatterns(RewritePatternSet &patterns,
-                                            PatternBenefit baseBenefit = 1);
-
 /// Match and rewrite for the pattern:
 /// ```
 ///    %alloc = ...
 ///    [optional] %view = memref.view %alloc ...
 ///    %subView = subview %allocOrView ...
 ///    [optional] linalg.fill(%allocOrView, %cst) ...
 ///    ...
 ///    memref.copy(%in, %subView) ...
struct ExtractSliceOfPadTensorSwapPattern
 
   LogicalResult matchAndRewrite(tensor::ExtractSliceOp sliceOp,
                                 PatternRewriter &rewriter) const override;
 
 private:
   ControlFn controlFn;
 };
 
-/// Split Reduction options.
-struct SplitReductionOptions {
-  // Ratio used to split the reduction dimension.  If the ratio is <= 1,
-  // nothing will be done.
-  int64_t ratio = 0;
-  // Index where the extra dimension is added to the intermediate tensor
-  // shape.
-  unsigned index = 0;
-  // If the inner dimension after splitting is parallel or reduction.
-  bool innerParallel = false;
-};
+//===----------------------------------------------------------------------===//
+// Populate functions.
+//===----------------------------------------------------------------------===//
 
-/// Function signature to control reduction splitting. This returns
-/// `SplitReductionOptions`.
-// TODO: don't use unsigned unless doing bit manipulation.
-using ControlSplitReductionFn =
-    std::function<SplitReductionOptions(LinalgOp op)>;
+/// 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);
 
-/// Patterns to apply `splitReduction` below.
-void populateSplitReductionPattern(
+/// Linalg generalization patterns
+
+/// Populates `patterns` with patterns to convert spec-generated named ops to
+/// linalg.generic ops.
+void populateLinalgNamedOpsGeneralizationPatterns(RewritePatternSet &patterns);
+
+/// Linalg decompose convolutions patterns
+
+/// Populates patterns to decompose high-D convolution ops into low-D ones.
+/// This is a step in progressive lowering for convolution ops, afterwards we
+/// can vectorize the low-D convolution ops.
+void populateDecomposeConvolutionPatterns(RewritePatternSet &patterns,
+                                          PatternBenefit benefit = 1);
+
+void populatePadTensorTilingPatterns(RewritePatternSet &patterns,
+                                     const LinalgTilingOptions &options);
+
+/// Populates `patterns` with patterns that vectorize tensor.pad.
+/// These patterns are meant to apply in a complementary fashion. Benefits
+/// are used to encode a certain ordering of pattern application. To avoid
+/// scattering magic constants throughout the code base, the patterns must be
+/// added with this function. `baseBenefit` can be used to offset the benefit
+/// of all tensor::PadOp vectorization patterns by a certain value.
+void populatePadOpVectorizationPatterns(RewritePatternSet &patterns,
+                                        PatternBenefit baseBenefit = 1);
+
+void populateExtractOpVectorizationPatterns(RewritePatternSet &patterns,
+                                            PatternBenefit baseBenefit = 1);
+
+/// Populate patterns for splitting a `LinalgOp` with multiple statements within
+/// its payload into multiple `GenericOp` that have a single statement.
+/// The option `removeDeadArgsAndResults` adds patterns to remove dead arguments
+/// and results from the generated decomposed ops. This is default `true` since
+/// the core decomposition patterns relies on these clean up patterns. It is set
+/// to false only for testing purposes.
+void populateDecomposeLinalgOpsPattern(RewritePatternSet &patterns,
+                                       bool removeDeadArgsAndResults = true);
+
+/// Populate patterns that convert non-destination-style ops to destination
+/// style ops.
+void populateConvertToDestinationStylePatterns(RewritePatternSet &patterns);
+
+/// 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);
+
+/// Populate patterns that are only useful in the context of sparse tensors.
+void populateSparseTensorRewriting(RewritePatternSet &patterns);
+
+/// Function type which is used to control when to stop fusion. It is expected
+/// that OpOperand is not modified in the callback. The OpOperand is not marked
+/// as const to allow callers to use non-const methods.
+using ControlFusionFn = std::function<bool(OpOperand *fusedOperand)>;
+
+/// Patterns for fusing linalg operation on tensors.
+
+/// Pattern to fuse `linalg.generic` -> `linalg.generic` operations
+/// when both operations are fusable elementwise operations.
+void populateElementwiseOpsFusionPatterns(
     RewritePatternSet &patterns,
-    const ControlSplitReductionFn &controlSplitReductionFn,
-    bool useAlloc = false);
+    const ControlFusionFn &controlElementwiseOpFusion);
 
-/// Apply transformation to split the single linalg op reduction into a
-/// parallel and reduction dimension. Then create a new linalg.generic op
-/// doing the rest of the reduction. Return the new linalg op with an extra
-/// parallel dimension or failure if the transformation didn't happen.
-///
-/// Example:
-/// ```
-///  %r = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>,
-///                                        affine_map<(d0) -> ()>],
-///       iterator_types = ["reduction"]}
-///  ins(%in : tensor<32xf32>)
-///  outs(%out : tensor<f32>) {
-///  ^bb0(%arg1: f32, %arg2: f32):
-///    %y = arith.addf %arg1, %arg2 : f32
-///    linalg.yield %y : f32
-///  } -> tensor<f32>
-/// ```
-/// To:
-/// ```
-///  %cst = arith.constant 0.000000e+00 : f32
-///  %0 = tensor.expand_shape %in [[0, 1]] : tensor<32xf32> into
-///  tensor<4x8xf32> %1 = tensor.empty [4] : tensor<4xf32> %2 = linalg.fill
-///  ins(%cst : f32) outs(%1 : tensor<4xf32>) -> tensor<4xf32> %3 =
-///  linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
-///                                        affine_map<(d0, d1) -> (d0)>],
-///    iterator_types = ["parallel", "reduction"]}
-///    ins(%0 : tensor<4x8xf32>) outs(%2 : tensor<4xf32>) {
-///    ^bb0(%arg3: f32, %arg5: f32):
-///    %5 = arith.addf %arg3, %arg4 : f32
-///    linalg.yield %5 : f32
-///  } -> tensor<4xf32>
-/// %r = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>,
-///                                       affine_map<(d0) -> ()>],
-///   iterator_types = ["reduction"]}
-///   ins(%3 : tensor<4xf32>) outs(%out : tensor<f32>) {
-///   ^bb0(%arg3: f32, %arg4: f32):
-///   %5 = arith.addf %arg3, %arg4 : f32
-///   linalg.yield %5 : f32
-/// } -> tensor<f32>
-/// ```
-struct SplitReductionResult {
-  Operation *initOrAlloc;
-  FillOp fillOp;
-  LinalgOp splitLinalgOp;
-  LinalgOp resultCombiningLinalgOp;
-};
-FailureOr<SplitReductionResult>
-splitReduction(PatternRewriter &b, LinalgOp op,
-               const ControlSplitReductionFn &controlSplitReductionFn,
-               bool useAlloc = false);
+/// Patterns to bubble up or down data layout ops across other operations.
+void populateDataLayoutPropagationPatterns(RewritePatternSet &patterns);
 
-/// Scaling-based implementation of the split reduction transformation.
-/// Instead of introducing an ExpandShapeOp, this rewrites a reduction
-/// dimension `k` into `k * scale + kk`.
-///
-/// Example:
-/// ```
-///  %0 = linalg.matmul ins(%A, %B: tensor<16x256xf32>, tensor<256x32xf32>)
-///    outs(%C: tensor<16x32xf32>) -> tensor<16x32xf32>
-/// ```
-///
-/// Is transformed to:
-///
-/// ```
-///  #map0 = affine_map<(d0, d1, d2, d3) -> (d0, d2 * 4 + d3)>
-///  #map1 = affine_map<(d0, d1, d2, d3) -> (d2 * 4 + d3, d1)>
-///  #map2 = affine_map<(d0, d1, d2, d3) -> (d2, d3)>
-///  #map3 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
-///  #map4 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
-///  #map5 = affine_map<(d0, d1, d2) -> (d0, d1)>
-///  %0 = tensor.empty [16, 32, 64] : tensor<16x32x64xf32>
-///  %cst = arith.constant 0.000000e+00 : f32
-///  %1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<16x32x64xf32>) ->
-///     tensor<16x32x64xf32>
-///  %2 = tensor.empty [64, 4] : tensor<64x4xi1>
-///
-///  %3 = linalg.generic {indexing_maps = [#map0, #map1, #map2, #map3],
-///    iterator_types = ["parallel", "parallel", "parallel", "reduction"]}
-///    ins(%A, %B, %2 : tensor<16x256xf32>, tensor<256x32xf32>,
-///    tensor<64x4xi1>)
-///   outs(%1 : tensor<16x32x64xf32>) {
-///      ^bb0(%arg3: f32, %arg4: f32, %arg5: i1, %arg6: f32):
-///        %5 = arith.mulf %arg3, %arg4 : f32
-///        %6 = arith.addf %arg6, %5 : f32
-///        linalg.yield %6 : f32
-///  } -> tensor<16x32x64xf32>
-///
-///  %4 = linalg.generic {indexing_maps = [#map4, #map5],
-///    iterator_types = ["parallel", "parallel", "reduction"]}
-//     ins(%3 : tensor<16x32x64xf32>)
-///    outs(%C : tensor<16x32xf32>) {
-///      ^bb0(%arg3: f32, %arg4: f32):
-///        %5 = arith.addf %arg3, %arg4 : f32
-///        linalg.yield %5 : f32
-///  } -> tensor<16x32xf32>
-///
-///  return %4 : tensor<16x32xf32>
-/// ```
-FailureOr<SplitReductionResult>
-splitReductionByScaling(PatternRewriter &b, LinalgOp op,
-                        const ControlSplitReductionFn &controlSplitReductionFn,
-                        bool useAlloc = false);
+/// Pattern to remove dead operands and results of `linalg.generic` operations.
+/// This is effectively DCE for a linalg op.
+void populateEraseUnusedOperandsAndResultsPatterns(RewritePatternSet &patterns);
 
-/// Collapses dimensions of linalg.generic operation. It also collapses inputs
-/// before the op and expands outputs after the op.
-FailureOr<SmallVector<Value>> collapseGenericOpIterationDims(
-    GenericOp genericOp, ArrayRef<ReassociationIndices> foldedIterationDims,
-    RewriterBase &rewriter);
+/// Patterns to promote inputs to outputs and remove unused inputs of
+/// `linalg.generic` ops.
+void populateEraseUnnecessaryInputsPatterns(RewritePatternSet &patterns);
 
-/// Struct to hold the result of a `pack` call.
-struct PackResult {
-  SmallVector<tensor::PackOp> packOps;
-  linalg::LinalgOp packedLinalgOp;
-  SmallVector<tensor::UnPackOp> unPackOps;
-};
-/// Implement packing of a single LinalgOp by `packedSizes`.
-/// There must be one packedSizes entry per `linalgOp` iterator.
-/// Return the packed Linalg op on success, failure otherwise.
-FailureOr<PackResult> pack(RewriterBase &rewriter, linalg::LinalgOp linalgOp,
-                           ArrayRef<OpFoldResult> packedSizes);
+/// Function type to control generic op dimension collapsing. It is expected
+/// to return an array of `ReassociationIndices` representing dimensions that
+/// should be merged.
+using GetCollapsableDimensionsFn =
+    std::function<SmallVector<ReassociationIndices>(linalg::GenericOp)>;
 
-/// Struct to hold the result of a `packTranspose` call.
-struct PackTransposeResult {
-  tensor::PackOp transposedPackOp;
-  linalg::LinalgOp transposedLinalgOp;
-  tensor::UnPackOp transposedUnPackOp;
-};
-/// Transpose a single PackOp -> LinalgOp -> UnPackOp chain and return the
-/// transposed PackOp -> LinalgOp -> UnPackOp chain after replacements.
-/// Return failure if either:
-///   1. the `packOp` does not have the `linalgOp` as its unique use.
-///   2. the `maybeUnPackOp`, if specified must be a consumer of the result tied
-///      to the unique `packOp` use.
-///   3. `outerPerm` (resp. `innerPerm`) must be valid permutations of
-///      `packOp.getOuterDimsPerm` (resp. `packOp.getInnerDimsPerm`) or empty.
-FailureOr<PackTransposeResult>
-packTranspose(RewriterBase &rewriter, tensor::PackOp packOp,
-              linalg::LinalgOp linalgOp, tensor::UnPackOp maybeUnPackOp,
-              ArrayRef<int64_t> outerPerm, ArrayRef<int64_t> innerPerm);
+/// Pattern to collapse dimensions in a linalg.generic op. This will collapse
+/// tensor operands when needed and expand back the result tensors.
+void populateCollapseDimensions(
+    RewritePatternSet &patterns,
+    const GetCollapsableDimensionsFn &controlCollapseDimensions);
 
-/// Rewrite tensor.from_elements to linalg.generic.
-FailureOr<Operation *>
-rewriteInDestinationPassingStyle(RewriterBase &rewriter,
-                                 tensor::FromElementsOp fromElementsOp);
+/// Patterns to fold an expanding (collapsing) tensor_reshape operation with its
+/// producer (consumer) generic operation by expanding the dimensionality of the
+/// loop in the generic op.
+void populateFoldReshapeOpsByExpansionPatterns(
+    RewritePatternSet &patterns, const ControlFusionFn &controlFoldingReshapes);
 
-/// Rewrite tensor.generate to linalg.generic.
-FailureOr<Operation *>
-rewriteInDestinationPassingStyle(RewriterBase &rewriter,
-                                 tensor::GenerateOp generateOp);
+/// Patterns to fold an expanding tensor.expand_shape operation with its
+/// producer generic operation by collapsing the dimensions of the generic op.
+void populateFoldReshapeOpsByCollapsingPatterns(
+    RewritePatternSet &patterns, const ControlFusionFn &controlFoldingReshapes);
 
-/// Rewrite tensor.pad to linalg.generic + tensor.insert_slice.
-FailureOr<Operation *> rewriteInDestinationPassingStyle(RewriterBase &rewriter,
-                                                        tensor::PadOp padOp);
+/// Patterns to constant fold Linalg operations.
+void populateConstantFoldLinalgOperations(RewritePatternSet &patterns,
+                                          const ControlFusionFn &controlFn);
+
+/// Pattern to fuse a `tensor.pad` operation with the producer of its source,
+/// if the producer is a `linalg` operation with all parallel iterator types.
+void populateFuseTensorPadWithProducerLinalgOpPatterns(
+    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);
+
+/// Patterns to fold unit-extent dimensions in operands/results of linalg ops on
+/// tensors via reassociative reshape ops.
+void populateFoldUnitExtentDimsViaReshapesPatterns(RewritePatternSet &patterns);
+
+/// Patterns to fold unit-extent dimensions in operands/results of linalg ops on
+/// tensors via rank-reducing slices.
+void populateFoldUnitExtentDimsViaSlicesPatterns(RewritePatternSet &patterns);
+
+/// A pattern that converts init operands to input operands.
+void populateMoveInitOperandsToInputPattern(RewritePatternSet &patterns);
+
+/// Patterns that are used to inline constant operands into linalg generic ops.
+void populateInlineConstantOperandsPatterns(RewritePatternSet &patterns);
+
+/// Patterns that are used to bubble up extract slice op above linalg op.
+void populateBubbleUpExtractSliceOpPatterns(RewritePatternSet &patterns);
+
+/// Adds patterns that waps tensor.extract_slice(linalg.fill(%cst, %init)) into
+/// linalg.fill(%cst, tensor.extract_slice(%init)).
+void populateSwapExtractSliceWithFillPatterns(RewritePatternSet &patterns);
+
+/// Patterns to apply `splitReduction` below.
+void populateSplitReductionPattern(
+    RewritePatternSet &patterns,
+    const ControlSplitReductionFn &controlSplitReductionFn,
+    bool useAlloc = false);
 
 } // namespace linalg
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_TRANSFORMS_TRANSFORMS_H

mlir/include/mlir/Dialect/SCF/Transforms/Transforms.h

 /// scf.for %iv = %c0 to %newUb step %c4 {
 ///   (loop body)
 /// }
 /// scf.for %iv2 = %newUb to %ub {
 ///   (loop body)
 /// }
 /// ```
 ///
-/// After loop peeling, this function tries to simplify/canonicalize affine.min
-/// and affine.max ops in the body of the peeled loop and in the body of the
-/// partial iteration loop, taking advantage of the fact that the peeled loop
-/// has only "full" iterations. This canonicalization is expected to enable
-/// further canonicalization opportunities through other patterns.
+/// After loop peeling, this function tries to simplify affine.min and
+/// affine.max ops in the body of the peeled loop and in the body of the partial
+/// iteration loop, taking advantage of the fact that the peeled loop has only
+/// "full" iterations. This simplification is expected to enable further
+/// canonicalization opportunities through other patterns.
 ///
 /// The return value indicates whether the loop was rewritten or not. Loops are
 /// not rewritten if:
 /// * Loop step size is 1 or
 /// * Loop bounds and step size are static, and step already divides the
 ///   iteration space evenly.
 ///
 /// Note: This function rewrites the given scf.for loop in-place and creates a
 /// new scf.for operation for the last iteration. It replaces all uses of the
 /// unpeeled loop with the results of the newly generated scf.for.
-LogicalResult peelAndCanonicalizeForLoop(RewriterBase &rewriter, ForOp forOp,
+LogicalResult peelForLoopAndSimplifyBounds(RewriterBase &rewriter, ForOp forOp,
                                            scf::ForOp &partialIteration);
 
 /// Tile a parallel loop of the form
 ///   scf.parallel (%i0, %i1) = (%arg0, %arg1) to (%arg2, %arg3)
 ///                                             step (%arg4, %arg5)
 ///
 /// into
 ///   scf.parallel (%i0, %i1) = (%arg0, %arg1) to (%arg2, %arg3)
 ///                                             step (%arg4*tileSize[0],

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

transform::PadOp::applyToOne(LinalgOp target,
   LinalgPaddingOptions paddingOptions;
   paddingOptions.setPaddingValues(paddingValues);
   paddingOptions.setPaddingDimensions(
       extractFromI64ArrayAttr(getPaddingDimensions()));
   paddingOptions.setPackPaddings(packPaddings);
   paddingOptions.setHoistPaddings(extractFromI64ArrayAttr(getHoistPaddings()));
   paddingOptions.setTransposePaddings(transposePaddings);
 
-  FailureOr<LinalgOp> result =
-      tryApply<LinalgPaddingPattern>(target, paddingOptions);
+  IRRewriter rewriter(target->getContext());
+  FailureOr<LinalgOp> result = padLinalgOp(rewriter, target, paddingOptions);
   if (succeeded(result)) {
     results.push_back(result->getOperation());
     return DiagnosedSilenceableFailure::success();
   }
 
   return emitDefaultSilenceableFailure(target);
 }
 

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

 #define DEBUG_TYPE "linalg-transforms"
 
 using namespace mlir;
 using namespace mlir::linalg;
 
 #define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE << "]: ")
 #define DBGSNL() (llvm::dbgs() << "\n")
 
-//===----------------------------------------------------------------------===//
-// Transformations exposed as rewrite patterns.
-//===----------------------------------------------------------------------===//
-
-LinalgTilingOptions &
-mlir::linalg::LinalgTilingOptions::setTileSizes(ArrayRef<int64_t> ts) {
-  assert(!tileSizeComputationFunction && "tile sizes already set");
-  SmallVector<int64_t, 4> tileSizes(ts.begin(), ts.end());
-  tileSizeComputationFunction = [tileSizes](OpBuilder &b, Operation *op) {
-    OpBuilder::InsertionGuard guard(b);
-    b.setInsertionPointToStart(
-        &op->getParentOfType<func::FuncOp>().getBody().front());
-    return llvm::to_vector<4>(map_range(tileSizes, [&](int64_t s) {
-      Value v = b.create<arith::ConstantIndexOp>(op->getLoc(), s);
-      return v;
-    }));
-  };
-  return *this;
-}
-
 /// Pad the `opOperand` in the `paddingDimensions` using the padding value and
 /// the nofold flag found in `paddingValues` and `packPaddings`, respectively.
 /// Exit early and return the `opOperand` value if the shape dimensions that
 /// match `paddingDimensions` have a static size and the nofold flag is not set.
 /// Otherwise, try to pad the shape dimensions that match the iterator
 /// dimensions `paddingDimensions` and return the tensor::PadOp result if
 /// padding succeeds or failure otherwise.
 static FailureOr<Value> padOperandToSmallestStaticBoundingBox(
static FailureOr<Value> padOperandToSmallestStaticBoundingBox(
 
   // Pad the operand to the bounding box defined by `paddedShape`.
   auto paddedTensorType = RankedTensorType::get(
       paddedShape, getElementTypeOrSelf(opOperand->get()));
   return makeComposedPadHighOp(b, opToPad->getLoc(), paddedTensorType,
                                opOperand->get(), paddingValue, nofold);
 }
 
+static SmallVector<utils::IteratorType>
+getNParallelLoopsAttrs(unsigned nParallelLoops) {
+  return SmallVector<utils::IteratorType>(nParallelLoops,
+                                          utils::IteratorType::parallel);
+}
+
+//===----------------------------------------------------------------------===//
+// Transformations exposed as functional-style API calls.
+//===----------------------------------------------------------------------===//
+
+//===----------------------------------------------------------------------===//
+// rewriteAsPaddedOp transformation.
+//===----------------------------------------------------------------------===//
 FailureOr<SmallVector<Value>>
 linalg::rewriteAsPaddedOp(OpBuilder &b, LinalgOp opToPad,
                           ArrayRef<int64_t> paddingDimensions,
                           ArrayRef<Attribute> paddingValues,
                           ArrayRef<bool> packPaddings, LinalgOp &paddedOp) {
   Location loc = opToPad->getLoc();
 
   // TODO: there are cases where we may still want to pad to larger sizes.
for (Value v : reifiedResultShapes[resultNumber])
       sizes.push_back(getAsOpFoldResult(v));
     SmallVector<OpFoldResult> strides(rank, b.getIndexAttr(1));
     paddedSubviewResults.push_back(b.create<tensor::ExtractSliceOp>(
         loc, paddedResult, offsets, sizes, strides));
   }
   return paddedSubviewResults;
 }
 
-/// Try to peel a loop `op` and return the new result.
+//===----------------------------------------------------------------------===//
+// peelLoop transformation.
+//===----------------------------------------------------------------------===//
+
+/// Try to peel and canonicalize loop `op` and return the new result.
+/// Also applies affine_min/max bounds simplification on the fly where relevant.
 // TODO: Add support for scf.parallel and affine.for loops.
 SmallVector<Value> mlir::linalg::peelLoop(RewriterBase &rewriter,
                                           Operation *op) {
   return llvm::TypeSwitch<Operation *, SmallVector<Value, 4>>(op)
       .Case<scf::ForOp>([&](scf::ForOp forOp) {
         scf::ForOp partialIteration;
-        if (succeeded(scf::peelAndCanonicalizeForLoop(rewriter, forOp,
+        if (succeeded(scf::peelForLoopAndSimplifyBounds(rewriter, forOp,
                                                         partialIteration)))
           return partialIteration->getResults();
         assert(!partialIteration && "expected that loop was not peeled");
         return forOp->getResults();
       })
       .Default([&](Operation *op) { return op->getResults(); });
 }
 
-/// Peel and canonicalize 'loops'.
+/// Peel 'loops' and applies affine_min/max bounds simplification on the fly
+/// where relevant.
 void mlir::linalg::peelLoops(RewriterBase &rewriter,
                              ArrayRef<scf::ForOp> loops) {
   for (auto loopOp : loops)
     peelLoop(rewriter, loopOp);
 }
 
-/// Linalg padding pattern.
-mlir::linalg::LinalgPaddingPattern::LinalgPaddingPattern(
-    MLIRContext *context, LinalgPaddingOptions options, PatternBenefit benefit)
-    : OpInterfaceRewritePattern<LinalgOp>(context, benefit),
-      options(std::move(options)) {}
+//===----------------------------------------------------------------------===//
+// pad transformation.
+//===----------------------------------------------------------------------===//
 
-FailureOr<LinalgOp>
-mlir::linalg::LinalgPaddingPattern::returningMatchAndRewrite(
-    LinalgOp linalgOp, PatternRewriter &rewriter) const {
+FailureOr<LinalgOp> mlir::linalg::padLinalgOp(RewriterBase &rewriter,
+                                              LinalgOp linalgOp,
+                                              LinalgPaddingOptions options) {
   if (!linalgOp.hasTensorSemantics())
-    return failure();
+    return rewriter.notifyMatchFailure(
+        linalgOp, "only applies to Linalg ops with tensor semantics");
 
   // Pad the operation.
   LinalgOp paddedOp;
   FailureOr<SmallVector<Value>> newResults =
       rewriteAsPaddedOp(rewriter, linalgOp, options.paddingDimensions,
                         options.paddingValues, options.packPaddings, paddedOp);
   if (failed(newResults))
-    return failure();
+    return rewriter.notifyMatchFailure(linalgOp,
+                                       "failed to rewrite as a padded op");
 
   // Hoist the padding.
   for (const auto &en : enumerate(options.hoistPaddings)) {
     if (static_cast<int64_t>(en.index()) >= paddedOp->getNumOperands())
       break;
     OpOperand &opOperand = paddedOp->getOpOperand(en.index());
     auto padOp = opOperand.get().getDefiningOp<tensor::PadOp>();
-    if (!padOp || en.value() == 0)
+    if (!padOp || en.value() == 0) {
+      (void)rewriter.notifyMatchFailure(linalgOp, "not a tensor.pad -- skip");
       continue;
+    }
 
     // Fail hoisting if the operand shape is not fully static.
-    if (llvm::any_of(paddedOp.getShape(&opOperand), ShapedType::isDynamic))
-      return failure();
+    if (llvm::any_of(paddedOp.getShape(&opOperand), ShapedType::isDynamic)) {
+      (void)rewriter.notifyMatchFailure(linalgOp,
+                                        "non static padding shape -- skip");
+      continue;
+    }
 
     tensor::PadOp hoistedOp;
     SmallVector<GenericOp> transposeOps;
     SmallVector<int64_t> transposeVector =
         en.index() < options.transposePaddings.size()
             ? options.transposePaddings[en.index()]
             : SmallVector<int64_t>{};
 
     FailureOr<Value> newResult = hoistPaddingOnTensors(
         padOp, en.value(), transposeVector, hoistedOp, transposeOps);
-    if (failed(newResult))
+    if (failed(newResult)) {
+      (void)rewriter.notifyMatchFailure(linalgOp,
+                                        "failed to apply hoistPadding");
       continue;
+    }
     rewriter.replaceOp(padOp, *newResult);
   }
 
   // Replace the original operation to pad.
   rewriter.replaceOp(linalgOp, *newResults);
 
   return paddedOp;
 }
 
+//===----------------------------------------------------------------------===//
+// pack transformation.
+//===----------------------------------------------------------------------===//
+
+#ifndef NDEBUG
+/// Return true if `map` has 0 or 1 result function of AffineDimExpr(dim).
+static bool hasAtMostOneResultFunctionOfDim(AffineMap map, int64_t dim) {
+  bool found = false;
+  for (AffineExpr e : map.getResults()) {
+    if (!e.isFunctionOfDim(dim))
+      continue;
+    if (found)
+      return false;
+    found = true;
+  }
+  return true;
+}
+#endif // NDEBUG
+
+/// Return the index of the first result of `map` that is a function of
+/// AffineDimExpr(dim), std::nullopt otherwise.
+static std::optional<int64_t> getFirstResultIndexFunctionOf(AffineMap map,
+                                                            int64_t dim) {
+  for (int64_t i = 0, e = map.getNumResults(); i < e; ++i) {
+    AffineExpr expr = map.getResult(i);
+    if (!expr.isFunctionOfDim(dim))
+      continue;
+    return i;
+  }
+  return std::nullopt;
+}
+
+/// Perform one step of packing of a LinalgOp's metadata along `dim` into the
+/// `newDim` at `iteratorTypes.size()` by:
+///   1. Appending `iteratorTypes[newDim]`, equal to `iteratorTypes[dim]`.
+///   2. Appending a `newDim` to the domain of every indexing map.
+///   3. For each operand (i.e. for each map in `indexingMaps`), perform packing
+///      by potentially adding a `newDim` result to `map`.
+/// The preserved invariant is that `iteratorTypes.size()` is always equal to
+/// `map.getNumDims()` for every map in `indexingMaps`.
+///
+/// Update `indexingMaps` and `iteratorTypes` inplace as one step of the update.
+/// Return a vector that records the optional packing for each operand.
+/// Return failure if the packed indexing cannot be represented with a LinalgOp.
+///
+/// Further details:
+/// ================
+/// The current implementation of packing (i.e. data tiling) consists of
+/// rewriting a linearized strip-mined form into a higher-dimensional access.
+/// e.g. consider an access `A[I][f(j, k, l)]` and packing by 4; we rewrite
+/// `I` into `4 * i + ii`, where `0 <= ii < 4`.
+/// The access is further rewritten as `A[i][f(j, k, l)][ii]`.
+///
+/// This rewrite into higher dimensional access is not possible for general
+/// AffineExpr in Linalg atm, it is restricted to an AffineDimExpr:
+/// e.g. consider an access `A[I + J][f(j, k, l)]` and packing by 4; we
+/// rewrite `I + J` into `4 * i + ii + J`, where `0 <= ii < 4`.
+/// The rewrite of the access would be a form not representable in Linalg:
+///   `A[i + (ii + J) / 4][f(j, k, l)][(ii + J) % 4]`.
+/// Note however that as `J` and `ii` iterate, the accesses do not have a
+/// particular alignment, so packing does not achieve alignment in this case
+///
+/// In the future, we may want to consider a mixed-form that allows some
+/// alignment in the presence of multiple accesses:
+///   `A[I][f(j, k, l)]` and `B[I + J][f(j, k, l)]`
+/// And would rewrite accesses as:
+///   `A[i][f(j, k, l)][ii]` and `B[4 * i + ii + J][f(j, k, l)]`
+static FailureOr<SmallVector<std::optional<int64_t>>>
+packLinalgMetadataOnce(SmallVectorImpl<AffineMap> &indexingMaps,
+                       SmallVectorImpl<utils::IteratorType> &iteratorTypes,
+                       int64_t dim) {
+  int64_t newDim = iteratorTypes.size();
+  iteratorTypes.push_back(iteratorTypes[dim]);
+
+  SmallVector<std::optional<int64_t>> packedDimPerIndexingMap(
+      indexingMaps.size(), std::nullopt);
+  SmallVector<AffineMap> newMaps;
+  for (int64_t operandIdx = 0, e = indexingMaps.size(); operandIdx < e;
+       ++operandIdx) {
+    AffineMap map = indexingMaps[operandIdx];
+
+    // Add the `newDim` to map whatever the case.
+    assert(map.getNumDims() == newDim && "num dims invariant violation");
+    map = map.shiftDims(1, newDim);
+
+    // Get the at-most-1 index of the result that is a function of `dim`.
+    // If we can find one, we insert `AffineDimExpr(newDim)` to the map, which
+    // logically chunks dimension `dim` into `K * dim + newDim`, where the
+    // packing factor `K` is specified separately.
+    assert(hasAtMostOneResultFunctionOfDim(map, dim) &&
+           "num results invariant violation");
+    auto maybeOperandDimensionToPack = getFirstResultIndexFunctionOf(map, dim);
+    if (!maybeOperandDimensionToPack.has_value()) {
+      newMaps.push_back(map);
+      continue;
+    }
+
+    // We can only pack AffineDimExpr atm.
+    if (!map.getResult(maybeOperandDimensionToPack.value())
+             .isa<AffineDimExpr>())
+      return failure();
+
+    // Add `newDim` to the results of the map.
+    map = map.insertResult(Builder(map.getContext()).getAffineDimExpr(newDim),
+                           map.getNumResults());
+    newMaps.push_back(map);
+
+    // Record the that `operandIdx` is packed.
+    packedDimPerIndexingMap[operandIdx] = maybeOperandDimensionToPack;
+  }
+  indexingMaps = newMaps;
+
+  return packedDimPerIndexingMap;
+}
+
+namespace {
+
+/// Helper struct to encode packing along one dimension of a LinalgOp.
+struct PackedOperandsDim {
+  OpFoldResult packedSize;
+  SmallVector<std::optional<int64_t>> packedDimForEachOperand;
+};
+
+/// Helper struct to encode packing along all dimensions of a LinalgOp.
+struct PackedOperandsDimList {
+  void push_back(PackedOperandsDim &&packedOperandsDims) {
+    spec.emplace_back(packedOperandsDims);
+  }
+  /// Return all the dims that have been packed for operand @ `operandPos`.
+  SmallVector<int64_t> extractPackedDimsForOperand(int64_t operandPos);
+  /// Return all the pack sizes by which an operand @ `operandPos` is packed.
+  SmallVector<OpFoldResult> extractPackSizesForOperand(int64_t operandPos);
+
+private:
+  SmallVector<PackedOperandsDim> spec;
+};
+
+} // namespace
+
+SmallVector<int64_t>
+PackedOperandsDimList::extractPackedDimsForOperand(int64_t operandPos) {
+  SmallVector<int64_t> res;
+  for (int64_t i = 0, e = spec.size(); i < e; ++i) {
+    if (!spec[i].packedDimForEachOperand[operandPos].has_value())
+      continue;
+    res.push_back(spec[i].packedDimForEachOperand[operandPos].value());
+  }
+  return res;
+}
+
+SmallVector<OpFoldResult>
+PackedOperandsDimList::extractPackSizesForOperand(int64_t operandPos) {
+  SmallVector<OpFoldResult> res;
+  for (int64_t i = 0, e = spec.size(); i < e; ++i) {
+    if (!spec[i].packedDimForEachOperand[operandPos].has_value())
+      continue;
+    res.push_back(spec[i].packedSize);
+  }
+  return res;
+}
+
+/// Implement packing of a single LinalgOp by performing packing by
+/// `packedSizes`. There must be one packedSizes entry per `linalgOp` iterator.
+/// Return the packed Linalg op on success, failure otherwise.
+FailureOr<PackResult> linalg::pack(RewriterBase &rewriter,
+                                   linalg::LinalgOp linalgOp,
+                                   ArrayRef<OpFoldResult> packedSizes) {
+  if (packedSizes.size() != linalgOp.getNumLoops()) {
+    return rewriter.notifyMatchFailure(linalgOp,
+                                       "incorrect number of pack sizes");
+  }
+
+  Location loc = linalgOp->getLoc();
+  SmallVector<AffineMap> indexingMaps = linalgOp.getIndexingMapsArray();
+  SmallVector<utils::IteratorType> iteratorTypes =
+      linalgOp.getIteratorTypesArray();
+  LLVM_DEBUG(DBGS() << "Start packing: " << linalgOp << "\n";
+             llvm::interleaveComma(indexingMaps, DBGS() << "maps: "); DBGSNL();
+             llvm::interleaveComma(iteratorTypes, DBGS() << "iterators: ");
+             DBGSNL(););
+
+  SmallVector<tensor::PackOp> packOps;
+  SmallVector<tensor::UnPackOp> unPackOps;
+  // Step 1. Pack each dim of the LinalgOp metadata by packedSizes[i].
+  PackedOperandsDimList listOfPackedOperandsDim;
+  for (int64_t i = 0, e = packedSizes.size(); i < e; ++i) {
+    std::optional<int64_t> maybeConstant = getConstantIntValue(packedSizes[i]);
+    // Skip tile sizes explicitly set to 0.
+    if (maybeConstant.has_value() && maybeConstant.value() == 0)
+      continue;
+
+    PackedOperandsDim packedOperandsDims;
+    packedOperandsDims.packedSize = packedSizes[i];
+    FailureOr<SmallVector<std::optional<int64_t>>>
+        maybePackedDimForEachOperand =
+            packLinalgMetadataOnce(indexingMaps, iteratorTypes, i);
+    if (failed(maybePackedDimForEachOperand))
+      return failure();
+    packedOperandsDims.packedDimForEachOperand = *maybePackedDimForEachOperand;
+    listOfPackedOperandsDim.push_back(std::move(packedOperandsDims));
+
+    LLVM_DEBUG(
+        DBGS() << "++++ After pack size #" << i << ": " << packedSizes[i]
+               << "\n";
+        llvm::interleaveComma(indexingMaps, DBGS() << "maps: "); DBGSNL();
+        llvm::interleaveComma(iteratorTypes, DBGS() << "iterators: "); DBGSNL();
+        llvm::interleaveComma(packedOperandsDims.packedDimForEachOperand,
+                              DBGS() << "packedDimForEachOperand: ");
+        DBGSNL(););
+  }
+
+  // Step 2. Propagate packing to all LinalgOp operands.
+  SmallVector<Value> inputsAndInits, results;
+  for (auto operandsList :
+       {linalgOp.getDpsInputOperands(), linalgOp.getDpsInitOperands()}) {
+    for (OpOperand *opOperandPtr : operandsList) {
+      int64_t pos = opOperandPtr->getOperandNumber();
+      Value operand = opOperandPtr->get();
+      SmallVector<int64_t> innerPos =
+          listOfPackedOperandsDim.extractPackedDimsForOperand(pos);
+      SmallVector<OpFoldResult> innerPackSizes =
+          listOfPackedOperandsDim.extractPackSizesForOperand(pos);
+      LLVM_DEBUG(
+          DBGS() << "operand: " << operand << "\n";
+          llvm::interleaveComma(innerPos, DBGS() << "innerPos: "); DBGSNL();
+          llvm::interleaveComma(innerPackSizes, DBGS() << "innerPackSizes: ");
+          DBGSNL(););
+      if (innerPackSizes.empty()) {
+        inputsAndInits.push_back(operand);
+        continue;
+      }
+      Value dest = tensor::PackOp::createDestinationTensor(
+          rewriter, loc, operand, innerPackSizes, innerPos,
+          /*outerDimsPerm=*/{});
+      // TODO: value of the padding attribute should be determined by consumers.
+      Attribute zeroAttr =
+          rewriter.getZeroAttr(getElementTypeOrSelf(dest.getType()));
+      Value zero = rewriter.create<arith::ConstantOp>(loc, zeroAttr);
+      packOps.push_back(rewriter.create<tensor::PackOp>(
+          loc, operand, dest, innerPos, innerPackSizes, zero));
+      inputsAndInits.push_back(packOps.back());
+    }
+  }
+
+  // Step 3. Build the packed op, use the type of `inits` as result types.
+  ValueRange inputs =
+      ValueRange{inputsAndInits}.take_front(linalgOp.getNumDpsInputs());
+  ValueRange inits =
+      ValueRange{inputsAndInits}.take_back(linalgOp.getNumDpsInits());
+  auto packedLinalgOp = rewriter.create<linalg::GenericOp>(
+      linalgOp.getLoc(), inits.getTypes(), inputs, inits, indexingMaps,
+      iteratorTypes);
+  packedLinalgOp.getRegion().takeBody(linalgOp->getRegion(0));
+
+  // Step 4. Propagate packing to all the op results.
+  for (OpResult result : packedLinalgOp->getResults()) {
+    int64_t resultNum = result.getResultNumber();
+    tensor::PackOp maybePackedInit =
+        inits[resultNum].getDefiningOp<tensor::PackOp>();
+    if (!maybePackedInit) {
+      results.push_back(result);
+      continue;
+    }
+    // Build the symmetrical UnPackOp to the existing PackOp.
+    unPackOps.push_back(rewriter.create<tensor::UnPackOp>(
+        packedLinalgOp->getLoc(), result, maybePackedInit.getSource(),
+        maybePackedInit.getInnerDimsPos(), maybePackedInit.getMixedTiles()));
+    results.push_back(unPackOps.back());
+  }
+
+  // Step 5. Replace `linalgOp`.
+  rewriter.replaceOp(linalgOp, results);
+
+  // Return packedLinalgOp.
+  return PackResult{packOps,
+                    cast<linalg::LinalgOp>(packedLinalgOp.getOperation()),
+                    unPackOps};
+}
+
+//===----------------------------------------------------------------------===//
+// packTranspose transformation.
+//===----------------------------------------------------------------------===//
+
+/// Return a copy of `tensorType` after permutation by `permutationVector`.
+// Note: Should be a new method in of MemRef/RankedTensor/VectorType::Builder
+// but this would introduce a dependence on Dialect in IR.
+// TODO: Restructure.
+static RankedTensorType permuteShape(RankedTensorType tensorType,
+                                     ArrayRef<int64_t> permutationVector) {
+  SmallVector<int64_t> shape(tensorType.getShape());
+  applyPermutationToVector(shape, permutationVector);
+  return RankedTensorType::Builder(tensorType).setShape(shape);
+}
+
+/// Return a new GenericOp obtained by transposing opOperand by the permutation
+/// vector:
+///   - the corresponding indexing map is transposed by `permutation`
+///   - the corresponding operand value is replaced by `transposedValue`
+/// `linalgOp` is replaced by the return op in the process.
+/// Asserts that `transposedValue` is of the proper transposed ShapedType.
+static LinalgOp transposeOneLinalgOperandAndReplace(
+    RewriterBase &rewriter, LinalgOp linalgOp, OpOperand &opOperand,
+    ArrayRef<int64_t> permutation, Value transposedValue) {
+  // Sanity check the operand.
+  assert(linalgOp == opOperand.getOwner() && "linalg op must own the operand");
+
+  // Sanity check of the expected transposed tensor type.
+  auto tensorType = permuteShape(
+      opOperand.get().getType().cast<RankedTensorType>(), permutation);
+  (void)tensorType;
+  assert(tensorType == transposedValue.getType() &&
+         "expected tensor type mismatch");
+
+  // Compute the transposed indexing map.
+  // Sigh unsigned pollution.
+  SmallVector<unsigned> tmpTransposition = llvm::to_vector(
+      llvm::map_range(permutation, [](int64_t i) -> unsigned { return i; }));
+  AffineMap permutationMap =
+      AffineMap::getPermutationMap(tmpTransposition, rewriter.getContext());
+  AffineMap transposedMap =
+      permutationMap.compose(linalgOp.getMatchingIndexingMap(&opOperand));
+
+  // Set the transposed indexing map in the proper position.
+  SmallVector<AffineMap> indexingMaps = linalgOp.getIndexingMapsArray();
+  indexingMaps[linalgOp.getIndexingMapIndex(&opOperand)] = transposedMap;
+  // Set the transposedValue in the proper operand position.
+  SmallVector<Value> operands = linalgOp->getOperands();
+  operands[opOperand.getOperandNumber()] = transposedValue;
+
+  ValueRange operandsRef(operands);
+  auto transposedGenericOp = rewriter.create<linalg::GenericOp>(
+      /*location=*/linalgOp->getLoc(),
+      /*resultTensorTypes=*/
+      operandsRef.drop_front(linalgOp.getNumDpsInputs()).getTypes(),
+      /*inputs=*/operandsRef.take_front(linalgOp.getNumDpsInputs()),
+      /*outputs=*/operandsRef.drop_front(linalgOp.getNumDpsInputs()),
+      /*indexingMaps=*/indexingMaps,
+      /*iteratorTypes=*/linalgOp.getIteratorTypesArray());
+  transposedGenericOp.getRegion().takeBody(linalgOp->getRegion(0));
+  rewriter.replaceOp(linalgOp, transposedGenericOp->getResults());
+
+  return cast<linalg::LinalgOp>(transposedGenericOp.getOperation());
+}
+
+FailureOr<PackTransposeResult>
+linalg::packTranspose(RewriterBase &rewriter, tensor::PackOp packOp,
+                      linalg::LinalgOp linalgOp, tensor::UnPackOp maybeUnPackOp,
+                      ArrayRef<int64_t> outerPerm,
+                      ArrayRef<int64_t> innerPerm) {
+  Location loc = linalgOp.getLoc();
+
+  // Step 1. Transpose packOp.
+  rewriter.setInsertionPoint(packOp);
+  tensor::PackOp transposedPackOp =
+      packOp.createTransposedClone(rewriter, loc, innerPerm, outerPerm);
+
+  if (!packOp.getResult().hasOneUse())
+    return rewriter.notifyMatchFailure(linalgOp, "expect single pack use");
+
+  OpOperand &packUse = *packOp->getUses().begin();
+  if (packUse.getOwner() != linalgOp) {
+    return rewriter.notifyMatchFailure(
+        linalgOp, "not a single use by the LinalgOp target");
+  }
+  if (maybeUnPackOp &&
+      (!linalgOp.isDpsInit(&packUse) ||
+       maybeUnPackOp.getSource() != linalgOp.getTiedOpResult(&packUse))) {
+    return rewriter.notifyMatchFailure(linalgOp,
+                                       "not produced by the LinalgOp target");
+  }
+
+  // Step 2. Transpose linalgOp.
+  // transposedPackOp.getOuterDimsPerm() may be empty, in which case it is the
+  // identity. Don't rely on it.
+  int64_t numLeadingDims = packOp.getSourceRank();
+  int64_t numTrailingDims = packOp.getInnerDimsPos().size();
+  // Step 2.a. Compute the permutation on the whole operand.
+  // Leading part just reuse the outerPerm.
+  SmallVector<int64_t> permutation(outerPerm);
+  if (permutation.empty())
+    llvm::append_range(permutation, llvm::seq<int64_t>(0, numLeadingDims));
+  // Trailing part needs to reindex positions by `numLeadingDims`.
+  if (innerPerm.empty()) {
+    llvm::append_range(
+        permutation,
+        llvm::seq<int64_t>(numLeadingDims, numLeadingDims + numTrailingDims));
+  } else {
+    llvm::append_range(permutation,
+                       llvm::map_range(innerPerm, [&](int64_t pos) {
+                         return numLeadingDims + pos;
+                       }));
+  }
+  if (!isPermutationVector(permutation))
+    return rewriter.notifyMatchFailure(linalgOp, "invalid permutation");
+
+  // Step 2.b. Save the transposedPackUse operand number in case we need to
+  // get the tied OpResult after `linalgOp` has been replaced.
+  int64_t packUseOperandNumber = packUse.getOperandNumber();
+  // Step 2.c. Actually perform the transposition.
+  rewriter.setInsertionPoint(linalgOp);
+  linalg::LinalgOp transposedLinalgOp = transposeOneLinalgOperandAndReplace(
+      rewriter, linalgOp, packUse, permutation, transposedPackOp.getResult());
+
+  // Step 3. Maybe transpose unPackOp.
+  tensor::UnPackOp transposedUnPackOp;
+  if (maybeUnPackOp) {
+    OpOperand &opOperand =
+        transposedLinalgOp->getOpOperand(packUseOperandNumber);
+    OpResult transposedResult = transposedLinalgOp.getTiedOpResult(&opOperand);
+    rewriter.setInsertionPoint(maybeUnPackOp);
+    transposedUnPackOp = maybeUnPackOp.createTransposedClone(
+        rewriter, loc, transposedResult, innerPerm, outerPerm);
+
+    rewriter.replaceOp(maybeUnPackOp, transposedUnPackOp->getResults());
+  }
+
+  // Step 4. Finally, replace packOp now that we don't need it anymore.
+  rewriter.replaceOp(packOp, transposedPackOp->getResults());
+
+  return PackTransposeResult{transposedPackOp, transposedLinalgOp,
+                             transposedUnPackOp};
+}
+
+//===----------------------------------------------------------------------===//
+// Transformations exposed as rewrite patterns.
+//===----------------------------------------------------------------------===//
+
+LinalgTilingOptions &
+mlir::linalg::LinalgTilingOptions::setTileSizes(ArrayRef<int64_t> ts) {
+  assert(!tileSizeComputationFunction && "tile sizes already set");
+  SmallVector<int64_t, 4> tileSizes(ts.begin(), ts.end());
+  tileSizeComputationFunction = [tileSizes](OpBuilder &b, Operation *op) {
+    OpBuilder::InsertionGuard guard(b);
+    b.setInsertionPointToStart(
+        &op->getParentOfType<func::FuncOp>().getBody().front());
+    return llvm::to_vector<4>(map_range(tileSizes, [&](int64_t s) {
+      Value v = b.create<arith::ConstantIndexOp>(op->getLoc(), s);
+      return v;
+    }));
+  };
+  return *this;
+}
+
+/// Linalg padding pattern.
+
+mlir::linalg::LinalgPaddingPattern::LinalgPaddingPattern(
+    MLIRContext *context, LinalgPaddingOptions options, PatternBenefit benefit)
+    : OpInterfaceRewritePattern<LinalgOp>(context, benefit),
+      options(std::move(options)) {}
+
+LogicalResult mlir::linalg::LinalgPaddingPattern::matchAndRewrite(
+    LinalgOp op, PatternRewriter &rewriter) const {
+  return padLinalgOp(rewriter, op, options);
+}
+
 LogicalResult mlir::linalg::CopyVectorizationPattern::matchAndRewrite(
     memref::CopyOp copyOp, PatternRewriter &rewriter) const {
   return vectorizeCopy(rewriter, copyOp);
 }
 
-static SmallVector<utils::IteratorType>
-getNParallelLoopsAttrs(unsigned nParallelLoops) {
-  return SmallVector<utils::IteratorType>(nParallelLoops,
-                                          utils::IteratorType::parallel);
-}
-
 /// Rewrite a tensor::PadOp into a sequence of EmptyOp, FillOp (to
 /// initialize with pad_val) and GenericOp (to copy contents).
 LogicalResult
 PadOpTransformationPattern::matchAndRewrite(tensor::PadOp padOp,
                                             PatternRewriter &rewriter) const {
 
   auto inputShapedType = padOp.getSource().getType().cast<ShapedType>();
   auto resultShapedType = padOp.getResult().getType().cast<ShapedType>();
patterns.add<
       DownscaleSizeOneWindowed2DConvolution<PoolingNhwcMaxUnsignedOp,
                                             PoolingNwcMaxUnsignedOp>,
       DownscaleSizeOneWindowed2DConvolution<PoolingNhwcMinOp, PoolingNwcMinOp>,
       DownscaleSizeOneWindowed2DConvolution<PoolingNhwcMinUnsignedOp,
                                             PoolingNwcMinUnsignedOp>,
       DownscaleSizeOneWindowed2DConvolution<PoolingNchwMaxOp, PoolingNcwMaxOp>>(
       patterns.getContext(), benefit);
 }
-
-//===----------------------------------------------------------------------===//
-// pack transformation.
-//===----------------------------------------------------------------------===//
-
-#ifndef NDEBUG
-/// Return true if `map` has 0 or 1 result function of AffineDimExpr(dim).
-static bool hasAtMostOneResultFunctionOfDim(AffineMap map, int64_t dim) {
-  bool found = false;
-  for (AffineExpr e : map.getResults()) {
-    if (!e.isFunctionOfDim(dim))
-      continue;
-    if (found)
-      return false;
-    found = true;
-  }
-  return true;
-}
-#endif // NDEBUG
-
-/// Return the index of the first result of `map` that is a function of
-/// AffineDimExpr(dim), std::nullopt otherwise.
-static std::optional<int64_t> getFirstResultIndexFunctionOf(AffineMap map,
-                                                            int64_t dim) {
-  for (int64_t i = 0, e = map.getNumResults(); i < e; ++i) {
-    AffineExpr expr = map.getResult(i);
-    if (!expr.isFunctionOfDim(dim))
-      continue;
-    return i;
-  }
-  return std::nullopt;
-}
-
-/// Perform one step of packing of a LinalgOp's metadata along `dim` into the
-/// `newDim` at `iteratorTypes.size()` by:
-///   1. Appending `iteratorTypes[newDim]`, equal to `iteratorTypes[dim]`.
-///   2. Appending a `newDim` to the domain of every indexing map.
-///   3. For each operand (i.e. for each map in `indexingMaps`), perform packing
-///      by potentially adding a `newDim` result to `map`.
-/// The preserved invariant is that `iteratorTypes.size()` is always equal to
-/// `map.getNumDims()` for every map in `indexingMaps`.
-///
-/// Update `indexingMaps` and `iteratorTypes` inplace as one step of the update.
-/// Return a vector that records the optional packing for each operand.
-/// Return failure if the packed indexing cannot be represented with a LinalgOp.
-///
-/// Further details:
-/// ================
-/// The current implementation of packing (i.e. data tiling) consists of
-/// rewriting a linearized strip-mined form into a higher-dimensional access.
-/// e.g. consider an access `A[I][f(j, k, l)]` and packing by 4; we rewrite
-/// `I` into `4 * i + ii`, where `0 <= ii < 4`.
-/// The access is further rewritten as `A[i][f(j, k, l)][ii]`.
-///
-/// This rewrite into higher dimensional access is not possible for general
-/// AffineExpr in Linalg atm, it is restricted to an AffineDimExpr:
-/// e.g. consider an access `A[I + J][f(j, k, l)]` and packing by 4; we
-/// rewrite `I + J` into `4 * i + ii + J`, where `0 <= ii < 4`.
-/// The rewrite of the access would be a form not representable in Linalg:
-///   `A[i + (ii + J) / 4][f(j, k, l)][(ii + J) % 4]`.
-/// Note however that as `J` and `ii` iterate, the accesses do not have a
-/// particular alignment, so packing does not achieve alignment in this case
-///
-/// In the future, we may want to consider a mixed-form that allows some
-/// alignment in the presence of multiple accesses:
-///   `A[I][f(j, k, l)]` and `B[I + J][f(j, k, l)]`
-/// And would rewrite accesses as:
-///   `A[i][f(j, k, l)][ii]` and `B[4 * i + ii + J][f(j, k, l)]`
-static FailureOr<SmallVector<std::optional<int64_t>>>
-packLinalgMetadataOnce(SmallVectorImpl<AffineMap> &indexingMaps,
-                       SmallVectorImpl<utils::IteratorType> &iteratorTypes,
-                       int64_t dim) {
-  int64_t newDim = iteratorTypes.size();
-  iteratorTypes.push_back(iteratorTypes[dim]);
-
-  SmallVector<std::optional<int64_t>> packedDimPerIndexingMap(
-      indexingMaps.size(), std::nullopt);
-  SmallVector<AffineMap> newMaps;
-  for (int64_t operandIdx = 0, e = indexingMaps.size(); operandIdx < e;
-       ++operandIdx) {
-    AffineMap map = indexingMaps[operandIdx];
-
-    // Add the `newDim` to map whatever the case.
-    assert(map.getNumDims() == newDim && "num dims invariant violation");
-    map = map.shiftDims(1, newDim);
-
-    // Get the at-most-1 index of the result that is a function of `dim`.
-    // If we can find one, we insert `AffineDimExpr(newDim)` to the map, which
-    // logically chunks dimension `dim` into `K * dim + newDim`, where the
-    // packing factor `K` is specified separately.
-    assert(hasAtMostOneResultFunctionOfDim(map, dim) &&
-           "num results invariant violation");
-    auto maybeOperandDimensionToPack = getFirstResultIndexFunctionOf(map, dim);
-    if (!maybeOperandDimensionToPack.has_value()) {
-      newMaps.push_back(map);
-      continue;
-    }
-
-    // We can only pack AffineDimExpr atm.
-    if (!map.getResult(maybeOperandDimensionToPack.value())
-             .isa<AffineDimExpr>())
-      return failure();
-
-    // Add `newDim` to the results of the map.
-    map = map.insertResult(Builder(map.getContext()).getAffineDimExpr(newDim),
-                           map.getNumResults());
-    newMaps.push_back(map);
-
-    // Record the that `operandIdx` is packed.
-    packedDimPerIndexingMap[operandIdx] = maybeOperandDimensionToPack;
-  }
-  indexingMaps = newMaps;
-
-  return packedDimPerIndexingMap;
-}
-
-namespace {
-
-/// Helper struct to encode packing along one dimension of a LinalgOp.
-struct PackedOperandsDim {
-  OpFoldResult packedSize;
-  SmallVector<std::optional<int64_t>> packedDimForEachOperand;
-};
-
-/// Helper struct to encode packing along all dimensions of a LinalgOp.
-struct PackedOperandsDimList {
-  void push_back(PackedOperandsDim &&packedOperandsDims) {
-    spec.emplace_back(packedOperandsDims);
-  }
-  /// Return all the dims that have been packed for operand @ `operandPos`.
-  SmallVector<int64_t> extractPackedDimsForOperand(int64_t operandPos);
-  /// Return all the pack sizes by which an operand @ `operandPos` is packed.
-  SmallVector<OpFoldResult> extractPackSizesForOperand(int64_t operandPos);
-
-private:
-  SmallVector<PackedOperandsDim> spec;
-};
-
-} // namespace
-
-SmallVector<int64_t>
-PackedOperandsDimList::extractPackedDimsForOperand(int64_t operandPos) {
-  SmallVector<int64_t> res;
-  for (int64_t i = 0, e = spec.size(); i < e; ++i) {
-    if (!spec[i].packedDimForEachOperand[operandPos].has_value())
-      continue;
-    res.push_back(spec[i].packedDimForEachOperand[operandPos].value());
-  }
-  return res;
-}
-
-SmallVector<OpFoldResult>
-PackedOperandsDimList::extractPackSizesForOperand(int64_t operandPos) {
-  SmallVector<OpFoldResult> res;
-  for (int64_t i = 0, e = spec.size(); i < e; ++i) {
-    if (!spec[i].packedDimForEachOperand[operandPos].has_value())
-      continue;
-    res.push_back(spec[i].packedSize);
-  }
-  return res;
-}
-
-/// Implement packing of a single LinalgOp by performing packing by
-/// `packedSizes`. There must be one packedSizes entry per `linalgOp` iterator.
-/// Return the packed Linalg op on success, failure otherwise.
-FailureOr<PackResult> linalg::pack(RewriterBase &rewriter,
-                                   linalg::LinalgOp linalgOp,
-                                   ArrayRef<OpFoldResult> packedSizes) {
-  if (packedSizes.size() != linalgOp.getNumLoops()) {
-    return rewriter.notifyMatchFailure(linalgOp,
-                                       "incorrect number of pack sizes");
-  }
-
-  Location loc = linalgOp->getLoc();
-  SmallVector<AffineMap> indexingMaps = linalgOp.getIndexingMapsArray();
-  SmallVector<utils::IteratorType> iteratorTypes =
-      linalgOp.getIteratorTypesArray();
-  LLVM_DEBUG(DBGS() << "Start packing: " << linalgOp << "\n";
-             llvm::interleaveComma(indexingMaps, DBGS() << "maps: "); DBGSNL();
-             llvm::interleaveComma(iteratorTypes, DBGS() << "iterators: ");
-             DBGSNL(););
-
-  SmallVector<tensor::PackOp> packOps;
-  SmallVector<tensor::UnPackOp> unPackOps;
-  // Step 1. Pack each dim of the LinalgOp metadata by packedSizes[i].
-  PackedOperandsDimList listOfPackedOperandsDim;
-  for (int64_t i = 0, e = packedSizes.size(); i < e; ++i) {
-    std::optional<int64_t> maybeConstant = getConstantIntValue(packedSizes[i]);
-    // Skip tile sizes explicitly set to 0.
-    if (maybeConstant.has_value() && maybeConstant.value() == 0)
-      continue;
-
-    PackedOperandsDim packedOperandsDims;
-    packedOperandsDims.packedSize = packedSizes[i];
-    FailureOr<SmallVector<std::optional<int64_t>>>
-        maybePackedDimForEachOperand =
-            packLinalgMetadataOnce(indexingMaps, iteratorTypes, i);
-    if (failed(maybePackedDimForEachOperand))
-      return failure();
-    packedOperandsDims.packedDimForEachOperand = *maybePackedDimForEachOperand;
-    listOfPackedOperandsDim.push_back(std::move(packedOperandsDims));
-
-    LLVM_DEBUG(
-        DBGS() << "++++ After pack size #" << i << ": " << packedSizes[i]
-               << "\n";
-        llvm::interleaveComma(indexingMaps, DBGS() << "maps: "); DBGSNL();
-        llvm::interleaveComma(iteratorTypes, DBGS() << "iterators: "); DBGSNL();
-        llvm::interleaveComma(packedOperandsDims.packedDimForEachOperand,
-                              DBGS() << "packedDimForEachOperand: ");
-        DBGSNL(););
-  }
-
-  // Step 2. Propagate packing to all LinalgOp operands.
-  SmallVector<Value> inputsAndInits, results;
-  for (auto operandsList :
-       {linalgOp.getDpsInputOperands(), linalgOp.getDpsInitOperands()}) {
-    for (OpOperand *opOperandPtr : operandsList) {
-      int64_t pos = opOperandPtr->getOperandNumber();
-      Value operand = opOperandPtr->get();
-      SmallVector<int64_t> innerPos =
-          listOfPackedOperandsDim.extractPackedDimsForOperand(pos);
-      SmallVector<OpFoldResult> innerPackSizes =
-          listOfPackedOperandsDim.extractPackSizesForOperand(pos);
-      LLVM_DEBUG(
-          DBGS() << "operand: " << operand << "\n";
-          llvm::interleaveComma(innerPos, DBGS() << "innerPos: "); DBGSNL();
-          llvm::interleaveComma(innerPackSizes, DBGS() << "innerPackSizes: ");
-          DBGSNL(););
-      if (innerPackSizes.empty()) {
-        inputsAndInits.push_back(operand);
-        continue;
-      }
-      Value dest = tensor::PackOp::createDestinationTensor(
-          rewriter, loc, operand, innerPackSizes, innerPos,
-          /*outerDimsPerm=*/{});
-      // TODO: value of the padding attribute should be determined by consumers.
-      Attribute zeroAttr =
-          rewriter.getZeroAttr(getElementTypeOrSelf(dest.getType()));
-      Value zero = rewriter.create<arith::ConstantOp>(loc, zeroAttr);
-      packOps.push_back(rewriter.create<tensor::PackOp>(
-          loc, operand, dest, innerPos, innerPackSizes, zero));
-      inputsAndInits.push_back(packOps.back());
-    }
-  }
-
-  // Step 3. Build the packed op, use the type of `inits` as result types.
-  ValueRange inputs =
-      ValueRange{inputsAndInits}.take_front(linalgOp.getNumDpsInputs());
-  ValueRange inits =
-      ValueRange{inputsAndInits}.take_back(linalgOp.getNumDpsInits());
-  auto packedLinalgOp = rewriter.create<linalg::GenericOp>(
-      linalgOp.getLoc(), inits.getTypes(), inputs, inits, indexingMaps,
-      iteratorTypes);
-  packedLinalgOp.getRegion().takeBody(linalgOp->getRegion(0));
-
-  // Step 4. Propagate packing to all the op results.
-  for (OpResult result : packedLinalgOp->getResults()) {
-    int64_t resultNum = result.getResultNumber();
-    tensor::PackOp maybePackedInit =
-        inits[resultNum].getDefiningOp<tensor::PackOp>();
-    if (!maybePackedInit) {
-      results.push_back(result);
-      continue;
-    }
-    // Build the symmetrical UnPackOp to the existing PackOp.
-    unPackOps.push_back(rewriter.create<tensor::UnPackOp>(
-        packedLinalgOp->getLoc(), result, maybePackedInit.getSource(),
-        maybePackedInit.getInnerDimsPos(), maybePackedInit.getMixedTiles()));
-    results.push_back(unPackOps.back());
-  }
-
-  // Step 5. Replace `linalgOp`.
-  rewriter.replaceOp(linalgOp, results);
-
-  // Return packedLinalgOp.
-  return PackResult{packOps,
-                    cast<linalg::LinalgOp>(packedLinalgOp.getOperation()),
-                    unPackOps};
-}
-
-//===----------------------------------------------------------------------===//
-// packTranspose transformation.
-//===----------------------------------------------------------------------===//
-
-/// Return a copy of `tensorType` after permutation by `permutationVector`.
-// Note: Should be a new method in of MemRef/RankedTensor/VectorType::Builder
-// but this would introduce a dependence on Dialect in IR.
-// TODO: Restructure.
-static RankedTensorType permuteShape(RankedTensorType tensorType,
-                                     ArrayRef<int64_t> permutationVector) {
-  SmallVector<int64_t> shape(tensorType.getShape());
-  applyPermutationToVector(shape, permutationVector);
-  return RankedTensorType::Builder(tensorType).setShape(shape);
-}
-
-/// Return a new GenericOp obtained by transposing opOperand by the permutation
-/// vector:
-///   - the corresponding indexing map is transposed by `permutation`
-///   - the corresponding operand value is replaced by `transposedValue`
-/// `linalgOp` is replaced by the return op in the process.
-/// Asserts that `transposedValue` is of the proper transposed ShapedType.
-static LinalgOp transposeOneLinalgOperandAndReplace(
-    RewriterBase &rewriter, LinalgOp linalgOp, OpOperand &opOperand,
-    ArrayRef<int64_t> permutation, Value transposedValue) {
-  // Sanity check the operand.
-  assert(linalgOp == opOperand.getOwner() && "linalg op must own the operand");
-
-  // Sanity check of the expected transposed tensor type.
-  auto tensorType = permuteShape(
-      opOperand.get().getType().cast<RankedTensorType>(), permutation);
-  (void)tensorType;
-  assert(tensorType == transposedValue.getType() &&
-         "expected tensor type mismatch");
-
-  // Compute the transposed indexing map.
-  // Sigh unsigned pollution.
-  SmallVector<unsigned> tmpTransposition = llvm::to_vector(
-      llvm::map_range(permutation, [](int64_t i) -> unsigned { return i; }));
-  AffineMap permutationMap =
-      AffineMap::getPermutationMap(tmpTransposition, rewriter.getContext());
-  AffineMap transposedMap =
-      permutationMap.compose(linalgOp.getMatchingIndexingMap(&opOperand));
-
-  // Set the transposed indexing map in the proper position.
-  SmallVector<AffineMap> indexingMaps = linalgOp.getIndexingMapsArray();
-  indexingMaps[linalgOp.getIndexingMapIndex(&opOperand)] = transposedMap;
-  // Set the transposedValue in the proper operand position.
-  SmallVector<Value> operands = linalgOp->getOperands();
-  operands[opOperand.getOperandNumber()] = transposedValue;
-
-  ValueRange operandsRef(operands);
-  auto transposedGenericOp = rewriter.create<linalg::GenericOp>(
-      /*location=*/linalgOp->getLoc(),
-      /*resultTensorTypes=*/
-      operandsRef.drop_front(linalgOp.getNumDpsInputs()).getTypes(),
-      /*inputs=*/operandsRef.take_front(linalgOp.getNumDpsInputs()),
-      /*outputs=*/operandsRef.drop_front(linalgOp.getNumDpsInputs()),
-      /*indexingMaps=*/indexingMaps,
-      /*iteratorTypes=*/linalgOp.getIteratorTypesArray());
-  transposedGenericOp.getRegion().takeBody(linalgOp->getRegion(0));
-  rewriter.replaceOp(linalgOp, transposedGenericOp->getResults());
-
-  return cast<linalg::LinalgOp>(transposedGenericOp.getOperation());
-}
-
-FailureOr<PackTransposeResult>
-linalg::packTranspose(RewriterBase &rewriter, tensor::PackOp packOp,
-                      linalg::LinalgOp linalgOp, tensor::UnPackOp maybeUnPackOp,
-                      ArrayRef<int64_t> outerPerm,
-                      ArrayRef<int64_t> innerPerm) {
-  Location loc = linalgOp.getLoc();
-
-  // Step 1. Transpose packOp.
-  rewriter.setInsertionPoint(packOp);
-  tensor::PackOp transposedPackOp =
-      packOp.createTransposedClone(rewriter, loc, innerPerm, outerPerm);
-
-  if (!packOp.getResult().hasOneUse())
-    return rewriter.notifyMatchFailure(linalgOp, "expect single pack use");
-
-  OpOperand &packUse = *packOp->getUses().begin();
-  if (packUse.getOwner() != linalgOp) {
-    return rewriter.notifyMatchFailure(
-        linalgOp, "not a single use by the LinalgOp target");
-  }
-  if (maybeUnPackOp &&
-      (!linalgOp.isDpsInit(&packUse) ||
-       maybeUnPackOp.getSource() != linalgOp.getTiedOpResult(&packUse))) {
-    return rewriter.notifyMatchFailure(linalgOp,
-                                       "not produced by the LinalgOp target");
-  }
-
-  // Step 2. Transpose linalgOp.
-  // transposedPackOp.getOuterDimsPerm() may be empty, in which case it is the
-  // identity. Don't rely on it.
-  int64_t numLeadingDims = packOp.getSourceRank();
-  int64_t numTrailingDims = packOp.getInnerDimsPos().size();
-  // Step 2.a. Compute the permutation on the whole operand.
-  // Leading part just reuse the outerPerm.
-  SmallVector<int64_t> permutation(outerPerm);
-  if (permutation.empty())
-    llvm::append_range(permutation, llvm::seq<int64_t>(0, numLeadingDims));
-  // Trailing part needs to reindex positions by `numLeadingDims`.
-  if (innerPerm.empty()) {
-    llvm::append_range(
-        permutation,
-        llvm::seq<int64_t>(numLeadingDims, numLeadingDims + numTrailingDims));
-  } else {
-    llvm::append_range(permutation,
-                       llvm::map_range(innerPerm, [&](int64_t pos) {
-                         return numLeadingDims + pos;
-                       }));
-  }
-  if (!isPermutationVector(permutation))
-    return rewriter.notifyMatchFailure(linalgOp, "invalid permutation");
-
-  // Step 2.b. Save the transposedPackUse operand number in case we need to
-  // get the tied OpResult after `linalgOp` has been replaced.
-  int64_t packUseOperandNumber = packUse.getOperandNumber();
-  // Step 2.c. Actually perform the transposition.
-  rewriter.setInsertionPoint(linalgOp);
-  linalg::LinalgOp transposedLinalgOp = transposeOneLinalgOperandAndReplace(
-      rewriter, linalgOp, packUse, permutation, transposedPackOp.getResult());
-
-  // Step 3. Maybe transpose unPackOp.
-  tensor::UnPackOp transposedUnPackOp;
-  if (maybeUnPackOp) {
-    OpOperand &opOperand =
-        transposedLinalgOp->getOpOperand(packUseOperandNumber);
-    OpResult transposedResult = transposedLinalgOp.getTiedOpResult(&opOperand);
-    rewriter.setInsertionPoint(maybeUnPackOp);
-    transposedUnPackOp = maybeUnPackOp.createTransposedClone(
-        rewriter, loc, transposedResult, innerPerm, outerPerm);
-
-    rewriter.replaceOp(maybeUnPackOp, transposedUnPackOp->getResults());
-  }
-
-  // Step 4. Finally, replace packOp now that we don't need it anymore.
-  rewriter.replaceOp(packOp, transposedPackOp->getResults());
-
-  return PackTransposeResult{transposedPackOp, transposedLinalgOp,
-                             transposedUnPackOp};
-}

mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp

transform::LoopPeelOp::applyToOne(scf::ForOp target,
                                   transform::TransformState &state) {
   scf::ForOp result;
   IRRewriter rewriter(target->getContext());
   // This helper returns failure when peeling does not occur (i.e. when the IR
   // is not modified). This is not a failure for the op as the postcondition:
   //    "the loop trip count is divisible by the step"
   // is valid.
   LogicalResult status =
-      scf::peelAndCanonicalizeForLoop(rewriter, target, result);
+      scf::peelForLoopAndSimplifyBounds(rewriter, target, result);
   // TODO: Return both the peeled loop and the remainder loop.
   results.push_back(failed(status) ? target : result);
   return DiagnosedSilenceableFailure::success();
 }
 
 //===----------------------------------------------------------------------===//
 // LoopPipelineOp
 //===----------------------------------------------------------------------===//

mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp

partialIteration.walk([&](Operation *affineOp) {
     if (!isa<AffineMinOp, AffineMaxOp>(affineOp))
       return WalkResult::advance();
     (void)scf::rewritePeeledMinMaxOp(rewriter, affineOp, partialIv, previousUb,
                                      step, /*insideLoop=*/false);
     return WalkResult::advance();
   });
 }
 
-LogicalResult mlir::scf::peelAndCanonicalizeForLoop(RewriterBase &rewriter,
+LogicalResult mlir::scf::peelForLoopAndSimplifyBounds(RewriterBase &rewriter,
                                                       ForOp forOp,
                                                       ForOp &partialIteration) {
   Value previousUb = forOp.getUpperBound();
   Value splitBound;
   if (failed(peelForLoop(rewriter, forOp, partialIteration, splitBound)))
     return failure();
 
   // Rewrite affine.min and affine.max ops.
   rewriteAffineOpAfterPeeling(rewriter, forOp, partialIteration, previousUb);
 
if (skipPartial) {
       Operation *op = forOp.getOperation();
       while ((op = op->getParentOfType<scf::ForOp>())) {
         if (op->hasAttr(kPartialIterationLabel))
           return failure();
       }
     }
     // Apply loop peeling.
     scf::ForOp partialIteration;
-    if (failed(peelAndCanonicalizeForLoop(rewriter, forOp, partialIteration)))
+    if (failed(peelForLoopAndSimplifyBounds(rewriter, forOp, partialIteration)))
       return failure();
     // Apply label, so that the same loop is not rewritten a second time.
     partialIteration->setAttr(kPeeledLoopLabel, rewriter.getUnitAttr());
     rewriter.updateRootInPlace(forOp, [&]() {
       forOp->setAttr(kPeeledLoopLabel, rewriter.getUnitAttr());
     });
     partialIteration->setAttr(kPartialIterationLabel, rewriter.getUnitAttr());
     return success();