PackAndUnpackPatterns.cpp

Go to the documentation of this file.

//===- FoldIntoPackAndUnpackPatterns.cpp ----------------------------------===//
//
// 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
//
//===----------------------------------------------------------------------===//
 
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/IR/BuiltinTypeInterfaces.h"
#include "mlir/IR/PatternMatch.h"
 
namespace mlir {
namespace linalg {
namespace {
 
/// Returns the number of shape sizes that is either dynamic or greater than 1.
static int64_t getNumGtOneDims(ArrayRef<int64_t> shape) {
  return llvm::count_if(
      shape, [](int64_t v) { return ShapedType::isDynamic(v) || v > 1; });
}
 
/// Returns the index of the first non-unit size in `sizes`. Returns -1 if
/// there are no non-unit sizes.
static int64_t getFirstNonUnitSizeIdx(ArrayRef<int64_t> sizes) {
  const auto *it = llvm::find_if(sizes, [](int64_t dim) { return dim != 1; });
  return (it != sizes.end()) ? std::distance(sizes.begin(), it) : -1;
}
 
/// Check whether `op` is effectively a 1D pack/unpack. Example:
///
///  %pack = linalg.pack %src
///     inner_dims_pos = [0, 1]
///     inner_tiles = [1, 2] into %dest
///   : tensor<1x32xf32> -> tensor<1x16x1x2xf32>
///
/// Returns success() if there is:
///   * only 1 non-unit dim in the un-packed domain,
///   * only 1 non-unit inner tile size, and
///   * the unique non-unit tile size is applied to the unique non-unit
///     un-packed dim.
template <typename PackOrUnpackOp>
static LogicalResult isPackOnEffectively1D(RewriterBase &rewriter,
                                           PackOrUnpackOp *op) {
  // Obtain the unpacked shape.
  auto pack = dyn_cast<linalg::PackOp>(op);
  auto unpack = dyn_cast<linalg::UnPackOp>(op);
 
  ArrayRef<int64_t> unpackedShape = pack ? pack->getSourceType().getShape()
                                         : unpack->getDestType().getShape();
 
  // Obtain the inner tile sizes.
  ArrayRef<int64_t> innerTileSizes = op->getStaticInnerTiles();
 
  // Make sure that there is exactly single non-unit unpacked dim.
  if (getNumGtOneDims(unpackedShape) != 1) {
    return rewriter.notifyMatchFailure(
        *op, "expects non-packed domain to have at most one non-unit dims");
  }
 
  // Make sure that there is at most one non-unit inner tile size.
  auto numNonUnitInnerTiles = getNumGtOneDims(innerTileSizes);
  if (numNonUnitInnerTiles > 1) {
    return rewriter.notifyMatchFailure(
        *op, "expects at most one non-unit inner tiles");
  }
 
  // If there are no non-unit tiles, there is nothing else to check.
  if (numNonUnitInnerTiles == 0)
    return success();
 
  // Get the index of the unique non-unit unpacked dim.
  int64_t nonUnitDimIdx = getFirstNonUnitSizeIdx(unpackedShape);
 
  // Get the index of the dim that the unique non-unit tile is applied to.
  int64_t nonUnitTileDestDimIdx = getFirstNonUnitSizeIdx(innerTileSizes);
 
  // Make sure that the unique non-unit tile is applied to the unique unit dim.
  if (nonUnitTileDestDimIdx != nonUnitDimIdx) {
    return rewriter.notifyMatchFailure(
        *op, "expects at most one non-unit inner tiles");
  }
 
  return success();
}
 
// If the `linalgOp` represents a transpose, return the permutation vector for
// the transpose. Otherwise, return failure.
static FailureOr<SmallVector<int64_t>>
getTransposeOpPermutation(linalg::LinalgOp linalgOp) {
  if (auto transposeOp = dyn_cast<linalg::TransposeOp>(linalgOp.getOperation()))
    return SmallVector<int64_t>(transposeOp.getPermutation());
  if (linalgOp.getNumParallelLoops() != linalgOp.getNumLoops())
    return failure();
 
  if (linalgOp.getNumDpsInputs() != 1 || linalgOp.getNumDpsInits() != 1)
    return failure();
  auto mapRange = linalgOp.getIndexingMapsArray();
  if (!mapRange.front().isPermutation() || !mapRange.back().isPermutation() ||
      mapRange.front() == mapRange.back()) {
    return failure();
  }
  if (!llvm::hasSingleElement(linalgOp.getBlock()->getOperations()))
    return failure();
  AffineMap outMap = mapRange.back();
  AffineMap inMap = mapRange.front();
  // To get the permutation, look at each output index and find which
  // dimension in the input we're reading from for that index.
  return llvm::map_to_vector(outMap.getResults(),
                             [&](AffineExpr expr) -> int64_t {
                               return *inMap.getResultPosition(expr);
                             });
}
 
/// Packing one-dimensional tensor can be expressed as an expand shape op.
struct SimplifyPackToExpandShape : public OpRewritePattern<PackOp> {
  using OpRewritePattern<PackOp>::OpRewritePattern;
 
