// Copyright 2025 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

#include "iree/compiler/Codegen/Common/CombineLayoutTransformation.h"
#include "iree/compiler/Codegen/Common/Transforms.h"
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h"
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenOps.h"
#include "iree/compiler/Dialect/LinalgExt/Transforms/Transforms.h"
#include "iree/compiler/Dialect/LinalgExt/Utils/Utils.h"
#include "iree/compiler/Dialect/Util/IR/UtilDialect.h"
#include "llvm/ADT/Repeated.h"
#include "llvm/Support/DebugLog.h"
#include "mlir/Analysis/SliceAnalysis.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Linalg/IR/LinalgInterfaces.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/UB/IR/UBOps.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/IR/IRMapping.h"
#include "mlir/Interfaces/DestinationStyleOpInterface.h"
#include "mlir/Interfaces/InferTypeOpInterface.h"
#include "mlir/Interfaces/ValueBoundsOpInterface.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"

#define DEBUG_TYPE "iree-codegen-combine-layout-transformation"

namespace mlir::iree_compiler {

#define GEN_PASS_DEF_COMBINERESULTLAYOUTTRANSFORMATIONPASS
#define GEN_PASS_DEF_COMBINESOURCELAYOUTTRANSFORMATIONPASS
#include "iree/compiler/Codegen/Common/Passes.h.inc"

using IREE::LinalgExt::MapLoadOp;
using IREE::LinalgExt::MapStoreOp;

//===----------------------------------------------------------------------===//
// Combining Layout Transformation Ops
//===----------------------------------------------------------------------===//

/// Folds an `op` that does not affect index computation into a `mapStoreOp`.
/// This is used for ops like `linalg::CopyOp`.
static MapStoreOp foldIdentityLikeOpIntoMapStore(RewriterBase &rewriter,
                                                 Operation *op,
                                                 MapStoreOp mapStoreOp) {
  assert(mapStoreOp.getInput() == op->getResult(0) &&
         "expected op to be the producer of mapStoreOp");
  rewriter.modifyOpInPlace(mapStoreOp, [&]() {
    mapStoreOp.getInputMutable().assign(op->getOperand(0));
  });
  return mapStoreOp;
}

/// Fold a transpose (given by `perm`) producer into a consumer `mapStoreOp`.
/// The `producerInput` is the tensor that feeds the transpose, and
/// `producerResult` is the transpose output (== mapStoreOp input).
static MapStoreOp foldTransposePermIntoMapStore(RewriterBase &rewriter,
                                                ArrayRef<int64_t> perm,
                                                Value producerInput,
                                                MapStoreOp mapStoreOp) {
  SmallVector<int64_t> permCopy(perm);
  rewriter.modifyOpInPlace(mapStoreOp, [&]() {
    mapStoreOp.insertTransformationAtStart(
        rewriter,
        [permCopy](ArrayRef<BlockArgument> indices) -> SmallVector<Value> {
          SmallVector<Value> indexValues(indices.begin(), indices.end());
          return applyPermutation(indexValues, permCopy);
        },
        permCopy.size());
    mapStoreOp.getInputMutable().assign(producerInput);
  });
  return mapStoreOp;
}

/// Fold a tensor::ExpandShapeOp or tensor::CollapseShapeOp into a consumer
/// `mapStoreOp` by composing the reshape's index transformation into the
/// front of the `mapStoreOp`'s.
///
/// Each reassociation group is folded independently: an `expand_shape` splits
/// one source index via a `delinearize` over just that group's result sizes,
/// and a `collapse_shape` merges a group of source indices via a `linearize`.
/// This keeps the split/merge factors per-group — typically static tile sizes
/// — rather than routing every index through one global `linearize` +
/// `delinearize` over the full (dynamic) shapes, which would force a dynamic
/// integer division per element when the transformation is later scalarized.
template <typename ReshapeOpTy>
static FailureOr<IREE::LinalgExt::MapStoreOp>
foldReshapeIntoMapStore(RewriterBase &rewriter, ReshapeOpTy reshapeOp,
                        IREE::LinalgExt::MapStoreOp mapStoreOp) {
  assert(mapStoreOp.getInput() == reshapeOp->getResult(0) &&
         "expected reshapeOp to be the producer of mapStoreOp");
  // Cannot fold if either side of the reshape is rank-0 (scalar). The source
  // being rank-0 would make the new map_store input rank-0, and the result
  // being rank-0 would produce an empty delinearization.
  if (cast<RankedTensorType>(reshapeOp.getSrc().getType()).getRank() == 0 ||
      cast<RankedTensorType>(reshapeOp.getResult().getType()).getRank() == 0) {
    return failure();
  }
  constexpr bool isExpand = std::is_same_v<ReshapeOpTy, tensor::ExpandShapeOp>;
  Location loc = reshapeOp->getLoc();
  OpBuilder::InsertionGuard g(rewriter);
  rewriter.setInsertionPointAfter(reshapeOp);
  SmallVector<OpFoldResult> srcDims =
      tensor::getMixedSizes(rewriter, loc, reshapeOp.getSrc());
  // There can be leftover tensor.dim ops consuming the result of the reshape,
  // but they are expected to be folded into some affine.apply ops on the source
  // sizes by later cleanup patterns.
  SmallVector<OpFoldResult> resultDims =
      tensor::getMixedSizes(rewriter, loc, reshapeOp.getResult());
  SmallVector<ReassociationIndices> reassociation =
      reshapeOp.getReassociationIndices();

  Location mapStoreLoc = mapStoreOp->getLoc();
  rewriter.modifyOpInPlace(mapStoreOp, [&]() {
    mapStoreOp.insertTransformationAtStart(
        rewriter,
        [&rewriter, mapStoreLoc, srcDims, resultDims, reassociation,
         isExpand](ArrayRef<BlockArgument> indices) -> SmallVector<Value> {
          // `indices` are in the reshape's source index space; produce the
          // reshape's result index space.
          SmallVector<Value> result(resultDims.size());
          for (auto [groupIdx, group] : llvm::enumerate(reassociation)) {
            if (isExpand) {
              // Source index `groupIdx` splits into result dims `group`.
              if (group.size() == 1) {
                result[group[0]] = indices[groupIdx];
                continue;
              }
              SmallVector<OpFoldResult> groupSizes;
              for (int64_t resultDim : group) {
                groupSizes.push_back(resultDims[resultDim]);
              }
              auto delinearizeOp = affine::AffineDelinearizeIndexOp::create(
                  rewriter, mapStoreLoc, indices[groupIdx], groupSizes,
                  /*hasOuterBound=*/true);
              for (auto [i, resultDim] : llvm::enumerate(group)) {
                result[resultDim] = delinearizeOp.getResult(i);
              }
            } else {
              // Source dims `group` merge into result index `groupIdx`.
              if (group.size() == 1) {
                result[groupIdx] = indices[group[0]];
                continue;
              }
              SmallVector<Value> groupIndices;
              SmallVector<OpFoldResult> groupSizes;
              for (int64_t srcDim : group) {
                groupIndices.push_back(indices[srcDim]);
                groupSizes.push_back(srcDims[srcDim]);
              }
              result[groupIdx] = affine::AffineLinearizeIndexOp::create(
                  rewriter, mapStoreLoc, groupIndices, groupSizes,
                  /*disjoint=*/true);
            }
          }
          return result;
        },
        srcDims.size());
    mapStoreOp.getInputMutable().assign(reshapeOp->getOperand(0));
  });
  return mapStoreOp;
}

/// Fold a tensor::ExpandShapeOp into a consumer `mapStoreOp`, by linearizing
/// and then delinearizing the source indices of the `mapStoreOp`s index
/// transformation.
static FailureOr<MapStoreOp>
foldExpandShapeIntoMapStore(RewriterBase &rewriter,
                            tensor::ExpandShapeOp expandShapeOp,
                            MapStoreOp mapStoreOp) {
  return foldReshapeIntoMapStore(rewriter, expandShapeOp, mapStoreOp);
}

/// Fold a tensor::CollapseShapeOp into a consumer `mapStoreOp`, by
/// linearizing and then delinearizing the source indices of the
/// `mapStoreOp`s index transformation.
static FailureOr<MapStoreOp>
foldCollapseShapeIntoMapStore(RewriterBase &rewriter,
                              tensor::CollapseShapeOp collapseShapeOp,
                              MapStoreOp mapStoreOp) {
  return foldReshapeIntoMapStore(rewriter, collapseShapeOp, mapStoreOp);
}

/// Fold an `extractSliceOp` into a consumer `mapStoreOp`. If the
/// extract_slice is an identity operation, the input is replaced directly.
/// Otherwise, a mask is applied based on the bounds of the extractSliceOp.
/// Currently, only zero offsets and unit strides are supported.
///
/// For a given `%mask` value, and slice sizes of `%size0`, `%size1`, ...,
/// `%sizeN`, the `%new_mask` becomes:
///
///   %bound0 = arith.cmpi ult, %idx0, %size0 : index
///   %mask0 = arith.andi %mask, %bound0 : i1
///   %bound1 = arith.cmpi ult, %idx1, %size1 : index
///   %mask1 = arith.andi %mask0, %bound1 : i1
///   ...
///   %boundN = arith.cmpi ult, %idxN, %sizeN : index
///   %new_mask = arith.andi %maskN-1, %boundN : i1
static FailureOr<MapStoreOp>
foldExtractSliceIntoMapStore(RewriterBase &rewriter,
                             tensor::ExtractSliceOp extractSliceOp,
                             MapStoreOp mapStoreOp) {
  assert(mapStoreOp.getInput() == extractSliceOp->getResult(0) &&
         "expected extractSliceOp to be the producer of mapStoreOp");
  // TODO(Max191): Support rank-reducing slices.
  if (extractSliceOp.getSourceType().getRank() !=
      extractSliceOp.getResultType().getRank()) {
    return rewriter.notifyMatchFailure(
        extractSliceOp, "rank reducing extract_slice op is not supported");
  }
  // TODO(Max191): Support non-zero offsets.
  if (!areAllConstantIntValue(extractSliceOp.getMixedOffsets(), 0)) {
    return rewriter.notifyMatchFailure(extractSliceOp,
                                       "non-zero offsets are not supported");
  }
  // TODO(Max191): Support non-unit strides.
  if (!areAllConstantIntValue(extractSliceOp.getMixedStrides(), 1)) {
    return rewriter.notifyMatchFailure(extractSliceOp,
                                       "non-unit strides are not supported");
  }

  // Check if this is an identity extract_slice (all sizes match the source).
  // If so, just replace the input without adding any masking.
  SmallVector<OpFoldResult> sliceSizes = extractSliceOp.getMixedSizes();
  Value source = extractSliceOp.getSource();
  bool isIdentity = true;
  for (auto [dim, sliceSize] : llvm::enumerate(sliceSizes)) {
    ValueBoundsConstraintSet::Variable sourceDimVar(source, dim);
    FailureOr<bool> areEqual =
        ValueBoundsConstraintSet::areEqual(sliceSize, sourceDimVar);
    if (!succeeded(areEqual) || !*areEqual) {
      isIdentity = false;
      break;
    }
  }

  if (isIdentity) {
    // Identity extract_slice: just replace the input without any masking.
    rewriter.modifyOpInPlace(
        mapStoreOp, [&]() { mapStoreOp.getInputMutable().assign(source); });
    return mapStoreOp;
  }

  SmallVector<OpFoldResult> bounds(extractSliceOp.getMixedSizes());
  Block &transformBody = mapStoreOp.getTransformationRegion().front();
  auto yieldOp = cast<IREE::LinalgExt::YieldOp>(transformBody.getTerminator());
  Value mask = yieldOp->getOperands().back();

  OpBuilder::InsertionGuard g(rewriter);
  rewriter.setInsertionPoint(yieldOp);
  Location loc = mapStoreOp->getLoc();
  ArrayRef<BlockArgument> srcIndices = transformBody.getArguments();
  for (auto [bound, srcIdx] : llvm::zip_equal(bounds, srcIndices)) {
    Value boundValue = getValueOrCreateConstantIndexOp(rewriter, loc, bound);
    auto isOutOfBounds =
        arith::CmpIOp::create(rewriter, loc, arith::CmpIPredicate::ult, srcIdx,
                              boundValue)
            ->getResult(0);
    mask = arith::AndIOp::create(rewriter, loc, mask, isOutOfBounds);
  }
  rewriter.modifyOpInPlace(yieldOp, [&]() {
    yieldOp->setOperand(yieldOp->getNumOperands() - 1, mask);
  });
  rewriter.modifyOpInPlace(mapStoreOp, [&]() {
    mapStoreOp.getInputMutable().assign(extractSliceOp.getSource());
  });
  return mapStoreOp;
}

static void buildNestedDistributionLoops(
    RewriterBase &rewriter, Location loc, int64_t distributionLevel,
    SmallVector<OpFoldResult> lbs, SmallVector<OpFoldResult> ubs,
    ArrayRef<DistributionConfig> distConfigs,
    function_ref<void(OpBuilder &, Location, ValueRange)> innerLoopBuilder) {
  DistributionConfig distConfig = distConfigs[distributionLevel];
  SmallVector<OpFoldResult> steps =
      getAsIndexOpFoldResult(rewriter.getContext(), distConfig.tileSizes);
  scf::ForallOp::create(
      rewriter, loc, lbs, ubs, steps, /*outputs=*/ValueRange(),
      rewriter.getArrayAttr(distConfig.mapping),
      /*bodyBuilder=*/[&](OpBuilder &b, Location nestedLoc, ValueRange ivs) {
        SmallVector<OpFoldResult> nestedLbs(ivs);
        SmallVector<OpFoldResult> nestedUbs;
        for (auto [start, step, ub] : llvm::zip_equal(nestedLbs, steps, ubs)) {
          auto tileEnd = IREE::LinalgExt::addOfrs(b, nestedLoc, start, step);
          auto minMap = AffineMap::get(
              2, 0, {b.getAffineDimExpr(0), b.getAffineDimExpr(1)},
              rewriter.getContext());
          auto min = affine::makeComposedFoldedAffineMin(b, nestedLoc, minMap,
                                                         {tileEnd, ub});
          nestedUbs.push_back(min);
        }
        // Continue distribution if there are more distribution levels.
        if (distributionLevel + 1 < distConfigs.size()) {
          buildNestedDistributionLoops(
              rewriter, nestedLoc, distributionLevel + 1, nestedLbs, nestedUbs,
              distConfigs, innerLoopBuilder);
          scf::InParallelOp::create(b, nestedLoc);
          return;
        }
        // Otherwise, tile to one, and generate the inner loop body.
        SmallVector<Value> nestedLbVals =
            getValueOrCreateConstantIndexOp(b, nestedLoc, nestedLbs);
        SmallVector<Value> nestedUbVals =
            getValueOrCreateConstantIndexOp(b, nestedLoc, nestedUbs);
        Value one = arith::ConstantIndexOp::create(rewriter, nestedLoc, 1);
        llvm::Repeated<Value> unitSteps(nestedLbs.size(), one);
        scf::buildLoopNest(rewriter, nestedLoc, nestedLbVals, nestedUbVals,
                           unitSteps, innerLoopBuilder);
        scf::InParallelOp::create(b, nestedLoc);
      });
}

FailureOr<MapStoreOp>
foldPadIntoMapStore(RewriterBase &rewriter, tensor::PadOp padOp,
                    MapStoreOp mapStoreOp,
                    PadDistributionConfigFn padDistributionConfigFn) {
  // Find the output buffer that the mapStoreOp is stored into.
  if (!mapStoreOp->hasOneUse()) {
    return rewriter.notifyMatchFailure(mapStoreOp,
                                       "map_store does not have a single user");
  }
  auto storeOp =
      dyn_cast<IREE::Codegen::StoreToBufferOp>(*mapStoreOp->getUsers().begin());
  if (!storeOp) {
    return rewriter.notifyMatchFailure(
        mapStoreOp, "map_store user is not an iree_codegen.store_to_buffer op");
  }

  rewriter.setInsertionPointAfter(storeOp);
  Location loc = padOp->getLoc();
  SmallVector<OpFoldResult> padSrcSizes =
      tensor::getMixedSizes(rewriter, loc, padOp.getSource());
  SmallVector<OpFoldResult> padResultSizes =
      tensor::getMixedSizes(rewriter, loc, padOp.getResult());

  // Write the padding values directly into the outputBuffer.
  Value outputBuffer = storeOp.getBuffer();
  auto innerLoopBuilder = [&](OpBuilder &b, Location loopLoc, ValueRange ivs) {
    // We need to scatter the padding values according to the existing
    // mapStoreOp transformation, so clone the transformation into the
    // loop nest.
    IRMapping mapping;
    Region &transformRegion = mapStoreOp.getTransformationRegion();
    Block *loopBody = rewriter.getInsertionBlock();
    transformRegion.cloneInto(loopBody->getParent(), mapping);
    Block *clonedTransformBody =
        mapping.lookup(&transformRegion.getBlocks().front());
    // Get a pointer to the YieldOp before inlining the Block.
    auto yieldOp =
        cast<IREE::LinalgExt::YieldOp>(clonedTransformBody->getTerminator());
    rewriter.inlineBlockBefore(clonedTransformBody, loopBody, loopBody->begin(),
                               ivs);
    OpBuilder::InsertionGuard g(b);
    b.setInsertionPointAfter(yieldOp);
    SmallVector<Value> storeIndices(yieldOp.getOperands());
    rewriter.eraseOp(yieldOp);
    Value mask = storeIndices.pop_back_val();
    // Create the store to the outputBuffer.
    auto thenBuilder = [&](OpBuilder &nestedBuilder, Location ifLoc) {
      memref::StoreOp::create(nestedBuilder, ifLoc,
                              padOp.getConstantPaddingValue(), outputBuffer,
                              storeIndices);
      scf::YieldOp::create(nestedBuilder, ifLoc);
    };
    scf::IfOp::create(b, loopLoc, mask, thenBuilder);
  };

  auto stopConditionFn = [](Value /*v*/, std::optional<int64_t> /*dim*/,
                            ValueBoundsConstraintSet & /*cstr*/) -> bool {
    return false;
  };

  // Prove that range [lb, ub) is empty (lb >= ub).
  //
  // For high padding, we have:
  //   lb = low + srcSize
  //   ub = resultSize
  // and for a well-formed pad, resultSize = srcSize + low + high, so:
  //   ub - lb = high
  // Therefore "lb >= ub" -> "high == 0". Expressing the check in terms of lb/ub
  // keeps it robust even when pad amount has been rewritten into a complex
  // form.
  auto proveEmptyRange = [&](RewriterBase &rewriter, Location loc, Value lb,
                             Value ub) -> bool {
    // Model diff = ub - lb  as an affine map (d0, d1) -> (d0 - d1)
    // and ask for a constant upper bound on diff.
    // If we can prove UB(diff) <= 0, then "diff <= 0" -> "ub - lb <= 0" ->
    // "lb >= ub".
    MLIRContext *ctx = rewriter.getContext();
    auto diffMap = AffineMap::get(
        /*dimCount=*/2, /*symbolCount=*/0,
        {getAffineDimExpr(0, ctx) - getAffineDimExpr(1, ctx)}, ctx);
    auto ubConstantBound = ValueBoundsConstraintSet::computeConstantBound(
        presburger::BoundType::UB,
        ValueBoundsConstraintSet::Variable(diffMap, ValueRange{ub, lb}),
        /*stopCondition=*/stopConditionFn);
    return succeeded(ubConstantBound) && *ubConstantBound <= 0;
  };

  // Distribute the padding of each dimension separately. This causes some
  // overlap of the iteration spaces across the loops, but simplifies the
  // implementation. The trade-off is expected to be okay because we expect
  // the padding to be small relative to the size of the output buffer, which
  // makes the overlap small. The below example shows what the overlap looks
  // like.
  //
  // 1st loop (low pad d0)        2nd loop (high pad d0)
  //        |———|                       |———————|
  //          |                             |
  //          v                             v
  //        +———+———————————————————————+———————+     ———
  //        | . | . . . . . . . . . . . | . . . | <--  |  3rd loop (low pad d1)
  //        +———+———————————————————————+———————+     ———
  //        | . | . . . . . . . . . . . | . . . |
  //        | . | . . . . . . . . . . . | . . . |
  //        | . | . . . . . . . . . . . | . . . |
  //        | . | . . . . . . . . . . . | . . . |
  //        +———+———————————————————————+———————+     ———
  //        | . | . . . . . . . . . . . | . . . | <--  |  4th loop (high pad d1)
  //        +———+———————————————————————+———————+     ———
  //
  // Each distributed loop performs the padding for a hyperplane spanning the
  // low or high pad of the corresponding dimension, and fully spanning all
  // other dimensions of the tensor. This means that there is overlap on the
  // intersection of high and low padding for each dimension.
  //
  // In most cases today, the low padding is generally either 0, or very small,
  // and the high padding will be small relative to the size of the full tensor,
  // since padding is usually done to the nearest multiple of some tile size, or
  // for padding an image for a convolution. Because the padding is small, the
  // overlap should also be small, and we can live with the small redundant
  // computation in order to have a more simplified kernel.
  for (auto [idx, low, high] :
       llvm::enumerate(padOp.getMixedLowPad(), padOp.getMixedHighPad())) {
    // Create a distributed loop for the low padding.
    if (!isConstantIntValue(low, 0)) {
      SmallVector<OpFoldResult> ubs(padResultSizes);
      SmallVector<OpFoldResult> lbs(ubs.size(), rewriter.getIndexAttr(0));
      SmallVector<int64_t> shape(padOp.getSourceType().getShape());
      shape[idx] = padOp.getStaticLow()[idx];
      ubs[idx] = low;
      SmallVector<DistributionConfig> distConfigs = padDistributionConfigFn(
          padOp.getSourceType().getShape(), rewriter.getContext());
      buildNestedDistributionLoops(rewriter, loc, /*distributionLevel=*/0, lbs,
                                   ubs, distConfigs, innerLoopBuilder);
    }
    // Create a distributed loop for the high padding.
    if (isConstantIntValue(high, 0)) {
      continue;
    }
    SmallVector<OpFoldResult> ubs(padResultSizes);
    SmallVector<OpFoldResult> lbs(ubs.size(), rewriter.getIndexAttr(0));
    SmallVector<int64_t> shape(padOp.getSourceType().getShape());
    shape[idx] = padOp.getStaticHigh()[idx];
    lbs[idx] = IREE::LinalgExt::addOfrs(rewriter, loc, low, padSrcSizes[idx]);
    // ubs[idx] is padResultSizes[idx] already (since ubs initialized from
    // padResultSizes)

    // Fallback: Even if we couldn't prove the pad amount is exactly 0, the
    // derived half-open iteration space [lb, ub) may still be empty. For
    // well-formed pads this typically means lb == ub (i.e., the high pad is 0),
    // but we phrase it as the general condition lb >= ub for half-open ranges.
    // If it's empty, skip generating this loop nest.
    {
      Value lbV = getValueOrCreateConstantIndexOp(rewriter, loc, lbs[idx]);
      Value ubV = getValueOrCreateConstantIndexOp(rewriter, loc, ubs[idx]);

      if (lbV == ubV) {
        continue;
      }

      if (proveEmptyRange(rewriter, loc, lbV, ubV)) {
        continue;
      }
    }

    SmallVector<DistributionConfig> distConfigs =
        padDistributionConfigFn(shape, rewriter.getContext());
    buildNestedDistributionLoops(rewriter, loc, /*distributionLevel=*/0, lbs,
                                 ubs, distConfigs, innerLoopBuilder);
  }

  // Now that the padding values are being written to the outputBuffer, the
  // padOp becomes a no-op with respect to the index transformation on the
  // non-padded values.
  rewriter.modifyOpInPlace(mapStoreOp, [&]() {
    mapStoreOp.getInputMutable().assign(padOp.getSource());
  });
  return mapStoreOp;
}

/// Decompose a `packOp` and fold all resulting ops into the consumer
/// `mapStoreOp` in one shot. lowerPack produces: source → [pad] →
/// expand_shape → transpose. We fold transpose then expand_shape into the
/// map_store. Identity pads are removed; non-identity pads are folded directly
/// using `foldPadIntoMapStore` when a `padDistributionConfigFn` is provided.
///
/// The behavior depends on whether a `padDistributionConfigFn` is supplied:
/// - If provided: always decompose and fold (non-identity pads are folded
///   directly via `foldPadIntoMapStore`).
/// - If not provided and the pack has no padding semantics: decompose and fold
///   (the pad from decomposition, if any, is a no-op and can be folded away).
/// - If not provided and the pack has padding semantics: return failure.
FailureOr<MapStoreOp>
foldPackIntoMapStore(RewriterBase &rewriter, linalg::PackOp packOp,
                     MapStoreOp mapStoreOp,
                     PadDistributionConfigFn padDistributionConfigFn) {
  assert(mapStoreOp.getInput() == packOp->getResult(0) &&
         "expected packOp to be the producer of mapStoreOp");
  if (!padDistributionConfigFn && packOp.getPaddingValue()) {
    return rewriter.notifyMatchFailure(
        packOp, "pack has padding semantics but no PadDistributionConfigFn "
                "was provided");
  }

  OpBuilder::InsertionGuard g(rewriter);
  rewriter.setInsertionPoint(packOp);
  auto result = linalg::lowerPack(rewriter, packOp,
                                  /*lowerPadLikeWithInsertSlice=*/false);
  if (failed(result)) {
    return failure();
  }

  // After lowerPack the chain is: source → [pad] → expand_shape → transpose.
  // The map_store input is now transpose.result.
  mapStoreOp = foldTransposePermIntoMapStore(
      rewriter, result->transposeOp.getPermutation(),
      result->transposeOp.getInput(), mapStoreOp);

  auto foldResult =
      foldExpandShapeIntoMapStore(rewriter, result->expandShapeOp, mapStoreOp);
  if (failed(foldResult)) {
    return failure();
  }
  mapStoreOp = *foldResult;

  if (result->padOp) {
    if (areAllConstantIntValue(result->padOp.getMixedLowPad(), 0) &&
        areAllConstantIntValue(result->padOp.getMixedHighPad(), 0)) {
      rewriter.replaceOp(result->padOp, result->padOp.getSource());
    } else {
      auto padFold = foldPadIntoMapStore(rewriter, result->padOp, mapStoreOp,
                                         padDistributionConfigFn);
      if (failed(padFold)) {
        return failure();
      }
      mapStoreOp = *padFold;
    }
  }

  return mapStoreOp;
}

/// Decompose an `unPackOp` and fold all resulting ops into the consumer
/// `mapStoreOp` in one shot. lowerUnPack (with lowerUnpadLikeWithExtractSlice
/// = false) produces: source → transpose → collapse_shape → extract_slice →
/// copy. We fold copy, extract_slice, collapse_shape, then transpose into the
/// map_store from innermost to outermost.
static FailureOr<MapStoreOp> foldUnpackIntoMapStore(RewriterBase &rewriter,
                                                    linalg::UnPackOp unPackOp,
                                                    MapStoreOp mapStoreOp) {
  assert(mapStoreOp.getInput() == unPackOp->getResult(0) &&
         "expected unPackOp to be the producer of mapStoreOp");
  OpBuilder::InsertionGuard g(rewriter);
  rewriter.setInsertionPoint(unPackOp);
  // lowerUnpadLikeWithExtractSlice=false ensures "like unpad" unpacks go
  // through the full transpose → collapse_shape → extract_slice chain instead
  // of being lowered to a single rank-reducing extract_slice, which
  // foldExtractSliceIntoMapStore does not support.
  auto result = linalg::lowerUnPack(rewriter, unPackOp,
                                    /*lowerUnpadLikeWithExtractSlice=*/false);
  if (failed(result)) {
    return failure();
  }

  // After lowerUnPack the chain is:
  //   source → transpose → collapse_shape → extract_slice → copy
  // The map_store input is now copy.result.
  // Fold from innermost (closest to map_store) to outermost.

  // 1. Fold copy (identity-like) into map_store.
  mapStoreOp =
      foldIdentityLikeOpIntoMapStore(rewriter, result->copyOp, mapStoreOp);

  // 2. Fold extract_slice into map_store.
  auto extractFold = foldExtractSliceIntoMapStore(
      rewriter, result->extractSliceOp, mapStoreOp);
  if (failed(extractFold)) {
    return failure();
  }
  mapStoreOp = *extractFold;

  // 3. Fold collapse_shape into map_store.
  auto collapseFold = foldCollapseShapeIntoMapStore(
      rewriter, result->collapseShapeOp, mapStoreOp);
  if (failed(collapseFold)) {
    return failure();
  }
  mapStoreOp = *collapseFold;

  // 4. Fold transpose into map_store.
  mapStoreOp = foldTransposePermIntoMapStore(
      rewriter, result->transposeOp.getPermutation(),
      result->transposeOp.getInput(), mapStoreOp);

  return mapStoreOp;
}

FailureOr<MapStoreOp> foldIntoMapStore(RewriterBase &rewriter, Operation *op,
                                       MapStoreOp mapStoreOp) {
  return llvm::TypeSwitch<Operation *, FailureOr<MapStoreOp>>(op)
      .Case([&](linalg::CopyOp copyOp) {
        return foldIdentityLikeOpIntoMapStore(rewriter, copyOp, mapStoreOp);
      })
      .Case([&](linalg::TransposeOp transposeOp) {
        FailureOr<MapStoreOp> res = foldTransposePermIntoMapStore(
            rewriter, transposeOp.getPermutation(), transposeOp.getInput(),
            mapStoreOp);
        return res;
      })
      .Case([&](tensor::ExpandShapeOp expandOp) {
        return foldExpandShapeIntoMapStore(rewriter, expandOp, mapStoreOp);
      })
      .Case([&](tensor::CollapseShapeOp collapseOp) {
        return foldCollapseShapeIntoMapStore(rewriter, collapseOp, mapStoreOp);
      })
      .Case([&](tensor::ExtractSliceOp extractSliceOp) {
        return foldExtractSliceIntoMapStore(rewriter, extractSliceOp,
                                            mapStoreOp);
      })
      .Case<linalg::PackOp>([&](linalg::PackOp packOp) {
        return foldPackIntoMapStore(rewriter, packOp, mapStoreOp,
                                    /*padDistributionConfigFn=*/nullptr);
      })
      .Case<linalg::UnPackOp>([&](linalg::UnPackOp unPackOp) {
        return foldUnpackIntoMapStore(rewriter, unPackOp, mapStoreOp);
      })
      .Case<linalg::GenericOp>([&](linalg::GenericOp genericOp) {
        if (auto perm = linalg::isaTransposeOpInterface(genericOp)) {
          FailureOr<MapStoreOp> res = foldTransposePermIntoMapStore(
              rewriter, perm.value(), genericOp.getDpsInputs()[0], mapStoreOp);
          return res;
        }
        return FailureOr<MapStoreOp>(failure());
      })
      .Default(failure());
}

// Insert identity map_store op after the root and replace all uses.
static MapStoreOp insertIdentityMapStore(RewriterBase &rewriter,
                                         OpResult root) {
  Location loc = root.getLoc();
  SetVector<OpOperand *> originalUses;
  for (OpOperand &use : root.getUses()) {
    originalUses.insert(&use);
  }
  OpBuilder::InsertionGuard g(rewriter);
  rewriter.setInsertionPointAfterValue(root);
  Type elementType = cast<RankedTensorType>(root.getType()).getElementType();
  SmallVector<OpFoldResult> sizes = tensor::getMixedSizes(rewriter, loc, root);
  Value mapStoreDest =
      tensor::EmptyOp::create(rewriter, loc, sizes, elementType);
  auto mapStoreOp =
      MapStoreOp::createIdentityMapStore(rewriter, loc, root, mapStoreDest);
  rewriter.replaceUsesWithIf(
      root, mapStoreOp.getResult(0),
      [&](OpOperand &use) { return originalUses.contains(&use); });
  LDBG() << "Created identity map_store:\n" << mapStoreOp;
  return mapStoreOp;
}

bool isSupportedSingleInputRelayoutOpForResult(Operation *op) {
  if (isa<tensor::ExpandShapeOp, tensor::CollapseShapeOp,
          tensor::ExtractSliceOp, tensor::PadOp, linalg::CopyOp,
          linalg::TransposeOp, linalg::PackOp, linalg::UnPackOp>(op)) {
    return true;
  }
  auto genericOp = dyn_cast<linalg::GenericOp>(op);
  return genericOp && linalg::isaTransposeOpInterface(genericOp).has_value();
}

bool isSupportedSingleInputRelayoutOpForSource(Operation *op) {
  if (isSupportedSingleInputRelayoutOpForResult(op)) {
    return true;
  }
  auto linalgOp = dyn_cast<linalg::LinalgOp>(op);
  return linalgOp && linalg::isaBroadcastOpInterface(linalgOp).has_value();
}

/// Returns true if `user` is a supported relayout op that uses `result` as a
/// valid source for extending the relayout chain.
///
/// We use different logic for DPS vs non-DPS ops so neither case needs to
/// consider operand numbers:
/// - DPS ops (e.g. linalg broadcast): Iterate getDpsInputOperands(). Only
///   consider uses as an *input* operand.
/// - Non-DPS ops (expand_shape, collapse_shape, transpose, etc.): Caller
///   iterates result.getUsers(), so we know user consumes result. These ops
///   have a single input.
static bool isRelayoutChainExtension(Operation *user, Value result) {
  if (!isSupportedSingleInputRelayoutOpForSource(user)) {
    return false;
  }
  auto dpsOp = dyn_cast<DestinationStyleOpInterface>(user);
  if (dpsOp) {
    for (OpOperand *input : dpsOp.getDpsInputOperands()) {
      if (input->get() == result) {
        return true;
      }
    }
    return false;
  }
  // Non-DPS: caller iterates result.getUsers(), so user consumes result.
  // Single-input relayout ops have only one place for it.
  return true;
}

/// Collects all relayout ops in the chain starting from `relayoutOp`
/// (inclusive). The chain extends through result->user edges where the user is
/// a supported relayout op.
static void collectRelayoutChain(Operation *relayoutOp,
                                 SmallPtrSetImpl<Operation *> &chain) {
  if (!chain.insert(relayoutOp).second) {
    return;
  }
  Value result = relayoutOp->getResult(0);
  for (Operation *user : result.getUsers()) {
    if (isRelayoutChainExtension(user, result)) {
      collectRelayoutChain(user, chain);
    }
  }
}

/// Returns true if the relayout op starts a "complex" chain.
/// A "complex" chain is one that is difficult for bufferization to handle. The
/// conditions are:
/// - The chain length must be >= 2.
/// - The chain must contain at least one reshape op (expand_shape or
/// collapse_shape).
/// - The chain must contain at least one op that is not extract_slice or a
/// reshape op.
/// - The chain must not contain any op that carries a `lowering_config`
/// attribute, since such ops are roots for downstream transformations (e.g.
/// tiling) and should not be folded into a MapLoad.
static bool isComplexRelayoutChain(Operation *relayoutOp) {
  assert(isSupportedSingleInputRelayoutOpForSource(relayoutOp) &&
         "expected a supported relayout op");
  SmallPtrSet<Operation *, 4> chain;
  collectRelayoutChain(relayoutOp, chain);
  if (chain.size() < 2) {
    return false;
  }
  bool hasReshape = llvm::any_of(
      chain, llvm::IsaPred<tensor::ExpandShapeOp, tensor::CollapseShapeOp>);
  // Need at least one op that is not reshape and not extract_slice (e.g. copy,
  // transpose, broadcast, pad).
  bool allReshapeOrExtractSlice = llvm::all_of(
      chain, llvm::IsaPred<tensor::ExpandShapeOp, tensor::CollapseShapeOp,
                           tensor::ExtractSliceOp>);
  if (!hasReshape || allReshapeOrExtractSlice) {
    return false;
  }
  if (llvm::any_of(chain,
                   [](Operation *op) { return getLoweringConfig(op); })) {
    return false;
  }
  return true;
}

/// Collects direct relayout op users of `loadResult` that start complex
/// relayout chains. When `combineNonComplexChains` is true it includes all
/// relayout op users regardless of complexity.
static SmallPtrSet<Operation *, 4>
getComplexChainRelayoutUsers(Value loadResult,
                             bool combineNonComplexChains = false) {
  SmallPtrSet<Operation *, 4> complexUsers;
  for (Operation *user : loadResult.getUsers()) {
    if (isRelayoutChainExtension(user, loadResult) &&
        (combineNonComplexChains || isComplexRelayoutChain(user))) {
      complexUsers.insert(user);
    }
  }
  return complexUsers;
}

// This is only desirable in the dispatch scope but not in the workgroup scope.
static bool
shouldDoReshapesByExpansion(IREE::Codegen::RelayoutCombinationScope scope) {
  if (scope == IREE::Codegen::RelayoutCombinationScope::Dispatch) {
    return true;
  }
  return false;
}

/// Insert identity map_store ops after the given operation if it is a valid
/// leaf op of a relayout op chain. A relayout op chain is a sequence of
/// relayout ops (defined by `isSupportedSingleInputRelayoutOpForResult`) for
/// which the only users of the ops in the chain are relayout ops, except for
/// the leaves of the chain. The leaves are simply relayout ops that have non
/// relayout op users. The `controlFn` is a callback on the leaf OpResult that
/// provides control over whether or not to insert a map_store op.
struct InsertMapStoreOpPattern : RewritePattern {
  InsertMapStoreOpPattern(MLIRContext *context,
                          CombineRelayoutOpsControlFnRef controlFn = nullptr,
                          PatternBenefit benefit = 1)
      : RewritePattern(MatchAnyOpTypeTag(), benefit, context),
        controlFn(controlFn) {}

