// Copyright 2023 The IREE Authors
//
// Licensed 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

//===- PropagateLinalgTranspose.cpp - Pass to propagate transposes ---------==//
//
// The pass is to propagate linalg.transpose operations through a restricted
// set of operations based on a set of local propagation decisions.
//
//===----------------------------------------------------------------------===//

#include "iree/compiler/Dialect/Flow/Conversion/TensorToFlow/Utils.h"
#include "iree/compiler/Dialect/Flow/Transforms/RegionOpUtils.h"
#include "iree/compiler/Dialect/LinalgExt/Transforms/Transforms.h"
#include "iree/compiler/Dialect/LinalgExt/Utils/Utils.h"
#include "iree/compiler/Dialect/Util/IR/UtilOps.h"
#include "iree/compiler/GlobalOptimization/Passes.h"
#include "llvm/Support/Debug.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
#include "mlir/Dialect/Linalg/Utils/Utils.h"
#include "mlir/Dialect/MemRef/Transforms/Transforms.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Tensor/Transforms/Transforms.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"

#define DEBUG_TYPE "iree-global-opt-propagate-linalg-transpose"

namespace mlir::iree_compiler::GlobalOptimization {

#define GEN_PASS_DEF_PROPAGATELINALGTRANSPOSEPASS
#include "iree/compiler/GlobalOptimization/Passes.h.inc"

//===----------------------------------------------------------------------===//
// Transpose permutation helpers
//===----------------------------------------------------------------------===//

static bool isIdentityPermutation(ArrayRef<int64_t> perm) {
  for (auto [index, dim] : llvm::enumerate(perm)) {
    if (index != dim) {
      return false;
    }
  }
  return true;
}

// Constructs a transpose of the given tensor and permutation.
static Value createTransposeInit(OpBuilder &builder, Value source,
                                 ArrayRef<int64_t> perm) {
  SmallVector<OpFoldResult> mixedSizes =
      tensor::getMixedSizes(builder, source.getLoc(), source);
  applyPermutationToVector(mixedSizes, perm);
  Type elemType = cast<RankedTensorType>(source.getType()).getElementType();
  Value empty =
      tensor::EmptyOp::create(builder, source.getLoc(), mixedSizes, elemType)
          .getResult();
  return empty;
}

// Constructs a transpose of the given tensor and permutation,
// or produces a transposed version of the producing tensor.empty op.
static Value createTranspose(OpBuilder &builder, Value source,
                             ArrayRef<int64_t> perm) {
  if (auto empty = source.getDefiningOp<tensor::EmptyOp>()) {
    Type elementType = empty.getType().getElementType();
    SmallVector<OpFoldResult> mixedSizes = empty.getMixedSizes();
    applyPermutationToVector(mixedSizes, perm);
    return tensor::EmptyOp::create(builder, empty.getLoc(), mixedSizes,
                                   elementType);
  }
  Value empty = createTransposeInit(builder, source, perm);
  return linalg::TransposeOp::create(builder, source.getLoc(), source, empty,
                                     perm)
      ->getResult(0);
}

static RankedTensorType getPermutedTensorType(RankedTensorType type,
                                              SmallVector<int64_t> perm) {
  SmallVector<int64_t> permutedShape = applyPermutation(type.getShape(), perm);
  return RankedTensorType::get(permutedShape, type.getElementType());
}

static bool isReshapeBlockingFusion(Operation *producer, Operation *consumer) {
  auto isFusableOp = [](Operation *op) {
    if (!op) {
      return false;
    }
    return isa_and_nonnull<linalg::LinalgDialect,
                           IREE::LinalgExt::IREELinalgExtDialect,
                           tensor::TensorDialect>(op->getDialect());
  };
  return isFusableOp(producer) && isFusableOp(consumer);
}

//===----------------------------------------------------------------------===//
// Transpose specialization
//===----------------------------------------------------------------------===//

// Specializes linalg.generic op to linalg.transpose if it is transposing a
// single input.
static void specializeGenericTransposeOp(RewriterBase &rewriter,
                                         linalg::GenericOp genericOp) {
  if (!IREE::LinalgExt::isaTransposeOpInterface(genericOp)) {
    return;
  }

  auto mapRange = genericOp.getIndexingMapsArray();
  AffineMap outMap = mapRange.back();
  AffineMap inMap = mapRange.front();
  SmallVector<int64_t> perm;
  // To get the permutation, look at each output index and find which
  // dimension in the input we're reading from for that index.
  for (AffineExpr expr : outMap.getResults()) {
    perm.push_back(*inMap.getResultPosition(expr));
  }
  rewriter.replaceOpWithNewOp<linalg::TransposeOp>(
      genericOp, genericOp.getDpsInputs()[0], genericOp.getDpsInits()[0], perm);
}

//===----------------------------------------------------------------------===//
// Other pattern helpers
//===----------------------------------------------------------------------===//

/// Returns the `op` if it is a linalg::GenericOp. If it is a named op and
/// `allowGeneralizing` is true, returns the generalized op. Otherwise, returns
/// failure.
/// TODO: Due to fragility around handling of convolutions, convolution
/// propagation is behind a flag.
static FailureOr<linalg::GenericOp>
getAllowedGenericOpOrGeneralizeNamedOp(RewriterBase &rewriter, Operation *op,
                                       bool allowGeneralizing, bool convProp) {
  auto linalgOp = dyn_cast<linalg::LinalgOp>(op);
  if (!linalgOp) {
    return failure();
  }

  if (!convProp && linalg::isaConvolutionOpInterface(linalgOp)) {
    return failure();
  }

  if (!isa<linalg::GenericOp>(linalgOp) &&
      !(allowGeneralizing && linalg::isaContractionOpInterface(linalgOp)) &&
      !(convProp && linalg::isaConvolutionOpInterface(linalgOp))) {
    return failure();
  }

  auto genericOp = dyn_cast<linalg::GenericOp>(op);
  if (genericOp) {
    return genericOp;
  }
  OpBuilder::InsertionGuard guard(rewriter);
  rewriter.setInsertionPoint(linalgOp);
  return linalg::generalizeNamedOp(rewriter, linalgOp);
}

//===----------------------------------------------------------------------===//
// Transpose Bubbling Patterns
//===----------------------------------------------------------------------===//

namespace {

// Fuses a transpose with the init of a linalg.generic op or contraction op.
// Contraction ops are generalized and then treated as a generic. For example,
//
// linalg.generic {
//   indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d0, d2)>,
//                    affine_map<(d0, d1, d2) -> (d0, d1, d2)>]
//   ins(%0 : tensor<2x7x5>) outs(%1 : tensor<7x2x5>)
//
// %2 = linalg.transpose ... permutation = [0, 2, 1] :
//                           tensor<7x2x5> -> tensor<7x5x2>
// Becomes
//
// linalg.generic {
//   indexing_maps = [affine_map<(d0, d1, d2) -> (d2, d0, d1)>,
//                    affine_map<(d0, d1, d2) -> (d0, d1, d2)>]
//   ins(%0 : tensor<2x7x5>) outs(%3 : tensor<7x5x2>)
class FuseTransposeWithProducerLinalgOp
    : public OpRewritePattern<linalg::TransposeOp> {
public:
  using Base::Base;
  FuseTransposeWithProducerLinalgOp(MLIRContext *ctx, bool aggressiveProp,
                                    bool convProp, PatternBenefit b = 1)
      : OpRewritePattern<linalg::TransposeOp>(ctx, b),
        allowGeneralizing(aggressiveProp), convProp(convProp) {}

  LogicalResult matchAndRewrite(linalg::TransposeOp transposeOp,
                                PatternRewriter &rewriter) const override {
    if (!IREE::Flow::isNonNullAndOutsideDispatch(transposeOp)) {
      return failure();
    }
    OpResult result = dyn_cast<OpResult>(transposeOp.getInput());
    if (!result) {
      return rewriter.notifyMatchFailure(
          transposeOp, "transpose input defined by block argument");
    }
    if (!result.hasOneUse()) {
      return rewriter.notifyMatchFailure(transposeOp,
                                         "multi use transpose input");
    }
    auto linalgOp = dyn_cast<linalg::LinalgOp>(result.getOwner());
    if (!linalgOp) {
      return rewriter.notifyMatchFailure(
          transposeOp, "non-linalg op producer for transpose input");
    }

    int64_t resultIndex = result.getResultNumber();
    auto maybeGenericOp = getAllowedGenericOpOrGeneralizeNamedOp(
        rewriter, result.getOwner(), allowGeneralizing, convProp);
    if (failed(maybeGenericOp)) {
      return rewriter.notifyMatchFailure(
          transposeOp, "linalg op producer is not generic or contraction");
    }

    auto genericOp = maybeGenericOp.value();
    result = genericOp->getOpResult(resultIndex);

    ArrayRef<int64_t> perm = transposeOp.getPermutation();
    auto invPerm = invertPermutationVector(perm);

    // 1. Get the transposed init of the generic.
    Value init = genericOp.getDpsInits()[resultIndex];
    SmallVector<Value> inits = genericOp.getDpsInits();
    Value newInit = createTranspose(rewriter, init, perm);
    inits[resultIndex] = newInit;

    SmallVector<Type> resultTypes(genericOp->getResultTypes());
    resultTypes[resultIndex] = newInit.getType();

    // 2. Update the indexing map of the transposed init operand by permuting
    // the results of the map.
    SmallVector<AffineMap> newIndexingMaps = genericOp.getIndexingMapsArray();
    AffineMap resultMap =
        newIndexingMaps[genericOp.getNumDpsInputs() + resultIndex];
    SmallVector<AffineExpr> newExprs =
        applyPermutation(resultMap.getResults(), perm);
    AffineMap transposedMap =
        AffineMap::get(resultMap.getNumDims(), resultMap.getNumSymbols(),
                       newExprs, rewriter.getContext());
    newIndexingMaps[genericOp.getNumDpsInputs() + resultIndex] = transposedMap;

    // 3. Create the new generic with the same iteration order.
    auto newGenericOp = linalg::GenericOp::create(
        rewriter, genericOp.getLoc(), resultTypes, genericOp.getDpsInputs(),
        newInit, newIndexingMaps, genericOp.getIteratorTypesArray(),
        /*bodyBuild=*/nullptr, linalg::getPrunedAttributeList(genericOp));
    rewriter.cloneRegionBefore(genericOp.getRegion(), newGenericOp.getRegion(),
                               newGenericOp.getRegion().begin());

    // 4. Remap iteration space of the generic to match the dimension order of
    // the output.
    if (newGenericOp.getNumResults() == 1) {
      SmallVector<unsigned int> interchange;
      int64_t permIdx = 0;
      for (int i = 0, e = transposedMap.getNumDims(); i < e; ++i) {
        if (transposedMap.isFunctionOfDim(i)) {
          interchange.push_back(
              cast<AffineDimExpr>(transposedMap.getResult(permIdx))
                  .getPosition());
          permIdx++;
          continue;
        }
        interchange.push_back(i);
      }
      auto interchangedGenericOp =
          linalg::interchangeGenericOp(rewriter, newGenericOp, interchange);
      // Interchange only fails if interchangeGenericOpPrecondition fails, which
      // only fails if the interchange vector is not invertible or doesn't match
      // the number of loops in the generic, both of which are guaranteed by
      // the fact that the output map must be a projection in the above
      // construction.
      assert(succeeded(interchangedGenericOp) &&
             "failed to interchange transposed generic");
      newGenericOp = *interchangedGenericOp;
    }

    // 5. Replace the result of the transpose with the transposed init.
    rewriter.replaceOp(transposeOp, newGenericOp->getResult(resultIndex));
    for (auto [oldRes, newRes] :
         llvm::zip_equal(genericOp.getResults(), newGenericOp->getResults())) {
      if (oldRes.getResultNumber() == resultIndex) {
        continue;
      }
      rewriter.replaceAllUsesWith(oldRes, newRes);
    }
    return success();
  }

private:
  bool allowGeneralizing = false;
  bool convProp = false;
};

// Bubbles a transpose through a tensor.collapse_shape.
class BubbleTransposeThroughCollapseShape
    : public OpRewritePattern<linalg::TransposeOp> {
public:
  using Base::Base;
  BubbleTransposeThroughCollapseShape(MLIRContext *ctx,
                                      bool enableEdgeReshapeProp,
                                      PatternBenefit b = 1)
      : OpRewritePattern<linalg::TransposeOp>(ctx, b),
        enableEdgeReshapePropagation(enableEdgeReshapeProp) {}

  LogicalResult matchAndRewrite(linalg::TransposeOp transposeOp,
                                PatternRewriter &rewriter) const override {
    if (!IREE::Flow::isNonNullAndOutsideDispatch(transposeOp)) {
      return failure();
    }
    Value source = transposeOp.getDpsInputOperand(0)->get();
    auto collapseOp = source.getDefiningOp<tensor::CollapseShapeOp>();
    // Do not propagate through reshapes if the transpose has multiple users, as
    // this could end up duplicating the transposes. We should only propagate
    // through reshape when it is free to do so.
    if (!collapseOp || !collapseOp->hasOneUse()) {
      return rewriter.notifyMatchFailure(
          transposeOp, "transpose input is not a single-use collapse shape");
    }

    if (!enableEdgeReshapePropagation &&
        !isReshapeBlockingFusion(transposeOp,
                                 collapseOp.getSrc().getDefiningOp())) {
      return rewriter.notifyMatchFailure(transposeOp,
                                         "transpose not blocking fusion");
    }

    SmallVector<ReassociationIndices> reassociations =
        collapseOp.getReassociationIndices();

    // Because we are doing transpose(collapse_shape), all expanded groups are
    // transposed together. As a result, to get the permutation of the new
    // transpose, we can just flatten the transposed reassociation indices.
    // For example,
    //
    // reassociation_map = [[0, 1, 2], [3], [4, 5]]
    // permutation = [1, 2, 0]
    //
    // Becomes
    //
    // permutation = [3, 4, 5, 0, 1, 2]
    // reassociation_map = [[0], [1, 2], [3, 4, 5]]
    applyPermutationToVector(reassociations, transposeOp.getPermutation());

    SmallVector<int64_t> newPerm;
    SmallVector<ReassociationIndices> newReassociations;
    int64_t expandedDim = 0;
    for (auto reassoc : reassociations) {
      ReassociationIndices newReassoc;
      for (auto dim : reassoc) {
        newPerm.push_back(dim);
        newReassoc.push_back(expandedDim++);
      }
      newReassociations.push_back(newReassoc);
    }

    Value newTranspose =
        createTranspose(rewriter, collapseOp.getSrc(), newPerm);
    Value newReshape = tensor::CollapseShapeOp::create(
        rewriter, collapseOp.getLoc(), transposeOp.getResultTypes()[0],
        newTranspose, newReassociations);
    rewriter.replaceOp(transposeOp, newReshape);
    return success();
  }

private:
  bool enableEdgeReshapePropagation = true;
};

} // namespace

//===----------------------------------------------------------------------===//
// Transpose Sinking Patterns
//===----------------------------------------------------------------------===//

namespace {

// Combines two transposes into one. This shouldn't be strictly necessary as
// fusion should cancel inverse transposes, but doing this here can open up
// new propagation opportunities and eases the analysis in fusion/later passes.
class ComposeTransposes : public OpRewritePattern<linalg::TransposeOp> {
public:
  using Base::Base;