  FailureOr<Value>
  insertExpand(RewriterBase &rewriter, Location loc, Value operand,
               Type newOperandType,
               ArrayRef<ReassociationIndices> reassociation) const {
    if (operand.getType() == newOperandType)
      return operand;
    return tensor::ExpandShapeOp::create(rewriter, loc, newOperandType, operand,
                                         reassociation)
        .getResult();
  }
 
  /// Returns success() if it is only packing on the innermost dimension.
  LogicalResult isPackOnInnerMostDim(RewriterBase &rewriter,
                                     PackOp packOp) const {
    auto outerDimsPerm = packOp.getOuterDimsPerm();
    if (!outerDimsPerm.empty() && !isIdentityPermutation(outerDimsPerm)) {
      return rewriter.notifyMatchFailure(
          packOp,
          "expects outer_dims_perm is empty or an identity permutation");
    }
 
    int64_t srcRank = packOp.getSourceRank();
    ArrayRef<int64_t> dimsPos = packOp.getInnerDimsPos();
    if (dimsPos.size() != 1 || (dimsPos[0] + 1 != srcRank)) {
      return rewriter.notifyMatchFailure(
          packOp, "expects packing at the innermost dimension");
    }
    return success();
  }
 
  LogicalResult matchAndRewrite(PackOp packOp,
                                PatternRewriter &rewriter) const override {
    if (packOp.getPaddingValue())
      return rewriter.notifyMatchFailure(packOp, "expects no padding value");
    // TODO: Support Memref PackOp. Temporarily return failure.
    if (!packOp.hasPureTensorSemantics())
      return failure();
 
    ShapedType sourceType = packOp.getSourceType();
    if (failed(isPackOnInnerMostDim(rewriter, packOp)) &&
        failed(isPackOnEffectively1D(rewriter, &packOp)) &&
        !packOp.isLikePad()) {
      return failure();
    }
 
    ShapedType destType = packOp.getDestType();
    auto reassociation =
        getReassociationIndicesForReshape(sourceType, destType);
    if (!reassociation)
      return failure();
    FailureOr<Value> expanded =
        insertExpand(rewriter, packOp.getLoc(), packOp.getSource(), destType,
                     *reassociation);
    if (failed(expanded)) {
      return rewriter.notifyMatchFailure(
          packOp, "unable to expand source of tensor.pack");
    }
    rewriter.replaceOp(packOp, *expanded);
    return success();
  }
};
 
struct SimplifyUnPackToCollapseShape : public OpRewritePattern<UnPackOp> {
  using OpRewritePattern<UnPackOp>::OpRewritePattern;
 
  Value insertCollapse(RewriterBase &rewriter, Location loc, Value operand,
                       Type newOperandType, ArrayAttr reassociation) const {
    if (operand.getType() == newOperandType)
      return operand;
    return tensor::CollapseShapeOp::create(rewriter, loc, newOperandType,
                                           operand, reassociation);
  }
 
  /// Returns success() if it is unpacking on the innermost dimension.
  LogicalResult isUnpackOnInnerMostDim(RewriterBase &rewriter,
                                       UnPackOp unpackOp) const {
    auto outerDimsPerm = unpackOp.getOuterDimsPerm();
    if (!outerDimsPerm.empty() && !isIdentityPermutation(outerDimsPerm)) {
      return rewriter.notifyMatchFailure(
          unpackOp,
          "expects outer_dims_perm is empty or an identity permutation");
    }
 
    ShapedType sourceType = unpackOp.getSourceType();
    ShapedType destType = unpackOp.getDestType();
    if (!sourceType.hasStaticShape() || !destType.hasStaticShape())
      return rewriter.notifyMatchFailure(unpackOp, "expects static shapes");
 
    ArrayRef<int64_t> dimsPos = unpackOp.getInnerDimsPos();
    if (dimsPos.size() != 1 || (dimsPos[0] + 1 != destType.getRank())) {
      return rewriter.notifyMatchFailure(
          unpackOp, "expects unpacking on the innermost dimension");
    }
 
    return success();
  }
 
