[mlir][Linalg] Refactor transform.structured.pad to separate out hoisting

[nicolasvasilache]

Summary

Depends on: D144717

[springerm]

mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td

771–772

Does this also compute an upper bound for tensor dimensions? What happens if the low/high padding and/or tensor source depends on the loop IV?

780

result

791

A loop handle may compose better with the remaining transform dialect ops.

[nicolasvasilache]

Rebase

[ftynse]

mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td

772

Nit: please describe the format of the transform attrbute.

780

I see this text is being copy-pasted, but the logic of apply-to-each doesn't just drop uncompatible ops, it emits a silenceable failure.

781

Nit: it's not a PDLOperation.

793

We don't want overconstrained result types to avoid proliferation of casts, see doc

[nicolasvasilache]

mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td

771–772

Does this also compute an upper bound for tensor dimensions?

It does and that has been around for a while; we want to evolve that and merge with your WIP PR that goes towards generalization.
But first, separation of concerns and making room for generalization.

791

Maybe, OTOH getting explicit handles is more cumbersome and makes it harder to e.g. devise tuning scripts.
I could also see us returning the top packing loop and the loop hoisted above for composition, TBD.
Let's keep this as a possible future improvement.

793

thanks!

[qcolombet]

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

414

Nit: If we don't want to give any special meaning to negative numLoops (e.g., -1 means all of them) maybe we should use unsigned.

Note: I realize that numLoops is used in different APIs with the same type so we can do this clean-up in a separate PR.

[qcolombet]

mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td

771

Nit: "given number of loops" => "at most the given number of loops"?

We may even want to add wording that the actual number depends on the number of enclosing loops and the given number of loops, whichever is the smallest.

It wasn't clear to me what would happen if we give a number of loops greater than the actual number of enclosing loops. I had to look at the implementation to check.

[nicolasvasilache]

Address comments and rebase.

[nicolasvasilache]

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

414

Moved to int64_t.

I am generally wary of using unsigned to represent non-negative.. unsigned should only be used in very specific limited cases in C++ (e.g. bit manipulation and overflow behavior).
I can dig back the references from some cpp con panel between Stroustrup, Carruth and others if needed, but I prefer to purge unsigned as much as possible.

[qcolombet]

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

414

Yeah, it has some unfortunate codegen implications.

Ideally we would want something that carries the semantic that this number is never negative, while not generating a bunch of codegen/perf overhead.

mlir/include/mlir/Analysis/SliceAnalysis.h

 /// Assuming all local orders match the numbering order:
 /// 1. after getting back to the root getForwardSlice, `forwardSlice` may
 ///    contain:
 ///      {9, 7, 8, 5, 1, 2, 6, 3, 4}
 /// 2. reversing the result of 1. gives:
 ///      {4, 3, 6, 2, 1, 5, 8, 7, 9}
 ///
 void getForwardSlice(Operation *op, SetVector<Operation *> *forwardSlice,
-                     TransitiveFilter filter = nullptr /* pass-through*/);
+                     TransitiveFilter filter = nullptr /* pass-through*/,
+                     bool inclusive = false);
 
 /// Value-rooted version of `getForwardSlice`. Return the union of all forward
 /// slices for the uses of the value `root`.
 void getForwardSlice(Value root, SetVector<Operation *> *forwardSlice,
-                     TransitiveFilter filter = nullptr /* pass-through*/);
+                     TransitiveFilter filter = nullptr /* pass-through*/,
+                     bool inclusive = false);
 
 /// Fills `backwardSlice` with the computed backward slice (i.e.
 /// all the transitive defs of op), **without** including that operation.
 ///
 /// This additionally takes a TransitiveFilter which acts as a frontier:
 /// when looking at defs transitively, an operation that does not pass the
 /// filter is never propagated through. This allows in particular to carve out
 /// the scope within a ForOp or the scope within an IfOp.

 ///      |_______________|
 ///              |
 ///              9
 ///
 /// Assuming all local orders match the numbering order:
 ///    {1, 2, 5, 3, 4, 6}
 ///
 void getBackwardSlice(Operation *op, SetVector<Operation *> *backwardSlice,
-                      TransitiveFilter filter = nullptr /* pass-through*/);
+                      TransitiveFilter filter = nullptr /* pass-through*/,
+                      bool inclusive = false);
 
 /// Value-rooted version of `getBackwardSlice`. Return the union of all backward
 /// slices for the op defining or owning the value `root`.
 void getBackwardSlice(Value root, SetVector<Operation *> *backwardSlice,
-                      TransitiveFilter filter = nullptr /* pass-through*/);
+                      TransitiveFilter filter = nullptr /* pass-through*/,
+                      bool inclusive = false);
 
 /// Iteratively computes backward slices and forward slices until
 /// a fixed point is reached. Returns an `SetVector<Operation *>` which
 /// **includes** the original operation.
 ///
 /// This allows building a slice (i.e. multi-root DAG where everything
 /// that is reachable from an Value in forward and backward direction is
 /// contained in the slice).

 /// Intuitively, we should be able to recurse like:
 ///   preorder(defs) - op - postorder(uses)
 /// and keep things ordered but this is still hand-wavy and not worth the
 /// trouble for now: punt to a simple worklist-based solution.
 ///
 SetVector<Operation *>
 getSlice(Operation *op,
          TransitiveFilter backwardFilter = nullptr /* pass-through*/,
-         TransitiveFilter forwardFilter = nullptr /* pass-through*/);
+         TransitiveFilter forwardFilter = nullptr /* pass-through*/,
+         bool inclusive = false);
 
 /// Multi-root DAG topological sort.
 /// Performs a topological sort of the Operation in the `toSort` SetVector.
 /// Returns a topologically sorted SetVector.
 SetVector<Operation *> topologicalSort(const SetVector<Operation *> &toSort);
 
 /// Utility to match a generic reduction given a list of iteration-carried
 /// arguments, `iterCarriedArgs` and the position of the potential reduction

mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.h

 
 namespace linalg {
 class GenericOp;
 class LinalgOp;
 } // namespace linalg
 
 namespace tensor {
 class PackOp;
+class PadOp;
 class UnPackOp;
 } // namespace tensor
 
 namespace transform {
 class TransformHandleTypeInterface;
 // Types needed for builders.
 struct TileSizesSpec {};
 struct NumThreadsSpec {};

mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td

let description = [{
     padded operations, which can be empty.
   }];
 
   let arguments =
     (ins PDL_Operation:$target,
          DefaultValuedAttr<ArrayAttr, "{}">:$padding_values,
          DefaultValuedAttr<I64ArrayAttr, "{}">:$padding_dimensions,
          DefaultValuedAttr<I64ArrayAttr, "{}">:$pack_paddings,
-         DefaultValuedAttr<I64ArrayAttr, "{}">:$hoist_paddings,
          DefaultValuedAttr<
           TypedArrayAttrBase<I64ArrayAttr, "array of arrays of i64">,
           "{}">:$transpose_paddings);
   let results = (outs PDL_Operation:$transformed);
 
   let assemblyFormat = "$target attr-dict";
   let hasVerifier = 1;
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::linalg::LinalgOp target,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 //===----------------------------------------------------------------------===//
+// HoistPadOp
+//===----------------------------------------------------------------------===//
+
+def HoistPadOp : Op<Transform_Dialect, "structured.hoist_pad",
+    [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
+     TransformOpInterface, TransformEachOpTrait]> {
+  let description = [{
+    Hoist the tensor.pad target operation by at most the given number of loops.
+    Optionally apply the transpose attribute to the inner dimensions.
+
+    TODO: In the future, we should consider rewriting as a tensor.pack after 
+    hoisting since this abstraction is now available.
+    TODO: Maybe also return the linalg.generic transpose created at some point.
+
+    #### Return modes
+
+    This operation ignores non-tensor.pad ops and drops them in the result.
+    If any non-tensor.pad is passed, the transform emits a silenceable failure.
+
+    If all the operations referred to by the `target` handle padproperly, the
+    transform succeeds. Otherwise the transform silently fails.
+
+    The return handle points to only the subset of successfully hoisted 
+    tensor.pad operations, which can be empty.
+  }];
+
+  // Also allow any !pdl.operation for simpler composition. Non-tensor.pad ops
+  // will be dropped from the results.
+  let arguments =
+    (ins TransformHandleTypeInterface:$target,
+         I64Attr:$num_loops,
+         DefaultValuedAttr<DenseI64ArrayAttr, "{}">:$transpose);
+  let results = (outs TransformHandleTypeInterface:$transformed);
+
+  let assemblyFormat = [{
+    $target 
+    `by` $num_loops `loops` 
+    (`,` `transpose` `by` $transpose^)? 
+    attr-dict
+    `:` functional-type(operands, results)
+  }];
+  let hasVerifier = 1;
+
+  let extraClassDeclaration = [{
+    ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::tensor::PadOp,
+        ::mlir::transform::ApplyToEachResultList &results,
+        ::mlir::transform::TransformState &state);
+  }];
+}
+
+//===----------------------------------------------------------------------===//
 // PromoteOp
 //===----------------------------------------------------------------------===//
 
 
 def PromoteOp : Op<Transform_Dialect, "structured.promote",
     [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
     TransformOpInterface, TransformEachOpTrait]> {
   let description = [{

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

-//===- HoistPadding.h - Hoisting for tensor::PadOp -*- C++ --------------*-===//
-//
-// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
-// See https://llvm.org/LICENSE.txt for license information.
-// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-//
-//===----------------------------------------------------------------------===//
-
-#ifndef MLIR_DIALECT_LINALG_TRANSFORMS_HOISTPADDING_H
-#define MLIR_DIALECT_LINALG_TRANSFORMS_HOISTPADDING_H
-
-#include "mlir/Support/LogicalResult.h"
-
-namespace mlir {
-class Value;
-
-namespace tensor {
-class PadOp;
-} // namespace tensor
-
-namespace linalg {
-class GenericOp;
-
-/// Mechanically hoist padding operations on tensors by `numLoops` into a new,
-/// generally larger tensor. This achieves packing of multiple padding ops into
-/// a larger tensor. On success, `opToHoist` is replaced by the cloned version
-/// in the packing loop so the caller can continue reasoning about the padding
-/// operation. If `transposeVector` is non-empty, hoist padding introduces a
-/// GenericOp to transpose the padded tensor before inserting it into the packed
-/// tensor. A `transposeVector` can change the storage order of the padded
-/// tensor but does not change the order of the pack or compute loops.
-///
-///
-/// Example in pseudo-mlir:
-/// =======================
-///
-/// If hoistPaddingOnTensors is called with `nLoops` = 2 on the following IR.
-/// ```
-///    scf.for (%i, %j, %k)
-///      %st0 = tensor.extract_slice f(%i, %k) : ... to tensor<?x?xf32>
-///      %0 = tensor.pad %st0 low[0, 0] high[...] {
-///      ^bb0( ... ):
-///        linalg.yield %pad
-///      } : tensor<?x?xf32> to tensor<4x8xf32>
-///      compute(%0)
-/// ```
-///
-/// IR resembling the following is produced:
-///
-/// ```
-///    scf.for (%i) {
-///      %packed_init = tensor.empty range(%j) : tensor<?x4x8xf32>
-///      %packed = scf.for (%k) iter_args(%p : %packed_init) {
-///        %st0 = tensor.extract_slice f(%i, %k) : ... to tensor<?x?xf32>
-///        %0 = tensor.pad %st0 low[0, 0] high[...] {
-///        ^bb0( ... ):
-///          linalg.yield %pad
-///        } : tensor<?x?xf32> to tensor<4x8xf32>
-///        %1 = tensor.insert_slice %0 ...
-///            : tensor<4x8xf32> to tensor<?x4x8xf32>
-///        scf.yield %1: tensor<?x4x8xf32>
-///      } -> tensor<?x4x8xf32>
-///      scf.for (%j, %k) {
-///        %st0 = tensor.extract_slice %packed [%k, 0, 0][1, 4, 8][1, 1, 1] :
-///                 tensor<?x4x8xf32> to tensor<4x8xf32>
-///        compute(%st0)
-///      }
-///    }
-/// ```
-FailureOr<Value> hoistPaddingOnTensors(
-    tensor::PadOp opToHoist, int numLoops, ArrayRef<int64_t> transposeVector,
-    tensor::PadOp &hoistedOp, SmallVectorImpl<GenericOp> &transposeOps);
-
-} // namespace linalg
-} // namespace mlir
-
-#endif // MLIR_DIALECT_LINALG_TRANSFORMS_HOISTPADDING_H

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

 
 /// 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,
+rewriteAsPaddedOp(RewriterBase &rewriter, 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,
+/// Mechanically hoist padding operations on tensors by `numLoops` into a new,
+/// generally larger tensor. This achieves packing of multiple padding ops into
+/// a larger tensor. On success, `opToHoist` is replaced by the cloned version
+/// in the packing loop so the caller can continue reasoning about the padding
+/// operation. If `transposeVector` is non-empty, hoist padding introduces a
+/// GenericOp to transpose the padded tensor before inserting it into the packed
+/// tensor. A `transposeVector` can change the storage order of the padded
+/// tensor but does not change the order of the pack or compute loops.
+///
+/// TODO: In the future, we should consider rewriting as a tensor.pack after
+/// hoisting since this abstraction is now available.
+///
+/// Example in pseudo-mlir:
+/// =======================
+///
+/// If hoistPaddingOnTensors is called with `nLoops` = 2 on the following IR.
+/// ```
+///    scf.for (%i, %j, %k)
+///      %st0 = tensor.extract_slice f(%i, %k) : ... to tensor<?x?xf32>
+///      %0 = tensor.pad %st0 low[0, 0] high[...] {
+///      ^bb0( ... ):
+///        linalg.yield %pad
+///      } : tensor<?x?xf32> to tensor<4x8xf32>
+///      compute(%0)
+/// ```
+///
+/// IR resembling the following is produced:
+///
+/// ```
+///    scf.for (%i) {
+///      %packed_init = tensor.empty range(%j) : tensor<?x4x8xf32>
+///      %packed = scf.for (%k) iter_args(%p : %packed_init) {
+///        %st0 = tensor.extract_slice f(%i, %k) : ... to tensor<?x?xf32>
+///        %0 = tensor.pad %st0 low[0, 0] high[...] {
+///        ^bb0( ... ):
+///          linalg.yield %pad
+///        } : tensor<?x?xf32> to tensor<4x8xf32>
+///        %1 = tensor.insert_slice %0 ...
+///            : tensor<4x8xf32> to tensor<?x4x8xf32>
+///        scf.yield %1: tensor<?x4x8xf32>
+///      } -> tensor<?x4x8xf32>
+///      scf.for (%j, %k) {
+///        %st0 = tensor.extract_slice %packed [%k, 0, 0][1, 4, 8][1, 1, 1] :
+///                 tensor<?x4x8xf32> to tensor<4x8xf32>
+///        compute(%st0)
+///      }
+///    }
+/// ```
+FailureOr<Value>
+hoistPaddingOnTensors(tensor::PadOp opToHoist, int64_t numLoops,
+                      ArrayRef<int64_t> transposeVector,
+                      tensor::PadOp &hoistedOp,
+                      SmallVectorImpl<GenericOp> &transposeOps);
+
+/// Apply padding and hoisting to `linalgOp` according to the configuration
+/// specified in `options`.
+FailureOr<LinalgOp> padAndHoistLinalgOp(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:

mlir/lib/Analysis/SliceAnalysis.cpp

for (Operation *userOp : result.getUsers())
       if (forwardSlice->count(userOp) == 0)
         getForwardSliceImpl(userOp, forwardSlice, filter);
   }
 
   forwardSlice->insert(op);
 }
 
 void mlir::getForwardSlice(Operation *op, SetVector<Operation *> *forwardSlice,
-                           TransitiveFilter filter) {
+                           TransitiveFilter filter, bool inclusive) {
   getForwardSliceImpl(op, forwardSlice, filter);
+  if (!inclusive) {
     // Don't insert the top level operation, we just queried on it and don't
     // want it in the results.
     forwardSlice->remove(op);
+  }
 
   // Reverse to get back the actual topological order.
   // std::reverse does not work out of the box on SetVector and I want an
   // in-place swap based thing (the real std::reverse, not the LLVM adapter).
   std::vector<Operation *> v(forwardSlice->takeVector());
   forwardSlice->insert(v.rbegin(), v.rend());
 }
 
 void mlir::getForwardSlice(Value root, SetVector<Operation *> *forwardSlice,
-                           TransitiveFilter filter) {
+                           TransitiveFilter filter, bool inclusive) {
   for (Operation *user : root.getUsers())
     getForwardSliceImpl(user, forwardSlice, filter);
 
   // Reverse to get back the actual topological order.
   // std::reverse does not work out of the box on SetVector and I want an
   // in-place swap based thing (the real std::reverse, not the LLVM adapter).
   std::vector<Operation *> v(forwardSlice->takeVector());
   forwardSlice->insert(v.rbegin(), v.rend());
for (const auto &en : llvm::enumerate(op->getOperands())) {
     }
   }
 
   backwardSlice->insert(op);
 }
 
 void mlir::getBackwardSlice(Operation *op,
                             SetVector<Operation *> *backwardSlice,
-                            TransitiveFilter filter) {
+                            TransitiveFilter filter, bool inclusive) {
   getBackwardSliceImpl(op, backwardSlice, filter);
 
+  if (!inclusive) {
     // Don't insert the top level operation, we just queried on it and don't
     // want it in the results.
     backwardSlice->remove(op);
   }
+}
 
 void mlir::getBackwardSlice(Value root, SetVector<Operation *> *backwardSlice,
-                            TransitiveFilter filter) {
+                            TransitiveFilter filter, bool inclusive) {
   if (Operation *definingOp = root.getDefiningOp()) {
-    getBackwardSlice(definingOp, backwardSlice, filter);
+    getBackwardSlice(definingOp, backwardSlice, filter, inclusive);
     return;
   }
   Operation *bbAargOwner = root.cast<BlockArgument>().getOwner()->getParentOp();
-  getBackwardSlice(bbAargOwner, backwardSlice, filter);
+  getBackwardSlice(bbAargOwner, backwardSlice, filter, inclusive);
 }
 
 SetVector<Operation *> mlir::getSlice(Operation *op,
                                       TransitiveFilter backwardFilter,
-                                      TransitiveFilter forwardFilter) {
+                                      TransitiveFilter forwardFilter,
+                                      bool inclusive) {
   SetVector<Operation *> slice;
   slice.insert(op);
 
   unsigned currentIndex = 0;
   SetVector<Operation *> backwardSlice;
   SetVector<Operation *> forwardSlice;
   while (currentIndex != slice.size()) {
     auto *currentOp = (slice)[currentIndex];
     // Compute and insert the backwardSlice starting from currentOp.
     backwardSlice.clear();
-    getBackwardSlice(currentOp, &backwardSlice, backwardFilter);
+    getBackwardSlice(currentOp, &backwardSlice, backwardFilter, inclusive);
     slice.insert(backwardSlice.begin(), backwardSlice.end());
 
     // Compute and insert the forwardSlice starting from currentOp.
     forwardSlice.clear();
-    getForwardSlice(currentOp, &forwardSlice, forwardFilter);
+    getForwardSlice(currentOp, &forwardSlice, forwardFilter, inclusive);
     slice.insert(forwardSlice.begin(), forwardSlice.end());
     ++currentIndex;
   }
   return topologicalSort(slice);
 }
 
 namespace {
 /// DFS post-order implementation that maintains a global count to work across

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

transform::PadOp::applyToOne(LinalgOp target,
   }
 
   // Extract the transpose vectors.
   SmallVector<SmallVector<int64_t>> transposePaddings;
   for (Attribute transposeVector : getTransposePaddings().cast<ArrayAttr>())
     transposePaddings.push_back(
         extractFromI64ArrayAttr(transposeVector.cast<ArrayAttr>()));
 
-  LinalgPaddingOptions paddingOptions;
-  paddingOptions.setPaddingValues(paddingValues);
-  paddingOptions.setPaddingDimensions(
-      extractFromI64ArrayAttr(getPaddingDimensions()));
-  paddingOptions.setPackPaddings(packPaddings);
-  paddingOptions.setHoistPaddings(extractFromI64ArrayAttr(getHoistPaddings()));
-  paddingOptions.setTransposePaddings(transposePaddings);
-
   IRRewriter rewriter(target->getContext());
-  FailureOr<LinalgOp> result = padLinalgOp(rewriter, target, paddingOptions);
+  LinalgOp paddedOp;
+  FailureOr<SmallVector<Value>> result = rewriteAsPaddedOp(
+      rewriter, target, extractFromI64ArrayAttr(getPaddingDimensions()),
+      paddingValues, packPaddings, paddedOp);
   if (succeeded(result)) {
-    results.push_back(result->getOperation());
+    // We need to perform our own replacement here because this API is still
+    // used in patterns that "pad and hoist", for which the replacement values
+    // need to be different.
+    // TODO: clean this up and stop "pad and hoist" behavior more globally now
+    // that we have more composable abstractions.
+    rewriter.replaceOp(target, *result);
+    results.push_back(paddedOp);
     return DiagnosedSilenceableFailure::success();
   }
 
   return emitDefaultSilenceableFailure(target);
 }
 
 LogicalResult transform::PadOp::verify() {
   SmallVector<int64_t> packPaddings =
SmallVector<int64_t> paddingDimensions =
       extractFromI64ArrayAttr(getPaddingDimensions());
   if (any_of(paddingDimensions,
              [](int64_t paddingDimension) { return paddingDimension < 0; })) {
     return emitOpError() << "expects padding_dimensions to contain positive "
                             "integers, found "
                          << getPaddingDimensions();
   }
 
-  SmallVector<int64_t> hoistPaddings =
-      extractFromI64ArrayAttr(getHoistPaddings());
-  if (any_of(hoistPaddings,
-             [](int64_t hoistPadding) { return hoistPadding < 0; })) {
-    return emitOpError()
-           << "expects hoist_paddings to contain positive integers, found "
-           << getHoistPaddings();
-  }
-
   ArrayAttr transposes = getTransposePaddings();
   for (Attribute attr : transposes) {
     SmallVector<int64_t> transpose = extractFromI64ArrayAttr(attr);
     auto sequence = llvm::to_vector(llvm::seq<int64_t>(0, transpose.size()));
     if (!std::is_permutation(sequence.begin(), sequence.end(),
                              transpose.begin(), transpose.end())) {
       return emitOpError()
              << "expects transpose_paddings to be a permutation, found "
              << attr;
     }
   }
   return success();
 }
 
+//===---------------------------------------------------------------------===//
+// HoistPadOp
+//===---------------------------------------------------------------------===//
+
+DiagnosedSilenceableFailure
+transform::HoistPadOp::applyToOne(tensor::PadOp target,
+                                  transform::ApplyToEachResultList &results,
+                                  transform::TransformState &state) {
+  IRRewriter rewriter(target->getContext());
+  tensor::PadOp hoistedPadOp;
+  SmallVector<GenericOp> transposeOps;
+  // TODO: Pass rewriter down to hoistPaddingOnTensors, in a followup commit.
+  FailureOr<Value> result = hoistPaddingOnTensors(
+      target, getNumLoops(), getTranspose(), hoistedPadOp, transposeOps);
+  if (succeeded(result)) {
+    // We need to perform our own replacement here because this API is still
+    // used in patterns that "pad and hoist", for which the replacement values
+    // need to be different.
+    // TODO: clean this up and stop "pad and hoist" behavior more globally now
+    // that we have more composable abstractions.
+    rewriter.replaceOp(target, *result);
+    results.push_back(hoistedPadOp);
+    return DiagnosedSilenceableFailure::success();
+  }
+  return emitDefaultSilenceableFailure(target);
+}
+
+LogicalResult transform::HoistPadOp::verify() {
+  ArrayRef<int64_t> transpose = getTranspose();
+  auto sequence = llvm::to_vector(llvm::seq<int64_t>(0, transpose.size()));
+  if (!std::is_permutation(sequence.begin(), sequence.end(), transpose.begin(),
+                           transpose.end())) {
+    return emitOpError() << "expects transpose to be a permutation, found "
+                         << getTranspose();
+  }
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // PromoteOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::PromoteOp::applyToOne(LinalgOp target,
                                  transform::ApplyToEachResultList &results,
                                  transform::TransformState &state) {

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

 //===- HoistPadding.cpp - Hoisting for tensor::PadOp ----------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements functions concerned with hoisting padding operations.
 //
 //===----------------------------------------------------------------------===//
 
-#include "mlir/Dialect/Linalg/Transforms/HoistPadding.h"
 #include "mlir/Analysis/SliceAnalysis.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/SCF/Utils/Utils.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Utils/IndexingUtils.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/Dialect/Vector/Utils/VectorUtils.h"
 #include "mlir/IR/AsmState.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Dominance.h"
+
 #include "mlir/IR/Matchers.h"
+#include "mlir/Transforms/RegionUtils.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Debug.h"
 
 using llvm::dbgs;
 
 #define DEBUG_TYPE "hoist-padding"
 
 #define DBGS() (dbgs() << '[' << DEBUG_TYPE << "] ")
computeTransposedType(RankedTensorType rankedTensorType,
   applyPermutationToVector(transposedShape, transposeVector);
 
   using RTTBuilder = RankedTensorType::Builder;
   RankedTensorType transposedTensorType =
       RTTBuilder(rankedTensorType).setShape(transposedShape);
   return transposedTensorType;
 }
 
+// Get all the ops in the backwards slice starting from `padOp` and that
+// are dominated by the outermost enclosing loop.
+// This also requires tracking ops defining values used in the region but
+// defined above.
+static void computeBackwardSlice(tensor::PadOp padOp,
+                                 scf::ForOp outermostEnclosingForOp,
+                                 SetVector<Operation *> &backwardSlice) {
+  DominanceInfo domInfo(outermostEnclosingForOp);
+  auto filter = [&](Operation *op) {
+    return domInfo.dominates(outermostEnclosingForOp, op) &&
+           !padOp->isProperAncestor(op);
+  };
+  // First, add the ops required to compute the region to the backwardSlice.
+  SetVector<Value> valuesDefinedAbove;
+  getUsedValuesDefinedAbove(padOp.getRegion(), padOp.getRegion(),
+                            valuesDefinedAbove);
+  for (Value v : valuesDefinedAbove) {
+    getBackwardSlice(v, &backwardSlice, filter, /*inclusive=*/true);
+  }
+  // Then, add the backward slice from padOp itself.
+  getBackwardSlice(padOp.getOperation(), &backwardSlice, filter,
+                   /*inclusive=*/true);
+}
+
 HoistingAnalysis::HoistingAnalysis(tensor::PadOp padOp, int numLoops) {
   valid = false;
 
   // Bail on any use that isn't an input of a LinalgOp.
   // Hoisting of inplace updates happens after vectorization.
   if (!isOnlyUsedAsInputOfLinalgOp(padOp))
     return;
 
HoistingAnalysis::HoistingAnalysis(tensor::PadOp padOp, int numLoops) {
   // dependencies captured by the padding region.
   Value paddingValue = padOp.getConstantPaddingValue();
   if (!paddingValue ||
       !isa_and_nonnull<arith::ConstantOp>(paddingValue.getDefiningOp())) {
     LLVM_DEBUG(DBGS() << "Cannot find constant padding value -> skip\n");
     return;
   }
 
-  // Get all the ops in the backwards slice starting from `padOp` and that
-  // are dominated by the outermost enclosing loop.
-  DominanceInfo domInfo(outermostEnclosingForOp);
-  getBackwardSlice(padOp.getOperation(), &backwardSlice, [&](Operation *op) {
-    return domInfo.dominates(outermostEnclosingForOp, op);
-  });
-  if (backwardSlice.empty())
+  computeBackwardSlice(padOp, outermostEnclosingForOp, backwardSlice);
+  if (backwardSlice.size() <= 1)
     return;
-  // Add `padOp` itself to the backward slice.
-  backwardSlice.insert(padOp.getOperation());
 
   // Remove all ops in the backward slice that are not used to index the padded
   // tensor. In particular, keep `padOp`, `sliceOp`, and the loop and
   // affine operations used for the index computation.
   if (failed(dropNonIndexDependencies(padOp, sliceOp)))
     return;
 
   // Add only the loops part of the filtered `backwardSlice` to the packing
if (!isDefinedOutsideOrConstant(outer, forOp.getLowerBound()) ||
     return Value();
   Value ivVal = forOp.getInductionVar(), lbVal = forOp.getLowerBound(),
         stepVal = forOp.getStep();
   auto loc = forOp->getLoc();
   return b.createOrFold<AffineApplyOp>(loc, (iv - lb).ceilDiv(step),
                                        ValueRange{ivVal, lbVal, stepVal});
 }
 
-FailureOr<Value> mlir::linalg::hoistPaddingOnTensors(
-    tensor::PadOp opToHoist, int numLoops, ArrayRef<int64_t> transposeVector,
-    tensor::PadOp &hoistedOp, SmallVectorImpl<GenericOp> &transposeOps) {
+FailureOr<Value>
+mlir::linalg::hoistPaddingOnTensors(tensor::PadOp opToHoist, int64_t numLoops,
+                                    ArrayRef<int64_t> transposeVector,
+                                    tensor::PadOp &hoistedOp,
+                                    SmallVectorImpl<GenericOp> &transposeOps) {
   LLVM_DEBUG(DBGS() << "Try to hoist " << *(opToHoist) << " by " << numLoops
                     << " loops\n");
   HoistingAnalysis analysis(opToHoist, numLoops);
   if (!analysis.isValid()) {
     LLVM_DEBUG(DBGS() << "Analysis failed -> Skip\n");
     return failure();
   }
 

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

 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
-#include "mlir/Dialect/Linalg/Transforms/HoistPadding.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/SCF/Transforms/Transforms.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tensor/IR/TensorTilingInterfaceImpl.h"
 #include "mlir/Dialect/Tensor/Utils/Utils.h"
 #include "mlir/Dialect/Utils/IndexingUtils.h"
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"

 /// 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(
-    OpBuilder &b, linalg::LinalgOp opToPad, OpOperand *opOperand,
+    RewriterBase &rewriter, linalg::LinalgOp opToPad, OpOperand *opOperand,
     ArrayRef<int64_t> paddingDimensions, ArrayRef<Attribute> paddingValues,
     ArrayRef<bool> packPaddings) {
   AffineMap indexingMap = opToPad.getMatchingIndexingMap(opOperand);
   ArrayRef<int64_t> shape = opToPad.getShape(opOperand);
 
   // Collect the shape dimension that are a function of the `paddingDimensions`.
   llvm::SmallDenseSet<int64_t> shapeDimsToPad;
   for (int64_t dim : paddingDimensions)
     for (const auto &en : enumerate(indexingMap.getResults()))
       if (en.value().isFunctionOfDim(dim))
         shapeDimsToPad.insert(en.index());
 
   // Return the unpadded operand if padding to a static shape is not needed and
   // if the nofold flag is not set.
   bool nofold = opOperand->getOperandNumber() < packPaddings.size()
                     ? packPaddings[opOperand->getOperandNumber()]
                     : false;
   bool hasStaticShape = llvm::none_of(shapeDimsToPad, [&](int64_t dim) {
     return ShapedType::isDynamic(shape[dim]);
   });
   if (!nofold && hasStaticShape)
     return opOperand->get();
 
   // Fail if `paddingValues` specifies no padding value.
-  if (opOperand->getOperandNumber() >= paddingValues.size())
-    return failure();
+  if (opOperand->getOperandNumber() >= paddingValues.size()) {
+    return rewriter.notifyMatchFailure(opToPad, "no padding value specified");
+  }
   Attribute paddingAttr = paddingValues[opOperand->getOperandNumber()];
-  Type paddingType = b.getType<NoneType>();
+  Type paddingType = rewriter.getType<NoneType>();
   if (auto typedAttr = paddingAttr.dyn_cast<TypedAttr>())
     paddingType = typedAttr.getType();
-  Value paddingValue =
-      b.create<arith::ConstantOp>(opToPad.getLoc(), paddingType, paddingAttr);
+  Value paddingValue = rewriter.create<arith::ConstantOp>(
+      opToPad.getLoc(), paddingType, paddingAttr);
 
   // Follow the use-def chain if `currOpOperand` is defined by a LinalgOp.
   OpOperand *currOpOperand = opOperand;
   while (auto linalgOp = currOpOperand->get().getDefiningOp<LinalgOp>()) {
     OpResult result = currOpOperand->get().cast<OpResult>();
     currOpOperand = linalgOp.getDpsInitOperand(result.getResultNumber());
   }
 
   // Fail if `currOpOperand` is not defined by an ExtractSliceOp or EmptyOp.
   auto sliceOp = currOpOperand->get().getDefiningOp<tensor::ExtractSliceOp>();
   auto emptyOp = currOpOperand->get().getDefiningOp<tensor::EmptyOp>();
-  if (!sliceOp && !emptyOp)
-    return failure();
+  if (!sliceOp && !emptyOp) {
+    // TODO: may want to add support for going through loop iter args.
+    // This is not strictly necessary as we can pad before hoisting but it would
+    // make the system overall more resilient to minor transformation
+    // reorderings.
+    return rewriter.notifyMatchFailure(
+        opToPad, "not defined by an extractSlice or emptyOp");
+  }
 
   llvm::SmallBitVector droppedDims;
   SmallVector<OpFoldResult> mixedSizes;
   if (sliceOp) {
     // Compute the dropped dimensions if `sliceOp` is ranke-reducing.
     droppedDims = sliceOp.getDroppedDims();
     mixedSizes = sliceOp.getMixedSizes();
   }
if (en.value().is<Attribute>()) {
       paddedShape[shapeIdx++] =
           en.value().get<Attribute>().dyn_cast<IntegerAttr>().getInt();
       continue;
     }
     // Otherwise, try to compute a constant upper bound for the size value.
     FailureOr<int64_t> upperBound =
         getConstantUpperBoundForIndex(en.value().get<Value>());
     if (failed(upperBound)) {
-      LLVM_DEBUG(DBGS() << "No constant bounding box can be found for padding");
-      return failure();
+      LLVM_DEBUG(DBGS() << "count not compute a bonding box for padding");
+      return rewriter.notifyMatchFailure(
+          opToPad, "count not compute a bonding box for padding");
     }
     paddedShape[shapeIdx++] = *upperBound;
   }
   assert(shapeIdx == static_cast<int64_t>(shape.size()) &&
          "expect the dynamic and static ranks to match");
 
   // Pad the operand to the bounding box defined by `paddedShape`.
   auto paddedTensorType = RankedTensorType::get(
       paddedShape, getElementTypeOrSelf(opOperand->get()));
-  return makeComposedPadHighOp(b, opToPad->getLoc(), paddedTensorType,
+  return makeComposedPadHighOp(rewriter, 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,
+linalg::rewriteAsPaddedOp(RewriterBase &rewriter, 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.
   assert(opToPad.hasTensorSemantics() &&
          "expected operation to have tensor semantics");
 
-  OpBuilder::InsertionGuard g(b);
+  OpBuilder::InsertionGuard g(rewriter);
   // Set IP after op because we also take the dims of the original output.
-  b.setInsertionPointAfter(opToPad);
+  rewriter.setInsertionPointAfter(opToPad);
+
   // Make a copy of the shaped operands and update it.
   SmallVector<Value> newOperands;
   newOperands.reserve(opToPad->getNumOperands());
   for (OpOperand &opOperand : opToPad->getOpOperands()) {
     FailureOr<Value> paddedOperand = padOperandToSmallestStaticBoundingBox(
-        b, opToPad, &opOperand, paddingDimensions, paddingValues, packPaddings);
+        rewriter, opToPad, &opOperand, paddingDimensions, paddingValues,
+        packPaddings);
     // Exit if `paddingDimensions` cannot be bounded statically.
     if (failed(paddedOperand))
       return failure();
     newOperands.push_back(*paddedOperand);
   }
 
   SmallVector<SmallVector<Value>> reifiedResultShapes;
   if (failed(cast<ReifyRankedShapedTypeOpInterface>(opToPad.getOperation())
-                 .reifyResultShapes(b, reifiedResultShapes)))
+                 .reifyResultShapes(rewriter, reifiedResultShapes)))
     return failure();
   assert(reifiedResultShapes.size() == opToPad->getNumResults() &&
          "expected same number of results");
 
   // Clone `opToPad` to operate on the statically padded shapes.
   auto resultTensorTypes =
       ValueRange(newOperands).take_back(opToPad.getNumDpsInits()).getTypes();
-  paddedOp = clone(b, opToPad, resultTensorTypes, newOperands);
+  // clone **should** properly notify the rewriter.
+  paddedOp = clone(rewriter, opToPad, resultTensorTypes, newOperands);
 
   // Recover the slice out of the new static results. This keeps the original
   // linalg op around because it uses the dims of the original results.
-  SmallVector<Value> paddedSubviewResults;
-  paddedSubviewResults.reserve(opToPad->getNumResults());
+  SmallVector<Value> paddedSubtensorResults;
+  paddedSubtensorResults.reserve(opToPad->getNumResults());
   for (const auto &en : llvm::enumerate(paddedOp->getResults())) {
     Value paddedResult = en.value();
     int64_t resultNumber = en.index();
     int64_t rank = paddedResult.getType().cast<RankedTensorType>().getRank();
-    SmallVector<OpFoldResult> offsets(rank, b.getIndexAttr(0));
+    SmallVector<OpFoldResult> offsets(rank, rewriter.getIndexAttr(0));
     SmallVector<OpFoldResult> 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>(
+    SmallVector<OpFoldResult> strides(rank, rewriter.getIndexAttr(1));
+    paddedSubtensorResults.push_back(rewriter.create<tensor::ExtractSliceOp>(
         loc, paddedResult, offsets, sizes, strides));
   }
-  return paddedSubviewResults;
+  return paddedSubtensorResults;
 }
 
 //===----------------------------------------------------------------------===//
 // 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.
void mlir::linalg::peelLoops(RewriterBase &rewriter,
   for (auto loopOp : loops)
     peelLoop(rewriter, loopOp);
 }
 
 //===----------------------------------------------------------------------===//
 // pad transformation.
 //===----------------------------------------------------------------------===//
 
-FailureOr<LinalgOp> mlir::linalg::padLinalgOp(RewriterBase &rewriter,
-                                              LinalgOp linalgOp,
-                                              LinalgPaddingOptions options) {
+FailureOr<LinalgOp>
+mlir::linalg::padAndHoistLinalgOp(RewriterBase &rewriter, LinalgOp linalgOp,
   if (!linalgOp.hasTensorSemantics())
     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,
tileSizeComputationFunction = [tileSizes](OpBuilder &b, Operation *op) {
     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.
-
+///
+/// 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);
+  return padAndHoistLinalgOp(rewriter, op, options);
 }
 
 LogicalResult mlir::linalg::CopyVectorizationPattern::matchAndRewrite(
     memref::CopyOp copyOp, PatternRewriter &rewriter) const {
   return vectorizeCopy(rewriter, copyOp);
 }
 
-/// Rewrite a tensor::PadOp into a sequence of EmptyOp, FillOp (to
+///
+/// Pattern to 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>();
 
   // Bail on non-static shapes.

mlir/python/mlir/dialects/_structured_transform_ops_ext.py

class PadOp:
 
   def __init__(self,
                target: Union[Operation, Value],
                *,
                padding_values: Optional[Union[ArrayAttr,
                                               Sequence[Attribute]]] = None,
                padding_dimensions: OptionalIntList = None,
                pack_paddings: OptionalIntList = None,
-               hoist_paddings: OptionalIntList = None,
                transpose_paddings: Optional[Union[ArrayAttr, Sequence[Union[
                    ArrayAttr, IntOrAttrList]]]] = None,
                loc=None,
                ip=None):
     pdl_operation_type = pdl.OperationType.get()
     padding_values_attr = _get_array_attr(padding_values)
     padding_dimensions_attr = _get_int_array_attr(padding_dimensions)
     pack_paddings_attr = _get_int_array_attr(pack_paddings)
-    hoist_paddings_attr = _get_int_array_attr(hoist_paddings)
     transpose_paddings_attr = _get_int_int_array_attr(transpose_paddings)
     super().__init__(
         pdl_operation_type,
         _get_op_result_or_value(target),
         padding_values=padding_values_attr,
         padding_dimensions=padding_dimensions_attr,
         pack_paddings=pack_paddings_attr,
-        hoist_paddings=hoist_paddings_attr,
         transpose_paddings=transpose_paddings_attr,
         loc=loc,
         ip=ip)
 
 
 class ScalarizeOp:
   """Specialization for ScalarizeOp class."""
 

mlir/test/Dialect/Linalg/transform-op-hoist-pad.mlir

+// RUN: mlir-opt --test-transform-dialect-interpreter -canonicalize -split-input-file --verify-diagnostics %s | FileCheck %s
+
+func.func @pad_and_hoist_rhs(
+  %arg0: tensor<24x12xf32>, %arg1: tensor<12x25xf32>, %arg2: tensor<24x25xf32>)
+     -> tensor<24x25xf32> 
+{
+  // expected-note @below {{payload operation}}
+  %0 = linalg.matmul ins(%arg0, %arg1 : tensor<24x12xf32>, tensor<12x25xf32>) outs(%arg2 : tensor<24x25xf32>) -> tensor<24x25xf32>
+  func.return %0 : tensor<24x25xf32>
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg1 
+    : (!pdl.operation) -> !pdl.operation
+
+  
+  %matmul_l1, %loops_l1 = transform.structured.tile_to_scf_for %matmul [5]
+
+  %matmul_padded = transform.structured.pad %matmul_l1 {
+    padding_values=[0.0: f32, 0.0 : f32, 0.0 : f32],
+    padding_dimensions=[0, 1, 2]
+  }
+
+  // In this case, the pad op is actually empty: we only tile the first dimension
+  // and it does not have an impact on the RHS operand.
+  // expected-error @below {{incompatible payload operation name}}
+  %pad = transform.get_producer_of_operand %matmul_padded[1]
+    : (!pdl.operation) -> !transform.op<"tensor.pad">
+
+  // We do not even reach this transform op.
+  transform.structured.hoist_pad %pad by 1 loops
+     : (!transform.op<"tensor.pad">) -> !pdl.operation
+}
+
+// -----
+
+func.func @pad_and_hoist_init(
+  %arg0: tensor<24x12xf32>, %arg1: tensor<12x25xf32>, %arg2: tensor<24x25xf32>)
+     -> tensor<24x25xf32> 
+{
+  // expected-note @below {{when applied to this op}}
+  %0 = linalg.matmul ins(%arg0, %arg1 : tensor<24x12xf32>, tensor<12x25xf32>) outs(%arg2 : tensor<24x25xf32>) -> tensor<24x25xf32>
+  func.return %0 : tensor<24x25xf32>
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg1 
+    : (!pdl.operation) -> !pdl.operation
+
+  
+  %matmul_l1, %loops_l1 = transform.structured.tile_to_scf_for %matmul [5]
+
+  %matmul_padded = transform.structured.pad %matmul_l1 {
+    padding_values=[0.0: f32, 0.0 : f32, 0.0 : f32],
+    padding_dimensions=[0, 1, 2]
+  }
+
+  %pad = transform.get_producer_of_operand %matmul_padded[2]
+    : (!pdl.operation) -> !transform.op<"tensor.pad">
+
+  // We do not know yet how to hoist the init.
+  // expected-error @below {{transform.structured.hoist_pad failed to apply}}
+  transform.structured.hoist_pad %pad by 1 loops
+     : (!transform.op<"tensor.pad">) -> !pdl.operation
+}
+
+// -----
+
+//     CHECK-LABEL: pad_and_hoist_lhs
+func.func @pad_and_hoist_lhs(
+  %arg0: tensor<24x12xf32>, %arg1: tensor<12x25xf32>, %arg2: tensor<24x25xf32>)
+     -> tensor<24x25xf32> 
+{
+  //     CHECK: %[[PACKED:.*]] = scf.for %{{.*}} -> (tensor<5x5x12xf32>) {
+  //     CHECK:   tensor.pad %{{.*}} 
+  //     CHECK:     : tensor<?x12xf32> to tensor<5x12xf32>
+  //     CHECK:   tensor.insert_slice %{{.*}} into %{{.*}}[%{{.*}}, 0, 0] [1, 5, 12] [1, 1, 1] 
+  // CHECK-SAME:   : tensor<5x12xf32> into tensor<5x5x12xf32>
+  //     CHECK: scf.for %{{.*}} -> (tensor<24x25xf32>) {
+  //     CHECK:   %[[PADDED:.*]] = tensor.extract_slice %[[PACKED]][%{{.*}}, 0, 0] [1, 5, 12] [1, 1, 1] 
+  // CHECK-SAME:    : tensor<5x5x12xf32> to tensor<5x12xf32>
+  //     CHECK:   linalg.matmul ins(%[[PADDED]]
+  %0 = linalg.matmul ins(%arg0, %arg1 : tensor<24x12xf32>, tensor<12x25xf32>) outs(%arg2 : tensor<24x25xf32>) -> tensor<24x25xf32>
+  func.return %0 : tensor<24x25xf32>
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg1 
+    : (!pdl.operation) -> !pdl.operation
+
+  
+  %matmul_l1, %loops_l1 = transform.structured.tile_to_scf_for %matmul [5]
+
+  %matmul_padded = transform.structured.pad %matmul_l1 {
+    padding_values=[0.0: f32, 0.0 : f32, 0.0 : f32],
+    padding_dimensions=[0, 1, 2]
+  }
+
+  %pad = transform.get_producer_of_operand %matmul_padded[0]
+    : (!pdl.operation) -> !pdl.operation
+
+  transform.structured.hoist_pad %pad by 1 loops
+     : (!pdl.operation) -> !pdl.operation
+}
+
+// -----
+
+//     CHECK-LABEL: pad_and_hoist_lhs_transpose
+func.func @pad_and_hoist_lhs_transpose(
+  %arg0: tensor<24x12xf32>, %arg1: tensor<12x25xf32>, %arg2: tensor<24x25xf32>)
+     -> tensor<24x25xf32> 
+{
+  //     CHECK: %[[PACKED:.*]] = scf.for %{{.*}} -> (tensor<5x12x5xf32>) {
+  //     CHECK:   tensor.pad %{{.*}} 
+  //     CHECK:     : tensor<?x12xf32> to tensor<5x12xf32>
+  //     CHECK:   linalg.generic
+  //     CHECK:     -> tensor<12x5xf32>
+  //     CHECK:   tensor.insert_slice %{{.*}} into %{{.*}}[%{{.*}}, 0, 0] [1, 12, 5] [1, 1, 1] 
+  // CHECK-SAME:   : tensor<12x5xf32> into tensor<5x12x5xf32>
+  //     CHECK: scf.for %{{.*}} -> (tensor<24x25xf32>) {
+  //     CHECK:   %[[PADDED:.*]] = tensor.extract_slice %[[PACKED]][%{{.*}}, 0, 0] [1, 12, 5] [1, 1, 1] 
+  // CHECK-SAME:    : tensor<5x12x5xf32> to tensor<12x5xf32>
+  //     CHECK:   %[[TRANSPOSED:.*]] = linalg.generic
+  //     CHECK:     -> tensor<5x12xf32>
+  //     CHECK:   linalg.matmul ins(%[[TRANSPOSED]]
+  %0 = linalg.matmul ins(%arg0, %arg1 : tensor<24x12xf32>, tensor<12x25xf32>) outs(%arg2 : tensor<24x25xf32>) -> tensor<24x25xf32>
+  func.return %0 : tensor<24x25xf32>
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg1 
+    : (!pdl.operation) -> !pdl.operation
+
+  
+  %matmul_l1, %loops_l1 = transform.structured.tile_to_scf_for %matmul [5]
+
+  %matmul_padded = transform.structured.pad %matmul_l1 {
+    padding_values=[0.0: f32, 0.0 : f32, 0.0 : f32],
+    padding_dimensions=[0, 1, 2]
+  }
+
+  %pad = transform.get_producer_of_operand %matmul_padded[0]
+    : (!pdl.operation) -> !pdl.operation
+
+  transform.structured.hoist_pad %pad by 1 loops, transpose by [1, 0]
+     : (!pdl.operation) -> !pdl.operation
+}

mlir/test/Dialect/Linalg/transform-op-pad.mlir

func.func @static_sizes_output_divisible(%arg0: tensor<24x12xf32>,
   %4 = linalg.matmul ins(%1, %2 : tensor<4x?xf32>, tensor<?x5xf32>) outs(%3 : tensor<4x5xf32>) -> tensor<4x5xf32>
   %5 = tensor.insert_slice %4 into %arg2[%iv0, %iv1] [4, 5] [1, 1] : tensor<4x5xf32> into tensor<24x25xf32>
   func.return %5 : tensor<24x25xf32>
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!pdl.operation) -> !pdl.operation
-  %1 = transform.structured.pad %0 {padding_values=[0.0 : f32, 0.0 : f32, 0.0 : f32], padding_dimensions=[0, 1, 2], pack_paddings=[1, 1, 0]}
+  %1 = transform.structured.pad %0 {
+    padding_values=[0.0 : f32, 0.0 : f32, 0.0 : f32], 
+    padding_dimensions=[0, 1, 2], 
+    pack_paddings=[1, 1, 0]
+  }
 }
 
 // -----
 
 #map = affine_map<()[s0] -> (-s0 + 12, 7)>
 
 // CHECK-LABEL: @static_sizes_output_divisible_on_empty_op
 func.func @static_sizes_output_divisible_on_empty_op(%arg0: tensor<24x12xf32>,
func.func @static_sizes_output_divisible_on_empty_op(%arg0: tensor<24x12xf32>,
   %4 = linalg.matmul ins(%1, %2 : tensor<4x?xf32>, tensor<?x5xf32>) outs(%3 : tensor<4x5xf32>) -> tensor<4x5xf32>
   %5 = tensor.insert_slice %4 into %arg2[%iv0, %iv1] [4, 5] [1, 1] : tensor<4x5xf32> into tensor<24x25xf32>
   func.return %5 : tensor<24x25xf32>
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!pdl.operation) -> !pdl.operation
-  %1 = transform.structured.pad %0 {padding_values=[0.0 : f32, 0.0 : f32, 0.0 : f32], padding_dimensions=[0, 1, 2], pack_paddings=[1, 1, 0]}
+  %1 = transform.structured.pad %0 {
+    padding_values=[0.0 : f32, 0.0 : f32, 0.0 : f32],
+    padding_dimensions=[0, 1, 2],
+    pack_paddings=[1, 1, 0]
+  }
 }
 
 // -----
 
 func.func @pad(%arg0: tensor<24x12xf32>,
                %arg1: tensor<12x25xf32>,
                %arg2: tensor<24x25xf32>) -> tensor<24x25xf32> {
   // expected-note @below {{when applied to this op}}
   %0 = linalg.matmul ins(%arg0, %arg1 : tensor<24x12xf32>, tensor<12x25xf32>) outs(%arg2 : tensor<24x25xf32>) -> tensor<24x25xf32>
   func.return %0 : tensor<24x25xf32>
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!pdl.operation) -> !pdl.operation
   // expected-error @below {{op expects a padding value of type 'f32', got 0 : i32}}
-  %1 = transform.structured.pad %0 {padding_values=[0: i32, 0.0 : f32, 0.0 : f32], padding_dimensions=[0, 1, 2], pack_paddings=[1, 1, 0]}
+  %1 = transform.structured.pad %0 {
+    padding_values=[0: i32, 0.0 : f32, 0.0 : f32],
+    padding_dimensions=[0, 1, 2],
+    pack_paddings=[1, 1, 0]
+  }
 }
 
 // -----
 
 func.func @pad(%arg0: tensor<24x12xf32>,
                %arg1: tensor<12x25xf32>,
                %arg2: tensor<24x25xf32>) -> tensor<24x25xf32> {
   // expected-note @below {{when applied to this op}}
   %0 = linalg.matmul ins(%arg0, %arg1 : tensor<24x12xf32>, tensor<12x25xf32>) outs(%arg2 : tensor<24x25xf32>) -> tensor<24x25xf32>
   func.return %0 : tensor<24x25xf32>
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!pdl.operation) -> !pdl.operation
   // expected-error @below {{expects a padding that parses to 'f32', got "foo"}}
-  %1 = transform.structured.pad %0 {padding_values=["foo", 0.0 : f32, 0.0 : f32], padding_dimensions=[0, 1, 2], pack_paddings=[1, 1, 0]}
+  %1 = transform.structured.pad %0 {
+    padding_values=["foo", 0.0 : f32, 0.0 : f32],
+    padding_dimensions=[0, 1, 2],
+    pack_paddings=[1, 1, 0]
+  }
 }
 
 // -----
 
 func.func @pad(%arg0: tensor<24x12xf32>,
                %arg1: tensor<12x25xf32>,
                %arg2: tensor<24x25xf32>) -> tensor<24x25xf32> {
   // This is attached to an error that is silenceable and is not reported by this transform
   //   {{when applied to this op}}
   %0 = linalg.matmul ins(%arg0, %arg1 : tensor<24x12xf32>, tensor<12x25xf32>) outs(%arg2 : tensor<24x25xf32>) -> tensor<24x25xf32>
   func.return %0 : tensor<24x25xf32>
 }
 
 transform.sequence failures(suppress) {
 ^bb1(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!pdl.operation) -> !pdl.operation
   // This error is silenceable and is not reported by this transform
   //   {{transform.structured.pad failed to apply}}
-  %1 = transform.structured.pad %0 {padding_values=[0.0 : f32, 0.0 : f32, 0.0 : f32], padding_dimensions=[0, 1, 2], pack_paddings=[1, 1, 0]}
+  %1 = transform.structured.pad %0 {
+    padding_values=[0.0 : f32, 0.0 : f32, 0.0 : f32],
+    padding_dimensions=[0, 1, 2],
+    pack_paddings=[1, 1, 0]
+  }
 }

mlir/test/Dialect/Linalg/transform-ops-invalid.mlir

^bb0(%arg0: !pdl.operation):
   // expected-error@below {{expects pack_paddings to contain booleans (0/1), found [1, 7]}}
   transform.structured.pad %arg0 {pack_paddings=[1, 7]}
 }
 
 // -----
 
 transform.sequence failures(propagate) {
 ^bb0(%arg0: !pdl.operation):
-  // expected-error@below {{expects hoist_paddings to contain positive integers, found [1, -7]}}
-  transform.structured.pad %arg0 {hoist_paddings=[1, -7]}
-}
-
-// -----
-
-transform.sequence failures(propagate) {
-^bb0(%arg0: !pdl.operation):
   // expected-error@below {{expects transpose_paddings to be a permutation, found [1, 1]}}
   transform.structured.pad %arg0 {transpose_paddings=[[1, 1]]}
 }
 
 // -----
 
 transform.sequence failures(propagate) {
 ^bb0(%arg0: !pdl.operation):

mlir/test/lib/Dialect/Linalg/TestLinalgTransforms.cpp

 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Passes.h"
-#include "mlir/Dialect/Linalg/Transforms/HoistPadding.h"
 #include "mlir/Dialect/Linalg/Transforms/Hoisting.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 #include "llvm/ADT/SmallVector.h"

mlir/test/python/dialects/transform_structured_ext.py

structured.PadOp(
         transpose_paddings=[[1, 0]])
     transform.YieldOp()
   # CHECK-LABEL: TEST: testPad
   # CHECK: transform.sequence
   # CHECK: transform.structured.pad
   # CHECK-DAG: padding_values = [4.200000e+01 : f32]
   # CHECK-DAG: padding_dimensions = [1]
   # CHECK-DAG: transpose_paddings = {{\[}}[1, 0]]
-  # (hoist_paddings and pack_paddings have default values)
+  # (pack_paddings has default values)
 
 @run
 def testScalarize():
   sequence = transform.SequenceOp(transform.FailurePropagationMode.PROPAGATE, [], pdl.OperationType.get())
   with InsertionPoint(sequence.body):
     structured.ScalarizeOp(sequence.bodyTarget)
     transform.YieldOp()
   # CHECK-LABEL: TEST: testScalarize