  LogicalResult matchAndRewrite(linalg::TransposeOp consumer,
                                PatternRewriter &rewriter) const override {
    if (!IREE::Flow::isNonNullAndOutsideDispatch(consumer)) {
      return failure();
    }
    Value input = consumer.getInput();
    auto producer = input.getDefiningOp<linalg::TransposeOp>();
    if (!producer) {
      return failure();
    }

    ArrayRef<int64_t> producerPerm = producer.getPermutation();
    ArrayRef<int64_t> consumerPerm = consumer.getPermutation();
    SmallVector<int64_t> composedPerm =
        applyPermutation(producerPerm, consumerPerm);

    Value transposedSource = producer.getInput();
    if (!isIdentityPermutation(composedPerm)) {
      transposedSource =
          createTranspose(rewriter, transposedSource, composedPerm);
    }
    rewriter.replaceOp(consumer, transposedSource);
    return success();
  }
};

// Sinks a transpose through a tensor.extract_slice iff the transpose turns
// the extracted slice into a contiguous slice.
class SinkTransposeThroughExtractSlice
    : public OpRewritePattern<tensor::ExtractSliceOp> {
public:
  using Base::Base;

  LogicalResult matchAndRewrite(tensor::ExtractSliceOp extractOp,
                                PatternRewriter &rewriter) const override {
    if (!IREE::Flow::isNonNullAndOutsideDispatch(extractOp)) {
      return failure();
    }
    Value source = extractOp.getSource();
    auto transposeOp = source.getDefiningOp<linalg::TransposeOp>();
    if (!transposeOp) {
      return failure();
    }

    // Applying `perm` takes a list from the pre-transpose ordering to the
    // post-transpose ordering.
    ArrayRef<int64_t> perm = transposeOp.getPermutation();
    // Applying `invPerm` takes a list from the post-transpose ordering to the
    // pre-transpose ordering. Sinking a transpose through an op is largely a
    // matter of rewriting it in pre-transpose space, and thus just applying
    // the inverse permutation.
    auto invPerm = invertPermutationVector(perm);

    SmallVector<OpFoldResult> offsets = extractOp.getMixedOffsets();
    SmallVector<OpFoldResult> sizes = extractOp.getMixedSizes();
    SmallVector<OpFoldResult> strides = extractOp.getMixedStrides();
    ArrayRef<int64_t> srcShape = extractOp.getSourceType().getShape();

    // Permute the offsets, sizes, and strides to pre-transpose ordering.
    applyPermutationToVector(offsets, invPerm);
    applyPermutationToVector(sizes, invPerm);
    applyPermutationToVector(strides, invPerm);
    SmallVector<int64_t> baseShape = applyPermutation(srcShape, invPerm);

    // Check if the resulting offsets, sizes, and strides correspond to a
    // contiguous slice and can thus be mappable to a `flow.tensor.update` op.
    // This should always be worth doing because this can remove a dispatch for
    // the slice, and the transpose is on the slice rather than the full tensor.
    if (!IREE::Flow::isOffsetSizeAndStrideMappableToFlow(offsets, sizes,
                                                         strides, baseShape)) {
      return rewriter.notifyMatchFailure(
          extractOp, "transposed slice not mappable to flow ops");
    }

    ArrayRef<int64_t> staticSizes = extractOp.getStaticSizes();
    ArrayRef<int64_t> sliceShape = extractOp.getResultType().getShape();
    std::optional<llvm::SmallDenseSet<unsigned>> maybeRankReducingMask =
        mlir::computeRankReductionMask(staticSizes, sliceShape);
    if (!maybeRankReducingMask) {
      return rewriter.notifyMatchFailure(
          extractOp, "failed to compute rank reducing mask");
    }
    llvm::SmallDenseSet<unsigned> rankReducingMask = *maybeRankReducingMask;

    // Find rank reducing map in the pre-transposed domain.
    int64_t dim = 0;
    llvm::SmallDenseMap<int64_t, int64_t> rankReducedMap;
    // Since `dim` is in the pre-transposed domain, and is incrementing each
    // iteration, `idx` must also be in the pre-transposed domain.
    for (int64_t idx = 0, e = invPerm.size(); idx < e; ++idx) {
      // Get index in the transposed domain, since `rankReducingMask` is in
      // the transposed domain.
      if (!rankReducingMask.contains(invPerm[idx])) {
        // Domain of `rankReducedMap` is in pre-transposed domain.
        rankReducedMap[idx] = dim++;
      }
    }

    // Compute the new permutation by dropping all rank-reduced dimensions.
    SmallVector<int64_t> rankReducedPerm;
    for (int64_t i : perm) {
      if (rankReducedMap.contains(i)) {
        rankReducedPerm.push_back(rankReducedMap[i]);
      }
    }

    auto rankReducedInvPerm = invertPermutationVector(rankReducedPerm);

    RankedTensorType sliceType = getPermutedTensorType(
        cast<RankedTensorType>(extractOp.getType()), rankReducedInvPerm);
    Value slice = tensor::ExtractSliceOp::create(
        rewriter, extractOp.getLoc(), sliceType, transposeOp.getInput(),
        offsets, sizes, strides);
    // Transpose back to the original slice.
    if (!isIdentityPermutation(rankReducedPerm)) {
      slice = createTranspose(rewriter, slice, rankReducedPerm);
    }
    rewriter.replaceOp(extractOp, slice);
    return success();
  }
};

// Sinks a transpose through a tensor.expand_shape.
class SinkTransposeThroughExpandShape
    : public OpRewritePattern<tensor::ExpandShapeOp> {
public:
  using Base::Base;
  SinkTransposeThroughExpandShape(MLIRContext *ctx, bool enableEdgeReshapeProp,
                                  PatternBenefit b = 1)
      : OpRewritePattern<tensor::ExpandShapeOp>(ctx, b),
        enableEdgeReshapePropagation(enableEdgeReshapeProp) {}

  LogicalResult matchAndRewrite(tensor::ExpandShapeOp expandOp,
                                PatternRewriter &rewriter) const override {
    if (!IREE::Flow::isNonNullAndOutsideDispatch(expandOp)) {
      return failure();
    }
    Value source = expandOp.getSrc();
    auto transposeOp = source.getDefiningOp<linalg::TransposeOp>();
    // Do not propagate through reshapes if the transpose has multiple users, as
    // this could end up duplicating the transposes. We should only propagate
    // through reshape when it is free to do so.
    if (!transposeOp || !transposeOp->hasOneUse()) {
      return rewriter.notifyMatchFailure(
          expandOp, "expand shape input is not a single-use transpose");
    }

    if (!enableEdgeReshapePropagation &&
        llvm::none_of(expandOp->getUsers(), [&](Operation *consumer) {
          return isReshapeBlockingFusion(transposeOp, consumer);
        })) {
      return rewriter.notifyMatchFailure(transposeOp,
                                         "transpose not blocking fusion");
    }

    auto invPerm = invertPermutationVector(transposeOp.getPermutation());
    SmallVector<ReassociationIndices> reassociations =
        expandOp.getReassociationIndices();

    // Because we are doing expand_shape(transpose), all expanded groups are
    // transposed together. As a result, to get the permutation of the new
    // transpose, we can just flatten the transposed reassociation indices.
    // For example,
    //
    // permutation = [0, 2, 1]
    // reassociation_map = [[0, 1, 2], [3], [4, 5]]
    //
    // Becomes
    //
    // reassociation_map = [[0, 1, 2], [3, 4], [5]]
    // permutation = [0, 1, 2, 4, 5, 3]
    applyPermutationToVector(reassociations, invPerm);

    SmallVector<int64_t> newInvPerm;
    SmallVector<ReassociationIndices> newReassociations;
    int64_t expandedDim = 0;
    for (auto reassoc : reassociations) {
      ReassociationIndices newReassoc;
      for (auto dim : reassoc) {
        newInvPerm.push_back(dim);
        newReassoc.push_back(expandedDim++);
      }
      newReassociations.push_back(newReassoc);
    }

    auto newPerm = invertPermutationVector(newInvPerm);

    RankedTensorType expandedType = getPermutedTensorType(
        cast<RankedTensorType>(expandOp.getType()), newInvPerm);
    // Permute the output shape to match the permuted type.
    SmallVector<OpFoldResult> permutedOutputShape =
        expandOp.getMixedOutputShape();
    applyPermutationToVector(permutedOutputShape, newInvPerm);
    Value transposedReshape = tensor::ExpandShapeOp::create(
        rewriter, expandOp.getLoc(), expandedType, transposeOp.getInput(),
        newReassociations, permutedOutputShape);
    Value originalReshape =
        createTranspose(rewriter, transposedReshape, newPerm);
    rewriter.replaceOp(expandOp, originalReshape);
    return success();
  }

private:
  bool enableEdgeReshapePropagation = true;
};

// Sinks a transpose through a tensor.pad.
class SinkTransposeThroughPad : public OpRewritePattern<tensor::PadOp> {
public:
  using Base::Base;