  LogicalResult matchAndRewrite(UnPackOp unpackOp,
                                PatternRewriter &rewriter) const override {
    // TODO: Support Memref UnPackOp. Temporarily return failure.
    if (!unpackOp.hasPureTensorSemantics())
      return failure();
 
    ShapedType destType = unpackOp.getDestType();
    if (failed(isUnpackOnInnerMostDim(rewriter, unpackOp)) &&
        failed(isPackOnEffectively1D(rewriter, &unpackOp)) &&
        !unpackOp.isLikeUnPad()) {
      return failure();
    }
 
    ShapedType sourceType = unpackOp.getSourceType();
    auto reassociation =
        getReassociationIndicesForReshape(sourceType, destType);
    if (!reassociation)
      return failure();
    Value collapsed = insertCollapse(
        rewriter, unpackOp.getLoc(), unpackOp.getSource(), destType,
        getReassociationIndicesAttribute(rewriter, *reassociation));
    rewriter.replaceOp(unpackOp, collapsed);
    return success();
  }
};
 
/// Fold a `pad` -> `pack` into `pack` if they have the same padding values and
/// the pad op has zero low paddings, or if `pack` has no padding values.
struct FoldPadWithPackOp : public OpRewritePattern<PackOp> {
public:
  FoldPadWithPackOp(MLIRContext *context, ControlFoldIntoPackUnpackFn controlFn)
      : OpRewritePattern<PackOp>(context), controlFn(std::move(controlFn)) {}
 
