[mlir][TransformDialect] Simplify the lowering of pack/unpack when these are just pad/unpad

[qcolombet]

Summary

This patch recognizes when tensor.pack/unpack operations are simple tensor.pad/unpad (a.k.a. tensor.extract_slice) and lowers them in a simpler sequence of instruction.

For pack, instead of doing:

pad
expand_shape
transpose

we do

pad
insert_slice

For unpack, instead of doing:

transpose
collapse_shape
extract_slice

we do

extract_slice

Note: returning nullptr for the transform dialect is fine. The related handles are just ignored by the following transformation.

[nicolasvasilache]

can we add a test with dynamic shape ?
(maybe we can't if some op does not support it, but I *think* we can)

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

833

This assumes static, we really want createOrFold<DimOp> here; we should have some helpers around that return the vector<OpFoldResult> with the right quantities.
Not sure where to look offhand for those helpers ..

839

3 newlines ? :)

840

can we early exit here and avoid the else nesting ?

864

with early exit, I believe we can dropp the trailing assert.

933

same re dynamic sizes

998

same re early exit and dropping the trailing assert.

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

3722

Something more idiomatic with ArrayRef == llvm::seq ?

3741

Something like return llvm::all_of(..., [](){ == 1; }) ?

[mravishankar]

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

41

I wondering if we even want to have an insert_slice here? we could just do a reshape here. In which case the only difference is whether there is a transpose or not? If we can teach bufferization to handle this reshape without introducing additional buffers, that'd be very useful (we used to have a memref.reshape to just change the indexing, wonder if we still have that).

[nicolasvasilache]

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

833

Ok, after quick offline chat, I see this is not strictly needed given that the shape is asserted to be static.
This is because tensor.expand_shape is not yet powerful enough re dynamic sizes, but this is a longer endeavor.
I still pulled out the relevant functionality here: https://reviews.llvm.org/D148201 for better reuse.

While it will do the same thing for now using an idiomatic API, it is also future proof.
Could you update to using getMixedDimensions ?

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

41

In principle we could but I'd like to see how these things connect in practice.

At the graph level, reshapes tend to simplify reasonably well.
At the codegen level, as soon as we start tiling and taking slices, reshapes don't behave properly atm.

I'd say this is dependent on where we apply the lowering and we will likely want this to be configurable in the future, based on data.

[qcolombet]

Forgot to submit these comments!

mlir/include/mlir/Dialect/Utils/IndexingUtils.h

208

I ended up not needing this refactoring but I wanted to show it nevertheless.

I'll remove them from that commit after I get your feedbacks.
Do you think it makes sense to expose these in the utility functions (in that case I'll make a separate NFC commit) or should I just drop them all together?

mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h

535

Ditto, not needed but I wanted to show it to you.

[nicolasvasilache]

mlir/include/mlir/Dialect/Utils/IndexingUtils.h

208

I think these are useful utils to expose in general, fine as is IMO.

mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h

535

I think these are useful utils to expose in general, fine as is IMO.

[qcolombet]

• Use all_of and llvm::seq where appropriate
• Use getMixedDimensions to future proof the code w.r.t. dynamic dimensions
• Early-exit to avoid else indentation

[qcolombet]

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

839

It was like that in the exiting code :).
Removed it.

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

41

Alright I'll land as is and we'll revisit later then.

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

let extraClassDeclaration = commonExtraClassDeclaration # [{
     ///   - 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);
+
+    /// Check if this PackOp is like a simple pad operation.
+    /// In other words, this operation:
+    /// 1. adds useless dimensions (dimension of size 1),
+    /// 2. pads the other ones, and
+    /// 3. doesn't shuffle the dimensions
+    bool isLikePad();
   }];
 
   let hasCanonicalizeMethod = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // UnPackOp
 //===----------------------------------------------------------------------===//
let extraClassDeclaration = commonExtraClassDeclaration # [{
     ///     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);
+
+    /// Check if this UnPackOp is like a simple unpad operation.
+    /// In other words, this operation:
+    /// 1. drops useless dimensions (dimension of size 1), and
+    /// 2. reduces dimensions in place (i.e., no tranpose.)
+    bool isLikeUnPad();
   }];
 
   let hasCanonicalizeMethod = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // YieldOp
 //===----------------------------------------------------------------------===//

mlir/include/mlir/Dialect/Utils/IndexingUtils.h

 }
 
 /// Helper method to apply to inverse a permutation.
 SmallVector<int64_t> invertPermutationVector(ArrayRef<int64_t> permutation);
 
 /// Method to check if an interchange vector is a permutation.
 bool isPermutationVector(ArrayRef<int64_t> interchange);
 