  LogicalResult matchAndRewrite(tensor::PadOp padOp,
                                PatternRewriter &rewriter) const override {
    if (!IREE::Flow::isNonNullAndOutsideDispatch(padOp)) {
      return failure();
    }
    Value source = padOp.getSource();
    auto transposeOp = source.getDefiningOp<linalg::TransposeOp>();
    if (!transposeOp) {
      return failure();
    }

    Block &block = padOp.getRegion().front();
    if (llvm::any_of(block.getArguments(), [](BlockArgument blockArg) {
          return blockArg.getNumUses();
        })) {
      return failure();
    }

    auto invPerm = invertPermutationVector(transposeOp.getPermutation());
    SmallVector<OpFoldResult> lowSizes = padOp.getMixedLowPad();
    SmallVector<OpFoldResult> highSizes = padOp.getMixedHighPad();
    applyPermutationToVector(lowSizes, invPerm);
    applyPermutationToVector(highSizes, invPerm);

    RankedTensorType oldPaddedType = cast<RankedTensorType>(padOp.getType());
    RankedTensorType newPaddedType = oldPaddedType.clone(
        applyPermutation(oldPaddedType.getShape(), invPerm));

    auto newPadOp = tensor::PadOp::create(
        rewriter, padOp.getLoc(), newPaddedType, transposeOp.getInput(),
        lowSizes, highSizes, padOp.getNofold());
    rewriter.cloneRegionBefore(padOp.getRegion(), newPadOp.getRegion(),
                               newPadOp.getRegion().begin());

    Value newTransposeOp =
        createTranspose(rewriter, newPadOp, transposeOp.getPermutation());
    rewriter.replaceOp(padOp, newTransposeOp);
    return success();
  }
};

// Fuses a transpose with the input of a linalg.generic op or contraction op.
// Contraction ops are generalized and then treated as a generic. For example,
//
// %0 = linalg.transpose ... permutation = [0, 2, 1] :
//                           tensor<2x5x7> -> tensor<2x7x5>
// linalg.generic {
//   indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d0, d2)>,
//                    affine_map<(d0, d1, d2) -> (d0, d1, d2)>]
//   ins(%0 : tensor<2x7x5>)
//
// Becomes
//
// linalg.generic {
//   indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2, d0)>,
//                    affine_map<(d0, d1, d2) -> (d0, d1, d2)>]
//   ins(%0 : tensor<2x5x7>)
//
// This is considered just one way to model transpose propagation to generics,
// another option would be to interpret the transpose on the iterators of the
// generic, thus producing a transpose on the output and any other inputs to
// the generic. This has the potential introduce more transposes/data movement
// and isn't the way this pass is modeled. Global data layout transformations
// like that are better suited for pack/unpack propagation rooted on specific
// operations.
//
// TODO: Rewrite this to use elementwise op fusion patterns.
class FuseTransposeWithLinalgOpConsumer
    : public OpInterfaceRewritePattern<linalg::LinalgOp> {
public:
  using OpInterfaceRewritePattern<linalg::LinalgOp>::OpInterfaceRewritePattern;
  FuseTransposeWithLinalgOpConsumer(MLIRContext *ctx, bool aggressiveProp,
                                    bool convProp, PatternBenefit b = 1)
      : OpInterfaceRewritePattern<linalg::LinalgOp>(ctx, b),
        allowGeneralizing(aggressiveProp), convProp(convProp) {}

  LogicalResult matchAndRewrite(linalg::LinalgOp linalgOp,
                                PatternRewriter &rewriter) const override {
    if (!IREE::Flow::isNonNullAndOutsideDispatch(linalgOp)) {
      return failure();
    }
    OpOperand *transposeOperand = nullptr;
    linalg::TransposeOp transposeOp;
    for (OpOperand *input : linalgOp.getDpsInputOperands()) {
      auto maybeTransposeOp = input->get().getDefiningOp<linalg::TransposeOp>();
      if (maybeTransposeOp && maybeTransposeOp->hasOneUse()) {
        transposeOp = maybeTransposeOp;
        transposeOperand = input;
        break;
      }
    }
    if (!transposeOperand) {
      return rewriter.notifyMatchFailure(linalgOp, "no transpose operand");
    }

    int64_t inputIndex = transposeOperand->getOperandNumber();
    ArrayRef<int64_t> perm = transposeOp.getPermutation();
    auto invPerm = invertPermutationVector(perm);

    // To do the fusion, we can simply apply the permutation of the transpose
    // to the results of the associated input's indexing map, and then forward
    // the input to the transpose to the consumer generic.
    auto maybeGenericOp = getAllowedGenericOpOrGeneralizeNamedOp(
        rewriter, linalgOp, allowGeneralizing, convProp);
    if (failed(maybeGenericOp)) {
      return failure();
    }
    auto genericOp = maybeGenericOp.value();
    transposeOperand = genericOp.getDpsInputOperand(inputIndex);
    rewriter.startOpModification(genericOp);

    SmallVector<AffineMap> newIndexingMaps = genericOp.getIndexingMapsArray();
    AffineMap inputMap = genericOp.getMatchingIndexingMap(transposeOperand);
    SmallVector<AffineExpr> newExprs =
        applyPermutation(inputMap.getResults(), invPerm);
    AffineMap transposedMap =
        AffineMap::get(inputMap.getNumDims(), inputMap.getNumSymbols(),
                       newExprs, rewriter.getContext());
    newIndexingMaps[inputIndex] = transposedMap;
    genericOp.setIndexingMapsAttr(
        rewriter.getAffineMapArrayAttr(newIndexingMaps));

    genericOp.setOperand(inputIndex, transposeOp.getInput());
    rewriter.finalizeOpModification(genericOp);
    return success();
  }

private:
  bool allowGeneralizing = false;
  bool convProp = false;
};

static bool isIndexingMapAffectedByTransposeMap(
    AffineMap indexingMap, ArrayRef<int64_t> iterationSpacePermutation) {
  int64_t prevIdx = -1;
  for (auto result : indexingMap.getResults()) {
    int64_t idx =
        iterationSpacePermutation[cast<AffineDimExpr>(result).getPosition()];
    // Verify that the relative ordering of indices in the map remain the same.
    // If not, then the transposition affects the access order for the given
    // map (and associated operand).
    if (idx <= prevIdx) {
      return true;
    }
    prevIdx = idx;
  }
  return false;
}

// Finds a single DPS input operand of the given |genericOp| that is affected by
// the |iterationSpacePermutation|. In other words, the permutation changes the
// relative ordering of any of the dimensions of that input operand.
//
// For example, with permutation [1, 0, 2], affine map (d0, d1, d2) -> (d0, d1)
// is affected by the permutation because the first two dimensions are iterated
// in a different order while (d0, d1, d2) -> (d0, d2) is unaffected.
//
// If no such operand is found or there is more than one such operation, nullptr
// is returned.
static OpOperand *
getSingleTransposedInputOperand(linalg::GenericOp genericOp,
                                ArrayRef<int64_t> iterationSpacePermutation) {
  OpOperand *operand = nullptr;
  for (auto input : genericOp.getDpsInputOperands()) {
    if (!isIndexingMapAffectedByTransposeMap(
            genericOp.getMatchingIndexingMap(input),
            iterationSpacePermutation)) {
      continue;
    }
    if (operand) {
      return nullptr;
    }
    operand = input;
  }
  return operand;
}

// Returns a new list of indexing maps that composes the iteration space
// permutation map |transposeMap| with all indexing maps of |genericOp| except
// for the |transposedInputIdx|'th operand. The unchanged operand is expected
// to have an explicit `linalg.transpose` op constructed for it so its map does
// not need to be updated.
static SmallVector<AffineMap>
getTransposedIndexingMaps(linalg::GenericOp genericOp,
                          int64_t transposedInputIdx, AffineMap transposeMap) {
  SmallVector<AffineMap> indexingMaps = genericOp.getIndexingMapsArray();
  for (unsigned i = 0, e = genericOp.getNumDpsInputs(); i < e; ++i) {
    if (i == transposedInputIdx) {
      continue;
    }
    indexingMaps[i] = indexingMaps[i].compose(transposeMap);
  }
  return indexingMaps;
}

// Sinks a transpose through the input of a elementwise operation where the
// transposition of the iteration space only affects a single input operand.
class SinkTransposeThroughUnaryElementwiseInput
    : public OpRewritePattern<linalg::GenericOp> {
public:
  using Base::Base;

  LogicalResult matchAndRewrite(linalg::GenericOp genericOp,
                                PatternRewriter &rewriter) const override {
    if (!IREE::Flow::isNonNullAndOutsideDispatch(genericOp)) {
      return rewriter.notifyMatchFailure(genericOp, "pre-formed dispatch");
    }

    if (!linalg::isElementwise(genericOp)) {
      return rewriter.notifyMatchFailure(genericOp, "non-elementwise generic");
    }

    if (genericOp.hasIndexSemantics()) {
      return rewriter.notifyMatchFailure(genericOp, "has index semantics");
    }

    if (genericOp.getNumDpsInits() != 1) {
      return rewriter.notifyMatchFailure(genericOp,
                                         "unimplemented: multiple results");
    }

    AffineMap resultMap =
        genericOp.getMatchingIndexingMap(genericOp.getDpsInitOperand(0));
    if (!resultMap.isIdentity()) {
      return rewriter.notifyMatchFailure(
          genericOp, "unimplemented: non-identity result map");
    }

    linalg::TransposeOp transposeOp;
    OpOperand *inputOperand;
    for (auto input : genericOp.getDpsInputOperands()) {
      // Skip broadcasted operands and transposed operands. If the input is
      // broadcasted then we would not want to propagate because that would
      // do the transpose on larger data, and if transposed we would rather
      // simply compose the transposes (handled in a separate pattern).
      if (genericOp.getMatchingIndexingMap(input) != resultMap) {
        continue;
      }

      auto maybeTransposeOp = input->get().getDefiningOp<linalg::TransposeOp>();
      // Skip multi-use transposes.
      if (!maybeTransposeOp || !maybeTransposeOp->hasOneUse()) {
        continue;
      }

      auto transposableInputOperand = getSingleTransposedInputOperand(
          genericOp, maybeTransposeOp.getPermutation());
      // Skip if more than one operand is affected by the transpose.
      if (transposableInputOperand != input) {
        continue;
      }

      transposeOp = maybeTransposeOp;
      inputOperand = transposableInputOperand;
      break;
    }

    if (!transposeOp) {
      return rewriter.notifyMatchFailure(genericOp,
                                         "no single use transpose operand");
    }

    ArrayRef<int64_t> perm = transposeOp.getPermutation();
    auto invPerm = invertPermutationVector(perm);

    // Create a new empty init for the transposed generic.
    Value newInit =
        createTransposeInit(rewriter, genericOp.getDpsInits()[0], invPerm);

    // We do not need to update iterator types because this is an elementwise
    // op. We just need to update the indexing maps of all other input operands
    // by composing the transpose map.
    AffineMap transposeMap =
        AffineMap::getPermutationMap(perm, rewriter.getContext());
    SmallVector<AffineMap> indexingMaps = getTransposedIndexingMaps(
        genericOp, inputOperand->getOperandNumber(), transposeMap);

    SmallVector<Value> newOperands = genericOp->getOperands();
    newOperands[inputOperand->getOperandNumber()] = transposeOp.getInput();
    newOperands[genericOp.getDpsInitOperand(0)->getOperandNumber()] = newInit;

    auto newGenericOp =
        mlir::clone(rewriter, genericOp, newInit.getType(), newOperands);
    newGenericOp.setIndexingMapsAttr(
        rewriter.getAffineMapArrayAttr(indexingMaps));
    rewriter.replaceOp(
        genericOp, createTranspose(rewriter, newGenericOp->getResult(0), perm));
    return success();
  }
};

// Bubbles a transpose through the init of a elementwise operation where the
// transposition of the iteration space only affects a single input operand.
class BubbleTransposeThroughUnaryElementwiseDpsInit
    : public OpRewritePattern<linalg::TransposeOp> {
public:
  using Base::Base;

  LogicalResult matchAndRewrite(linalg::TransposeOp transposeOp,
                                PatternRewriter &rewriter) const override {
    auto genericOp = transposeOp.getInput().getDefiningOp<linalg::GenericOp>();
    if (!genericOp) {
      return rewriter.notifyMatchFailure(transposeOp, "non-generic producer");
    }

    if (genericOp.getNumDpsInits() != 1) {
      return rewriter.notifyMatchFailure(transposeOp,
                                         "unimplemented: multiple results");
    }

    if (!IREE::Flow::isNonNullAndOutsideDispatch({genericOp, transposeOp})) {
      return failure();
    }

    if (!linalg::isElementwise(genericOp) ||
        !genericOp.getMatchingIndexingMap(genericOp.getDpsInitOperand(0))
             .isIdentity()) {
      return rewriter.notifyMatchFailure(transposeOp, "not elementwise");
    }

    if (genericOp.hasIndexSemantics()) {
      return rewriter.notifyMatchFailure(genericOp, "has index semantics");
    }

    if (!genericOp->hasOneUse()) {
      return rewriter.notifyMatchFailure(transposeOp, "not single user");
    }

    ArrayRef<int64_t> perm = transposeOp.getPermutation();
    auto invPerm = invertPermutationVector(perm);

    auto inputOperand = getSingleTransposedInputOperand(genericOp, invPerm);
    if (!inputOperand ||
        !genericOp.getMatchingIndexingMap(inputOperand).isIdentity()) {
      return rewriter.notifyMatchFailure(
          genericOp, "no single transposable input operand");
    }

    Value newTranspose = createTranspose(rewriter, inputOperand->get(), perm);

    // Create a new empty init for the transposed generic.
    Value newInit =
        createTransposeInit(rewriter, genericOp.getDpsInits()[0], perm);

    SmallVector<Value> newOperands = genericOp->getOperands();
    newOperands[inputOperand->getOperandNumber()] = newTranspose;
    newOperands[genericOp.getDpsInitOperand(0)->getOperandNumber()] = newInit;

    AffineMap transposeMap =
        AffineMap::getPermutationMap(invPerm, rewriter.getContext());

    // We do not need to update iterator types because this is an elementwise
    // op. We just need to update the indexing maps of all other input operands
    // by composing the transpose map.
    SmallVector<AffineMap> indexingMaps = getTransposedIndexingMaps(
        genericOp, inputOperand->getOperandNumber(), transposeMap);

    // We do not need to update indexing maps because this is an elementwise
    // op where the input and output maps are the same.
    // Just replace the operands with transposed variants.
    auto newGenericOp =
        mlir::clone(rewriter, genericOp, newInit.getType(), newOperands);
    newGenericOp.setIndexingMapsAttr(
        rewriter.getAffineMapArrayAttr(indexingMaps));
    rewriter.replaceOp(transposeOp, newGenericOp);
    return success();
  }
};

} // namespace

//===----------------------------------------------------------------------===//
// Linalg Named Op -> Named Op Conversions
//===----------------------------------------------------------------------===//

namespace {

template <typename OpTy, int64_t inputIdx>
class NamedOpConversion : public OpRewritePattern<OpTy> {
public:
  using OpRewritePattern<OpTy>::OpRewritePattern;
  NamedOpConversion(MLIRContext *ctx, SmallVector<int64_t> perm,
                    PatternBenefit b = 1)
      : OpRewritePattern<OpTy>(ctx, b), permutation(perm) {}