  LogicalResult matchAndRewrite(PackOp packOp,
                                PatternRewriter &rewriter) const override {
    auto padOp = packOp.getSource().getDefiningOp<tensor::PadOp>();
 
    if (!padOp || padOp.getNofold() || !padOp.hasZeroLowPad())
      return failure();
 
    // User controlled folding function.
    if (controlFn && !controlFn(&packOp.getSourceMutable()))
      return failure();
 
    Value constantPaddingValue = padOp.getConstantPaddingValue();
    if (!constantPaddingValue)
      return failure();
 
    if (auto paddingValue = packOp.getPaddingValue())
      if (!isEqualConstantIntOrValue(paddingValue, constantPaddingValue))
        return failure();
 
    // Folding is not allowed if it were to introduce artificial padding.
    // Folding is also disabled in the case of dynamic dimensions and/or tile
    // sizes - that is because it would be impossible to compute the padding
    // size and hence to establish whether "artificial" padding would be
    // created.
    ShapedType unpackedType = packOp.getSourceType();
    SmallVector<int64_t> outerShapeWithoutTranspose =
        getPackedOuterShapeWithoutTransposition(packOp);
    for (auto [pos, tileSize, high] :
         llvm::zip_equal(packOp.getInnerDimsPos(), packOp.getStaticInnerTiles(),
                         padOp.getMixedHighPad())) {
      if (unpackedType.isDynamicDim(pos))
        return failure();
      if (ShapedType::isDynamic(outerShapeWithoutTranspose[pos]))
        return failure();
      if (ShapedType::isDynamic(tileSize))
        return failure();
      std::optional<int64_t> cstHigh = getConstantIntValue(high);
      if (!cstHigh)
        return failure();
      int64_t paddingSize = outerShapeWithoutTranspose[pos] * tileSize -
                            unpackedType.getDimSize(pos);
      // Do not fold the op if it requires artificial padding.
      if (paddingSize + cstHigh.value() >= tileSize)
        return failure();
    }
 
    rewriter.replaceOpWithNewOp<PackOp>(
        packOp, padOp.getSource(), packOp.getDest(), packOp.getInnerDimsPos(),
        packOp.getMixedTiles(), constantPaddingValue,
        packOp.getOuterDimsPerm());
    return success();
  }
 
private:
  ControlFoldIntoPackUnpackFn controlFn;
};
 
/// Fold a `unpack` -> `extract_slice` into the `unpack` since it already
/// has extract_slice semantics.
struct FoldUnpackWithExtractSliceOp
    : public OpRewritePattern<tensor::ExtractSliceOp> {
public:
  FoldUnpackWithExtractSliceOp(MLIRContext *context,
                               ControlFoldIntoPackUnpackFn controlFn)
      : OpRewritePattern<tensor::ExtractSliceOp>(context),
        controlFn(std::move(controlFn)) {}
 
  LogicalResult matchAndRewrite(tensor::ExtractSliceOp sliceOp,
                                PatternRewriter &rewriter) const override {
    auto unpackOp = sliceOp.getSource().getDefiningOp<UnPackOp>();
    if (!unpackOp)
      return failure();
 
    // TODO: Support Memref UnPackOp. Temporarily return failure.
    if (!unpackOp.hasPureTensorSemantics())
      return failure();
 
    // User controlled folding function.
    if (controlFn && !controlFn(&sliceOp.getSourceMutable()))
      return failure();
 
    if (!unpackOp.canFoldSliceOp(sliceOp))
      return failure();
 
    // Create a new empty output tensor.
    Type elementType = unpackOp.getDestType().getElementType();
    Value output = tensor::EmptyOp::create(
        rewriter, sliceOp.getLoc(), sliceOp.getMixedSizes(), elementType);
    rewriter.replaceOpWithNewOp<UnPackOp>(
        sliceOp, unpackOp.getSource(), output, unpackOp.getInnerDimsPos(),
        unpackOp.getMixedTiles(), unpackOp.getOuterDimsPerm());
    return success();
  }
 
private:
  ControlFoldIntoPackUnpackFn controlFn;
};
 
// Applies 'permutation' on 'inVec' and stores the result in resVec.
// 'inVec' may be empty, in that case it's one-to-one mapping with permutation.
// `rank` sets the boundary for permutation i.e., the permutation dim can't be
// greater than the rank specified. If it's so then return false.
// For e.g., permutation {1, 0, 3, 2} with rank 2 is allowed since the values in
// permutation[:rank] doesn't exceed rank, whereas, permutation {1, 3, 0, 2} is
// not allowed since `3` exceeds the value of the rank in the given range.
static bool checkAndPermute(ArrayRef<int64_t> permutation,
                            ArrayRef<int64_t> inVec,
                            SmallVectorImpl<int64_t> &resVec, int64_t rank) {
 
  for (unsigned int i = 0; i < rank; ++i) {
    int64_t remappedPosition = permutation[i];
    if (remappedPosition >= rank)
      return false;
    if (!inVec.empty())
      remappedPosition = inVec[remappedPosition];
    resVec.push_back(remappedPosition);
  }
 
  return true;
}
 
/// Fold 'pack' -> 'transpose' into 'pack' since 'pack' already has transpose
/// semantics.
struct FoldProducerPackWithConsumerLinalgTransposeOp
    : public OpInterfaceRewritePattern<linalg::LinalgOp> {
 
public:
  FoldProducerPackWithConsumerLinalgTransposeOp(
      MLIRContext *context, ControlFoldIntoPackUnpackFn controlFn)
      : OpInterfaceRewritePattern<linalg::LinalgOp>(context),
        controlFn(std::move(controlFn)) {}
 
  LogicalResult matchAndRewrite(linalg::LinalgOp linalgOp,
                                PatternRewriter &rewriter) const override {
    auto packOp = linalgOp->getOperand(0).getDefiningOp<PackOp>();
 
    if (!packOp)
      return failure();
 
    // TODO: Support Memref PackOp. Temporarily return failure.
    if (!packOp.hasPureTensorSemantics())
      return failure();
 
    // User controlled folding function.
    if (controlFn && !controlFn(&linalgOp->getOpOperand(0)))
      return failure();
 
    FailureOr<SmallVector<int64_t>> maybePerm =
        getTransposeOpPermutation(linalgOp);
    if (failed(maybePerm))
      return failure();
 
    auto innerDimsPos = packOp.getInnerDimsPos();
    auto mixedInnerTiles = packOp.getMixedTiles();
    auto outerDimsPerm = packOp.getOuterDimsPerm();
    const auto &transposePerm = maybePerm.value();
    SmallVector<int64_t> newOuterDimsPermVec;
    SmallVector<int64_t> newInnerDimsPosVec;
    SmallVector<OpFoldResult> newMixedInnerTilesVec;
    int64_t srcRank = packOp.getSourceRank();
 
    if (!checkAndPermute(transposePerm, outerDimsPerm, newOuterDimsPermVec,
                         srcRank))
      return rewriter.notifyMatchFailure(
          linalgOp,
          "Cannot fold in tensor.pack if a tile dimension was transposed "
          "with a non-tile dimension in linalg.transpose.");
 
    // Process transpose operation for tiled inner dimensions
    for (unsigned int i = srcRank; i < transposePerm.size(); ++i) {
      int64_t remappedPosition = transposePerm[i] - srcRank;
      newMixedInnerTilesVec.push_back(mixedInnerTiles[remappedPosition]);
      newInnerDimsPosVec.push_back(innerDimsPos[remappedPosition]);
    }
 
    Value output = packOp.createDestinationTensor(
        rewriter, linalgOp.getLoc(), packOp.getSource(), newMixedInnerTilesVec,
        newInnerDimsPosVec, newOuterDimsPermVec);
 
    rewriter.replaceOpWithNewOp<PackOp>(
        linalgOp, packOp.getSource(), output, newInnerDimsPosVec,
        newMixedInnerTilesVec, packOp.getPaddingValue(), newOuterDimsPermVec);
 
    return success();
  }
 
private:
  ControlFoldIntoPackUnpackFn controlFn;
};
 
/// Fold 'transpose' -> 'pack' into 'pack' since 'pack' already has transpose
/// semantics.
struct FoldConsumerPackWithProducerLinalgTransposeOp
    : public OpRewritePattern<PackOp> {
 
public:
  FoldConsumerPackWithProducerLinalgTransposeOp(
      MLIRContext *context, ControlFoldIntoPackUnpackFn controlFn)
      : OpRewritePattern<PackOp>(context), controlFn(std::move(controlFn)) {}
 
  LogicalResult matchAndRewrite(PackOp packOp,
                                PatternRewriter &rewriter) const override {
    // TODO: Support Memref PackOp. Temporarily return failure.
    if (!packOp.hasPureTensorSemantics())
      return failure();
 
    auto linalgOp = packOp.getSource().getDefiningOp<linalg::LinalgOp>();
    if (!linalgOp)
      return failure();
 
    // User controlled folding function.
    if (controlFn && !controlFn(&packOp.getSourceMutable()))
      return failure();
 
    FailureOr<SmallVector<int64_t>> maybePerm =
        getTransposeOpPermutation(linalgOp);
    if (failed(maybePerm))
      return failure();
 
    auto transposePermutation = maybePerm.value();
    auto outerDimsPerm = packOp.getOuterDimsPerm();
    auto innerDimsPos = packOp.getInnerDimsPos();
    SmallVector<int64_t> newInnerDimsPosVec;
    SmallVector<int64_t> newOuterDimsPermVec =
        llvm::to_vector(transposePermutation);
 
    if (!outerDimsPerm.empty())
      applyPermutationToVector(newOuterDimsPermVec, outerDimsPerm);
 
    // Can't use applyPermutationToVector for newInnerDimsPosVec since input and
    // permutation rank won't necessarily be equal in all cases.
    for (auto dim : innerDimsPos)
      newInnerDimsPosVec.push_back(transposePermutation[dim]);
 
    Value output = packOp.createDestinationTensor(
        rewriter, packOp.getLoc(), linalgOp->getOperand(0),
        packOp.getMixedTiles(), newInnerDimsPosVec, newOuterDimsPermVec);
 
    rewriter.replaceOpWithNewOp<PackOp>(
        packOp, linalgOp->getOperand(0), output, newInnerDimsPosVec,
        packOp.getMixedTiles(), packOp.getPaddingValue(), newOuterDimsPermVec);
 
    return success();
  }
 
private:
  ControlFoldIntoPackUnpackFn controlFn;
};
 
/// Fold 'unpack' -> 'transpose' into 'unpack' since 'unpack' already has
/// transpose semantics.
struct FoldProducerUnPackWithConsumerLinalgTransposeOp
    : public OpInterfaceRewritePattern<linalg::LinalgOp> {
 
public:
  FoldProducerUnPackWithConsumerLinalgTransposeOp(
      MLIRContext *context, ControlFoldIntoPackUnpackFn controlFn)
      : OpInterfaceRewritePattern<linalg::LinalgOp>(context),
        controlFn(std::move(controlFn)) {}
 
  LogicalResult matchAndRewrite(linalg::LinalgOp linalgOp,
                                PatternRewriter &rewriter) const override {
    auto unPackOp = linalgOp->getOperand(0).getDefiningOp<UnPackOp>();
 
    if (!unPackOp)
      return failure();
 
    // TODO: Support Memref UnPackOp. Temporarily return failure.
    if (!unPackOp.hasPureTensorSemantics())
      return failure();
 
    // User controlled folding function.
    if (controlFn && !controlFn(&linalgOp->getOpOperand(0)))
      return failure();
 
    FailureOr<SmallVector<int64_t>> maybePerm =
        getTransposeOpPermutation(linalgOp);
    if (failed(maybePerm))
      return failure();
 
    auto outerDimsPerm = unPackOp.getOuterDimsPerm();
    auto innerDimsPos = unPackOp.getInnerDimsPos();
    SmallVector<int64_t> newInnerDimsPosVec;
    SmallVector<int64_t> newOuterDimsPermVec =
        invertPermutationVector(maybePerm.value());
 
    // Can't use applyPermutationToVector for newInnerDimsPosVec since input and
    // permutation rank won't necessarily be equal in all cases.
    for (auto dim : innerDimsPos)
      newInnerDimsPosVec.push_back(newOuterDimsPermVec[dim]);
 
    if (!outerDimsPerm.empty())
      applyPermutationToVector(newOuterDimsPermVec, outerDimsPerm);
 
    // Reuse the destination of the transpose op.
    rewriter.replaceOpWithNewOp<UnPackOp>(
        linalgOp, unPackOp.getSource(), linalgOp.getDpsInits()[0],
        newInnerDimsPosVec, unPackOp.getMixedTiles(), newOuterDimsPermVec);
 
    return success();
  }
 
private:
  ControlFoldIntoPackUnpackFn controlFn;
};
 
/// Fold 'transpose' -> 'unpack' into 'unpack' since 'unpack' already has
/// transpose semantics.
struct FoldConsumerUnPackWithProducerLinalgTransposeOp
    : public OpRewritePattern<UnPackOp> {
  using OpRewritePattern<UnPackOp>::OpRewritePattern;
 
public:
  FoldConsumerUnPackWithProducerLinalgTransposeOp(
      MLIRContext *context, ControlFoldIntoPackUnpackFn controlFn)
      : OpRewritePattern<UnPackOp>(context), controlFn(std::move(controlFn)) {}
 
  LogicalResult matchAndRewrite(UnPackOp unPackOp,
                                PatternRewriter &rewriter) const override {
    // TODO: Support Memref UnPackOp. Temporarily return failure.
    if (!unPackOp.hasPureTensorSemantics())
      return failure();
 
    auto linalgOp = unPackOp.getSource().getDefiningOp<linalg::LinalgOp>();
    if (!linalgOp)
      return failure();
 
    // User controlled folding function.
    if (controlFn && !controlFn(&unPackOp.getSourceMutable()))
      return failure();
 
    FailureOr<SmallVector<int64_t>> maybePerm =
        getTransposeOpPermutation(linalgOp);
    if (failed(maybePerm))
      return failure();
 
    SmallVector<SmallVector<OpFoldResult>> unpackOpResultDims;
    if (failed(reifyResultShapes(rewriter, unPackOp, unpackOpResultDims))) {
      return failure();
    }
 
    SmallVector<int64_t> inverseTransposePerm =
        invertPermutationVector(maybePerm.value());
    auto outerDimsPerm = unPackOp.getOuterDimsPerm();
    auto innerDimsPos = unPackOp.getInnerDimsPos();
    int64_t destRank = unPackOp.getSourceRank() - innerDimsPos.size();
    auto mixedInnerTilesVec = unPackOp.getMixedTiles();
    SmallVector<int64_t> newOuterDimsPermVec;
    SmallVector<int64_t> newInnerDimsPosVec;
    SmallVector<OpFoldResult> newMixedInnerTilesVec;
    if (!checkAndPermute(inverseTransposePerm, outerDimsPerm,
                         newOuterDimsPermVec, destRank))
      return rewriter.notifyMatchFailure(
          unPackOp,
          "Cannot fold in tensor.unpack if a tile dimension was transposed "
          "with a non-tile dimension in linalg.transpose.");
 
    // Process transpose operation for tiled inner dimensions
    for (unsigned int i = destRank; i < inverseTransposePerm.size(); ++i) {
      int64_t remappedPosition = inverseTransposePerm[i] - destRank;
      newMixedInnerTilesVec.push_back(mixedInnerTilesVec[remappedPosition]);
      newInnerDimsPosVec.push_back(innerDimsPos[remappedPosition]);
    }
 
    auto elemType =
        cast<ShapedType>(unPackOp->getResultTypes()[0]).getElementType();
    Value output = tensor::EmptyOp::create(rewriter, unPackOp->getLoc(),
                                           unpackOpResultDims[0], elemType);
 
    rewriter.replaceOpWithNewOp<UnPackOp>(
        unPackOp, linalgOp->getOperand(0), output, newInnerDimsPosVec,
        newMixedInnerTilesVec, newOuterDimsPermVec);
 
    return success();
  }
 
private:
  ControlFoldIntoPackUnpackFn controlFn;
};
 
/// tensor.empty does not define any tensor contents, so an unpadded pack
/// can be folded away.
struct FoldEmptyTensorWithPackOp : public OpRewritePattern<PackOp> {
  using OpRewritePattern<PackOp>::OpRewritePattern;
 