+/// Return a permutation vector of size permSize that would result in moving
+/// positions into desiredPositions.
+///
+/// For example, permSize == 5, positions = {2, 4}, desiredPositions = {1, 0}
+/// would result in a {4, 2, 0, 1, 3} permutation vector.
+SmallVector<int64_t>
+computePermutationVector(int64_t permSize, ArrayRef<int64_t> positions,
+                         ArrayRef<int64_t> desiredPositions);
+
 /// Helper to return a subset of `arrayAttr` as a vector of int64_t.
 // TODO: Port everything relevant to DenseArrayAttr and drop this util.
 SmallVector<int64_t> getI64SubArray(ArrayAttr arrayAttr, unsigned dropFront = 0,
                                     unsigned dropBack = 0);
 
 } // namespace mlir
 
 #endif // MLIR_DIALECT_UTILS_INDEXINGUTILS_H

mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h

 ///                                   [1, 1, 1]
 ///    : tensor<?x1x30xf32> to tensor<?x30xf32>
 /// ```
 FailureOr<CollapseShapeRankReducingSliceSimplificationInfo>
 getSimplifyCollapseShapeWithRankReducingSliceInfo(
     RankedTensorType sourceType,
     ArrayRef<ReassociationIndices> reassociationIndices);
 
+struct PackingMetadata {
+  SmallVector<int64_t> insertPositions;
+  SmallVector<ReassociationIndices> reassociations;
+};
+
+/// Given a vector of `positions` indices representing desired packing insertion
+/// points into a target vector (i.e. pack/unpack.inner_dim_pos), compute the
+/// final positions in the target shape as well as the reshape reassociations.
+// Note: This should not be called with a large positions array (or the
+// implementation needs to be updated to use an N.log N sort instead of
+// repeated N^2 counts).
+PackingMetadata computePackingMetadata(int64_t packedRank,
+                                       ArrayRef<int64_t> innerDimPos);
 } // namespace mlir
 
 #endif // MLIR_DIALECT_UTILS_RESHAPEOPSUTILS_H

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

if (op->getNumResults() != 1 || !op->getResult(0).getType().isIndex()) {
       return diag;
     }
     result.push_back(op->getResult(0));
   }
 
   return DiagnosedSilenceableFailure::success();
 }
 
-/// Return a permutation vector of size permSize that would result in moving
-/// positions into desiredPositions.
-///
-/// For example, permSize == 5, positions = {2, 4}, desiredPositions = {1, 0}
-/// would result in a {4, 2, 0, 1, 3} permutation vector.
-static SmallVector<int64_t>
-computePermutationVector(int64_t permSize, ArrayRef<int64_t> positions,
-                         ArrayRef<int64_t> desiredPositions) {
-  SmallVector<int64_t> res(permSize, -1);
-  DenseSet<int64_t> seen;
-  for (auto [pos, desiredPos] : llvm::zip_equal(positions, desiredPositions)) {
-    res[desiredPos] = pos;
-    seen.insert(pos);
-  }
-  int64_t nextPos = 0;
-  for (int64_t &entry : res) {
-    if (entry != -1)
-      continue;
-    while (seen.contains(nextPos))
-      ++nextPos;
-    entry = nextPos;
-    ++nextPos;
-  }
-  return res;
-}
-
-struct PackingMetadata {
-  SmallVector<int64_t> insertPositions;
-  SmallVector<ReassociationIndices> reassociations;
-};
-/// Given a vector of `positions` indices representing desired packing insertion
-/// points into a target vector (i.e. pack/unpack.inner_dim_pos), compute the
-/// final positions in the target shape as well as the reshape reassociations.
-// Note: This should not be called with a large positions array (or the
-// implementation needs to be updated to use an N.log N sort instead of
-// repeated N^2 counts).
-static PackingMetadata computePackingMetadata(int64_t packedRank,
-                                              ArrayRef<int64_t> innerDimPos) {
-  PackingMetadata res;
-  res.insertPositions.reserve(innerDimPos.size());
-  // The pack insert position is the position + the number of previously
-  // inserted positions + offset.
-  // The offset controls whether the packing dimension is the first or last.
-  //
-  // Example
-  // =======
-  // Consider packing from a hypothetical ABCD layout to ABCDba whose
-  // pack.inner_dims is [1, 0]. The first step consists in undoing the
-  // permutation and producing AaBbCD. This is achieved purely by computing the
-  // insert positions of `b` and `a` into `ABCD`, starting from [1, 0]. One
-  // possibility, is to produce insert positions [2, 0], this would result in an
-  // aAbBCD layout (i.e. offset 0). The other possibility, is to produce insert
-  // positions [3, 1], this would result in an AaBbCD layout (i.e. offset 1).
-  // The latter is what we expect from packing.
-  int64_t offset = 1;
-  for (int64_t pos : innerDimPos) {
-    int64_t numInsertedBefore = llvm::count_if(
-        innerDimPos, [&pos](int64_t pos2) { return pos > pos2; });
-    res.insertPositions.push_back(pos + numInsertedBefore + offset);
-  }
-
-  DenseSet<int64_t> posSet(res.insertPositions.begin(),
-                           res.insertPositions.end());
-  res.reassociations.reserve(packedRank);
-  for (int64_t i = 1; i <= packedRank; ++i) {
-    if (!posSet.contains(i)) {
-      res.reassociations.push_back(ReassociationIndices{i - 1});
-      continue;
-    }
-    res.reassociations.push_back(ReassociationIndices{i - 1, i});
-    ++i;
-  }
-  return res;
-}
-
 //===----------------------------------------------------------------------===//
 // BufferizeToAllocationOp
 //===----------------------------------------------------------------------===//
 DiagnosedSilenceableFailure
 transform::BufferizeToAllocationOp::apply(transform::TransformResults &results,
                                           transform::TransformState &state) {
   Attribute memorySpace =
       getMemorySpace().has_value() ? getMemorySpace().value() : Attribute();
LLVM_DEBUG(
           packingMetadata.reassociations, DBGS() << "reassociations: ",
           [&](ReassociationIndices ri) {
             llvm::interleaveComma(ri, llvm::dbgs() << "|");
           });
       DBGSNL();
       llvm::interleaveComma(stripMinedShape, DBGS() << "stripMinedShape: ");
       DBGSNL(); DBGS() << "collapsed type: " << collapsed; DBGSNL(););
 