  LogicalResult matchAndRewrite(Operation *op,
                                PatternRewriter &rewriter) const override {
    if (!isSupportedSingleInputRelayoutOpForResult(op)) {
      return failure();
    }
    // Relayout ops with only relayout op users are not leaves.
    auto isDimOrSupportedRelayoutOp = [](Operation *op) {
      return isSupportedSingleInputRelayoutOpForResult(op) ||
             isa<tensor::DimOp>(op);
    };
    if (llvm::all_of(op->getUsers(), isDimOrSupportedRelayoutOp)) {
      return failure();
    }
    // All relayout ops have a single result.
    OpResult leaf = op->getResult(0);
    if (controlFn && !controlFn(leaf)) {
      return failure();
    }
    (void)insertIdentityMapStore(rewriter, leaf);
    return success();
  }

private:
  CombineRelayoutOpsControlFnRef controlFn;
};

LogicalResult
combineLayoutTransformation(MLIRContext *ctx, FunctionOpInterface funcOp,
                            PadDistributionConfigFn padDistributionConfigFn,
                            bool doReshapeByExpansion,
                            CombineRelayoutOpsControlFnRef controlFn) {
  // Sink relayout operations to the end of the funcOp.
  RewritePatternSet propagationPatterns(ctx);
  tensor::populateFoldTensorEmptyPatterns(propagationPatterns);
  tensor::ExpandShapeOp::getCanonicalizationPatterns(propagationPatterns, ctx);
  tensor::CollapseShapeOp::getCanonicalizationPatterns(propagationPatterns,
                                                       ctx);
  if (doReshapeByExpansion) {
    // Only sink reshape ops, so bail if the consumer operation is a reshape.
    auto controlSinkReshapesFn = [](OpOperand *operand) -> bool {
      Operation *consumer = operand->getOwner();
      return !isa<tensor::ExpandShapeOp, tensor::CollapseShapeOp>(consumer);
    };
    linalg::populateFoldReshapeOpsByExpansionPatterns(propagationPatterns,
                                                      controlSinkReshapesFn);
  }
  // Only sink unpack ops, so bail if the producer operation is not an unpack.
  // Also only sink unpack ops when new pack operations will not be created.
  // This means the consumer op must have at most one additional destination
  // operand, and it must come from an empty tensor.
  auto controlPropagationFn = [](OpOperand *operand) -> bool {
    Operation *producer = operand->get().getDefiningOp();
    Operation *consumer = operand->getOwner();
    if (!isa_and_nonnull<linalg::UnPackOp>(producer)) {
      return false;
    }
    // Pads and reshapes will not produce extra pack ops.
    if (isa<tensor::PadOp, tensor::ExpandShapeOp>(consumer)) {
      return true;
    }
    // Otherwise, the consumer must be a GenericOp with all of its `outs`
    // operands coming from tensor.empty ops, and the `operand` must be the
    // sole `ins` operand of the generic op. This ensures that no additional
    // linalg.pack ops will be created on other inputs of the generic op.
    // TODO(Max191): Remove the restriction of not creating new pack ops once
    // codegen can handle it.
    auto genericConsumer = dyn_cast<linalg::GenericOp>(consumer);
    if (!genericConsumer || genericConsumer.getNumDpsInputs() != 1 ||
        *genericConsumer.getDpsInputOperand(0) != *operand) {
      return false;
    }
    return llvm::all_of(
        genericConsumer.getDpsInits(), [&](Value consumerOperand) -> bool {
          return consumerOperand.getDefiningOp<tensor::EmptyOp>();
        });
  };
  linalg::populateDataLayoutPropagationPatterns(
      propagationPatterns, controlPropagationFn, /*PoisonPaddingOk=*/true);
  // TODO(Max191): The propagation patterns could be applied at the same time as
  // relayout ops are folded into the map_store, which may enable even more
  // folding. This requires the relayout op folding to be done as pattern
  // rewrites, and also direct foldings for pack and unpack ops.
  if (failed(applyPatternsGreedily(funcOp, std::move(propagationPatterns)))) {
    return failure();
  }

  // Resolve dims aggressively and canonicalize to simplify the IR.
  // This improves pattern matching and makes it easier for subsequent
  // rewrites/proofs to recognize identity/no-op cases.
  {
    RewritePatternSet patterns(ctx);
    memref::populateResolveRankedShapedTypeResultDimsPatterns(patterns);
    IREE::Util::AssumeIntOp::getCanonicalizationPatterns(patterns, ctx);

    if (failed(applyPatternsGreedily(funcOp, std::move(patterns)))) {
      return failure();
    }
  }

  // Combine relayout operations into new the map_store ops.
  RewritePatternSet relayoutCombinationPatterns(ctx);
  relayoutCombinationPatterns.add<InsertMapStoreOpPattern>(ctx, controlFn);
  populateCombineRelayoutOpPatterns(relayoutCombinationPatterns,
                                    padDistributionConfigFn);
  memref::populateResolveRankedShapedTypeResultDimsPatterns(
      relayoutCombinationPatterns);
  if (failed(applyPatternsGreedily(funcOp,
                                   std::move(relayoutCombinationPatterns)))) {
    return failure();
  }

  // Clean up any identity map_store ops after combining.
  IRRewriter rewriter(ctx);
  funcOp->walk([&](MapStoreOp mapStoreOp) {
    if (mapStoreOp.isIdentity()) {
      rewriter.replaceOp(mapStoreOp, mapStoreOp.getInput());
    }
  });
  return success();
}

/// TODO(#20530): Improve heuristic for tile size selection.
static SmallVector<DistributionConfig>
defaultPadWorkgroupDistributionConfigFn(ArrayRef<int64_t> iterationBounds,
                                        MLIRContext *ctx) {
  DistributionConfig workgroupDistributionConfig;
  workgroupDistributionConfig.tileSizes.assign(iterationBounds.size(), 1);
  workgroupDistributionConfig.tileSizes.back() = 64;
  workgroupDistributionConfig.mapping = llvm::map_to_vector(
      llvm::seq<int64_t>(iterationBounds.size()),
      [&](int64_t dim) -> Attribute {
        switch (dim) {
        case 0:
        case 1:
        case 2:
          return IREE::Codegen::WorkgroupMappingAttr::get(
              ctx, IREE::Codegen::symbolizeWorkgroupId(dim).value());
        default:
          return IREE::Codegen::WorkgroupMappingAttr::get(
              ctx, IREE::Codegen::WorkgroupId::IdZ, dim - 2);
        }
      });
  std::reverse(workgroupDistributionConfig.mapping.begin(),
               workgroupDistributionConfig.mapping.end());
  return {workgroupDistributionConfig};
}

CombineRelayoutOpsControlFn
getCombineRelayoutOpsControlFn(IREE::Codegen::RelayoutCombinationScope scope) {
  CombineRelayoutOpsControlFn controlFn;
  switch (scope) {
  // Control function for Dispatch scope. Filters to only relayout ops with
  // a single iree_codegen.store_to_buffer user.
  case IREE::Codegen::RelayoutCombinationScope::Dispatch:
    controlFn = [](OpResult leaf) {
      if (leaf.getNumUses() != 1) {
        return false;
      }
      return isa<IREE::Codegen::StoreToBufferOp>(*leaf.getUsers().begin());
    };
    break;
  // Control function for DispatchReshape scope. Like Dispatch, but additionally
  // requires the backward relayout slice to contain a `tensor.expand_shape` or
  // `tensor.collapse_shape` — the non-tileable reshape whose presence is what
  // makes folding into a `map_store` necessary (see the enum doc comment).
  case IREE::Codegen::RelayoutCombinationScope::DispatchReshape:
    controlFn = [](OpResult leaf) {
      if (leaf.getNumUses() != 1) {
        return false;
      }
      if (!isa<IREE::Codegen::StoreToBufferOp>(*leaf.getUsers().begin())) {
        return false;
      }
      llvm::SetVector<Operation *> slice;
      BackwardSliceOptions options;
      options.filter = isSupportedSingleInputRelayoutOpForResult;
      options.inclusive = true;
      if (failed(getBackwardSlice(leaf, &slice, options))) {
        return false;
      }
      return llvm::any_of(
          slice, llvm::IsaPred<tensor::CollapseShapeOp, tensor::ExpandShapeOp>);
    };
    break;
  // Control function for Workgroup scope. Filters to only relayout ops with
  // a single tensor.parallel_insert_slice user inside of a workgroup
  // scf.forall op. Relayout chains of only reshapes are also filtered out,
  // because these chains can usually be handled by bufferization.
  case IREE::Codegen::RelayoutCombinationScope::Workgroup:
    controlFn = [](OpResult leaf) {
      if (leaf.getNumUses() != 1) {
        return false;
      }
      auto parallelInsertOp =
          dyn_cast<tensor::ParallelInsertSliceOp>(*leaf.getUsers().begin());
      if (!parallelInsertOp) {
        return false;
      }
      auto forallOp = parallelInsertOp->getParentOfType<scf::ForallOp>();
      if (!forallOp || !forallOp.getMapping() ||
          !llvm::all_of(forallOp.getMapping().value(),
                        llvm::IsaPred<IREE::Codegen::WorkgroupMappingAttr>)) {
        return false;
      }
      auto bbArg = dyn_cast<BlockArgument>(parallelInsertOp.getDest());
      if (!bbArg || forallOp.getCombiningOps(bbArg).size() != 1) {
        return false;
      }
      // If there are only reshape ops, then bufferization can usually handle
      // it, so don't introduce map_store.
      llvm::SetVector<Operation *> slice;
      BackwardSliceOptions options;
      options.filter = isSupportedSingleInputRelayoutOpForResult;
      options.inclusive = true;
      LogicalResult result =
          getBackwardSlice(parallelInsertOp.getSource(), &slice, options);
      if (failed(result)) {
        return false;
      }
      return !llvm::all_of(
          slice, llvm::IsaPred<tensor::CollapseShapeOp, tensor::ExpandShapeOp>);
    };
    break;
  }
  return controlFn;
}

namespace {

struct CombineResultLayoutTransformationPass final
    : impl::CombineResultLayoutTransformationPassBase<
          CombineResultLayoutTransformationPass> {
  using Base::Base;

  void runOnOperation() override {
    CombineRelayoutOpsControlFn controlFn =
        getCombineRelayoutOpsControlFn(this->scope);
    bool doReshapesByExpansion = shouldDoReshapesByExpansion(this->scope);
    if (failed(
            combineLayoutTransformation(&getContext(), getOperation(),
                                        defaultPadWorkgroupDistributionConfigFn,
                                        doReshapesByExpansion, controlFn))) {
      return signalPassFailure();
    }

    MLIRContext *context = &getContext();
    FunctionOpInterface funcOp = getOperation();
    {
      RewritePatternSet patterns(context);
      populateFuseTilableForallConsumersPattern(patterns);
      scf::ForallOp::getCanonicalizationPatterns(patterns, context);
      tensor::populateFoldTensorEmptyPatterns(patterns);
      if (failed(applyPatternsGreedily(funcOp, std::move(patterns)))) {
        return signalPassFailure();
      }
    }
  }
};

} // namespace

//===----------------------------------------------------------------------===//
// Combining Layout Transformation Ops into MapLoadOp
//===----------------------------------------------------------------------===//

/// Folds `consumerOp` into the producer `mapLoadOp` using `rewriter`.
/// The `indexTransformBuilder` lambda defines how output indices of the new
/// map_load map to input indices of the original map_load. Optionally,
/// `newFillValue` can be provided to override the fill value (e.g., for pad).
///
/// This handles the shared boilerplate:
/// 1. Create new dest tensor from consumer result shape
/// 2. Clone map_load with new dest and region
/// 3. Apply index transformation via insertTransformationAtStart
/// 4. Optionally update fill value
/// 5. Replace consumer with new map_load result
static FailureOr<MapLoadOp> foldConsumerIntoMapLoadImpl(
    RewriterBase &rewriter, Operation *consumerOp, MapLoadOp mapLoadOp,
    function_ref<SmallVector<Value>(ArrayRef<BlockArgument>)>
        indexTransformBuilder,
    std::optional<Value> newFillValue = std::nullopt) {
  Location loc = consumerOp->getLoc();
  OpBuilder::InsertionGuard g(rewriter);
  rewriter.setInsertionPoint(consumerOp);

  // Reify may emit small helper ops at the insertion point. They dominate the
  // new `tensor.empty` / `map_load`.
  Value consumerResult = consumerOp->getResult(0);
  Type elementType = getElementTypeOrSelf(consumerResult.getType());
  FailureOr<SmallVector<OpFoldResult>> newSizes =
      reifyShapeOfResult(rewriter, consumerOp, /*resultIndex=*/0);
  if (failed(newSizes)) {
    return rewriter.notifyMatchFailure(
        consumerOp,
        "consumer does not implement ReifyRankedShapedTypeOpInterface");
  }
  Value newDest =
      tensor::EmptyOp::create(rewriter, loc, *newSizes, elementType);

  // Clone the map_load with new dest.
  auto newMapLoad = MapLoadOp::create(rewriter, loc, newDest.getType(),
                                      mapLoadOp.getSource(), newDest);
  rewriter.cloneRegionBefore(mapLoadOp.getTransformationRegion(),
                             newMapLoad.getTransformationRegion(),
                             newMapLoad.getTransformationRegion().begin());

  // Prepend the index transformation.
  int64_t newRank = cast<ShapedType>(consumerResult.getType()).getRank();
  newMapLoad.insertTransformationAtStart(rewriter, indexTransformBuilder,
                                         newRank);

  // Optionally update fill value (for pad ops).
  if (newFillValue.has_value()) {
    Block &transformBody = newMapLoad.getTransformationRegion().front();
    auto yieldOp =
        cast<IREE::LinalgExt::YieldOp>(transformBody.getTerminator());
    rewriter.modifyOpInPlace(yieldOp, [&]() {
      yieldOp.setOperand(yieldOp.getNumOperands() - 1, *newFillValue);
    });
  }

  rewriter.replaceOp(consumerOp, newMapLoad.getResult(0));
  LDBG() << "Folded consumer " << consumerOp->getName().getStringRef()
         << " into map_load:\n"
         << newMapLoad;
  return newMapLoad;
}

/// Fold an `op` that does not affect index computation into a `mapLoadOp`.
/// This is used for ops like `linalg::CopyOp`.
static FailureOr<MapLoadOp>
foldIdentityLikeOpIntoMapLoad(RewriterBase &rewriter, Operation *op,
                              MapLoadOp mapLoadOp) {
  assert(op->getOperand(0) == mapLoadOp.getResult(0) &&
         "expected mapLoadOp to be the producer of op input");
  rewriter.replaceOp(op, mapLoadOp.getResult(0));
  LDBG() << "Folded consumer " << op->getName().getStringRef()
         << " into map_load (identity fold)";
  return mapLoadOp;
}

/// Fold a transpose (given by `perm`) consumer into a producer `mapLoadOp`.
/// For consumer folding, we apply the INVERSE permutation (output indices →
/// input indices). `consumerOp` is the op being folded (used for replacement).
static FailureOr<MapLoadOp> foldTransposePermIntoMapLoad(RewriterBase &rewriter,
                                                         Operation *consumerOp,
                                                         ArrayRef<int64_t> perm,
                                                         MapLoadOp mapLoadOp) {
  SmallVector<int64_t> inversePerm = invertPermutationVector(perm);
  return foldConsumerIntoMapLoadImpl(
      rewriter, consumerOp, mapLoadOp,
      [inversePerm](ArrayRef<BlockArgument> indices) -> SmallVector<Value> {
        SmallVector<Value> indexValues(indices.begin(), indices.end());
        return applyPermutation(indexValues, inversePerm);
      });
}

/// Fold a consumer `transposeOp` into a producer `mapLoadOp`.
static FailureOr<MapLoadOp>
foldTransposeIntoMapLoad(RewriterBase &rewriter,
                         linalg::TransposeOp transposeOp, MapLoadOp mapLoadOp) {
  assert(transposeOp.getInput() == mapLoadOp.getResult(0) &&
         "expected mapLoadOp to be the producer of transposeOp input");
  return foldTransposePermIntoMapLoad(rewriter, transposeOp,
                                      transposeOp.getPermutation(), mapLoadOp);
}

/// Fold a consumer reshape op (expand_shape or collapse_shape) into a producer
/// `mapLoadOp`. Index transformation: linearize in result space, delinearize
/// in source space.
template <typename ReshapeOpTy>
static FailureOr<MapLoadOp> foldReshapeIntoMapLoad(RewriterBase &rewriter,
                                                   ReshapeOpTy reshapeOp,
                                                   MapLoadOp mapLoadOp) {
  assert(reshapeOp.getSrc() == mapLoadOp.getResult(0) &&
         "expected mapLoadOp to be the producer of reshapeOp input");
  // Cannot fold if either side of the reshape is rank-0 (scalar). The result
  // being rank-0 would make the new map_load output rank-0, and the source
  // being rank-0 would produce an empty delinearization.
  if (cast<RankedTensorType>(reshapeOp.getSrc().getType()).getRank() == 0 ||
      cast<RankedTensorType>(reshapeOp.getResult().getType()).getRank() == 0) {
    return failure();
  }
  Location loc = reshapeOp->getLoc();
  OpBuilder::InsertionGuard g(rewriter);
  rewriter.setInsertionPoint(reshapeOp);
  SmallVector<OpFoldResult> srcDims =
      tensor::getMixedSizes(rewriter, loc, mapLoadOp.getOutput());
  FailureOr<SmallVector<OpFoldResult>> resultDims =
      reifyShapeOfResult(rewriter, reshapeOp, /*resultIndex=*/0);
  if (failed(resultDims)) {
    return rewriter.notifyMatchFailure(
        reshapeOp,
        "reshape does not implement ReifyRankedShapedTypeOpInterface");
  }

  return foldConsumerIntoMapLoadImpl(
      rewriter, reshapeOp, mapLoadOp,
      [&rewriter, loc, resultDims = *resultDims,
       srcDims](ArrayRef<BlockArgument> indices) -> SmallVector<Value> {
        OpBuilder::InsertionGuard g(rewriter);
        SmallVector<Value> indexValues(indices.begin(), indices.end());
        auto linearizeIndexOp = affine::AffineLinearizeIndexOp::create(
            rewriter, loc, indexValues, resultDims, /*disjoint=*/true);
        auto delinearizeIndexOp = affine::AffineDelinearizeIndexOp::create(
            rewriter, loc, linearizeIndexOp.getResult(), srcDims,
            /*hasOuterBound=*/true);
        return delinearizeIndexOp->getResults();
      });
}

/// Fold a consumer tensor::ExpandShapeOp into a producer `mapLoadOp`.
static FailureOr<MapLoadOp>
foldExpandShapeIntoMapLoad(RewriterBase &rewriter,
                           tensor::ExpandShapeOp expandShapeOp,
                           MapLoadOp mapLoadOp) {
  return foldReshapeIntoMapLoad(rewriter, expandShapeOp, mapLoadOp);
}

/// Fold a consumer tensor::CollapseShapeOp into a producer `mapLoadOp`.
static FailureOr<MapLoadOp>
foldCollapseShapeIntoMapLoad(RewriterBase &rewriter,
                             tensor::CollapseShapeOp collapseShapeOp,
                             MapLoadOp mapLoadOp) {
  return foldReshapeIntoMapLoad(rewriter, collapseShapeOp, mapLoadOp);
}

/// Fold a consumer `extractSliceOp` into a producer `mapLoadOp`.
/// Index transformation: original_idx = offset + new_idx * stride
static FailureOr<MapLoadOp>
foldExtractSliceIntoMapLoad(RewriterBase &rewriter,
                            tensor::ExtractSliceOp extractSliceOp,
                            MapLoadOp mapLoadOp) {
  assert(extractSliceOp.getSource() == mapLoadOp.getResult(0) &&
         "expected mapLoadOp to be the producer of extractSliceOp input");

  if (extractSliceOp.getSourceType().getRank() !=
      extractSliceOp.getResultType().getRank()) {
    return rewriter.notifyMatchFailure(
        extractSliceOp, "rank reducing extract_slice op is not supported");
  }

  Location loc = extractSliceOp->getLoc();
  SmallVector<OpFoldResult> offsets = extractSliceOp.getMixedOffsets();
  SmallVector<OpFoldResult> strides = extractSliceOp.getMixedStrides();

  auto indexTransformBuilder =
      [&](ArrayRef<BlockArgument> newOutputIndices) -> SmallVector<Value> {
    SmallVector<Value> originalIndices;
    for (auto [idx, offset, stride] :
         llvm::zip_equal(newOutputIndices, offsets, strides)) {
      Value offsetValue =
          getValueOrCreateConstantIndexOp(rewriter, loc, offset);
      Value strideValue =
          getValueOrCreateConstantIndexOp(rewriter, loc, stride);
      // original_idx = offset + idx * stride
      Value scaledIdx = arith::MulIOp::create(rewriter, loc, idx, strideValue,
                                              arith::IntegerOverflowFlags::nsw);
      Value originalIdx =
          arith::AddIOp::create(rewriter, loc, offsetValue, scaledIdx,
                                arith::IntegerOverflowFlags::nsw);
      originalIndices.push_back(originalIdx);
    }
    return originalIndices;
  };

  return foldConsumerIntoMapLoadImpl(rewriter, extractSliceOp, mapLoadOp,
                                     indexTransformBuilder);
}

/// Fold a consumer broadcast op (named or generic) into a producer `map_load`.
static FailureOr<MapLoadOp>
foldBroadcastIntoMapLoad(RewriterBase &rewriter, linalg::LinalgOp broadcastOp,
                         MapLoadOp mapLoadOp) {
  if (!linalg::isaBroadcastOpInterface(broadcastOp).has_value()) {
    return rewriter.notifyMatchFailure(broadcastOp.getOperation(),
                                       "op is not a broadcast");
  }
  assert(broadcastOp.getDpsInputs()[0] == mapLoadOp.getResult(0) &&
         "expected map_load to be the producer of broadcast input");

  AffineMap inputMap = broadcastOp.getIndexingMapsArray()[0];
  return foldConsumerIntoMapLoadImpl(
      rewriter, broadcastOp.getOperation(), mapLoadOp,
      [inputMap](ArrayRef<BlockArgument> indices) -> SmallVector<Value> {
        SmallVector<Value> sourceIndices;
        sourceIndices.reserve(inputMap.getNumResults());
        for (AffineExpr expr : inputMap.getResults()) {
          unsigned pos = cast<AffineDimExpr>(expr).getPosition();
          sourceIndices.push_back(indices[pos]);
        }
        return sourceIndices;
      });
}

/// Fold a consumer `padOp` into a producer `mapLoadOp`.
/// Index transformation: source_idx = new_idx - low_pad
/// Fill value is set to the pad value.
static FailureOr<MapLoadOp> foldPadIntoMapLoad(RewriterBase &rewriter,
                                               tensor::PadOp padOp,
                                               MapLoadOp mapLoadOp) {
  assert(padOp.getSource() == mapLoadOp.getResult(0) &&
         "expected mapLoadOp to be the producer of padOp input");

  // Only support constant pad values for now.
  Value padValue = padOp.getConstantPaddingValue();
  if (!padValue) {
    return rewriter.notifyMatchFailure(
        padOp, "non-constant padding value is not supported");
  }

  // Check if the map_load already has a real (non-poison) padding value.
  // If so, we cannot safely replace it as that would lose the existing padding.
  Value currentPadValue = mapLoadOp.getPaddingValue();
  if (!currentPadValue.getDefiningOp<ub::PoisonOp>()) {
    return rewriter.notifyMatchFailure(
        padOp, "map_load already has a non-poison padding value; folding "
               "another pad would overwrite it");
  }

  Location loc = padOp->getLoc();
  SmallVector<OpFoldResult> lowPad = padOp.getMixedLowPad();

  auto indexTransformBuilder =
      [&](ArrayRef<BlockArgument> newOutputIndices) -> SmallVector<Value> {
    SmallVector<Value> sourceIndices;
    for (auto [idx, low] : llvm::zip_equal(newOutputIndices, lowPad)) {
      Value lowValue = getValueOrCreateConstantIndexOp(rewriter, loc, low);
      // source_idx = idx - low_pad
      Value sourceIdx = arith::SubIOp::create(rewriter, loc, idx, lowValue,
                                              arith::IntegerOverflowFlags::nsw);
      sourceIndices.push_back(sourceIdx);
    }
    return sourceIndices;
  };

  return foldConsumerIntoMapLoadImpl(rewriter, padOp, mapLoadOp,
                                     indexTransformBuilder, padValue);
}

/// Decompose a consumer `packOp` and fold all resulting ops into the producer
/// `mapLoadOp` in one shot. lowerPack produces: source → [pad] →
/// expand_shape → transpose. We fold pad, expand_shape, then transpose into
/// the map_load from outermost to innermost.
static FailureOr<MapLoadOp> foldPackIntoMapLoad(RewriterBase &rewriter,
                                                linalg::PackOp packOp,
                                                MapLoadOp mapLoadOp) {
  assert(packOp.getSource() == mapLoadOp.getResult(0) &&
         "expected mapLoadOp to be the producer of packOp input");
  OpBuilder::InsertionGuard g(rewriter);
  rewriter.setInsertionPoint(packOp);
  auto result = linalg::lowerPack(rewriter, packOp,
                                  /*lowerPadLikeWithInsertSlice=*/false);
  if (failed(result)) {
    return failure();
  }

  // After lowerPack the chain is: map_load → [pad] → expand_shape → transpose.
  // Fold from outermost (closest to map_load) to innermost.

  // 1. Handle pad.
  if (result->padOp) {
    if (areAllConstantIntValue(result->padOp.getMixedLowPad(), 0) &&
        areAllConstantIntValue(result->padOp.getMixedHighPad(), 0)) {
      rewriter.replaceOp(result->padOp, result->padOp.getSource());
    } else {
      auto padFold = foldPadIntoMapLoad(rewriter, result->padOp, mapLoadOp);
      if (failed(padFold)) {
        return failure();
      }
      mapLoadOp = *padFold;
    }
  }

  // 2. Fold expand_shape into map_load.
  auto expandFold =
      foldExpandShapeIntoMapLoad(rewriter, result->expandShapeOp, mapLoadOp);
  if (failed(expandFold)) {
    return failure();
  }
  mapLoadOp = *expandFold;

  // 3. Fold transpose into map_load.
  auto transposeFold = foldTransposePermIntoMapLoad(
      rewriter, result->transposeOp, result->transposeOp.getPermutation(),
      mapLoadOp);
  if (failed(transposeFold)) {
    return failure();
  }
  mapLoadOp = *transposeFold;
  return mapLoadOp;
}

/// Decompose a consumer `unPackOp` and fold all resulting ops into the producer
/// `mapLoadOp` in one shot. lowerUnPack (with lowerUnpadLikeWithExtractSlice
/// = false) produces: source → transpose → collapse_shape → extract_slice →
/// copy. We fold from outermost (closest to map_load) to innermost.
static FailureOr<MapLoadOp> foldUnpackIntoMapLoad(RewriterBase &rewriter,
                                                  linalg::UnPackOp unPackOp,
                                                  MapLoadOp mapLoadOp) {
  assert(unPackOp.getSource() == mapLoadOp.getResult(0) &&
         "expected mapLoadOp to be the producer of unPackOp input");
  OpBuilder::InsertionGuard g(rewriter);
  rewriter.setInsertionPoint(unPackOp);
  // lowerUnpadLikeWithExtractSlice=false ensures "like unpad" unpacks go
  // through the full transpose → collapse_shape → extract_slice chain instead
  // of being lowered to a single rank-reducing extract_slice, which
  // foldExtractSliceIntoMapLoad does not support.
  auto result = linalg::lowerUnPack(rewriter, unPackOp,
                                    /*lowerUnpadLikeWithExtractSlice=*/false);
  if (failed(result)) {
    return failure();
  }

  // After lowerUnPack the chain is:
  //   map_load → transpose → collapse_shape → extract_slice → copy
  // Fold from outermost (closest to map_load) to innermost.

  // 1. Fold transpose into map_load.
  auto transposeFold = foldTransposePermIntoMapLoad(
      rewriter, result->transposeOp, result->transposeOp.getPermutation(),
      mapLoadOp);
  if (failed(transposeFold)) {
    return failure();
  }
  mapLoadOp = *transposeFold;

  // 2. Fold collapse_shape into map_load.
  auto collapseFold = foldCollapseShapeIntoMapLoad(
      rewriter, result->collapseShapeOp, mapLoadOp);
  if (failed(collapseFold)) {
    return failure();
  }
  mapLoadOp = *collapseFold;

  // 3. Fold extract_slice into map_load.
  auto extractFold =
      foldExtractSliceIntoMapLoad(rewriter, result->extractSliceOp, mapLoadOp);
  if (failed(extractFold)) {
    return failure();
  }
  mapLoadOp = *extractFold;

  // 4. Fold copy (identity-like) into map_load.
  auto copyFold =
      foldIdentityLikeOpIntoMapLoad(rewriter, result->copyOp, mapLoadOp);
  if (failed(copyFold)) {
    return failure();
  }
  mapLoadOp = *copyFold;
  return mapLoadOp;
}

/// Fold a consumer relayout op into a producer map_load.
FailureOr<MapLoadOp> foldIntoMapLoad(RewriterBase &rewriter, Operation *op,
                                     MapLoadOp mapLoadOp) {
  return llvm::TypeSwitch<Operation *, FailureOr<MapLoadOp>>(op)
      .Case<linalg::CopyOp>([&](linalg::CopyOp copyOp) {
        return foldIdentityLikeOpIntoMapLoad(rewriter, copyOp, mapLoadOp);
      })
      .Case<linalg::TransposeOp>([&](linalg::TransposeOp transposeOp) {
        return foldTransposePermIntoMapLoad(
            rewriter, transposeOp, transposeOp.getPermutation(), mapLoadOp);
      })
      .Case<tensor::ExpandShapeOp>([&](tensor::ExpandShapeOp expandOp) {
        return foldExpandShapeIntoMapLoad(rewriter, expandOp, mapLoadOp);
      })
      .Case<tensor::CollapseShapeOp>([&](tensor::CollapseShapeOp collapseOp) {
        return foldCollapseShapeIntoMapLoad(rewriter, collapseOp, mapLoadOp);
      })
      .Case<tensor::ExtractSliceOp>([&](tensor::ExtractSliceOp extractSliceOp) {
        return foldExtractSliceIntoMapLoad(rewriter, extractSliceOp, mapLoadOp);
      })
      .Case<tensor::PadOp>([&](tensor::PadOp padOp) {
        return foldPadIntoMapLoad(rewriter, padOp, mapLoadOp);
      })
      .Case<linalg::PackOp>([&](linalg::PackOp packOp) {
        return foldPackIntoMapLoad(rewriter, packOp, mapLoadOp);
      })
      .Case<linalg::UnPackOp>([&](linalg::UnPackOp unPackOp) {
        return foldUnpackIntoMapLoad(rewriter, unPackOp, mapLoadOp);
      })
      .Case<linalg::GenericOp>([&](linalg::GenericOp genericOp) {
        if (auto perm = linalg::isaTransposeOpInterface(genericOp)) {
          return foldTransposePermIntoMapLoad(rewriter, genericOp, perm.value(),
                                              mapLoadOp);
        }
        if (linalg::isaBroadcastOpInterface(genericOp)) {
          return foldBroadcastIntoMapLoad(
              rewriter, cast<linalg::LinalgOp>(genericOp.getOperation()),
              mapLoadOp);
        }
        return FailureOr<MapLoadOp>(failure());
      })
      .Case<linalg::LinalgOp>([&](linalg::LinalgOp linalgOp) {
        return foldBroadcastIntoMapLoad(rewriter, linalgOp, mapLoadOp);
      })
      .Default([](Operation *) { return failure(); });
}

/// Pattern to fold consumer relayout ops into a producer map_load.
struct FoldConsumerRelayoutIntoMapLoadPattern
    : OpRewritePattern<IREE::LinalgExt::MapLoadOp> {
  using Base::Base;

  LogicalResult matchAndRewrite(IREE::LinalgExt::MapLoadOp mapLoadOp,
                                PatternRewriter &rewriter) const override {
    // Find a consumer relayout op (one that uses map_load result as its input).
    Operation *consumerOp = nullptr;
    Value mapLoadResult = mapLoadOp.getResult(0);
    for (Operation *user : mapLoadOp->getUsers()) {
      if (isRelayoutChainExtension(user, mapLoadResult)) {
        consumerOp = user;
        break;
      }
    }
    if (!consumerOp) {
      return failure();
    }
    if (failed(foldIntoMapLoad(rewriter, consumerOp, mapLoadOp))) {
      return failure();
    }
    return success();
  }
};

// Insert identity map_load op after the root and replace only uses whose
// owner is in `complexChainUsers` (i.e. uses that are part of a complex
// relayout chain). Other uses keep using the load/root directly.
static MapLoadOp
insertIdentityMapLoad(RewriterBase &rewriter, OpResult root,
                      const SmallPtrSetImpl<Operation *> &complexChainUsers) {
  Location loc = root.getLoc();
  SetVector<OpOperand *> originalUses;
  for (OpOperand &use : root.getUses()) {
    if (complexChainUsers.contains(use.getOwner())) {
      originalUses.insert(&use);
    }
  }
  OpBuilder::InsertionGuard g(rewriter);
  rewriter.setInsertionPointAfterValue(root);
  Type elementType = getElementTypeOrSelf(root.getType());
  SmallVector<OpFoldResult> sizes = tensor::getMixedSizes(rewriter, loc, root);
  Value mapLoadDest =
      tensor::EmptyOp::create(rewriter, loc, sizes, elementType);
  auto mapLoadOp =
      MapLoadOp::createIdentityMapLoad(rewriter, loc, root, mapLoadDest);
  rewriter.replaceUsesWithIf(root, mapLoadOp.getResult(0), [&](OpOperand &use) {
    return originalUses.contains(&use);
  });
  LDBG() << "Created identity map_load:\n" << mapLoadOp;
  return mapLoadOp;
}

/// Insert identity map_load op after a LoadFromBufferOp if it has relayout
/// op consumers. The identity map_load can then be used to fold consumer
/// relayout ops into it iteratively.
struct InsertMapLoadOpPattern
    : OpRewritePattern<IREE::Codegen::LoadFromBufferOp> {
  InsertMapLoadOpPattern(MLIRContext *context, bool combineNonComplexChains)
      : OpRewritePattern(context),
        combineNonComplexChains(combineNonComplexChains) {}

  LogicalResult matchAndRewrite(IREE::Codegen::LoadFromBufferOp loadOp,
                                PatternRewriter &rewriter) const override {
    Value loadResult = loadOp.getResult();
    SmallPtrSet<Operation *, 4> complexChainUsers =
        getComplexChainRelayoutUsers(loadResult, combineNonComplexChains);
    // Only introduce map_load when there is at least one complex chain.
    if (complexChainUsers.empty()) {
      return failure();
    }
    // Check that the load doesn't already have a map_load user (avoid
    // infinite loop).
    bool hasMapLoadUser = llvm::any_of(loadOp->getUsers(), [](Operation *user) {
      return isa<IREE::LinalgExt::MapLoadOp>(user);
    });
    if (hasMapLoadUser) {
      return failure();
    }
    (void)insertIdentityMapLoad(rewriter, cast<OpResult>(loadResult),
                                complexChainUsers);
    return success();
  }

private:
  bool combineNonComplexChains;
};

namespace {

struct CombineSourceLayoutTransformationPass final
    : impl::CombineSourceLayoutTransformationPassBase<
          CombineSourceLayoutTransformationPass> {
  using Base::Base;

  void runOnOperation() override {
    MLIRContext *context = &getContext();
    FunctionOpInterface funcOp = getOperation();

    // Insert identity map_load ops after load_from_buffer ops and fold
    // consumer relayout ops into them.
    RewritePatternSet patterns(context);
    patterns.add<InsertMapLoadOpPattern>(context, testCombineNonComplexChains);
    patterns.add<FoldConsumerRelayoutIntoMapLoadPattern>(context);
    memref::populateResolveRankedShapedTypeResultDimsPatterns(patterns);
    if (failed(applyPatternsGreedily(funcOp, std::move(patterns)))) {
      return signalPassFailure();
    }
  }
};

} // namespace
} // namespace mlir::iree_compiler
