[mlir][Linalg] Add a structured.pack_transpose transform op

[nicolasvasilache]

Summary

This transform is complementary to the structured.pack op which
allows packing a whole op but does not allow transposes on the individual
operands.

structured.pack_transpose allows transposing single operands connected to
pack or unpack ops after the fact.

This makes the system overall more composable than e.g. a giant transform
op with all permutation specified at once.

[nicolasvasilache]

Still needs tests but I'd appreciate early feedback on the logic and descriptions.

[nicolasvasilache]

Bug fixes + add first test.

[nicolasvasilache]

More bugfixes and add second test.

[chelini]

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

1159

Mind improving this comment please? The comment reflects only the first few lines of this method. If this should not be part of the doc we can use //

1183

use //?

1218

debug leftover. Same below.

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

3242

nit: unfinished comment.

[nicolasvasilache]

Address comments.

[nicolasvasilache]

Mark addressed.

[nicolasvasilache]

Update commit message.

[ftynse]

I didn't check the logic of transposition in great detail, deferring to @chelini for that.

mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td

433–451

Do we expect concrete ops to ever reimplement this behavior? Otherwise, this looks like a good fit for extraSharedClassDeclaration instead of being an interface method. This lets the body live in C++ rather than string, allows for faster dispatch (no pseudo-vtable), and avoids funny discrepancies like the difference between $_op and this->getOperation().

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

439

480