+  Operation *replacementOp = nullptr;
+  tensor::ExpandShapeOp reshapeOp = nullptr;
+  linalg::TransposeOp transposeOp = nullptr;
+  if (packOp.isLikePad()) {
+    // This pack is just a plain pad.
+    // Just insert the pad in the higher ranked tensor.
+    ArrayRef<int64_t> origShape = collapsed.getShape();
+    auto emptyOp =
+        rewriter.create<tensor::EmptyOp>(loc, packedTensorType, ValueRange{});
+    // offsets.
+    SmallVector<OpFoldResult> zeros(packedRank, rewriter.getIndexAttr(0));
+    // Strides.
+    OpFoldResult one = rewriter.getIndexAttr(1);
+    SmallVector<OpFoldResult> ones(packedRank, one);
+    // The inner dimensions stay the same, but the outer ones are additional 1s.
+    SmallVector<OpFoldResult> sizes(packedRank - origShape.size(), one);
+    for (int64_t dstSize : origShape)
+      sizes.push_back(rewriter.getIndexAttr(dstSize));
+
+    replacementOp = rewriter.create<tensor::InsertSliceOp>(
+        loc, /*source=*/padOp, /*dest=*/emptyOp,
+        /*offsets=*/zeros, sizes,
+        /*strides=*/ones);
+    LLVM_DEBUG(DBGSNL(); DBGSNL(); DBGSNL();
+               DBGS() << "insert_slice op: " << replacementOp; DBGSNL(););
+  } else {
     // 5. Expand from the padded result to the stripMinedShape.
-  auto reshapeOp = rewriter.create<tensor::ExpandShapeOp>(
+    reshapeOp = rewriter.create<tensor::ExpandShapeOp>(
         loc,
         RankedTensorType::Builder(packedTensorType).setShape(stripMinedShape),
         padOp.getResult(), packingMetadata.reassociations);
 
     // 6. Transpose stripMinedShape to packedShape.
-  SmallVector<int64_t> insertPositionsToLastDimsPerm = computePermutationVector(
-      packedRank, packingMetadata.insertPositions, lastDims);
-  auto transposeOp = rewriter.create<linalg::TransposeOp>(
+    SmallVector<int64_t> insertPositionsToLastDimsPerm =
+        computePermutationVector(packedRank, packingMetadata.insertPositions,
+                                 lastDims);
+    transposeOp = rewriter.create<linalg::TransposeOp>(
         loc, reshapeOp.getResult(), packOp.getDest(),
         insertPositionsToLastDimsPerm);
 
-  LLVM_DEBUG(DBGSNL(); DBGSNL(); DBGSNL();
-             DBGS() << "reshape op: " << reshapeOp; DBGSNL();
-             llvm::interleaveComma(insertPositionsToLastDimsPerm,
-                                   DBGS() << "insertPositionsToLastDimsPerm: ");
-             DBGSNL(); DBGS() << "transpose op: " << transposeOp; DBGSNL(););
-
-  // 7. Replace packOp by transposeOp.
-  rewriter.replaceOp(packOp, transposeOp->getResults());
+    LLVM_DEBUG(
+        DBGSNL(); DBGSNL(); DBGSNL(); DBGS() << "reshape op: " << reshapeOp;
+        DBGSNL();
+    replacementOp = transposeOp;
+  }
+  assert(padOp && replacementOp &&
+         (packOp.isLikePad() || (reshapeOp && transposeOp)) &&
+         "If pack is not a pad, all intermediates steps should happen");
+  // Replace packOp by the final replacement op.
+  rewriter.replaceOp(packOp, replacementOp->getResults());
 
   return LowerPackResult{padOp, reshapeOp, transposeOp};
 }
 
 DiagnosedSilenceableFailure transform::LowerPackOp::applyToOne(
     tensor::PackOp target, transform::ApplyToEachResultList &transformResults,
     transform::TransformState &state) {
   IRRewriter rewriter(target->getContext());
return rewriter.notifyMatchFailure(
         unPackOp,
         "non-static shape NYI, needs a more powerful tensor.expand_shape op");
   }
 
   Location loc = unPackOp->getLoc();
   OpBuilder::InsertionGuard g(rewriter);
   rewriter.setInsertionPoint(unPackOp);
 
-  // 2. Compute the permutation vector to move the last `numPackedDims` into the
-  // `innerPosDims` of a shape of rank `packedRank`.
-  int64_t numPackedDims = unPackOp.getInnerDimsPos().size();
   int64_t packedRank = packedTensorType.getRank();
+
+  tensor::EmptyOp emptyOp = nullptr;
+  linalg::TransposeOp transposeOp = nullptr;
+  tensor::CollapseShapeOp reshapeOp = nullptr;
+  tensor::ExtractSliceOp extractSliceOp = nullptr;
+
+  OpFoldResult zero = rewriter.getIndexAttr(0), one = rewriter.getIndexAttr(1);
+  auto destTensorType = unPackOp.getDest().getType().cast<RankedTensorType>();
+  if (unPackOp.isLikeUnPad()) {
+    // This unpack is just a plain pad.
+    // Just extract the pad from the higher ranked tensor.
+    ArrayRef<int64_t> destShape = destTensorType.getShape();
+    SmallVector<OpFoldResult> sizes(packedRank - destShape.size(), one);
+    for (int64_t dstSize : destShape) {
+      sizes.push_back(rewriter.getIndexAttr(dstSize));
+    }
+    extractSliceOp = rewriter.create<tensor::ExtractSliceOp>(
+        loc, destTensorType, unPackOp.getSource(),
+        SmallVector<OpFoldResult>(packedRank, zero), sizes,
+        SmallVector<OpFoldResult>(packedRank, one));
+  } else {
+    // 2. Compute the permutation vector to move the last `numPackedDims` into
+    // the `innerPosDims` of a shape of rank `packedRank`.
+    int64_t numPackedDims = unPackOp.getInnerDimsPos().size();
     auto lastDims = llvm::to_vector(
         llvm::seq<int64_t>(packedRank - numPackedDims, packedRank));
     PackingMetadata packingMetadata =
         computePackingMetadata(packedRank, unPackOp.getInnerDimsPos());
-  SmallVector<int64_t> lastDimsToInsertPositionsPerm = computePermutationVector(
-      packedRank, lastDims, packingMetadata.insertPositions);
+    SmallVector<int64_t> lastDimsToInsertPositionsPerm =
+        computePermutationVector(packedRank, lastDims,
+                                 packingMetadata.insertPositions);
 
-  // 3. Compute the stripMinedShape: this is the packed shape without outer and
-  // inner permutations.
+    // 3. Compute the stripMinedShape: this is the packed shape without outer
+    // and inner permutations.
     SmallVector<int64_t> stripMinedShape(packedTensorType.getShape());
     applyPermutationToVector(stripMinedShape, lastDimsToInsertPositionsPerm);
 
     // 4. Transpose packedShape to stripMinedShape.
     RankedTensorType stripMinedTensorType =
         RankedTensorType::Builder(packedTensorType).setShape(stripMinedShape);
-  RankedTensorType collapsedType = tensor::CollapseShapeOp::inferCollapsedType(
+    RankedTensorType collapsedType =
+        tensor::CollapseShapeOp::inferCollapsedType(
             stripMinedTensorType, packingMetadata.reassociations);
-  auto emptyOp =
-      rewriter.create<tensor::EmptyOp>(loc, stripMinedTensorType, ValueRange{});
-  auto transposeOp = rewriter.create<linalg::TransposeOp>(
+    emptyOp = rewriter.create<tensor::EmptyOp>(loc, stripMinedTensorType,
+                                               ValueRange{});
+    transposeOp = rewriter.create<linalg::TransposeOp>(
         loc, unPackOp.getSource(), emptyOp, lastDimsToInsertPositionsPerm);
 
     LLVM_DEBUG(
-      DBGSNL(); DBGSNL(); llvm::interleaveComma(packingMetadata.insertPositions,
-                                                DBGS() << "insertPositions: ");
+        DBGSNL(); DBGSNL(); llvm::interleaveComma(
+            packingMetadata.insertPositions, DBGS() << "insertPositions: ");
         DBGSNL(); llvm::interleaveComma(packedTensorType.getShape(),
                                         DBGS() << "packedShape: ");
         DBGSNL();
         llvm::interleaveComma(lastDimsToInsertPositionsPerm,
                               DBGS() << "lastDimsToInsertPositionsPerm: ");
         DBGSNL(); llvm::interleaveComma(
             packingMetadata.reassociations, DBGS() << "reassociations: ",
             [&](ReassociationIndices ri) {
               llvm::interleaveComma(ri, llvm::dbgs() << "|");
             });
         DBGSNL();
         llvm::interleaveComma(stripMinedShape, DBGS() << "stripMinedShape: ");
         DBGSNL(); DBGS() << "collapsed type: " << collapsedType; DBGSNL(););
 
     // 5. Collapse from the stripMinedShape to the padded result.
-  auto reshapeOp = rewriter.create<tensor::CollapseShapeOp>(
+    reshapeOp = rewriter.create<tensor::CollapseShapeOp>(
         loc, collapsedType, transposeOp->getResult(0),
         packingMetadata.reassociations);
 
     // 6. ExtractSlice
-  auto destTensorType = unPackOp.getDest().getType().cast<RankedTensorType>();
     int64_t destRank = destTensorType.getRank();
-  OpFoldResult zero = rewriter.getIndexAttr(0), one = rewriter.getIndexAttr(1);
-  auto extractSliceOp = rewriter.create<tensor::ExtractSliceOp>(
+    extractSliceOp = rewriter.create<tensor::ExtractSliceOp>(
         loc, destTensorType, reshapeOp->getResult(0),
         SmallVector<OpFoldResult>(destRank, zero),
         tensor::getMixedSizes(rewriter, loc, unPackOp->getResult(0)),
         SmallVector<OpFoldResult>(destRank, one));
+  }
 
-  // 7. Replace unPackOp by transposeOp.
+  assert(extractSliceOp &&
+         (unPackOp.isLikeUnPad() || (emptyOp && transposeOp && reshapeOp)) &&
+         "If unPackOp is not a pad, all intermediate steps should happen");
+  // Replace unPackOp by extractSliceOp.
   rewriter.replaceOp(unPackOp, extractSliceOp->getResults());
 
   return LowerUnPackOpResult{emptyOp, transposeOp, reshapeOp, extractSliceOp};
 }
 
 DiagnosedSilenceableFailure transform::LowerUnPackOp::applyToOne(
     tensor::UnPackOp target, transform::ApplyToEachResultList &transformResults,
     transform::TransformState &state) {

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

LogicalResult PackOp::canonicalize(PackOp packOp, PatternRewriter &rewriter) {
   if (packOp.getPaddingValue() ||
       !hasSameInnerOuterAttribute(packOp, unPackOp) ||
       !haveSameTiles(packOp, unPackOp))
     return failure();
   rewriter.replaceOp(packOp, unPackOp.getSource());
   return success();
 }
 
+template <typename PackOrUnpackOp>
+static bool isLikePadUnPad(PackOrUnpackOp packOp,
+                           RankedTensorType packedTensorType) {
+  static_assert(std::is_same<PackOrUnpackOp, tensor::PackOp>::value ||
+                    std::is_same<PackOrUnpackOp, tensor::UnPackOp>::value,
+                "Function meant for pack/unpack");
+  // This is a pad if packing only adds ones and we don't transpose dimensions.
+
+  // Check that we are not transposing any dimensions.
+  int64_t expectedDimIdx = 0;
+  ArrayRef<int64_t> innerDimsPos = packOp.getInnerDimsPos();
+  for (int64_t dimIdx : innerDimsPos) {
+    if (dimIdx != expectedDimIdx++) {
+      // This dimension doesn't happen in order.
+      return false;
+    }
+  }
+  int64_t numPackedDims = innerDimsPos.size();
+  ArrayRef<int64_t> packedShape = packedTensorType.getShape();
+  int64_t packedRank = packedTensorType.getRank();
+  // At this point we know that we are taking numPackedDims outer
+  // dimensions and pushing them all the way as the inner most dimensions.
+  // What's left on the outer most dimensions is in this order:
+  // - the factor of the packed dimensions, then
+  // - the untouched dimensions
+  // This shifting inward of dimensions is a no-op (as opposed to a transpose)
+  // if all the dimensions that bubble outerward are ones.
+  // Therefore check that all the dimensions but the numPackedDims inner most
+  // ones are ones.
+  for (int i = 0; i != packedRank - numPackedDims; ++i) {
+    if (packedShape[i] != 1)
+      return false;
+  }
+  return true;
+}
+
+bool PackOp::isLikePad() {
+  auto packedTensorType =
+      (*this)->getResultTypes().front().cast<RankedTensorType>();
+  return isLikePadUnPad(*this, packedTensorType);
+}
+
 //===----------------------------------------------------------------------===//
 // UnPackOp
 //===----------------------------------------------------------------------===//
 
 void UnPackOp::getAsmResultNames(
     function_ref<void(Value, StringRef)> setNameFn) {
   setNameFn(getResult(), "unpack");
 }
LogicalResult UnPackOp::canonicalize(UnPackOp unPackOp,
   if (packOp.getPaddingValue() ||
       !hasSameInnerOuterAttribute(packOp, unPackOp) ||
       !haveSameTiles(packOp, unPackOp))
     return failure();
   rewriter.replaceOp(unPackOp, packOp.getSource());
   return success();
 }
 
+bool UnPackOp::isLikeUnPad() {
+  RankedTensorType packedTensorType = getSourceType();
+  return isLikePadUnPad(*this, packedTensorType);
+}
 //===----------------------------------------------------------------------===//
 // Common Canonicalizers and Folders.
 //===----------------------------------------------------------------------===//
 
 /// Folds a tensor.cast op into a consuming DestinationStyleOpInterface op if
 /// the `tensor.cast` has source that is more static than the consuming op.
 ///
 /// Example:

mlir/lib/Dialect/Utils/IndexingUtils.cpp

bool mlir::isPermutationVector(ArrayRef<int64_t> interchange) {
   for (auto val : interchange) {
     if (seenVals.count(val))
       return false;
     seenVals.insert(val);
   }
   return seenVals.size() == interchange.size();
 }
 
+SmallVector<int64_t>
+mlir::computePermutationVector(int64_t permSize, ArrayRef<int64_t> positions,
+                               ArrayRef<int64_t> desiredPositions) {
+  SmallVector<int64_t> res(permSize, -1);
+  DenseSet<int64_t> seen;
+  for (auto [pos, desiredPos] : llvm::zip_equal(positions, desiredPositions)) {
+    res[desiredPos] = pos;
+    seen.insert(pos);
+  }
+  int64_t nextPos = 0;
+  for (int64_t &entry : res) {
+    if (entry != -1)
+      continue;
+    while (seen.contains(nextPos))
+      ++nextPos;
+    entry = nextPos;
+    ++nextPos;
+  }
+  return res;
+}
+
 SmallVector<int64_t> mlir::getI64SubArray(ArrayAttr arrayAttr,
                                           unsigned dropFront,
                                           unsigned dropBack) {
   assert(arrayAttr.size() > dropFront + dropBack && "Out of bounds");
   auto range = arrayAttr.getAsRange<IntegerAttr>();
   SmallVector<int64_t> res;
   res.reserve(arrayAttr.size() - dropFront - dropBack);
   for (auto it = range.begin() + dropFront, eit = range.end() - dropBack;
        it != eit; ++it)
     res.push_back((*it).getValue().getSExtValue());
   return res;
 }

mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp

if ((*trivialSegments)[groupIdx] ||
       reassociation.clear();
       groupIdx++;
     }
   }
 
   return CollapseShapeRankReducingSliceSimplificationInfo{
       sliceType, newReassociationIndices};
 }
+
+PackingMetadata mlir::computePackingMetadata(int64_t packedRank,
+                                             ArrayRef<int64_t> innerDimPos) {
+  PackingMetadata res;
+  res.insertPositions.reserve(innerDimPos.size());
+  // The pack insert position is the position + the number of previously
+  // inserted positions + offset.
+  // The offset controls whether the packing dimension is the first or last.
+  //
+  // Example
+  // =======
+  // Consider packing from a hypothetical ABCD layout to ABCDba whose
+  // pack.inner_dims is [1, 0]. The first step consists in undoing the
+  // permutation and producing AaBbCD. This is achieved purely by computing the
+  // insert positions of `b` and `a` into `ABCD`, starting from [1, 0]. One
+  // possibility, is to produce insert positions [2, 0], this would result in an
+  // aAbBCD layout (i.e. offset 0). The other possibility, is to produce insert
+  // positions [3, 1], this would result in an AaBbCD layout (i.e. offset 1).
+  // The latter is what we expect from packing.
+  int64_t offset = 1;
+  for (int64_t pos : innerDimPos) {
+    int64_t numInsertedBefore = llvm::count_if(
+        innerDimPos, [&pos](int64_t pos2) { return pos > pos2; });
+    res.insertPositions.push_back(pos + numInsertedBefore + offset);
+  }
+
+  DenseSet<int64_t> posSet(res.insertPositions.begin(),
+                           res.insertPositions.end());
+  res.reassociations.reserve(packedRank);
+  for (int64_t i = 1; i <= packedRank; ++i) {
+    if (!posSet.contains(i)) {
+      res.reassociations.push_back(ReassociationIndices{i - 1});
+      continue;
+    }
+    res.reassociations.push_back(ReassociationIndices{i - 1, i});
+    ++i;
+  }
+  return res;
+}

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

func.func @pack(%arg0: tensor<129x47x16x16xf32>, %arg1: tensor<17x2x16x16x32x8xf32>) -> tensor<17x2x16x16x32x8xf32> {
   // CHECK-SAME:   permutation = [0, 2, 4, 5, 3, 1]
   %pack = tensor.pack %arg0 padding_value(%cst_0 : f32) inner_dims_pos = [1, 0] inner_tiles = [32, 8] into %arg1
     : tensor<129x47x16x16xf32> -> tensor<17x2x16x16x32x8xf32>
   return %pack : tensor<17x2x16x16x32x8xf32>
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%module_op: !pdl.operation):
-  %pack = transform.structured.match ops{["tensor.pack"]} in %module_op 
+  %pack = transform.structured.match ops{["tensor.pack"]} in %module_op
     : (!pdl.operation) -> !transform.op<"tensor.pack">
-  transform.structured.lower_pack %pack : (!transform.op<"tensor.pack">) 
+  transform.structured.lower_pack %pack : (!transform.op<"tensor.pack">)
     -> (!transform.op<"tensor.pad">, !transform.op<"tensor.expand_shape">, !transform.op<"linalg.transpose">)
 }
 
 // -----
 
+// CHECK-LABEL: func.func @pack_as_pad(
+func.func @pack_as_pad(%arg0: tensor<129x47x16x16xf32>, %arg1: tensor<1x1x1x1x136x64x16x16xf32>) -> tensor<1x1x1x1x136x64x16x16xf32> {
+  %cst_0 = arith.constant 0.0 : f32
+
+  // tensor.pack is lowered to tensor.pad + tensor.insert_slice
+  //      CHECK: %[[C0:.*]] = arith.constant 0 : index
+  //      CHECK: %[[PAD:.*]] = tensor.pad {{.*}} low[%[[C0]], %[[C0]], %[[C0]], %[[C0]]]
+  //      CHECK:   : tensor<129x47x16x16xf32> to tensor<136x64x16x16xf32>
+  //      CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<1x1x1x1x136x64x16x16xf32>
+  //      CHECK: %[[RES:.*]] = tensor.insert_slice %[[PAD]] into %[[EMPTY]]
+  // offsets.
+  // CHECK-SAME:   [0, 0, 0, 0, 0, 0, 0, 0]
+  // sizes.
+  // CHECK-SAME:   [1, 1, 1, 1, 136, 64, 16, 16]
+  // strides multipliers.
+  // CHECK-SAME:   [1, 1, 1, 1, 1, 1, 1, 1]
+  // CHECK-SAME:   : tensor<136x64x16x16xf32> into tensor<1x1x1x1x136x64x16x16xf32>
+  //      CHECK: return %[[RES]]
+  %pack = tensor.pack %arg0 padding_value(%cst_0 : f32) inner_dims_pos = [0, 1, 2, 3] inner_tiles = [136, 64, 16, 16] into %arg1
+    : tensor<129x47x16x16xf32> -> tensor<1x1x1x1x136x64x16x16xf32>
+  return %pack :  tensor<1x1x1x1x136x64x16x16xf32>
+}
+
+transform.sequence failures(propagate) {
+^bb1(%module_op: !pdl.operation):
+  %pack = transform.structured.match ops{["tensor.pack"]} in %module_op
+    : (!pdl.operation) -> !transform.op<"tensor.pack">
+  transform.structured.lower_pack %pack : (!transform.op<"tensor.pack">)
+    -> (!transform.op<"tensor.pad">, !transform.op<"tensor.expand_shape">, !transform.op<"linalg.transpose">)
+}
+
+// -----
+
+// Check that we don't lower the following pack as a pad.
+// Although all the outer most dimensions in the resulting shape are 1s,
+// some of the original dimensions are not part of the inner_dims_pos, hence
+// some transpose needs to happen.
+// CHECK-LABEL: func.func @pack_not_a_pad(
+func.func @pack_not_a_pad(%arg0: tensor<129x47x16x16xf32>, %arg1: tensor<1x1x16x16x136x64xf32>) -> tensor<1x1x16x16x136x64xf32> {
+  %cst_0 = arith.constant 0.0 : f32
+
+  //      CHECK: %[[C0:.*]] = arith.constant 0 : index
+  //      CHECK: tensor.pad {{.*}} low[%[[C0]], %[[C0]], %[[C0]], %[[C0]]]
+  //      CHECK:   : tensor<129x47x16x16xf32> to tensor<136x64x16x16xf32>
+  //      CHECK: tensor.expand_shape %{{.*}} [{{.*}}[0, 1], [2, 3], [4], [5]]
+  // CHECK-SAME:   : tensor<136x64x16x16xf32> into tensor<1x136x1x64x16x16xf32>
+  //      CHECK: linalg.transpose
+  // CHECK-SAME:   ins(%{{.*}} : tensor<1x136x1x64x16x16xf32>)
+  // CHECK-SAME:   outs(%{{.*}} : tensor<1x1x16x16x136x64xf32>)
+  // CHECK-SAME:   permutation = [0, 2, 4, 5, 1, 3]
+
+  %pack = tensor.pack %arg0 padding_value(%cst_0 : f32) inner_dims_pos = [0, 1] inner_tiles = [136, 64] into %arg1
+    : tensor<129x47x16x16xf32> -> tensor<1x1x16x16x136x64xf32>
+  return %pack :  tensor<1x1x16x16x136x64xf32>
+}
+
+transform.sequence failures(propagate) {
+^bb1(%module_op: !pdl.operation):
+  %pack = transform.structured.match ops{["tensor.pack"]} in %module_op
+    : (!pdl.operation) -> !transform.op<"tensor.pack">
+  transform.structured.lower_pack %pack : (!transform.op<"tensor.pack">)
+    -> (!transform.op<"tensor.pad">, !transform.op<"tensor.expand_shape">, !transform.op<"linalg.transpose">)
+}
+
+// -----
 // CHECK-LABEL: func.func @unpack(
 func.func @unpack(%arg0: tensor<17x2x16x16x32x8xf32>, %arg1: tensor<129x47x16x16xf32>) -> tensor<129x47x16x16xf32> {
   %cst_0 = arith.constant 0.0 : f32
 
   //      CHECK: tensor.empty() : tensor<17x8x2x32x16x16xf32>
   //      CHECK: linalg.transpose
   // CHECK-SAME:    ins(%{{.*}} : tensor<17x2x16x16x32x8xf32>)
   // CHECK-SAME:   outs(%{{.*}} : tensor<17x8x2x32x16x16xf32>)
   // CHECK-SAME:   permutation = [0, 5, 1, 4, 2, 3]
-  //      CHECK: tensor.collapse_shape {{.*}}[0, 1], [2, 3], [4], [5]] 
+  //      CHECK: tensor.collapse_shape {{.*}}[0, 1], [2, 3], [4], [5]]
   // CHECK-SAME:   : tensor<17x8x2x32x16x16xf32> into tensor<136x64x16x16xf32>
-  //      CHECK: tensor.extract_slice %{{.*}}[0, 0, 0, 0] [129, 47, 16, 16] [1, 1, 1, 1]  
+  //      CHECK: tensor.extract_slice %{{.*}}[0, 0, 0, 0] [129, 47, 16, 16] [1, 1, 1, 1]
   // CHECK-SAME:   : tensor<136x64x16x16xf32> to tensor<129x47x16x16xf32>
   %pack = tensor.unpack %arg0 inner_dims_pos = [1, 0] inner_tiles = [32, 8] into %arg1
     : tensor<17x2x16x16x32x8xf32> -> tensor<129x47x16x16xf32>
   return %pack : tensor<129x47x16x16xf32>
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%module_op: !pdl.operation):
-  %unpack = transform.structured.match ops{["tensor.unpack"]} in %module_op 
+  %unpack = transform.structured.match ops{["tensor.unpack"]} in %module_op
     : (!pdl.operation) -> !transform.op<"tensor.unpack">
-  transform.structured.lower_unpack %unpack : (!transform.op<"tensor.unpack">) 
-    -> (!transform.op<"tensor.empty">, 
+  transform.structured.lower_unpack %unpack : (!transform.op<"tensor.unpack">)
+    -> (!transform.op<"tensor.empty">,
         !transform.op<"linalg.transpose">,
         !transform.op<"tensor.collapse_shape">,
         !transform.op<"tensor.extract_slice">)
 }
+
+// -----
+// When an unpack is a plain 'unpad', lower it to a simple extract_slice.
+// CHECK-LABEL: func.func @unpack_as_pad(
+func.func @unpack_as_pad(%arg0: tensor<1x1x1x1x136x64x16x16xf32>, %arg1: tensor<129x47x16x16xf32>) -> tensor<129x47x16x16xf32> {
+  %cst_0 = arith.constant 0.0 : f32
+
+  // CHECK-SAME: %[[ARG0:[^:]*]]: tensor<1x1x1x1x136x64x16x16xf32>
+  //      CHECK: %[[RES:.*]] = tensor.extract_slice %[[ARG0]]
+  // offsets.
+  // CHECK-SAME:   [0, 0, 0, 0, 0, 0, 0, 0]
+  // sizes.
+  // CHECK-SAME:   [1, 1, 1, 1, 129, 47, 16, 16]
+  // strides multiplers.
+  // CHECK-SAME:   [1, 1, 1, 1, 1, 1, 1, 1]
+  // CHECK-SAME:   : tensor<1x1x1x1x136x64x16x16xf32> to tensor<129x47x16x16xf32>
+  %pack = tensor.unpack %arg0 inner_dims_pos = [0, 1, 2, 3] inner_tiles = [136, 64, 16, 16] into %arg1
+    : tensor<1x1x1x1x136x64x16x16xf32> -> tensor<129x47x16x16xf32>
+  return %pack : tensor<129x47x16x16xf32>
+}
+
+transform.sequence failures(propagate) {
+^bb1(%module_op: !pdl.operation):
+  %unpack = transform.structured.match ops{["tensor.unpack"]} in %module_op
+    : (!pdl.operation) -> !transform.op<"tensor.unpack">
+  transform.structured.lower_unpack %unpack : (!transform.op<"tensor.unpack">)
+    -> (!transform.op<"tensor.empty">,
+        !transform.op<"linalg.transpose">,
+        !transform.op<"tensor.collapse_shape">,
+        !transform.op<"tensor.extract_slice">)
+}