  LogicalResult matchAndRewrite(OpTy namedOp,
                                PatternRewriter &rewriter) const override {
    if (!IREE::Flow::isNonNullAndOutsideDispatch(namedOp)) {
      return failure();
    }

    Value input = namedOp.getInputs()[inputIdx];
    auto transpose = input.getDefiningOp<linalg::TransposeOp>();
    if (!transpose) {
      return failure();
    }

    SmallVector<int64_t> transPerm(transpose.getPermutation());
    if (transPerm != permutation) {
      return rewriter.notifyMatchFailure(
          namedOp, "transpose permutation does not match target permutation");
    }
    SmallVector<NamedAttribute> attrs = getPrunedAttributeList(namedOp);
    SmallVector<Value> newInputs = namedOp.getInputs();
    newInputs[inputIdx] = transpose.getInput();

    auto replaceOp = [&](auto *typePtr) {
      rewriter.replaceOpWithNewOp<std::remove_pointer_t<decltype(typePtr)>>(
          namedOp, newInputs, namedOp.getDpsInits(), attrs);
    };

    Operation *op = namedOp.getOperation();
    if (isa<linalg::MatmulTransposeAOp>(op) && inputIdx == 0) {
      replaceOp(static_cast<linalg::MatmulOp *>(nullptr));
    } else if (isa<linalg::MatmulTransposeBOp>(op) && inputIdx == 1) {
      replaceOp(static_cast<linalg::MatmulOp *>(nullptr));
    } else if (IREE::LinalgExt::isPureMatmul(op) && inputIdx == 0) {
      replaceOp(static_cast<linalg::MatmulTransposeAOp *>(nullptr));
    } else if (IREE::LinalgExt::isPureMatmul(op) && inputIdx == 1) {
      replaceOp(static_cast<linalg::MatmulTransposeBOp *>(nullptr));
    } else if (isa<linalg::BatchMatmulTransposeAOp>(op) && inputIdx == 0) {
      replaceOp(static_cast<linalg::BatchMatmulOp *>(nullptr));
    } else if (isa<linalg::BatchMatmulTransposeBOp>(op) && inputIdx == 1) {
      replaceOp(static_cast<linalg::BatchMatmulOp *>(nullptr));
    } else if (IREE::LinalgExt::isPureBatchMatmul(op) && inputIdx == 0) {
      replaceOp(static_cast<linalg::BatchMatmulTransposeAOp *>(nullptr));
    } else if (IREE::LinalgExt::isPureBatchMatmul(op) && inputIdx == 1) {
      replaceOp(static_cast<linalg::BatchMatmulTransposeBOp *>(nullptr));
    } else {
      return failure();
    }
    return success();
  }

private:
  // Non-type literal array template parameters are a C++20 feature, so instead
  // all the named op patterns pass their permutation explicitly as a
  // SmallVector.
  SmallVector<int64_t> permutation;
};

// Fuses a transpose into a reduction linalg.generic by absorbing it into the
// indexing map.
class FuseTransposeThroughGenericReduction
    : public OpRewritePattern<linalg::GenericOp> {
public:
  using Base::Base;

  LogicalResult matchAndRewrite(linalg::GenericOp genericOp,
                                PatternRewriter &rewriter) const override {
    if (!IREE::Flow::isNonNullAndOutsideDispatch(genericOp)) {
      return failure();
    }

    // Only try to fuse when the generic has at least one reduction dimension.
    if (genericOp.getNumParallelLoops() == genericOp.getNumLoops()) {
      return rewriter.notifyMatchFailure(genericOp, "not a reduction");
    }

    // All maps must be projected permutations.
    if (!llvm::all_of(genericOp.getIndexingMapsArray(), [](AffineMap map) {
          return map.isProjectedPermutation();
        })) {
      return rewriter.notifyMatchFailure(genericOp,
                                         "not a projected permutation");
    }

    // Look for a transpose on any input.
    for (int64_t inputIdx = 0, endIdx = genericOp.getNumDpsInputs();
         inputIdx < endIdx; ++inputIdx) {
      auto transpose = genericOp.getDpsInputs()[inputIdx]
                           .getDefiningOp<linalg::TransposeOp>();
      if (!transpose) {
        continue;
      }

      // Update the indexing map according to the transpose.
      AffineMap inputMap = genericOp.getMatchingIndexingMap(
          genericOp.getDpsInputOperand(inputIdx));

      // Fuse by updating the indexing map to absorb the transpose.
      auto invPerm = invertPermutationVector(transpose.getPermutation());
      SmallVector<AffineExpr> newExprs =
          applyPermutation(inputMap.getResults(), invPerm);

      // Only fuse if dimension indices are in increasing order to maintain
      // efficient memory access patterns. This should offset the fact that the
      // transpose may have multiple uses.
      int64_t prevDim = -1;
      for (AffineExpr expr : newExprs) {
        const int64_t dim = cast<AffineDimExpr>(expr).getPosition();
        if (dim < prevDim) {
          return rewriter.notifyMatchFailure(genericOp,
                                             "newExprs are not contiguous");
        }
        prevDim = dim;
      }

      auto transposedMap =
          AffineMap::get(inputMap.getNumDims(), inputMap.getNumSymbols(),
                         newExprs, rewriter.getContext());

      SmallVector<AffineMap> newIndexingMaps = genericOp.getIndexingMapsArray();
      newIndexingMaps[inputIdx] = transposedMap;

      rewriter.startOpModification(genericOp);
      genericOp.setIndexingMapsAttr(
          rewriter.getAffineMapArrayAttr(newIndexingMaps));
      genericOp.setOperand(inputIdx, transpose.getInput());
      rewriter.finalizeOpModification(genericOp);
      return success();
    }

    return rewriter.notifyMatchFailure(genericOp,
                                       "no matching transpose found");
  }
};

} // namespace

//===----------------------------------------------------------------------===//
// Pass definition
//===----------------------------------------------------------------------===//

namespace {
struct PropagateLinalgTransposePass
    : impl::PropagateLinalgTransposePassBase<PropagateLinalgTransposePass> {
  using Base::Base;
  void getDependentDialects(DialectRegistry &registry) const override {
    registry.insert<linalg::LinalgDialect, tensor::TensorDialect>();
  }
  explicit PropagateLinalgTransposePass(bool enableAggressivePropagation) {
    this->enableAggressivePropagation = enableAggressivePropagation;
  }

  void runOnOperation() override;
};
} // namespace

static void populateMatmulSinkingPatterns(MLIRContext *context,
                                          RewritePatternSet &sinkingPatterns) {
  sinkingPatterns.insert<NamedOpConversion</*OpType=*/linalg::MatmulOp,
                                           /*inputIdx=*/1>>(
      context, SmallVector<int64_t>{1, 0});
  sinkingPatterns.insert<NamedOpConversion</*OpType=*/linalg::MatmulOp,
                                           /*inputIdx=*/0>>(
      context, SmallVector<int64_t>{1, 0});
  sinkingPatterns.insert<NamedOpConversion</*OpType=*/linalg::BatchMatmulOp,
                                           /*inputIdx=*/1>>(
      context, SmallVector<int64_t>{0, 2, 1});
  sinkingPatterns.insert<NamedOpConversion</*OpType=*/linalg::BatchMatmulOp,
                                           /*inputIdx=*/0>>(
      context, SmallVector<int64_t>{0, 2, 1});
  sinkingPatterns.insert<FuseTransposeThroughGenericReduction>(context);
}

static void
populateCommonCanonicalizationPatterns(MLIRContext *context,
                                       RewritePatternSet &patterns) {
  linalg::FillOp::getCanonicalizationPatterns(patterns, context);
  tensor::EmptyOp::getCanonicalizationPatterns(patterns, context);
  tensor::ExpandShapeOp::getCanonicalizationPatterns(patterns, context);
  tensor::CollapseShapeOp::getCanonicalizationPatterns(patterns, context);
  memref::populateResolveRankedShapedTypeResultDimsPatterns(patterns);
  tensor::populateFoldTensorEmptyPatterns(patterns,
                                          /*foldSingleUseOnly=*/false);
}

void PropagateLinalgTransposePass::runOnOperation() {
  MLIRContext *context = &getContext();
  mlir::FunctionOpInterface funcOp = getOperation();
  // First, specialize all transposes to `linalg.transpose`. This dramatically
  // simplifies all subsequent propagation patterns, both in matching and
  // rewriting.
  {
    SmallVector<linalg::GenericOp> genericCandidates;
    funcOp.walk([&](linalg::GenericOp genericOp) {
      if (IREE::Flow::isNonNullAndOutsideDispatch(genericOp)) {
        genericCandidates.push_back(genericOp);
      }
    });
    IRRewriter rewriter(&getContext());
    for (auto genericOp : genericCandidates) {
      rewriter.setInsertionPoint(genericOp);
      specializeGenericTransposeOp(rewriter, genericOp);
    }
  }

  LLVM_DEBUG({
    llvm::dbgs() << "\n--- After specializing transpose ops ---\n";
    funcOp->print(llvm::dbgs(), OpPrintingFlags().useLocalScope());
    llvm::dbgs() << "\n\n";
  });

  // First try to fuse transposes with some consumer linalg named ops before
  // any reshape propagation. Some transposes may be adjacent to named ops,
  // and it is more canonical if we can fuse the ops into a new named op.
  if (!testBubblingOnly) {
    RewritePatternSet sinkingPatterns(context);
    sinkingPatterns.insert<SinkTransposeThroughExtractSlice>(context);
    sinkingPatterns.insert<SinkTransposeThroughExpandShape>(
        context, enableEdgeReshapePropagation);
    populateMatmulSinkingPatterns(context, sinkingPatterns);
    populateCommonCanonicalizationPatterns(context, sinkingPatterns);
    sinkingPatterns.add<SinkTransposeThroughUnaryElementwiseInput>(
        context, /*benefit=*/2);
    if (failed(applyPatternsGreedily(funcOp, std::move(sinkingPatterns)))) {
      funcOp.emitError("Transpose initial sinking patterns failed");
      return signalPassFailure();
    }
  }

  LLVM_DEBUG({
    llvm::dbgs() << "\n--- After canonicalizing transpose in place ---\n";
    funcOp->print(llvm::dbgs(), OpPrintingFlags().useLocalScope());
    llvm::dbgs() << "\n\n";
  });

  // Propagate transposes upwards, and fuse with any producer generic ops. Also
  // propagate reshapes upwards to open up more transpose fusion opportunities.
  if (!testSinkingOnly) {
    linalg::ControlFusionFn reshapePropagationFn =
        [&](OpOperand *fusedOperand) {
          Operation *producer = fusedOperand->get().getDefiningOp();
          Operation *consumer = fusedOperand->getOwner();
          if (!IREE::Flow::isNonNullAndOutsideDispatch({producer, consumer})) {
            return false;
          }

          // Do not reshape producer linalg op if it has more than one user.
          auto producerLinalgOp = dyn_cast<linalg::LinalgOp>(producer);
          if (!producerLinalgOp || !producerLinalgOp->hasOneUse()) {
            return false;
          }
          // Only reshape generic ops, or any op if aggressive propagation is
          // enabled.
          if (!enableAggressivePropagation &&
              !isa<linalg::GenericOp>(producerLinalgOp)) {
            return false;
          }
          // Only propagate expand_shape ops up through producers because it
          // is always possible to bubble a transpose through an collapse_shape
          // and thus is handled separately.
          if (!isa<tensor::ExpandShapeOp>(consumer)) {
            return false;
          }

          if (!enableEdgeReshapePropagation &&
              llvm::none_of(
                  consumer->getUsers(), [&](Operation *expandConsumer) {
                    return isReshapeBlockingFusion(producer, expandConsumer);
                  })) {
            return false;
          }
          // Only propagate if the immediate consumer of the reshape is a
          // transpose.
          return consumer->hasOneUse() &&
                 isa<linalg::TransposeOp>(*(consumer->user_begin()));
        };
    RewritePatternSet bubblingPatterns(context);
    linalg::populateFoldReshapeOpsByExpansionPatterns(bubblingPatterns,
                                                      reshapePropagationFn);
    linalg::FillOp::getCanonicalizationPatterns(bubblingPatterns, context);

    if (enableAttentionVTranspose) {
      linalg::ControlFusionFn bubbleTransposeControlFn =
          [](OpOperand *fusedOperand) {
            Operation *producer = fusedOperand->get().getDefiningOp();
            Operation *consumer = fusedOperand->getOwner();

            return IREE::Flow::isNonNullAndOutsideDispatch(
                {producer, consumer});
          };
      IREE::LinalgExt::populateBubbleTransposeFromLinalgExtOps(
          bubblingPatterns, bubbleTransposeControlFn);
    }
    bubblingPatterns.insert<FuseTransposeWithProducerLinalgOp>(
        context, enableAggressivePropagation, enableConvolutionPropagation);
    bubblingPatterns.insert<BubbleTransposeThroughCollapseShape>(
        context, enableEdgeReshapePropagation);
    bubblingPatterns.add<BubbleTransposeThroughUnaryElementwiseDpsInit>(
        context, /*benefit=*/2);
    bubblingPatterns.insert<ComposeTransposes>(context);
    populateCommonCanonicalizationPatterns(context, bubblingPatterns);

    GreedyRewriteConfig config;
    config.setMaxIterations(GreedyRewriteConfig::kNoLimit);
    if (failed(applyPatternsGreedily(funcOp, std::move(bubblingPatterns),
                                     config))) {
      funcOp.emitError("Transpose bubbling patterns failed");
      return signalPassFailure();
    }
  }

  LLVM_DEBUG({
    llvm::dbgs() << "\n--- After bubbling transpose ops up ---\n";
    funcOp->print(llvm::dbgs(), OpPrintingFlags().useLocalScope());
    llvm::dbgs() << "\n\n";
  });

  // Propagate transposes downwards, and fuse with any non-unary generic ops
  // or linalg named ops. Also propagate reshapes downwards to open up more
  // transpose fusion opportunities.
  if (!testBubblingOnly) {
    RewritePatternSet sinkingPatterns(context);
    linalg::ControlFusionFn reshapePropagationFn =
        [&](OpOperand *fusedOperand) {
          Operation *producer = fusedOperand->get().getDefiningOp();
          Operation *consumer = fusedOperand->getOwner();
          if (!IREE::Flow::isNonNullAndOutsideDispatch({producer, consumer})) {
            return false;
          }
          auto consumerLinalgOp = dyn_cast<linalg::LinalgOp>(consumer);
          if (!consumerLinalgOp || consumerLinalgOp.getNumReductionLoops()) {
            return false;
          }
          // Only reshape generic ops.
          if (!enableAggressivePropagation &&
              !isa<linalg::GenericOp>(consumerLinalgOp)) {
            return false;
          }
          // Only propagate collapse_shape ops down through consumers because it
          // is always possible to sink a transpose through an expand_shape and
          // thus is handled separately.
          if (!isa<tensor::CollapseShapeOp>(producer)) {
            return false;
          }

          if (!enableEdgeReshapePropagation &&
              !isReshapeBlockingFusion(producer->getOperand(0).getDefiningOp(),
                                       consumer)) {
            return false;
          }

          // Require that the immediate producer of the reshape is a transpose.
          return isa_and_nonnull<linalg::TransposeOp>(
              producer->getOperand(0).getDefiningOp());
        };
    linalg::populateFoldReshapeOpsByExpansionPatterns(sinkingPatterns,
                                                      reshapePropagationFn);
    sinkingPatterns.insert<SinkTransposeThroughExtractSlice>(context);
    sinkingPatterns.insert<SinkTransposeThroughExpandShape>(
        context, enableEdgeReshapePropagation);
    if (enableSinkTransposeThroughPad) {
      sinkingPatterns.insert<SinkTransposeThroughPad>(context);
    }
    sinkingPatterns.insert<FuseTransposeWithLinalgOpConsumer>(
        context, enableAggressivePropagation, enableConvolutionPropagation);
    sinkingPatterns.insert<ComposeTransposes>(context);
    populateMatmulSinkingPatterns(context, sinkingPatterns);
    populateCommonCanonicalizationPatterns(context, sinkingPatterns);
    sinkingPatterns.add<SinkTransposeThroughUnaryElementwiseInput>(
        context, /*benefit=*/2);
    GreedyRewriteConfig config;
    // TODO: This is inefficient. Consider rewriting this pass to use a
    // worklist of just the transpose operations.
    config.setMaxIterations(GreedyRewriteConfig::kNoLimit);
    if (failed(applyPatternsGreedily(funcOp, std::move(sinkingPatterns),
                                     config))) {
      funcOp.emitError("Transpose sinking patterns failed");
      return signalPassFailure();
    }
  }

  LLVM_DEBUG({
    llvm::dbgs() << "\n--- After sinking transpose ops down ---\n";
    funcOp->print(llvm::dbgs(), OpPrintingFlags().useLocalScope());
    llvm::dbgs() << "\n\n";
  });
}

std::unique_ptr<InterfacePass<mlir::FunctionOpInterface>>
createPropagateLinalgTransposePass(bool enableAggressivePropagation) {
  return std::make_unique<PropagateLinalgTransposePass>(
      enableAggressivePropagation);
}

} // namespace mlir::iree_compiler::GlobalOptimization