Nit: this will generate getUnPackOp that would be inconsistent with getTargetPackOrUnpackOp (P not capitalized in "unpack).

490–494

Nit: you can just do DeclareOpInterfaceMethods<TransformOpInterface>` in the definition of the op class and this becomes unnecessary.

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

1139

Missed it the last time, but this also needs transform::modifiesPayloadIR effect.

1148

This also needs a transform::modifiesPayloadIR.

1148–1152

Looks like this could have had a FunctionalStyleTransformOpTrait.

1203–1204

Maybe we don't need to pass it then?

1239

This moves the body from linalgOp leaving it in potentially invalid state. I see that linalgOp is still necessary at the callsite so it cannot be immediately deleted here, but this needs to be at least documented on the function. Better, store everything that's needed from linalgOp at the callsite before calling this and actually call replaceOp here.

1291–1293

Incorrect error message.

1751–1754

One day you'll fix your editor to tell it 80 cols are fine and there's no need to do 78 on irrelevant parts of the file...

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

3261

The combination of name and type confused me. getInnerDimPos sounds like it returns one value (position), which would be treated as the size of the vector. It is not in fact the case, but looks confusing. Any clarification (comment, different code structure, different name) is welcome.

[nicolasvasilache]

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

1751–1754

I've been at 81 for a while now.
I am utterly unclear why this would happen here and not in other places..

2006

why would it reflow here and not on the code on l2003 above is beyond me ..

[ftynse]

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

1751–1754

Maybe some rouge clang-format pre-commit git hook?

[nicolasvasilache]

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

1751–1754

Sometimes I can comment / uncomment locally while debugging and things get auto-formatted and not reflowed back.
But this instance is still puzzling to me ..

[nicolasvasilache]

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

480

went for UnPack to be consistent with tensor::UnPackOp

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

1203–1204

I find the explicit chaining requirement to help with the legibility atm.
I'd prefer to remove later if/when we feel like its a burden.

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

3261

I am not seeing this, maybe I applied some cleanups in the meantime ?

[nicolasvasilache]

Rebase and address comments.

[nicolasvasilache]

Drop spurious formatting changes

[chelini]

Looks good to me, thanks.

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

1304

nit: UnPackOp

[nicolasvasilache]

Typo

mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td

InterfaceMethod<
         return *(indexingMaps.begin() +
                  cast<DestinationStyleOpInterface>(*this->getOperation())
                      .getNumDpsInputs() +
                  result.getResultNumber());
       }]
     >,
     InterfaceMethod<
       /*desc=*/[{
+        Return the index in the indexingMaps vector that corresponds to this `opOperand`
+      }],
+      /*retTy=*/"int64_t",
+      /*methodName=*/"getIndexingMapIndex",
+      /*args=*/(ins "OpOperand *":$opOperand),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        auto operandNumber = opOperand->getOperandNumber();
+        auto dpsIface = cast<DestinationStyleOpInterface>(*this->getOperation());
+        if (!dpsIface.isDpsInput(opOperand)) return operandNumber;
+        auto [start, end] = dpsIface.getDpsInitsPositionRange();
+        assert(!$_op.isDpsInit(opOperand));
+        // Account for potential inputs that are not DPS and may not appear in
+        // `indexingMaps`.
+        return cast<DestinationStyleOpInterface>(*this->getOperation())
+                     .getNumDpsInputs() +
+               operandNumber - start;
+      }]
+    >,
+    InterfaceMethod<
+      /*desc=*/[{
         Return the value yielded by the region corresponding to an output
         `opOperand`.
       }],
       /*retTy=*/"OpOperand *",
       /*methodName=*/"getMatchingYieldValue",
       /*args=*/(ins "OpOperand*":$opOperand),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{

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

def MultiTileSizesOp : Op<Transform_Dialect, "structured.multitile_sizes",
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::linalg::LinalgOp target,
         ::mlir::transform::ApplyToEachResultList &results,
         TransformState &state);
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// PackOp
+//===----------------------------------------------------------------------===//
 def PackOp : Op<Transform_Dialect, "structured.pack", [
                 TransformOpInterface,
                 DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,]> {
   let description = [{
     Pack a LinalgOp by applying a data tiling transformation on the op and
     packing the operands according to the `packed_sizes` specification.
     
     Iterator dimensions are tiled in their canonical order in the op spec.
::mlir::DiagnosedSilenceableFailure apply(
       transform::TransformResults &transformResults,
       transform::TransformState &state);
 
     ::llvm::SmallVector<::mlir::OpFoldResult> getMixedPackedSizes();
   }];
 }
 
 //===----------------------------------------------------------------------===//
+// PackTransposeOp
+//===----------------------------------------------------------------------===//
+def PackTransposeOp : Op<Transform_Dialect, "structured.pack_transpose", [
+                         TransformOpInterface,
+                         DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,]> {
+  let description = [{
+    Apply a transposition to a single `tensor.pack` (resp. `tensor.unpack`) and 
+    updates the `linalg.generic` op that consumes (resp. produces) the operation.
+
+    This transform allows composing a simple `structured.pack` with additional
+    transpositions to e.g. match the data format required by a specific library
+    call or ISA instruction.
+
+    The transpose spec must specify at least one of `outer_perm` or `inner_perm`
+    attributes, which will act upon the `outer_dims_perm` or `inner_dims_pos` of
+    the specified `tensor.pack` or `tensor.unpack` op.
+
+    If the `target` of this op is a `tensor.pack` then a new `tensor.empty` will
+    be created along with transposed versions of the `tensor.pack` and the 
+    consuming `linalg.generic`, which is expected to be the sole consumer.
+
+    If the `target` of this op is a `tensor.unpack` then the whole pack / compute
+    / unpack chain will be transposed and transposed clones of `tensor.pack`,
+    the consuming `linalg.generic` and the tail `tensor.pack` will be created.
+
+    #### Return modes
+
+    This operation targets a single `tensor.pack` / `tensor.unpack` op and a
+    single matching `linalg.generic` that consumes / produces the op. Otherwise,
+    it produces a silenceableFailure.
+
+    This operation may produce a silenceableFailure if the transpose spec is
+    ill-formed (i.e. `outer_perm` or `inner_perm` are not permutations of the
+    proper rank) or if the tranposition of all involved operations fails for any
+    reason.
+
+    This operation returns 3 handles, one to the transformed LinalgOp, one to
+    the transformed `tensor.pack` and one to the transformed `tensor.unpack`.
+    The last handle for `tensor.unpack` is empty if `target_pack_or_unpack_op` 
+    was not itself a `tensor.unpack`.
+  }];
+
+  let arguments = (ins TransformHandleTypeInterface:$target_pack_or_unpack_op,
+                       TransformHandleTypeInterface:$target_linalg_op,
+                       DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$outer_perm,
+                       DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$inner_perm);
+  let results = (outs TransformHandleTypeInterface:$packed_op,
+                      TransformHandleTypeInterface:$pack_op,
+                      TransformHandleTypeInterface:$un_pack_op);
+  let assemblyFormat = [{
+    $target_pack_or_unpack_op
+    `with_compute_op` `(` $target_linalg_op `)`
+    (`outer_perm` `=` $outer_perm^ )?
+    (`inner_perm` `=` $inner_perm^ )?
+    attr-dict
+    `:` functional-type(operands, results)
+  }];
+
+  let extraClassDeclaration = [{
+    ::mlir::DiagnosedSilenceableFailure apply(
+      transform::TransformResults &transformResults,
+      transform::TransformState &state);
+  }];
+
+  let hasVerifier = 1;
+}
+
+//===----------------------------------------------------------------------===//
 // PadOp
 //===----------------------------------------------------------------------===//
 
 def PadOp : Op<Transform_Dialect, "structured.pad",
     [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
      TransformOpInterface, TransformEachOpTrait]> {
   let description = [{
     Pads the operations pointed to by the target handle using the options

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

let extraClassDeclaration = commonExtraClassDeclaration # [{
     // outer loops (outerDimsPerm).
     static ShapedType inferPackedType(ShapedType sourceType,
         ArrayRef<int64_t> innerTileSizes, ArrayRef<int64_t> innerDimsPos,
         ArrayRef<int64_t> outerDimsPerm = {});
 
     static Value createDestinationTensor(OpBuilder &b, Location loc,
         Value source, ArrayRef<OpFoldResult> innerTileSizes,
         ArrayRef<int64_t> innerDimsPos, ArrayRef<int64_t> outerDimsPerm);
+
+    /// Build and return a new PackOp that is a clone of the current PackOp with
+    /// (innerDimsPos, innerTiles) (resp. outerDimsPerm) are permuted by 
+    /// innerPermutation (resp. outerPermutation).
+    /// A new `tensor.empty` of the proper shape is built in the process.
+    /// Asserts that:
+    ///   - At least one of innerPermutation or outerPermutation is non-empty.
+    ///   - If not empty, innerPermutation is a valid permutation of size
+    ///     matching innerDimPos.
+    ///   - If not empty, outerPermutation is a valid permutation of size 
+    ///     matching outerDimsPerm.
+    PackOp createTransposedClone(OpBuilder &b,
+                                 Location loc,
+                                 ArrayRef<int64_t> innerPermutation,
+                                 ArrayRef<int64_t> outerPermutation);
   }];
 
   let hasCanonicalizeMethod = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // UnPackOp
 //===----------------------------------------------------------------------===//
def Tensor_UnPackOp : Tensor_RelayoutOp<"unpack"> {
 
   let builders = [
     OpBuilder<(ins "Value":$source, "Value":$dest,
     "ArrayRef<int64_t>":$innerDimsPos,
     "ArrayRef<OpFoldResult>":$innerTiles,
     CArg<"ArrayRef<int64_t>", "{}">:$outerDimsPerm)>
   ];
 
-  let extraClassDeclaration = commonExtraClassDeclaration;
+  let extraClassDeclaration = commonExtraClassDeclaration # [{
+    /// Build and return a new UnPackOp that is a clone of the current UnPackOp
+    /// with (innerDimsPos, innerTiles) (resp. outerDimsPerm) are permuted by 
+    /// innerPermutation (resp. outerPermutation).
+    /// Asserts that:
+    ///   - At least one of innerPermutation or outerPermutation is non-empty.
+    ///   - If not empty, innerPermutation is a valid permutation of size
+    ///     matching innerDimPos.
+    ///   - If not empty, outerPermutation is a valid permutation of size 
+    ///     matching outerDimsPerm.
+    UnPackOp createTransposedClone(OpBuilder &b,
+                                   Location loc,
+                                   Value transposedSource,
+                                   ArrayRef<int64_t> innerPermutation,
+                                   ArrayRef<int64_t> outerPermutation);
+  }];
+
   let hasCanonicalizeMethod = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // YieldOp
 //===----------------------------------------------------------------------===//
 
 def Tensor_YieldOp : Tensor_Op<"yield",

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

 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/PDL/IR/PDL.h"
 #include "mlir/Dialect/PDL/IR/PDLTypes.h"
 #include "mlir/Dialect/SCF/Transforms/TileUsingInterface.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
 #include "mlir/Dialect/Transform/IR/TransformUtils.h"
 #include "mlir/Dialect/Transform/Utils/Utils.h"
+#include "mlir/Dialect/Utils/IndexingUtils.h"
+#include "mlir/IR/AffineMap.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Interfaces/TilingInterface.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/Support/Debug.h"
 
 using namespace mlir;
 using namespace mlir::linalg;
 using namespace mlir::transform;
 
transform::PackOp::apply(transform::TransformResults &transformResults,
   // If nothing to pack, propagate success.
   if (targetOps.empty()) {
     transformResults.set(getPackedOp().cast<OpResult>(), {});
     return DiagnosedSilenceableFailure::success();
   }
   // Fail on multi-op handles.
   auto linalgOp = dyn_cast<linalg::LinalgOp>(targetOps.front());
   if (targetOps.size() != 1 || !linalgOp) {
-    // TODO: remove this unnecessary set to empty once crashes are fixed.
+    // TODO: remove this unnecessary "set to empty" once crashes are fixed.
     transformResults.set(getPackedOp().cast<OpResult>(), {});
     return emitSilenceableError()
            << "requires target to map to exactly 1 LinalgOp (got "
            << targetOps.size() << ")";
   }
   // Fail on mismatched number of pack sizes.
   if (getMixedPackedSizes().size() != linalgOp.getNumLoops()) {
-    // TODO: remove this unnecessary set to empty once crashes are fixed.
+    // TODO: remove this unnecessary "set to empty" once crashes are fixed.
     transformResults.set(getPackedOp().cast<OpResult>(), {});
     return emitSilenceableError()
            << "requires number of packed sizes match the number of loops ("
            << getMixedPackedSizes().size() << " vs " << linalgOp.getNumLoops()
            << ")";
   }
 
   IRRewriter rewriter(linalgOp->getContext());
   rewriter.setInsertionPoint(linalgOp);
   FailureOr<LinalgOp> maybeResult =
       packOneLinalgOp(rewriter, state, *this, linalgOp, getMixedPackedSizes());
   if (failed(maybeResult))
     return emitDefiniteFailure("data tiling failed");
 
   transformResults.set(getPackedOp().cast<OpResult>(),
                        maybeResult->getOperation());
   return DiagnosedSilenceableFailure::success();
 }
 
 void transform::PackOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   transform::consumesHandle(getTarget(), effects);
   transform::onlyReadsHandle(getPackedSizes(), effects);
   transform::producesHandle(getPackedOp(), effects);
 }
 
 //===---------------------------------------------------------------------===//
+// PackTransposeOp
+//===---------------------------------------------------------------------===//
+
+void transform::PackTransposeOp::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  transform::consumesHandle(getTargetPackOrUnpackOp(), effects);
+  transform::consumesHandle(getTargetLinalgOp(), effects);
+  transform::producesHandle(getPackedOp(), effects);
+  transform::producesHandle(getPackOp(), effects);
+  transform::producesHandle(getUnPackOp(), effects);
+}
+
+namespace {
+enum class OuterOrInnerPerm { Outer = 0, Inner = 1 };
+} // namespace
+
+/// Return true if `permutation` is a valid permutation of the `outer_dims_perm`
+/// (case OuterOrInnerPerm::Outer) or `inner_dims_pos` (OuterOrInnerPerm::Inner)
+/// of the `tensor.pack` or `tensor.unpack` `op.
+/// This is the case when the `permutation` rank matches the rank expected by
+/// `op` and `permutation` is itself a permutation vector.
+/// Return true if either `op` or `permutation` are empty to allow a simpler
+/// polymorphic implementation.
+template <typename RelayoutOpTy>
+bool isValidPackingPermutation(
+    RelayoutOpTy op, ArrayRef<int64_t> permutation,
+    OuterOrInnerPerm outerOrInnerPerm = OuterOrInnerPerm::Outer) {
+  static_assert(
+      llvm::is_one_of<RelayoutOpTy, tensor::PackOp, tensor::UnPackOp>::value,
+      "applies to only pack or unpack operations");
+  if (!op || permutation.empty())
+    return true;
+  int64_t innerRank = op.getInnerDimsPos().size();
+  if (outerOrInnerPerm == OuterOrInnerPerm::Inner)
+    return permutation.size() == innerRank && isPermutationVector(permutation);
+  // op.getOuterDimsPerm() may be empty, in which case it is identity.
+  // Don't rely on it.
+  if (std::is_same<RelayoutOpTy, tensor::PackOp>::value) {
+    return permutation.size() == op.getSourceRank() &&
+           isPermutationVector(permutation);
+  }
+  return permutation.size() == op.getDestRank() &&
+         isPermutationVector(permutation);
+}
+
+/// 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`
+/// Asserts that `transposedValue` is of the proper transposed ShapedType.
+static LinalgOp transposeOneLinalgOperand(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);
+  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));
+
+  return cast<linalg::LinalgOp>(transposedGenericOp.getOperation());
+}
+
+LogicalResult transform::PackTransposeOp::verify() {
+  if (!isPermutationVector(getInnerPerm())) {
+    return emitOpError() << getInnerPermAttrName()
+                         << " is not a valid permutation";
+  }
+  if (!isPermutationVector(getOuterPerm())) {
+    return emitOpError() << getOuterPermAttrName()
+                         << " is not a valid permutation";
+  }
+  if (getInnerPerm().empty() && getOuterPerm().empty()) {
+    return emitOpError() << " at least one of " << getInnerPermAttrName()
+                         << " or " << getOuterPermAttrName()
+                         << " must be specified";
+  }
+  return success();
+}
+
+DiagnosedSilenceableFailure
+transform::PackTransposeOp::apply(transform::TransformResults &transformResults,
+                                  transform::TransformState &state) {
+  ArrayRef<Operation *> packOrUnpackOps =
+      state.getPayloadOps(getTargetPackOrUnpackOp());
+  ArrayRef<Operation *> linalgOps = state.getPayloadOps(getTargetLinalgOp());
+  // Step 1. If nothing to pack, propagate success.
+  if (packOrUnpackOps.empty()) {
+    transformResults.set(getPackedOp().cast<OpResult>(), {});
+    transformResults.set(getPackOp().cast<OpResult>(), {});
+    transformResults.set(getUnPackOp().cast<OpResult>(), {});
+    return DiagnosedSilenceableFailure::success();
+  }
+
+  // Step 2. Bunch of runtime sanity check and error messages.
+  // Fail on multi-op handles.
+  if (packOrUnpackOps.size() != 1 || linalgOps.size() != 1) {
+    return emitSilenceableError()
+           << "requires target to map to exactly 1 "
+              "tensor.pack or tensor.unpack and exactly 1 LinalgOp (got "
+           << packOrUnpackOps.size() << " and " << linalgOps.size() << ")";
+  }
+  // Fail on wrong type.
+  auto packOp = dyn_cast<tensor::PackOp>(packOrUnpackOps.front());
+  auto unPackOp = dyn_cast<tensor::UnPackOp>(packOrUnpackOps.front());
+  if ((!packOp && !unPackOp)) {
+    // TODO: remove this unnecessary "set to empty" once crashes are fixed.
+    transformResults.set(getPackedOp().cast<OpResult>(), {});
+    transformResults.set(getPackOp().cast<OpResult>(), {});
+    transformResults.set(getUnPackOp().cast<OpResult>(), {});
+    return emitSilenceableError() << "requires target to map to exactly 1 "
+                                     "tensor.pack or tensor.unpack (got "
+                                  << packOrUnpackOps.size() << ")";
+  }
+  // Fail if we can't get the producer / consumer Linalg op.
+  LinalgOp linalgOp;
+  if (packOp && packOp.getResult().hasOneUse())
+    linalgOp = dyn_cast<LinalgOp>(*(packOp.getResult().getUsers().begin()));
+  else if (unPackOp)
+    linalgOp = unPackOp.getSource().getDefiningOp<LinalgOp>();
+  if (!linalgOp || linalgOp != linalgOps.front()) {
+    // TODO: remove this unnecessary "set to empty" once crashes are fixed.
+    transformResults.set(getPackedOp().cast<OpResult>(), {});
+    transformResults.set(getPackOp().cast<OpResult>(), {});
+    transformResults.set(getUnPackOp().cast<OpResult>(), {});
+    if (packOp)
+      return emitSilenceableError()
+             << "not a single use by the specified LinalgOp";
+    return emitSilenceableError() << "not produced by the specified LinalgOp";
+  }
+  // If we have an UnpackOp, we also need to fetch the symmetrical PackOp.
+  if (unPackOp) {
+    assert(!packOp && "packOp must be null on entry when unPackOp is not null");
+    OpOperand *packUse = linalgOp.getDpsInitOperand(
+        unPackOp.getSource().cast<OpResult>().getResultNumber());
+    packOp = dyn_cast_or_null<tensor::PackOp>(packUse->get().getDefiningOp());
+    if (!packOp || !packOp.getResult().hasOneUse())
+      return emitSilenceableError() << "could not find matching pack op";
+  }
+  // Fail if any permutation does not validate.
+  for (auto permType : {OuterOrInnerPerm::Outer, OuterOrInnerPerm::Inner}) {
+    ArrayRef<int64_t> perm =
+        (permType == OuterOrInnerPerm::Outer) ? getOuterPerm() : getInnerPerm();
+    auto errorMsg = (permType == OuterOrInnerPerm::Outer)
+                        ? StringLiteral{"invalid outer_perm"}
+                        : StringLiteral{"invalid inner_perm"};
+    if (!isValidPackingPermutation(packOp, perm, permType) ||
+        !isValidPackingPermutation(unPackOp, perm, permType)) {
+      // TODO: remove this unnecessary "set to empty" once crashes are fixed.
+      transformResults.set(getPackedOp().cast<OpResult>(), {});
+      transformResults.set(getPackOp().cast<OpResult>(), {});
+      transformResults.set(getUnPackOp().cast<OpResult>(), {});
+      Operation *packOrUnpackOp =
+          unPackOp ? unPackOp.getOperation() : packOp.getOperation();
+      return emitSilenceableError() << errorMsg << ": " << *packOrUnpackOp;
+    }
+  }
+
+  // From here on, packOp and linalgOp are always present, unPackOp may or may
+  // not be present.
+  assert(packOp && linalgOp && "unexpected null op");
+
+  // Step 3. Actually transpose the ops.
+  IRRewriter rewriter(getContext());
+
+  // Step 3.a. Transpose packOp.
+  rewriter.setInsertionPoint(packOp);
+  tensor::PackOp transposedPackOp = packOp.createTransposedClone(
+      rewriter, packOp.getLoc(), getInnerPerm(), getOuterPerm());
+
+  // Step 3.b. Transpose linalgOp.
+  assert(packOp.getResult().hasOneUse() && "expect single use");
+  // packOp.getOuterDimsPerm() may be empty, in which case it is identity.
+  // Don't rely on it.
+  int64_t numLeadingDims = packOp.getSourceRank();
+  int64_t numTrailingDims = packOp.getInnerDimsPos().size();
+  // Step 3.b.i. Compute the permutation on the whole operand.
+  // Leading part just reuse the outerPerm.
+  SmallVector<int64_t> permutation(getOuterPerm());
+  if (permutation.empty())
+    llvm::append_range(permutation, llvm::seq<int64_t>(0, numLeadingDims));
+  // Trailing part needs to reindex positions by `numLeadingDims`.
+  if (getInnerPerm().empty()) {
+    llvm::append_range(
+        permutation,
+        llvm::seq<int64_t>(numLeadingDims, numLeadingDims + numTrailingDims));
+  } else {
+    llvm::append_range(permutation,
+                       llvm::map_range(getInnerPerm(), [&](int64_t pos) {
+                         return numLeadingDims + pos;
+                       }));
+  }
+  assert(isPermutationVector(permutation) && "invalid permutation");
+  // Step 3.b.ii. Actually perform the transposition.
+  OpOperand &packUse = *(packOp.getResult().getUses().begin());
+  rewriter.setInsertionPoint(linalgOp);
+  linalg::LinalgOp transposedLinalgOp = transposeOneLinalgOperand(
+      rewriter, linalgOp, packUse, permutation, transposedPackOp.getResult());
+
+  // Step 3.c. Maybe transpose unPackOp.
+  tensor::UnPackOp transposedUnPackOp;
+  if (unPackOp) {
+    OpResult transposedResult = linalgOp.getTiedOpResult(&packUse);
+    rewriter.setInsertionPoint(unPackOp);
+    transposedUnPackOp = unPackOp.createTransposedClone(
+        rewriter, packOp.getLoc(), transposedResult, getInnerPerm(),
+        getOuterPerm());
+  }
+
+  // Step 4. Replace and return results.
+  rewriter.replaceOp(packOp, transposedPackOp->getResults());
+  rewriter.replaceOp(linalgOp, transposedLinalgOp->getResults());
+  transformResults.set(getPackOp().cast<OpResult>(), {transposedPackOp});
+  transformResults.set(getPackedOp().cast<OpResult>(), {transposedLinalgOp});
+  if (unPackOp) {
+    rewriter.replaceOp(unPackOp, transposedUnPackOp->getResults());
+    transformResults.set(getUnPackOp().cast<OpResult>(), {transposedUnPackOp});
+  } else {
+    transformResults.set(getUnPackOp().cast<OpResult>(), {});
+  }
+  return DiagnosedSilenceableFailure::success();
+}
+
+//===---------------------------------------------------------------------===//
 // PadOp
 //===---------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::PadOp::applyToOne(linalg::LinalgOp target,
                              transform::ApplyToEachResultList &results,
                              transform::TransformState &state) {
   // Convert the integer packing flags to booleans.
transform::ReplaceOp::apply(TransformResults &transformResults,
 
   // Check for invalid targets.
   for (Operation *target : payload) {
     if (target->getNumOperands() > 0)
       return emitDefiniteFailure() << "expected target without operands";
     if (!target->hasTrait<OpTrait::IsIsolatedFromAbove>() &&
         target->getNumRegions() > 0)
       return emitDefiniteFailure()
-             << "expected target that is isloated from above";
+             << "expected target that is isolated from above";
   }
 
   // Clone and replace.
   IRRewriter rewriter(getContext());
   Operation *pattern = &getBodyRegion().front().front();
   SmallVector<Operation *> replacements;
   for (Operation *target : payload) {
     if (getOperation()->isAncestor(target))
if (!second.back()) {
       noSecondPart = target;
       second.pop_back();
     }
   }
 
   if (second.size() != first.size() && !second.empty()) {
     results.set(getFirst().cast<OpResult>(), {});
     results.set(getSecond().cast<OpResult>(), {});
-    auto diag =
-        emitSilenceableError()
-        << "splitting does not produce the second part for a subset of targets";
+    auto diag = emitSilenceableError()
+                << "splitting does not produce the second part for a subset "
+                   "of targets";
     diag.attachNote() << "expected splitting to produce the second part of all "
                          "or none of the targets";
     diag.attachNote(noSecondPart->getLoc())
         << "first target with no second part";
     return diag;
   }
 
   results.set(getFirst().cast<OpResult>(), first);
   results.set(getSecond().cast<OpResult>(), second);

 void transform::TileOp::build(OpBuilder &builder, OperationState &result,
                               Value target,
                               ArrayRef<OpFoldResult> mixedTileSizes,
                               ArrayRef<int64_t> interchange) {
   SmallVector<int64_t> staticTileSizes;
   SmallVector<Value> dynamicTileSizes;
   dispatchIndexOpFoldResults(mixedTileSizes, dynamicTileSizes, staticTileSizes);
   // Call the default builder which sets up the proper operands segment sizes
-  // attributes for multiple variadic operands. In the absence of this, horrible
-  // bugs ensue.
+  // attributes for multiple variadic operands. In the absence of this,
+  // horrible bugs ensue.
   MLIRContext *ctx = builder.getContext();
   auto operationType = pdl::OperationType::get(ctx);
   auto staticTileSizesAttr = builder.getDenseI64ArrayAttr(staticTileSizes);
   build(builder, result,
         /*resultTypes=*/TypeRange{operationType, operationType},
         /*target=*/target,
         /*dynamic_sizes=*/dynamicTileSizes,
         /*static_sizes=*/staticTileSizesAttr,

 void transform::TileToForeachThreadOp::build(
     OpBuilder &builder, OperationState &result, Value target,
     ArrayRef<OpFoldResult> mixedTileSizes, transform::TileSizesSpec,
     ArrayAttr mapping) {
   SmallVector<int64_t> staticTileSizes;
   SmallVector<Value> dynamicTileSizes;
   dispatchIndexOpFoldResults(mixedTileSizes, dynamicTileSizes, staticTileSizes);
   // Call the default builder which sets up the proper operands segment sizes
-  // attributes for multiple variadic operands. In the absence of this, horrible
-  // bugs ensue.
+  // attributes for multiple variadic operands. In the absence of this,
+  // horrible bugs ensue.
   MLIRContext *ctx = builder.getContext();
   auto operationType = pdl::OperationType::get(ctx);
   auto staticTileSizesAttr = builder.getDenseI64ArrayAttr(staticTileSizes);
   build(builder, result,
         /*resultTypes=*/TypeRange{operationType, operationType},
         /*target=*/target,
         /*num_threads=*/ValueRange{},
         /*tile_sizes=*/dynamicTileSizes,
void transform::TileToForeachThreadOp::build(
     OpBuilder &builder, OperationState &result, Value target,
     ArrayRef<OpFoldResult> mixedNumThreads, transform::NumThreadsSpec,
     ArrayAttr mapping) {
   SmallVector<int64_t> staticNumThreads;
   SmallVector<Value> dynamicNumThreads;
   dispatchIndexOpFoldResults(mixedNumThreads, dynamicNumThreads,
                              staticNumThreads);
   // Call the default builder which sets up the proper operands segment sizes
-  // attributes for multiple variadic operands. In the absence of this, horrible
-  // bugs ensue.
+  // attributes for multiple variadic operands. In the absence of this,
+  // horrible bugs ensue.
   MLIRContext *ctx = builder.getContext();
   auto operationType = pdl::OperationType::get(ctx);
   auto staticNumThreadsAttr = builder.getDenseI64ArrayAttr(staticNumThreads);
   build(builder, result,
         /*resultTypes=*/TypeRange{operationType, operationType},
         /*target=*/target,
         /*num_threads=*/dynamicNumThreads,
         /*tile_sizes=*/ValueRange{},
DiagnosedSilenceableFailure transform::tileToForeachThreadOpImpl(
     ArrayRef<OpFoldResult> mixedNumThreads,
     ArrayRef<OpFoldResult> mixedTileSizes, std::optional<ArrayAttr> mapping,
     SmallVector<Operation *> &tileOps, SmallVector<Operation *> &tiledOps) {
   if (targets.empty())
     return DiagnosedSilenceableFailure::success();
 
   // Transform all targets one by one.
   for (Operation *target : targets) {
-    auto tilableOp = dyn_cast<TilingInterface>(target);
-    if (!tilableOp) {
+    auto tileableOp = dyn_cast<TilingInterface>(target);
+    if (!tileableOp) {
       DiagnosedSilenceableFailure diag =
           transformOp.emitSilenceableError()
           << "only TilingInterface ops are supported";
       diag.attachNote(target->getLoc()) << "target op";
       return diag;
     }
-    rewriter.setInsertionPoint(tilableOp);
+    rewriter.setInsertionPoint(tileableOp);
     FailureOr<linalg::ForeachThreadTilingResult> tilingResult = failure();
     if (!mixedNumThreads.empty()) {
-      tilingResult = linalg::tileToForeachThreadOp(rewriter, tilableOp,
+      tilingResult = linalg::tileToForeachThreadOp(rewriter, tileableOp,
                                                    mixedNumThreads, mapping);
     } else {
       tilingResult = linalg::tileToForeachThreadOpUsingTileSizes(
-          rewriter, tilableOp, mixedTileSizes, mapping);
+          rewriter, tileableOp, mixedTileSizes, mapping);
     }
 
     if (failed(tilingResult))
-      return transformOp.emitDefaultSilenceableFailure(tilableOp);
-    rewriter.replaceOp(tilableOp, tilingResult->tileOp->getResults());
+      return transformOp.emitDefaultSilenceableFailure(tileableOp);
+    rewriter.replaceOp(tileableOp, tilingResult->tileOp->getResults());
 
     tileOps.push_back(tilingResult->tileOp);
     tiledOps.push_back(tilingResult->tiledOp);
   }
   return DiagnosedSilenceableFailure::success();
 }
 
 DiagnosedSilenceableFailure transform::TileToForeachThreadOp::apply(
if (numThreadsSpec > 1)
     return emitOpError(
         "num_threads and packed_num_threads are mutually exclusive");
   int tileSizesSpec = static_cast<int>(!getMixedTileSizes().empty()) +
                       static_cast<int>(getPackedTileSizes() != Value());
   if (tileSizesSpec > 1)
     return emitOpError(
         "tile_sizes and packed_tile_sizes are mutually exclusive");
   if (numThreadsSpec == 0 && tileSizesSpec == 0)
-    return emitOpError(
-        "either (packed_)num_threads or (packed_)tile_sizes must be specified");
+    return emitOpError("either (packed_)num_threads or (packed_)tile_sizes "
+                       "must be specified");
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // TileToScfForOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

if (!llvm::all_of(
             if (ShapedType::isDynamic(shape)) {
               // For the shape being dynamic when tile size is
               // specified, return true. In canonical form a constant
               // tile size should lead to constant shape of the tiled
               // dimension, but not needed for verification.
               return true;
             }
             return shape == constTileSize.value();
-
           })) {
     return op->emitError("mismatch in inner tile sizes specified and shaped of "
                          "tiled dimension in the packed type");
   }
   return success();
 }
 
+namespace {
+/// Subset of PackOp/UnPackOp fields used to compute the result of applying
+/// various permutations to the op.
+// TODO: Add linalg.transpose + pack/unpack folding patterns that just reuse
+// these. These may or may not become true foldings / canonicalizations
+// depending on how aggressive we want to be in automatically folding
+// transposes.
+struct PackOrUnPackTransposeResult {
+  SmallVector<int64_t> innerDimsPos;
+  SmallVector<OpFoldResult> innerTiles;
+  SmallVector<int64_t> outerDimsPerm;
+};
+} // namespace
+
+template <typename OpTy>
+static PackOrUnPackTransposeResult
+commonPermutationOfPackAndUnPackOp(OpTy packOrUnPackOp,
+                                   ArrayRef<int64_t> innerPermutation,
+                                   ArrayRef<int64_t> outerPermutation) {
+  static_assert(llvm::is_one_of<OpTy, PackOp, UnPackOp>::value,
+                "applies to only pack or unpack operations");
+  assert((!innerPermutation.empty() || !outerPermutation.empty()) &&
+         "some permutation must be non-empty");
+  PackOrUnPackTransposeResult metadata;
+  metadata.innerDimsPos =
+      SmallVector<int64_t>(packOrUnPackOp.getInnerDimsPos());
+  metadata.innerTiles =
+      SmallVector<OpFoldResult>(packOrUnPackOp.getMixedTiles());
+  int64_t numOuterDims = std::is_same<OpTy, PackOp>::value
+                             ? packOrUnPackOp.getSourceRank()
+                             : packOrUnPackOp.getDestRank();
+  metadata.outerDimsPerm =
+      packOrUnPackOp.getOuterDimsPerm().empty()
+          ? llvm::to_vector(llvm::seq<int64_t>(0, numOuterDims))
+          : SmallVector<int64_t>(packOrUnPackOp.getOuterDimsPerm());
+  if (!innerPermutation.empty()) {
+    assert(innerPermutation.size() == metadata.innerDimsPos.size() &&
+           isPermutationVector(innerPermutation) &&
+           "invalid inner permutation");
+    applyPermutationToVector(metadata.innerDimsPos, innerPermutation);
+    applyPermutationToVector(metadata.innerTiles, innerPermutation);
+  }
+  if (!outerPermutation.empty()) {
+    assert(outerPermutation.size() == metadata.outerDimsPerm.size() &&
+           isPermutationVector(outerPermutation) &&
+           "invalid outer permutation");
+    applyPermutationToVector(metadata.outerDimsPerm, outerPermutation);
+  }
+  return metadata;
+}
+
 //===----------------------------------------------------------------------===//
 // PackOp
 //===----------------------------------------------------------------------===//
 
 void PackOp::getAsmResultNames(function_ref<void(Value, StringRef)> setNameFn) {
   setNameFn(getResult(), "pack");
 }
 

 SmallVector<OpFoldResult> PackOp::getMixedTiles() {
   return getMixedTilesImpl(*this);
 }
 
 SmallVector<int64_t> PackOp::getStaticTiles() {
   return getStaticTilesImpl(*this);
 }
 
-/// Check if we have enough static information to catch undefined behavior when
-/// the tile size does not divide perfectly the dimension of the input tensor.
+/// Check if we have enough static information to catch undefined behavior
+/// when the tile size does not divide perfectly the dimension of the input
+/// tensor.
 static bool
 areNotFullTiles(ArrayRef<int64_t> inputShape,
                 DenseMap<int64_t, OpFoldResult> const &dimAndTileMapping) {
   int64_t rank = inputShape.size();
   for (int64_t dim = 0; dim < rank; dim++) {
     if (ShapedType::isDynamic(inputShape[dim]))
       continue;
     auto it = dimAndTileMapping.find(dim);
     if (it == dimAndTileMapping.end())
       continue;
     std::optional<int64_t> constantTile = getConstantIntValue(it->second);
     if (!constantTile)
       continue;
     if (inputShape[dim] % (*constantTile) != 0)
       return true;
   }
   return false;
 }
 
 LogicalResult PackOp::verify() {
   if (failed(commonVerifierPackAndUnPackOp(*this)))
     return failure();
 
   // Verify padding value, and bail out if the tile does not divide the
-  // dimension fully. In the case of dynamic tile factors or dimensions, having
-  // a partial tile is undefined behavior.
+  // dimension fully. In the case of dynamic tile factors or dimensions,
+  // having a partial tile is undefined behavior.
   auto paddingValue = getPaddingValue();
   if (paddingValue &&
       paddingValue.getType() != getSourceType().getElementType()) {
     return emitOpError("expected padding_value has ")
            << getSourceType().getElementType()
            << " but got: " << paddingValue.getType();
   }
 
   auto dimAndTileMapping = getDimAndTileMapping();
   if (!paddingValue &&
       areNotFullTiles(getSourceType().getShape(), dimAndTileMapping)) {
     return emitOpError("invalid tile factor provided. Only full tiles are "
                        "supported when padding_value is not set");
   }
   return success();
 }
 
-/// Get the expected packed type based on source type, tile factors, position of
-/// the inner tiles and permutation of the outer tiled loop.
+/// Get the expected packed type based on source type, tile factors, position
+/// of the inner tiles and permutation of the outer tiled loop.
 ShapedType PackOp::inferPackedType(ShapedType sourceType,
                                    ArrayRef<int64_t> innerTileSizes,
                                    ArrayRef<int64_t> innerDimsPos,
                                    ArrayRef<int64_t> outerDimsPerm) {
   SmallVector<int64_t> resultShape = llvm::to_vector(sourceType.getShape());
   for (const auto &tiledDim : llvm::enumerate(innerDimsPos)) {
     if (ShapedType::isDynamic(resultShape[tiledDim.value()]))
       continue;
Value PackOp::createDestinationTensor(OpBuilder &b, Location loc, Value source,
   if (!outerDimsPerm.empty())
     applyPermutationToVector<OpFoldResult>(mixedSizes, outerDimsPerm);
 
   mixedSizes.append(innerTileSizes.begin(), innerTileSizes.end());
   auto elemType = source.getType().cast<ShapedType>().getElementType();
   return b.create<tensor::EmptyOp>(loc, mixedSizes, elemType);
 }
 
+PackOp PackOp::createTransposedClone(OpBuilder &b, Location loc,
+                                     ArrayRef<int64_t> innerPermutation,
+                                     ArrayRef<int64_t> outerPermutation) {
+  PackOrUnPackTransposeResult metadata = commonPermutationOfPackAndUnPackOp(
+      *this, innerPermutation, outerPermutation);
+  Value transposedDest =
+      createDestinationTensor(b, loc, getSource(), metadata.innerTiles,
+                              metadata.innerDimsPos, metadata.outerDimsPerm);
+  return b.create<PackOp>(loc, getSource(), transposedDest,
+                          metadata.innerDimsPos, metadata.innerTiles,
+                          getPaddingValue(), metadata.outerDimsPerm);
+}
+
 /// Returns true if the tiles and the tiled dims are constant.
 template <typename OpTy>
 bool areTilesAndTiledDimsAllConstant(OpTy op) {
   static_assert(llvm::is_one_of<OpTy, PackOp, UnPackOp>::value,
                 "applies to only pack or unpack operations");
   ShapedType packedType = (std::is_same<OpTy, PackOp>::value)
                               ? op.getDestType()
                               : op.getSourceType();
   SmallVector<OpFoldResult> mixedTiles = op.getMixedTiles();
   for (auto [dimDest, tile] : llvm::zip(
            packedType.getShape().take_back(mixedTiles.size()), mixedTiles)) {
     std::optional<int64_t> constTileSize = getConstantIntValue(tile);
     if (!constTileSize || ShapedType::isDynamic(dimDest))
       return false;
   }
   return true;
 }
 
 Speculation::Speculatability PackOp::getSpeculatability() {
   if (auto paddingValue = getPaddingValue())
     return Speculation::Speculatable;
 
-  // The verifier rejects already operations if we can statically prove that the
-  // sizes of the tiles do not divide perfectly the dimension; thus, check only
-  // to have constant tiles and tiled inner dimensions.
+  // The verifier rejects already operations if we can statically prove that
+  // the sizes of the tiles do not divide perfectly the dimension; thus, check
+  // only to have constant tiles and tiled inner dimensions.
   if (!areTilesAndTiledDimsAllConstant(*this))
     return Speculation::NotSpeculatable;
 
   return Speculation::Speculatable;
 }
 
 // Return true if `inner_dims_pos` and `outer_dims_perm` target the same
 // dimensions for pack and unpack.
void UnPackOp::build(OpBuilder &builder, OperationState &state, Value source,
   dispatchIndexOpFoldResults(innerTiles, dynamicTileSizes, staticTileSizes);
   build(builder, state, dest.getType(), source, dest,
         outerDimsPerm.empty() ? nullptr
                               : builder.getDenseI64ArrayAttr(outerDimsPerm),
         builder.getDenseI64ArrayAttr(innerDimsPos), dynamicTileSizes,
         builder.getDenseI64ArrayAttr(staticTileSizes));
 }
 
+UnPackOp UnPackOp::createTransposedClone(OpBuilder &b, Location loc,
+                                         Value transposedSource,
+                                         ArrayRef<int64_t> innerPermutation,
+                                         ArrayRef<int64_t> outerPermutation) {
+  PackOrUnPackTransposeResult metadata = commonPermutationOfPackAndUnPackOp(
+      *this, innerPermutation, outerPermutation);
+  return b.create<UnPackOp>(loc, transposedSource, getDest(),
+                            metadata.innerDimsPos, metadata.innerTiles,
+                            metadata.outerDimsPerm);
+}
+
 /// pack(unpack(x)) -> x
 LogicalResult UnPackOp::canonicalize(UnPackOp unPackOp,
                                      PatternRewriter &rewriter) {
   PackOp packOp = unPackOp.getSource().getDefiningOp<tensor::PackOp>();
   if (!packOp || packOp.getDestType() != unPackOp.getSourceType())
     return failure();
   if (packOp.getPaddingValue() ||
       !hasSameInnerOuterAttribute(packOp, unPackOp) ||

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

 
 #map = affine_map<(d0, d1) -> (d0, d1)>
 #map1 = affine_map<(d0, d1) -> (d1)>
 #col_reduction_2d_trait = {
   indexing_maps = [#map, #map1],
   iterator_types = ["reduction", "parallel"]
 }
 
-//    CHECK-DAG: #[[$PACKED_MAP_0:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+//    CHECK-DAG: #[[$PACKED_MAP_0:.*]] = affine_map<(d0, d1, d2) -> (d1, d0, d2)>
 //    CHECK-DAG: #[[$PACKED_MAP_1:.*]] = affine_map<(d0, d1, d2) -> (d1)>
 
 //  CHECK-LABEL: @col_reduction_2d_static
 //   CHECK-SAME:   %[[T0:.+]]: tensor<7x3xf16>,
 //   CHECK-SAME:   %[[T1:.+]]: tensor<3xf16>
 func.func @col_reduction_2d_static(%t0: tensor<7x3xf16>, %t1: tensor<3xf16>) -> tensor<3xf16> {
-  //      CHECK:  %[[EMPTY:.*]] = tensor.empty() : tensor<2x3x4xf16>
+  //      CHECK:  %[[EMPTY:.*]] = tensor.empty() : tensor<3x2x4xf16>
   //      CHECK: %[[PACKED:.*]] = tensor.pack %[[T0]] padding_value(%{{.*}} : f16) 
-  // CHECK-SAME:   inner_dims_pos = [0] inner_tiles = [4] into %[[EMPTY]] : tensor<7x3xf16> -> tensor<2x3x4xf16>
+  // CHECK-SAME:   outer_dims_perm = [1, 0] inner_dims_pos = [0] inner_tiles = [4] into %[[EMPTY]] : tensor<7x3xf16> -> tensor<3x2x4xf16>
   //  CHECK-NOT: tensor.pack
   //      CHECK: linalg.generic 
   // CHECK-SAME:   indexing_maps = [#[[$PACKED_MAP_0]], #[[$PACKED_MAP_1]]]
   // CHECK-SAME:   iterator_types = ["reduction", "parallel", "reduction"]
-  // CHECK-SAME:   ins(%{{.*}} : tensor<2x3x4xf16>)
+  // CHECK-SAME:   ins(%{{.*}} : tensor<3x2x4xf16>)
   // CHECK-SAME:  outs(%{{.*}} : tensor<3xf16>)
   %2 = linalg.generic #col_reduction_2d_trait ins(%t0 : tensor<7x3xf16>) outs(%t1 : tensor<3xf16>) {
   ^bb0(%in: f16, %out: f16):
     %3 = arith.addf %in, %out : f16
     linalg.yield %3 : f16
   } -> tensor<3xf16>
 
   //  CHECK-NOT: tensor.unpack
   return %2 : tensor<3xf16>
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["linalg.generic"]} in %arg1
-  transform.structured.pack %0 packed_sizes = [4, 0]
+  %1 = transform.structured.pack %0 packed_sizes = [4, 0]
       : (!pdl.operation) -> (!transform.op<"linalg.generic">)
+  %pack = transform.get_producer_of_operand %1[0] 
+    : (!transform.op<"linalg.generic">) -> (!transform.op<"tensor.pack">)
+  %2, %pack_2, %empty_unpack_2 = 
+    transform.structured.pack_transpose %pack with_compute_op(%1) 
+    outer_perm = [1, 0]
+     : (!transform.op<"tensor.pack">, !transform.op<"linalg.generic">) 
+    -> (!transform.op<"linalg.generic">, !transform.op<"tensor.pack">, !pdl.operation)
 }
 
 // -----
 
 #map = affine_map<(d0, d1) -> (d0, d1)>
 #map1 = affine_map<(d0, d1) -> (d0)>
 #reduction_2d_trait = {
   indexing_maps = [#map, #map1],

 
 // -----
 
 //                                                M   N   K   m   n   k       M   K   m   k
 // CHECK-DAG: #[[$PACKED_MAP_0:.*]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d2, d3, d5)>
 //                                                                            K   N   n   k
 // CHECK-DAG: #[[$PACKED_MAP_1:.*]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d2, d1, d4, d5)>
 //                                                                            M   N   m   n
-// CHECK-DAG: #[[$PACKED_MAP_2:.*]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d3, d4)>
+// CHECK-DAG: #[[$PACKED_MAP_2:.*]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d1, d0, d4, d3)>
 
 // CHECK-LABEL: @matmul
 //  CHECK-SAME:   %[[A:[0-9a-zA-Z]+]]: tensor<?x?xf32>,
 //  CHECK-SAME:   %[[B:[0-9a-zA-Z]+]]: tensor<?x?xf32>,
 //  CHECK-SAME:   %[[C:[0-9a-zA-Z]+]]: tensor<?x?xf32>
 func.func @matmul(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>, %C: tensor<?x?xf32>)
     -> tensor<?x?xf32> {
 
   //      CHECK: %[[PACK_A:.*]] = tensor.pack %{{.*}} inner_dims_pos = [0, 1] inner_tiles = [2, 4]
   // CHECK-SAME:   : tensor<?x?xf32> -> tensor<?x?x2x4xf32>
   //      CHECK: %[[PACK_B:.*]] = tensor.pack %{{.*}} inner_dims_pos = [1, 0] inner_tiles = [3, 4]
   // CHECK-SAME:   : tensor<?x?xf32> -> tensor<?x?x3x4xf32>
-  //      CHECK: %[[PACK_C:.*]] = tensor.pack %{{.*}} inner_dims_pos = [0, 1] inner_tiles = [2, 3]
-  // CHECK-SAME:   : tensor<?x?xf32> -> tensor<?x?x2x3xf32>
+  //      CHECK: %[[PACK_C:.*]] = tensor.pack %{{.*}} outer_dims_perm = [1, 0] inner_dims_pos = [1, 0] inner_tiles = [3, 2]
+  // CHECK-SAME:   : tensor<?x?xf32> -> tensor<?x?x3x2xf32>
 
   //      CHECK: linalg.generic {indexing_maps = [#[[$PACKED_MAP_0]], #[[$PACKED_MAP_1]], #[[$PACKED_MAP_2]]]
   // CHECK-SAME:     iterator_types = ["parallel", "parallel", "reduction", "parallel", "parallel", "reduction"]} 
   // CHECK-SAME:  ins(%{{.*}} : tensor<?x?x2x4xf32>, tensor<?x?x3x4xf32>)
-  // CHECK-SAME: outs(%{{.*}} : tensor<?x?x2x3xf32>)
+  // CHECK-SAME: outs(%{{.*}} : tensor<?x?x3x2xf32>)
   %0 = linalg.matmul  ins(%A, %B: tensor<?x?xf32>, tensor<?x?xf32>)
                      outs(%C: tensor<?x?xf32>)
     -> tensor<?x?xf32>
 
-  //      CHECK: tensor.unpack %{{.*}} inner_dims_pos = [0, 1] inner_tiles = [2, 3]
-  // CHECK-SAME:   : tensor<?x?x2x3xf32> -> tensor<?x?xf32>
+  //      CHECK: tensor.unpack %{{.*}} outer_dims_perm = [1, 0] inner_dims_pos = [1, 0] inner_tiles = [3, 2]
+  // CHECK-SAME:   : tensor<?x?x3x2xf32> -> tensor<?x?xf32>
   return %0 : tensor<?x?xf32>
 }
 
 transform.sequence failures(propagate) {
   ^bb0(%arg1: !pdl.operation):
     %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
     //                                            M  N  K
     %1 = transform.structured.pack %0 packed_sizes = [2, 3, 4]
       : (!pdl.operation) -> (!transform.op<"linalg.generic">)
+
+    %unpack = transform.get_consumers_of_result %1[0] 
+      : (!transform.op<"linalg.generic">) -> (!transform.op<"tensor.unpack">)
+    %2, %pack_2, %unpack_2 = 
+      transform.structured.pack_transpose %unpack with_compute_op(%1) 
+      outer_perm = [1, 0] inner_perm = [1, 0]
+      : (!transform.op<"tensor.unpack">, !transform.op<"linalg.generic">) 
+      -> (!transform.op<"linalg.generic">, !transform.op<"tensor.pack">, !transform.op<"tensor.unpack">)
 }
 
 // -----
 
 //                                                N   F   H   W   C  KH  KW   f   c
 // CHECK-DAG: #[[$PACKED_MAP_0:.*]] = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d0, d4, d2 + d5, d3 + d6, d8)>
 // CHECK-DAG: #[[$PACKED_MAP_1:.*]] = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d1, d4, d5, d6, d7, d8)>
 // CHECK-DAG: #[[$PACKED_MAP_2:.*]] = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d0, d1, d2, d3, d7)>