  LogicalResult matchAndRewrite(PackOp packOp,
                                PatternRewriter &rewriter) const override {
    // TODO: Support Memref PackOp. Temporarily return failure.
    if (!packOp.hasPureTensorSemantics())
      return failure();
 
    // Check for tensor.empty source.
    auto emptyOp = packOp.getSource().getDefiningOp<tensor::EmptyOp>();
    if (!emptyOp)
      return failure();
 
    // Check for padding.
    // Packing with padding cannot be simply removed.
    if (packOp.getPaddingValue())
      return rewriter.notifyMatchFailure(packOp, "expects no padding value");
 
    // Replace the pack directly with its destination.
    rewriter.replaceOp(packOp, packOp.getDest());
 
    return success();
  }
};
 
/// tensor.empty does not define any tensor contents, so an unpack
/// can be folded away.
struct FoldEmptyTensorWithUnPackOp : public OpRewritePattern<UnPackOp> {
  using OpRewritePattern<UnPackOp>::OpRewritePattern;
 
  LogicalResult matchAndRewrite(UnPackOp unPackOp,
                                PatternRewriter &rewriter) const override {
    // TODO: Support Memref UnPackOp. Temporarily return failure.
    if (!unPackOp.hasPureTensorSemantics())
      return failure();
 
    // Check for tensor.empty source.
    auto emptyOp = unPackOp.getSource().getDefiningOp<tensor::EmptyOp>();
    if (!emptyOp)
      return failure();
 
    // Replace the unpack directly with its destination.
    rewriter.replaceOp(unPackOp, unPackOp.getDest());
 
    return success();
  }
};
 
} // namespace
 
void populateFoldIntoPackAndUnpackPatterns(
    RewritePatternSet &patterns, const ControlFoldIntoPackUnpackFn &controlFn) {
  patterns.insert<FoldUnpackWithExtractSliceOp, FoldPadWithPackOp,
                  FoldProducerPackWithConsumerLinalgTransposeOp,
                  FoldConsumerPackWithProducerLinalgTransposeOp,
                  FoldConsumerUnPackWithProducerLinalgTransposeOp,
                  FoldProducerUnPackWithConsumerLinalgTransposeOp>(
      patterns.getContext(), controlFn);
}
 
void populateSimplifyPackAndUnpackPatterns(RewritePatternSet &patterns) {
  patterns.add<SimplifyPackToExpandShape, SimplifyUnPackToCollapseShape>(
      patterns.getContext());
}
 
void populateFoldPackUnpackIntoTensorEmptyPatterns(
    RewritePatternSet &patterns) {
  patterns.add<FoldEmptyTensorWithPackOp, FoldEmptyTensorWithUnPackOp>(
      patterns.getContext());
}
 
} // namespace linalg
} // namespace mlir