// Copyright 2024 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
//
//
// The content of this file is adapted from linalg's ElemenwiseOpFusion.cpp and
// modified to work with LinalgExt ops.

#include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtInterfaces.h"
#include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.h"
#include "iree/compiler/Dialect/LinalgExt/Transforms/Transforms.h"
#include "iree/compiler/Dialect/LinalgExt/Utils/Utils.h"
#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Utils/ReshapeOpsUtils.h"
#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Transforms/RegionUtils.h"

#include <cstdint>
#include <optional>

namespace mlir::iree_compiler::IREE::LinalgExt {
namespace {

/// Represents the size of a dimension of some ShapedType value in the IR. This
/// is used instead of OpFoldResult when modifying the IR is illegal. This can
/// still be constructed from an OpFoldResult in cases where the value can be
/// obtained without IR modification.
class DimSize {
public:
  DimSize(TypedValue<ShapedType> val, int64_t dim)
      : ofr(nullptr), val(val), dim(dim) {}
  DimSize(OpFoldResult ofr) : ofr(ofr), val(nullptr), dim(-1) {}

  bool isStatic() const {
    if (ofr) {
      return getConstantIntValue(ofr).has_value();
    }
    return val.getType().isStaticDim(dim);
  }

  // Get an OpFoldResult by possibly inserting IR.
  OpFoldResult materialize(OpBuilder &b) const {
    if (ofr) {
      return ofr;
    }
    return getDim(b, val.getLoc(), val, dim);
  }

private:
  OpFoldResult ofr;
  TypedValue<ShapedType> val;
  int64_t dim;
};
} // namespace

static SmallVector<DimSize> getDimSizes(Value v) {
  auto shapedVal = cast<TypedValue<ShapedType>>(v);
  int64_t rank = shapedVal.getType().getRank();
  SmallVector<DimSize> sizes;
  for (int i = 0; i < rank; ++i) {
    sizes.emplace_back(shapedVal, i);
  }
  return sizes;
}

static bool
isIdentityReassoc(const SmallVector<ReassociationIndices> &indices) {
  for (auto &index : indices) {
    if (index.size() != 1) {
      return false;
    }
  }
  return true;
};

static SmallVector<ReassociationIndices>
computeReassocFromShapeMap(ArrayRef<SmallVector<DimSize>> shapeMap) {
  SmallVector<ReassociationIndices> reassoc;
  int64_t dimCount = 0;
  for (auto &shape : shapeMap) {
    reassoc.emplace_back(
        llvm::to_vector(llvm::seq<int64_t>(dimCount, dimCount + shape.size())));
    dimCount += shape.size();
  }
  return reassoc;
}

namespace {
/// Helper class that supports fusing reshapes with operands when not all of the
/// shape dims map to the iteration space.
struct ReshapeOperandInfo {
  static constexpr int64_t kNoMapping = -1;

  // Original shape of this operand.
  SmallVector<DimSize> originalShape;

  // Similar to the results of the operand's `AffineMap` except `kNoMapping` if
  // that dim doesn't map to the iteration space. For example, the indexed
  // dimensions in a LinalgExt::ScatterOp.
  SmallVector<int64_t> operandToIterationSpace;
};

/// Information needed to expand an operation to fold the reshape with
/// it.
class ExpansionInfo {
public:
  // Computes the mapping from original dimensions of the op to the dimensions
  // of the expanded op given the `indexingMap` of the fused operand/result of
  // the op, the `reassociationMaps` of the reshape op and the shape of
  // the expanded op.
  LogicalResult compute(SmallVector<ReshapeOperandInfo> infos,
                        SmallVector<int64_t> loopRanges,
                        OpOperand *fusableOpOperand,
                        ArrayRef<ReassociationIndices> operandReassoc,
                        ArrayRef<DimSize> expandedShape);

  std::optional<Value> getOrCreateExpanded(Location loc, OpOperand *operand,
                                           RewriterBase &rewriter) {
    auto shapeMap = this->getShapeMap(operand);
    auto reassoc = computeReassocFromShapeMap(shapeMap);
    if (isIdentityReassoc(reassoc)) {
      return operand->get();
    }
    SmallVector<OpFoldResult> outputShape;
    for (auto &shape : shapeMap) {
      llvm::append_range(
          outputShape, llvm::map_range(shape, [&rewriter](const DimSize &size) {
            return size.materialize(rewriter);
          }));
    }
    auto [staticShape, dynamicShape] = decomposeMixedValues(outputShape);
    (void)dynamicShape;
    auto oldType = cast<ShapedType>(operand->get().getType());
    auto newType = RankedTensorType::get(staticShape, oldType.getElementType());
    if (failed(reshapeLikeShapesAreCompatible(
            [&](const Twine &msg) {
              return rewriter.notifyMatchFailure(loc, msg);
            },
            oldType.getShape(), newType.getShape(), reassoc,
            /*isExpandingReshape=*/true))) {
      return {};
    }
    return tensor::ExpandShapeOp::create(rewriter, loc, newType, operand->get(),
                                         reassoc, outputShape);
  };

  /// Get the shape map for the operand.
  SmallVector<SmallVector<DimSize>> getShapeMap(OpOperand *operand) const {
    auto info = reshapeInfos[operand->getOperandNumber()];
    SmallVector<SmallVector<DimSize>> shapeMap;
    for (auto [operandIdx, loopIdx] :
         llvm::enumerate(info.operandToIterationSpace)) {
      if (loopIdx == ReshapeOperandInfo::kNoMapping) {
        shapeMap.push_back(
            SmallVector<DimSize>{info.originalShape[operandIdx]});
      } else {
        shapeMap.push_back(loopShapeMap[loopIdx]);
      }
    }
    return shapeMap;
  }

  SmallVector<ReassociationIndices> getReassoc(OpOperand *operand) const {
    auto shapeMap = this->getShapeMap(operand);
    return computeReassocFromShapeMap(shapeMap);
  }

  unsigned getOrigNumLoops() const { return loopReassoc.size(); }
  unsigned getExpandedNumLoops() const { return expandedOpNumDims; }
  ReassociationIndicesRef getExpandedLoops(unsigned i) const {
    return loopReassoc[i];
  }

private:
  SmallVector<ReassociationIndices> loopReassoc;
  /// Mapping from extent of loops in the original operation, to the extent of
  /// loops in the expanded operation.
  SmallVector<SmallVector<DimSize>> loopShapeMap;
  unsigned expandedOpNumDims;
  /// Info about the reassociation and original shape for each operand.
  SmallVector<ReshapeOperandInfo> reshapeInfos;
};

class CollapsingInfo {
public:
  LogicalResult initialize(unsigned origNumLoops,
                           ArrayRef<ReassociationIndices> foldedIterationDims);

  /// Return mapping from collapsed loop domain to original loop domain.
  ArrayRef<ReassociationIndices> getCollapsedOpToOrigOpMapping() const {
    return collapsedOpToOrigOpIterationDim;
  }

  /// Return mapping from original loop domain to collapsed loop domain. The
  /// mapping is a pair. First value is the dimension in the collapsed loop that
  /// the original loop is mapped to. Second is the relative position in folded
  /// list of this domain. For example if the original loop domain is 3D, and
  /// the collapsed loop domain is folding all of it, i.e.
  ///
  /// ```
  /// collapsedOpToOrigOpMapping = [[0, 1, 2] [3, 4]]`
  /// ```
  ///
  /// then
  ///
  /// ```
  ///  origOpToCollapsedOpMapping[0] = {0, 0};
  ///  origOpToCollapsedOpMapping[1] = {0, 1};
  ///  origOpToCollapsedOpMapping[2] = {0, 2};
  ///  origOpToCollapsedOpMapping[3] = {1, 0};
  ///  origOpToCollapsedOpMapping[4] = {1, 1};
  /// ```
  ///
  ArrayRef<std::pair<int64_t, unsigned>> getOrigOpToCollapsedOpMapping() const {
    return origOpToCollapsedOpIterationDim;
  }

  /// Return the collapsed op iteration domain rank.
  unsigned getCollapsedOpIterationRank() const {
    return collapsedOpToOrigOpIterationDim.size();
  }

private:
  /// Map from the iteration domain index in collapsed op to the iteration
  /// domain indices in the original op.
  SmallVector<ReassociationIndices> collapsedOpToOrigOpIterationDim;

  /// Map from iteration domain index in the original op to the iteration domain
  /// index in the collapsed op.
  SmallVector<std::pair<int64_t, unsigned>> origOpToCollapsedOpIterationDim;
};

} // namespace

LogicalResult ExpansionInfo::compute(
    SmallVector<ReshapeOperandInfo> infos, SmallVector<int64_t> loopRanges,
    OpOperand *fusableOpOperand, ArrayRef<ReassociationIndices> operandReassoc,
    ArrayRef<DimSize> expandedShape) {
  if (operandReassoc.empty()) {
    return failure();
  }

  // Check that the operand dim size matches the iteration space dim size. This
  // can fail when one is static and the other is dynamic.
  for (const ReshapeOperandInfo &info : infos) {
    for (auto [operandDim, iterDim] :
         llvm::enumerate(info.operandToIterationSpace)) {
      if (iterDim != ReshapeOperandInfo::kNoMapping &&
          ShapedType::isStatic(loopRanges[iterDim]) !=
              info.originalShape[operandDim].isStatic()) {
        return failure();
      }
    }
  }

  int64_t operandNum = fusableOpOperand->getOperandNumber();
  ReshapeOperandInfo &fusionOperandInfo = infos[operandNum];
  this->loopShapeMap.clear();
  this->loopShapeMap.resize(loopRanges.size());
  for (auto [operandIdx, loopIdx] :
       llvm::enumerate(fusionOperandInfo.operandToIterationSpace)) {
    if (loopIdx == ReshapeOperandInfo::kNoMapping) {
      continue;
    }

    // Compute the shape map at element `loopIdx`
    ReassociationIndicesRef indices = operandReassoc[operandIdx];
    for (auto [dimIdx, shapeIdx] : llvm::enumerate(indices)) {
      this->loopShapeMap[loopIdx].push_back(expandedShape[shapeIdx]);
    }
  }

  // Fill in the remaining elements.
  for (const ReshapeOperandInfo &info : infos) {
    for (auto [operandIdx, loopIdx] :
         llvm::enumerate(info.operandToIterationSpace)) {
      if (loopIdx == ReshapeOperandInfo::kNoMapping ||
          !this->loopShapeMap[loopIdx].empty()) {
        continue;
      }

      this->loopShapeMap[loopIdx] =
          SmallVector<DimSize>{info.originalShape[operandIdx]};
    }
  }

  this->expandedOpNumDims = 0;
  for (const auto &shapeMap : this->loopShapeMap) {
    this->expandedOpNumDims += shapeMap.size();
  }

  if (llvm::all_of(this->loopShapeMap,
                   [&](auto vec) { return vec.size() == 1; })) {
    return failure();
  }
  this->loopReassoc = computeReassocFromShapeMap(this->loopShapeMap);
  this->reshapeInfos = std::move(infos);
  return success();
}

LogicalResult
CollapsingInfo::initialize(unsigned origNumLoops,
                           ArrayRef<ReassociationIndices> foldedIterationDims) {
  llvm::SmallDenseSet<int64_t, 4> processedDims;
  // Find all the dims that are folded.
  for (ReassociationIndicesRef foldedIterationDim : foldedIterationDims) {
    if (foldedIterationDim.empty()) {
      continue;
    }
    // If the folded dims contain dims already folded, that's illegal
    // specification. Repetition within a list is also illegal.
    for (auto dim : foldedIterationDim) {
      if (dim >= origNumLoops) {
        return failure();
      }
      if (processedDims.contains(dim)) {
        return failure();
      }
      processedDims.insert(dim);
    }
    collapsedOpToOrigOpIterationDim.emplace_back(foldedIterationDim.begin(),
                                                 foldedIterationDim.end());
  }
  if (processedDims.size() > origNumLoops) {
    return failure();
  }

  // Add all the preserved dims of the original op as single
  // elements to `collapsedOpToOrigOpIterationDim`.
  for (auto dim : llvm::seq<int64_t>(0, origNumLoops)) {
    if (processedDims.contains(dim)) {
      continue;
    }
    collapsedOpToOrigOpIterationDim.emplace_back(ReassociationIndices{dim});
  }

  llvm::sort(collapsedOpToOrigOpIterationDim,
             [&](ReassociationIndicesRef lhs, ReassociationIndicesRef rhs) {
               return lhs[0] < rhs[0];
             });
  origOpToCollapsedOpIterationDim.resize(origNumLoops);
  for (const auto &foldedDims :
       llvm::enumerate(collapsedOpToOrigOpIterationDim)) {
    for (const auto &dim : enumerate(foldedDims.value())) {
      origOpToCollapsedOpIterationDim[dim.value()] =
          std::make_pair<int64_t, unsigned>(foldedDims.index(), dim.index());
    }
  }
  return success();
}

template <typename AttentionOpTy>
static SmallVector<ReshapeOperandInfo>
getAttentionLikeReshapeInfo(AttentionOpTy attentionOp) {
  return llvm::map_to_vector(
      attentionOp->getOpOperands(), [&](OpOperand &opOperand) {
        ReshapeOperandInfo operandInfo;
        auto operandType = dyn_cast<ShapedType>(opOperand.get().getType());
        if (!operandType) {
          assert(
              attentionOp.getMatchingIndexingMap(&opOperand).getNumResults() ==
                  0 &&
              "expected non-shaped type to have no results in indexing map");
          return operandInfo;
        }

        operandInfo.originalShape = getDimSizes(opOperand.get());
        for (auto result :
             attentionOp.getMatchingIndexingMap(&opOperand).getResults()) {
          operandInfo.operandToIterationSpace.push_back(
              cast<AffineDimExpr>(result).getPosition());
        }
        return operandInfo;
      });
}

static SmallVector<ReshapeOperandInfo>
getReshapeInfo(LinalgExt::AttentionOp attentionOp) {
  return getAttentionLikeReshapeInfo(attentionOp);
}

static SmallVector<ReshapeOperandInfo>
getReshapeInfo(LinalgExt::OnlineAttentionOp attentionOp) {
  return getAttentionLikeReshapeInfo(attentionOp);
}

static SmallVector<ReshapeOperandInfo>
getReshapeInfo(LinalgExt::ScatterOp scatterOp) {
  SmallVector<ReshapeOperandInfo> infos;
  auto rankOfContiguousSlice =
      scatterOp.getOriginalType().getRank() - scatterOp.getIndexDepth();
  auto updateRank = scatterOp.getUpdateType().getRank();

  ReshapeOperandInfo updateInfo;
  updateInfo.originalShape = getDimSizes(scatterOp.getUpdates());
  llvm::append_range(updateInfo.operandToIterationSpace,
                     llvm::seq<int64_t>(0, updateRank));
  infos.push_back(std::move(updateInfo));

  ReshapeOperandInfo indicesInfo;
  indicesInfo.originalShape = getDimSizes(scatterOp.getIndices());
  llvm::append_range(indicesInfo.operandToIterationSpace,
                     llvm::seq<int64_t>(scatterOp.getBatchRank()));
  if (scatterOp.getBatchRank() != scatterOp.getIndicesType().getRank()) {
    indicesInfo.operandToIterationSpace.push_back(
        ReshapeOperandInfo::kNoMapping);
  }
  infos.push_back(std::move(indicesInfo));

  if (Value mask = scatterOp.getMask()) {
    ReshapeOperandInfo maskInfo;
    maskInfo.originalShape = getDimSizes(mask);
    llvm::append_range(maskInfo.operandToIterationSpace,
                       llvm::seq<int64_t>(scatterOp.getBatchRank()));
    infos.push_back(std::move(maskInfo));
  }

  ReshapeOperandInfo originalInfo;
  originalInfo.originalShape = getDimSizes(scatterOp.getOriginal());
  originalInfo.operandToIterationSpace.append(scatterOp.getIndexDepth(),
                                              ReshapeOperandInfo::kNoMapping);
  llvm::append_range(originalInfo.operandToIterationSpace,
                     llvm::seq(updateRank - rankOfContiguousSlice, updateRank));
  infos.push_back(std::move(originalInfo));
  return infos;
};

static SmallVector<ReshapeOperandInfo>
getReshapeInfo(LinalgExt::GatherOp gatherOp) {
  SmallVector<ReshapeOperandInfo> infos;
  auto rankOfContiguousSlice = gatherOp.getOutputSliceRank();
  auto outputRank = gatherOp.getOutputType().getRank();

  ReshapeOperandInfo sourceInfo;
  sourceInfo.originalShape = getDimSizes(gatherOp.getSource());
  sourceInfo.operandToIterationSpace.append(gatherOp.getIndexDepth(),
                                            ReshapeOperandInfo::kNoMapping);
  llvm::append_range(sourceInfo.operandToIterationSpace,
                     llvm::seq(outputRank - rankOfContiguousSlice, outputRank));
  infos.push_back(std::move(sourceInfo));

  ReshapeOperandInfo indicesInfo;
  indicesInfo.originalShape = getDimSizes(gatherOp.getIndices());
  llvm::append_range(indicesInfo.operandToIterationSpace,
                     llvm::seq<int64_t>(gatherOp.getBatchRank()));
  if (gatherOp.getBatchRank() != gatherOp.getIndicesType().getRank()) {
    indicesInfo.operandToIterationSpace.push_back(
        ReshapeOperandInfo::kNoMapping);
  }
  infos.push_back(std::move(indicesInfo));

  if (Value mask = gatherOp.getMask()) {
    ReshapeOperandInfo maskInfo;
    maskInfo.originalShape = getDimSizes(mask);
    llvm::append_range(maskInfo.operandToIterationSpace,
                       llvm::seq<int64_t>(gatherOp.getBatchRank()));
    infos.push_back(std::move(maskInfo));
  }

  ReshapeOperandInfo outputInfo;
  outputInfo.originalShape = getDimSizes(gatherOp.getOutput());
  llvm::append_range(outputInfo.operandToIterationSpace,
                     llvm::seq<int64_t>(0, outputRank));
  infos.push_back(std::move(outputInfo));
  return infos;
};

static AffineMap
getIndexingMapInExpandedOp(OpBuilder &builder, AffineMap indexingMap,
                           const ExpansionInfo &expansionInfo) {
  SmallVector<AffineExpr> newExprs;
  for (AffineExpr expr : indexingMap.getResults()) {
    unsigned pos = cast<AffineDimExpr>(expr).getPosition();
    auto expandedExprs = llvm::to_vector_of<AffineExpr, 6>(
        llvm::map_range(expansionInfo.getExpandedLoops(pos), [&](int64_t v) {
          return builder.getAffineDimExpr(static_cast<unsigned>(v));
        }));
    newExprs.append(expandedExprs.begin(), expandedExprs.end());
  }
  return AffineMap::get(expansionInfo.getExpandedNumLoops(),
                        indexingMap.getNumSymbols(), newExprs,
                        builder.getContext());
}

template <typename OpTy>
static std::optional<SmallVector<Value>>
fuseWithReshapeByExpansion(OpTy op, Operation *reshapeOp,
                           OpOperand *fusableOpOperand,
                           PatternRewriter &rewriter) {
  Location loc = op.getLoc();
  // Check if reshape is expanding or collapsing.
  auto expandingReshapeOp = dyn_cast<tensor::ExpandShapeOp>(*reshapeOp);
  auto collapsingReshapeOp = dyn_cast<tensor::CollapseShapeOp>(*reshapeOp);
  bool isExpanding = (expandingReshapeOp != nullptr);
  Value expandedVal = isExpanding ? expandingReshapeOp.getResult()
                                  : collapsingReshapeOp.getSrc();
  SmallVector<DimSize> expandedSize;
  if (isExpanding) {
    // The SSA dims must dominate `op` in order to use them to create new
    // expand_shape ops.
    if (failed(moveValueDefinitions(rewriter,
                                    expandingReshapeOp.getOutputShape(), op))) {
      return std::nullopt;
    }
    llvm::append_range(expandedSize, expandingReshapeOp.getMixedOutputShape());
  } else {
    expandedSize = getDimSizes(expandedVal);
  }
  ExpansionInfo info;
  if (failed(info.compute(
          getReshapeInfo(op), op.getStaticLoopRanges(), fusableOpOperand,
          isExpanding ? expandingReshapeOp.getReassociationIndices()
                      : collapsingReshapeOp.getReassociationIndices(),
          expandedSize))) {
    return std::nullopt;
  }

  OpBuilder::InsertionGuard g(rewriter);
  rewriter.setInsertionPoint(op);

  IRMapping mapping;
  for (OpOperand &operand : op->getOpOperands()) {
    std::optional<Value> maybeNewOperand =
        info.getOrCreateExpanded(loc, &operand, rewriter);
    assert(maybeNewOperand.has_value());
    mapping.map(operand.get(), maybeNewOperand.value());
  }

  auto newOp = cast<OpTy>(rewriter.clone(*op, mapping));
  for (auto [i, initOperand] : llvm::enumerate(op.getDpsInitsMutable())) {
    newOp->getResult(i).setType(mapping.lookup(initOperand.get()).getType());
  }

  if constexpr (std::is_same_v<OpTy, AttentionOp> ||
                std::is_same_v<OpTy, OnlineAttentionOp>) {
    auto expandedOpIndexingMaps = llvm::to_vector_of<AffineMap, 6>(
        llvm::map_range(op.getIndexingMapsArray(), [&](AffineMap m) {
          return getIndexingMapInExpandedOp(rewriter, m, info);
        }));
    newOp.setIndexingMapsAttr(
        rewriter.getAffineMapArrayAttr(expandedOpIndexingMaps));
  }

  SmallVector<Value> replacements;
  replacements.reserve(op->getNumResults());
  for (auto [i, initOperand] : llvm::enumerate(op.getDpsInitsMutable())) {
    SmallVector<ReassociationIndices> originalReassoc =
        info.getReassoc(&initOperand);
    Value newResult = newOp->getResult(i);
    if (isIdentityReassoc(originalReassoc)) {
      replacements.push_back(newResult);
    } else {
      replacements.push_back(tensor::CollapseShapeOp::create(
          rewriter, loc, op.getResult(i).getType(), newResult,
          originalReassoc));
    }
  }
  return replacements;
}

//===----------------------------------------------------------------------===//
// Fuse By Expansion Patterns
//===----------------------------------------------------------------------===//

namespace {

struct DropScatterUnitIndexDepth final : OpRewritePattern<ScatterOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(ScatterOp scatterOp,
                                PatternRewriter &rewriter) const override {
    llvm::ArrayRef<int64_t> indicesShape =
        scatterOp.getIndicesType().getShape();
    int64_t batchRank = scatterOp.getBatchRank();
    if (indicesShape.size() == batchRank || indicesShape.back() != 1) {
      return failure();
    }
    SmallVector<ReassociationIndices> reassoc;
    reassoc.reserve(indicesShape.size());
    for (auto i : llvm::seq<int64_t>(0, batchRank - 1)) {
      reassoc.emplace_back(1, i);
    }
    reassoc.push_back(ReassociationIndices{batchRank - 1, batchRank});
    auto collapseOp = tensor::CollapseShapeOp::create(
        rewriter, scatterOp.getLoc(), scatterOp.getIndices(), reassoc);

    rewriter.modifyOpInPlace(scatterOp, [&]() {
      scatterOp.getIndicesMutable().set(collapseOp.getResult());
    });
    return success();
  }
};

FailureOr<Value> rankReduceOperand(RewriterBase &rewriter, Location loc,
                                   int64_t startDim, int64_t numDims,
                                   Value operand, ShapedType ty,
                                   const linalg::ControlDropUnitDims &options) {
  if (numDims == 0) {
    return failure();
  }
  // Find list of dims to drop and the target shape.
  ArrayRef<int64_t> shape = ty.getShape();
  SmallVector<bool> unitDims(shape.size(), false);
  for (auto [isUnitDim, size] : llvm::zip_equal(unitDims, shape)) {
    if (size == 1) {
      isUnitDim = true;
    }
  }
  // Because gather/scatter like to define special behavior to allow eliding
  // the coordinate dimension, we have to make sure we don't generate a 0-D
  // tensor.
  auto slice = MutableArrayRef(unitDims).slice(startDim, numDims);
  if (llvm::all_of(slice, [](bool x) { return x; })) {
    slice.back() = false;
  }
  SmallVector<int64_t> targetShape;
  targetShape.append(shape.begin(), shape.begin() + startDim);
  for (int i = startDim; i < numDims + startDim; ++i) {
    if (!unitDims[i]) {
      targetShape.push_back(shape[i]);
    }
  }
  targetShape.append(shape.begin() + startDim + numDims, shape.end());
  assert(targetShape.size() <= shape.size());
  // No unit dims dropped.
  if (targetShape.size() == shape.size()) {
    return failure();
  }
  // Drop unit dims using extract_slice.
  if (options.rankReductionStrategy ==
      linalg::ControlDropUnitDims::RankReductionStrategy::ExtractInsertSlice) {
    FailureOr<Value> rankReducingExtract =
        tensor::ExtractSliceOp::rankReduceIfNeeded(rewriter, loc, operand,
                                                   targetShape);
    assert(succeeded(rankReducingExtract) && "not a unit-extent collapse");
    return rankReducingExtract.value();
  } else if (options.rankReductionStrategy ==
             linalg::ControlDropUnitDims::RankReductionStrategy::
                 ReassociativeReshape) {
    std::optional<SmallVector<ReassociationIndices>> reassoc =
        getReassociationIndicesForCollapse(shape, targetShape);
    assert(reassoc.has_value());
    return tensor::CollapseShapeOp::create(rewriter, loc, ty.clone(targetShape),
                                           operand, reassoc.value())
        .getResult();
  }
  llvm_unreachable("unhandled rank reduction strategy");
};

// Expands the rank of `val` to match `destVal` (by adding unit dimensions)
// using either `expand_shape` or `insert_slice`.
Value rankExpandValue(RewriterBase &rewriter, Location loc, Value destVal,
                      Value val, const linalg::ControlDropUnitDims &options) {
  if (val.getType() == destVal.getType()) {
    return val;
  }
  if (options.rankReductionStrategy ==
      linalg::ControlDropUnitDims::RankReductionStrategy::ExtractInsertSlice) {
    int64_t rank = cast<ShapedType>(destVal.getType()).getRank();
    SmallVector<OpFoldResult> offsets(rank, rewriter.getIndexAttr(0));
    SmallVector<OpFoldResult> sizes =
        tensor::getMixedSizes(rewriter, loc, destVal);
    SmallVector<OpFoldResult> strides(rank, rewriter.getIndexAttr(1));
    return tensor::InsertSliceOp::create(rewriter, loc, val, destVal, offsets,
                                         sizes, strides);
  } else if (options.rankReductionStrategy ==
             linalg::ControlDropUnitDims::RankReductionStrategy::
                 ReassociativeReshape) {
    std::optional<SmallVector<ReassociationIndices>> reassoc =
        getReassociationIndicesForReshape(cast<ShapedType>(val.getType()),
                                          cast<ShapedType>(destVal.getType()));
    assert(reassoc.has_value());
    return tensor::ExpandShapeOp::create(rewriter, loc, destVal.getType(), val,
                                         reassoc.value());
  } else {
    llvm_unreachable("unhandled rank reduction strategy");
  }
}

struct DropMapStoreUnitDims final : OpRewritePattern<MapStoreOp> {
  using Base::Base;
  DropMapStoreUnitDims(MLIRContext *context,
                       linalg::ControlDropUnitDims options,
                       PatternBenefit benefit = 1)
      : OpRewritePattern<MapStoreOp>(context, benefit),
        options(std::move(options)) {}

  LogicalResult matchAndRewrite(MapStoreOp mapStoreOp,
                                PatternRewriter &rewriter) const override {
    auto inputType = dyn_cast<RankedTensorType>(mapStoreOp.getInputType());
    if (!inputType) {
      return failure();
    }
    Location loc = mapStoreOp.getLoc();
    FailureOr<Value> newInput = rankReduceOperand(
        rewriter, loc, /*startDim=*/0, /*numDims=*/mapStoreOp.getInputRank(),
        mapStoreOp.getInput(), mapStoreOp.getInputType(), options);
    if (failed(newInput)) {
      return failure();
    }

    auto newInputType = cast<ShapedType>(newInput->getType());
    auto unitFoldingBuilder = [&](ArrayRef<BlockArgument> nonUnitIndices) {
      Value zero = arith::ConstantIndexOp::create(rewriter, loc, 0);
      int nonUnitArgIdx = 0;
      return llvm::map_to_vector(
          llvm::seq<int64_t>(inputType.getRank()), [&](int64_t dim) -> Value {
            return inputType.getDimSize(dim) == 1
                       ? zero
                       : cast<Value>(nonUnitIndices[nonUnitArgIdx++]);
          });
    };
    // The map_store op is generally only used in Codegen, where it is the
    // last op in the dispatch, so for now, we don't bother collapsing the
    // result shape and inserting an expansion after the op.
    rewriter.modifyOpInPlace(mapStoreOp, [&]() {
      mapStoreOp.getInputMutable().assign(newInput.value());
      mapStoreOp.insertTransformationAtStart(
          rewriter, unitFoldingBuilder,
          /*numSourceIndices=*/newInputType.getRank());
    });
    return success();
  }

private:
  linalg::ControlDropUnitDims options;
};

struct DropGatherUnitDims final : OpRewritePattern<GatherOp> {
  DropGatherUnitDims(MLIRContext *context, linalg::ControlDropUnitDims options,
                     PatternBenefit benefit = 1)
      : OpRewritePattern<GatherOp>(context, benefit),
        options(std::move(options)) {}

  LogicalResult matchAndRewrite(GatherOp gatherOp,
                                PatternRewriter &rewriter) const override {
    Location loc = gatherOp.getLoc();
    if (!gatherOp.hasPureTensorSemantics()) {
      return rewriter.notifyMatchFailure(
          gatherOp, "dropping unit dims not implemented for buffer semantics");
    }

    bool changed = false;
    // Drop batch dimensions.
    Value reducedSource = gatherOp.getSource();
    Value reducedIndices = gatherOp.getIndices();
    Value reducedMask = gatherOp.getMask();
    Value reducedOutput = gatherOp.getOutput();
    if (gatherOp.getBatchRank() > 1) {
      // The only reaason we have to do these rank reductions separate is
      // because gather/scatter have special behavior for eliding the coordinate
      // dimension.
      // TODO: Do the rank reduction in one go after this behavior is changed.
      FailureOr<Value> newIndices = rankReduceOperand(
          rewriter, loc, /*startDim=*/0, /*numDims=*/gatherOp.getBatchRank(),
          gatherOp.getIndices(), gatherOp.getIndicesType(), options);
      FailureOr<Value> newOutput = rankReduceOperand(
          rewriter, loc, /*startDim=*/0, /*numDims=*/gatherOp.getBatchRank(),
          gatherOp.getOutput(), gatherOp.getOutputType(), options);
      FailureOr<Value> newMask = failure();
      if (reducedMask) {
        newMask =
            rankReduceOperand(rewriter, loc, /*startDim=*/0,
                              /*numDims=*/gatherOp.getBatchRank(), reducedMask,
                              cast<ShapedType>(reducedMask.getType()), options);
      }
      if (succeeded(newIndices) && succeeded(newOutput) &&
          (!reducedMask || succeeded(newMask))) {
        reducedIndices = newIndices.value();
        reducedOutput = newOutput.value();
        if (reducedMask) {
          reducedMask = newMask.value();
        }
        changed = true;
      }
    }
    // Drop slice dimensions.
    FailureOr<Value> newSource = rankReduceOperand(
        rewriter, loc, /*startDim=*/gatherOp.getIndexDepth(),
        /*numDims=*/gatherOp.getOutputSliceRank(), gatherOp.getSource(),
        gatherOp.getSourceType(), options);
    ShapedType newOutputTy = cast<ShapedType>(reducedOutput.getType());
    FailureOr<Value> newOutput = rankReduceOperand(
        rewriter, loc,
        /*startDim=*/newOutputTy.getRank() - gatherOp.getOutputSliceRank(),
        /*numDims=*/gatherOp.getOutputSliceRank(), reducedOutput, newOutputTy,
        options);
    if (succeeded(newSource) && succeeded(newOutput)) {
      reducedSource = newSource.value();
      reducedOutput = newOutput.value();
      changed = true;
    }

    if (!changed) {
      return failure();
    }

    auto newGather = GatherOp::create(
        rewriter, gatherOp.getLoc(), TypeRange{reducedOutput.getType()},
        /*source=*/reducedSource, /*indices=*/reducedIndices,
        /*mask=*/reducedMask, /*output=*/reducedOutput,
        gatherOp.getDimensionMap());
    rewriter.replaceOp(gatherOp,
                       rankExpandValue(rewriter, loc, gatherOp.getOutput(),
                                       newGather.getResult(0), options));
    return success();
  }

private:
  linalg::ControlDropUnitDims options;
};

struct DropScatterUnitDims final : OpRewritePattern<ScatterOp> {
  DropScatterUnitDims(MLIRContext *context, linalg::ControlDropUnitDims options,
                      PatternBenefit benefit = 1)
      : OpRewritePattern<ScatterOp>(context, benefit),
        options(std::move(options)) {}

  LogicalResult matchAndRewrite(ScatterOp scatterOp,
                                PatternRewriter &rewriter) const override {
    Location loc = scatterOp.getLoc();
    if (!scatterOp.hasPureTensorSemantics()) {
      return rewriter.notifyMatchFailure(
          scatterOp, "dropping unit dims not implemented for buffer semantics");
    }

    bool changed = false;
    // Drop batch dimensions.
    Value original = scatterOp.getOriginal();
    Value indices = scatterOp.getIndices();
    Value mask = scatterOp.getMask();
    Value updates = scatterOp.getUpdates();
    if (scatterOp.getBatchRank() > 1) {
      FailureOr<Value> newIndices = rankReduceOperand(
          rewriter, loc, /*startDim=*/0, /*numDims=*/scatterOp.getBatchRank(),
          indices, cast<ShapedType>(indices.getType()), options);
      FailureOr<Value> newOutput = rankReduceOperand(
          rewriter, loc, /*startDim=*/0, /*numDims=*/scatterOp.getBatchRank(),
          updates, cast<ShapedType>(updates.getType()), options);
      FailureOr<Value> newMask = failure();
      if (mask) {
        newMask = rankReduceOperand(rewriter, loc, /*startDim=*/0,
                                    /*numDims=*/scatterOp.getBatchRank(), mask,
                                    cast<ShapedType>(mask.getType()), options);
      }
      if (succeeded(newIndices) && succeeded(newOutput) &&
          (!mask || succeeded(newMask))) {
        indices = newIndices.value();
        updates = newOutput.value();
        if (mask) {
          mask = newMask.value();
        }
        changed = true;
      }
    }
    // Drop slice dimensions.
    FailureOr<Value> newSource =
        rankReduceOperand(rewriter, loc, /*startDim=*/scatterOp.getIndexDepth(),
                          /*numDims=*/scatterOp.getUpdateSliceRank(), original,
                          cast<ShapedType>(original.getType()), options);
    ShapedType newOutputTy = cast<ShapedType>(updates.getType());
    FailureOr<Value> newOutput = rankReduceOperand(
        rewriter, loc,
        /*startDim=*/newOutputTy.getRank() - scatterOp.getUpdateSliceRank(),
        /*numDims=*/scatterOp.getUpdateSliceRank(), updates, newOutputTy,
        options);
    if (succeeded(newSource) && succeeded(newOutput)) {
      original = newSource.value();
      updates = newOutput.value();
      changed = true;
    }

    if (!changed) {
      return failure();
    }

    auto newScatter = ScatterOp::create(
        rewriter, scatterOp.getLoc(), TypeRange{original.getType()}, updates,
        indices, mask, original, scatterOp.getDimensionMap(),
        scatterOp.getUniqueIndices());
    rewriter.inlineRegionBefore(scatterOp.getRegion(), newScatter.getRegion(),
                                newScatter.getRegion().begin());
    rewriter.replaceOp(scatterOp,
                       rankExpandValue(rewriter, loc, scatterOp.getOriginal(),
                                       newScatter.getResult(0), options));
    return success();
  }

private:
  linalg::ControlDropUnitDims options;
};

template <typename OpTy>
struct FoldWithProducerReshapeByExpansion final : OpRewritePattern<OpTy> {
  FoldWithProducerReshapeByExpansion(
      MLIRContext *context, linalg::ControlFusionFn controlFoldingReshapes,
      PatternBenefit benefit = 1)
      : OpRewritePattern<OpTy>(context, benefit),
        controlFoldingReshapes(std::move(controlFoldingReshapes)) {}

  LogicalResult matchAndRewrite(OpTy op,
                                PatternRewriter &rewriter) const override {
    if (!op.hasPureTensorSemantics()) {
      return failure();
    }

    for (OpOperand &opOperand : op->getOpOperands()) {
      tensor::CollapseShapeOp reshapeOp =
          opOperand.get().getDefiningOp<tensor::CollapseShapeOp>();
      if (!reshapeOp) {
        continue;
      }
      if (!controlFoldingReshapes(&opOperand)) {
        continue;
      }

      std::optional<SmallVector<Value>> replacementValues =
          fuseWithReshapeByExpansion(op, reshapeOp, &opOperand, rewriter);
      if (!replacementValues) {
        return failure();
      }
      rewriter.replaceOp(op, *replacementValues);
      return success();
    }
    return failure();
  }

  linalg::ControlFusionFn controlFoldingReshapes;
};

template <typename OpTy>
struct FoldWithConsumerReshapeByExpansion final
    : OpRewritePattern<tensor::ExpandShapeOp> {
  FoldWithConsumerReshapeByExpansion(
      MLIRContext *context, linalg::ControlFusionFn controlFoldingReshapes,
      PatternBenefit benefit = 1)
      : OpRewritePattern<tensor::ExpandShapeOp>(context, benefit),
        controlFoldingReshapes(std::move(controlFoldingReshapes)) {}

  LogicalResult matchAndRewrite(tensor::ExpandShapeOp expandOp,
                                PatternRewriter &rewriter) const override {
    auto producerResult = dyn_cast<OpResult>(expandOp.getSrc());
    if (!producerResult) {
      return rewriter.notifyMatchFailure(expandOp,
                                         "source not produced by an operation");
    }

    auto op = expandOp.getSrc().getDefiningOp<OpTy>();
    if (!op || !op.hasPureTensorSemantics()) {
      return failure();
    }

    if (!controlFoldingReshapes(&expandOp.getSrcMutable())) {
      return failure();
    }

    std::optional<SmallVector<Value>> replacementValues =
        fuseWithReshapeByExpansion(
            op, expandOp, op.getTiedOpOperand(producerResult), rewriter);
    if (!replacementValues) {
      return failure();
    }
    rewriter.replaceOp(op, *replacementValues);
    return success();
  }

  linalg::ControlFusionFn controlFoldingReshapes;
};

} // namespace

/// Return the `reassociation` indices to use to collapse the operand when the
/// iteration space of a generic op is collapsed.
static SmallVector<ReassociationIndices>
getOperandReassociation(AffineMap indexingMap,
                        const CollapsingInfo &collapsingInfo) {
  unsigned counter = 0;
  SmallVector<ReassociationIndices> operandReassociation;
  auto origOpToCollapsedOpMapping =
      collapsingInfo.getOrigOpToCollapsedOpMapping();
  auto collapsedOpToOrigOpMapping =
      collapsingInfo.getCollapsedOpToOrigOpMapping();
  while (counter < indexingMap.getNumResults()) {
    unsigned dim =
        cast<AffineDimExpr>(indexingMap.getResult(counter)).getPosition();
    // This is the start of a collapsed dimensions of the iteration that
    // is guaranteed to be preserved in the indexing map. The number of folded
    // dims is obtained from the collapsed op to original op mapping.
    unsigned numFoldedDims =
        collapsedOpToOrigOpMapping[origOpToCollapsedOpMapping[dim].first]
            .size();
    if (origOpToCollapsedOpMapping[dim].second == 0) {
      auto range = llvm::seq<unsigned>(counter, counter + numFoldedDims);
      operandReassociation.emplace_back(range.begin(), range.end());
    }
    counter += numFoldedDims;
  }
  return operandReassociation;
}

/// Get the new value to use for a given `OpOperand` in the collapsed operation.
template <typename AttentionOpTy>
static Value getCollapsedOpOperand(Location loc, AttentionOpTy op,
                                   OpOperand *opOperand,
                                   const CollapsingInfo &collapsingInfo,
                                   OpBuilder &builder) {
  AffineMap indexingMap = op.getMatchingIndexingMap(opOperand);
  SmallVector<ReassociationIndices> operandReassociation =
      getOperandReassociation(indexingMap, collapsingInfo);

  // If the number of entries in the reassociation for the operand is same as
  // the number of results of the indexing map, then nothing to do for this
  // operand.
  Value operand = opOperand->get();
  if (operandReassociation.size() == indexingMap.getNumResults()) {
    return operand;
  }

  // Insert a reshape to collapse the dimensions.
  if (isa<MemRefType>(operand.getType())) {
    return memref::CollapseShapeOp::create(builder, loc, operand,
                                           operandReassociation)
        .getResult();
  }
  return tensor::CollapseShapeOp::create(builder, loc, operand,
                                         operandReassociation)
      .getResult();
}

template <typename AttentionOpTy>
static void collapseOperandsAndResults(AttentionOpTy op,
                                       const CollapsingInfo &collapsingInfo,
                                       RewriterBase &rewriter,
                                       SmallVectorImpl<Value> &inputOperands,
                                       SmallVectorImpl<Value> &outputOperands,
                                       SmallVectorImpl<Type> &resultTypes) {
  Location loc = op->getLoc();
  inputOperands =
      llvm::map_to_vector(op.getDpsInputOperands(), [&](OpOperand *opOperand) {
        return getCollapsedOpOperand(loc, op, opOperand, collapsingInfo,
                                     rewriter);
      });

  // Get the output operands and result types.
  resultTypes.reserve(op.getNumDpsInits());
  outputOperands.reserve(op.getNumDpsInits());
  for (OpOperand &output : op.getDpsInitsMutable()) {
    Value newOutput =
        getCollapsedOpOperand(loc, op, &output, collapsingInfo, rewriter);
    outputOperands.push_back(newOutput);
    // If the op has "buffer semantics", then the init operands are ranked
    // memrefs and the op has no results.
    if (!op.hasPureBufferSemantics()) {
      resultTypes.push_back(newOutput.getType());
    }
  }
}

/// Compute the indexing map in the collapsed op that corresponds to the given
/// `indexingMap` of the original operation.
static AffineMap
getCollapsedOpIndexingMap(AffineMap indexingMap,
                          const CollapsingInfo &collapsingInfo) {
  MLIRContext *context = indexingMap.getContext();
  assert(indexingMap.isProjectedPermutation() &&
         "expected indexing map to be projected permutation");
  SmallVector<AffineExpr> resultExprs;
  auto origOpToCollapsedOpMapping =
      collapsingInfo.getOrigOpToCollapsedOpMapping();
  for (auto expr : indexingMap.getResults()) {
    unsigned dim = cast<AffineDimExpr>(expr).getPosition();
    // If the dim is not the first of the collapsed dim, do nothing.
    if (origOpToCollapsedOpMapping[dim].second != 0) {
      continue;
    }
    // The next n-dims are guaranteed to be collapsed. So just use the
    // iteration dimension of the collapsed op.
    resultExprs.push_back(
        getAffineDimExpr(origOpToCollapsedOpMapping[dim].first, context));
  }
  return AffineMap::get(collapsingInfo.getCollapsedOpIterationRank(), 0,
                        resultExprs, context);
}

/// Returns a copy of `origOp` with collapsed iteration dimensions.
template <typename AttentionOpTy>
static Operation *createCollapsedOp(AttentionOpTy origOp,
                                    const CollapsingInfo &collapsingInfo,
                                    RewriterBase &rewriter) {
  SmallVector<Value> inputOperands, outputOperands;
  SmallVector<Type> resultTypes;
  collapseOperandsAndResults(origOp, collapsingInfo, rewriter, inputOperands,
                             outputOperands, resultTypes);
  SmallVector<AffineMap> indexingMaps(
      llvm::map_range(origOp.getIndexingMapsArray(), [&](AffineMap map) {
        return getCollapsedOpIndexingMap(map, collapsingInfo);
      }));

  auto collapsedOp = AttentionOpTy::create(
      rewriter, origOp.getLoc(), resultTypes, inputOperands, outputOperands,
      rewriter.getAffineMapArrayAttr(indexingMaps));
  rewriter.inlineRegionBefore(origOp.getRegion(), collapsedOp.getRegion(),
                              collapsedOp.getRegion().begin());
  return collapsedOp;
}

template <typename AttentionOpTy>
static FailureOr<CollapseResult> collapseAttentionLikeOpIterationDims(
    AttentionOpTy op, ArrayRef<ReassociationIndices> foldedIterationDims,
    RewriterBase &rewriter) {
  if (op.getNumLoops() <= 1 || foldedIterationDims.empty() ||
      llvm::all_of(foldedIterationDims, [](ReassociationIndicesRef foldedDims) {
        return foldedDims.size() <= 1;
      })) {
    return failure();
  }

  CollapsingInfo collapsingInfo;
  if (failed(
          collapsingInfo.initialize(op.getNumLoops(), foldedIterationDims))) {
    return rewriter.notifyMatchFailure(
        op, "illegal to collapse specified dimensions");
  }

  Operation *collapsedOp = createCollapsedOp(op, collapsingInfo, rewriter);

  auto loc = op.getLoc();
  SmallVector<Value> results;
  for (const auto &originalResult : llvm::enumerate(op->getResults())) {
    Value collapsedOpResult = collapsedOp->getResult(originalResult.index());
    auto originalResultType =
        cast<ShapedType>(originalResult.value().getType());
    auto collapsedOpResultType = cast<ShapedType>(collapsedOpResult.getType());
    if (collapsedOpResultType.getRank() != originalResultType.getRank()) {
      AffineMap indexingMap =
          op.getIndexingMapMatchingResult(originalResult.value());
      SmallVector<ReassociationIndices> reassociation =
          getOperandReassociation(indexingMap, collapsingInfo);
      Value result;
      if (isa<MemRefType>(collapsedOpResult.getType())) {
        MemRefType expandShapeResultType = MemRefType::get(
            originalResultType.getShape(), originalResultType.getElementType());
        result =
            memref::ExpandShapeOp::create(rewriter, loc, expandShapeResultType,
                                          collapsedOpResult, reassociation);
      } else {
        result =
            tensor::ExpandShapeOp::create(rewriter, loc, originalResultType,
                                          collapsedOpResult, reassociation);
      }
      results.push_back(result);
    } else {
      results.push_back(collapsedOpResult);
    }
  }
  return CollapseResult{results, collapsedOp};
}

FailureOr<CollapseResult>
collapseOpIterationDims(AttentionOp op,
                        ArrayRef<ReassociationIndices> foldedIterationDims,
                        RewriterBase &rewriter) {
  return collapseAttentionLikeOpIterationDims(op, foldedIterationDims,
                                              rewriter);
}

FailureOr<CollapseResult>
collapseOpIterationDims(OnlineAttentionOp op,
                        ArrayRef<ReassociationIndices> foldedIterationDims,
                        RewriterBase &rewriter) {
  return collapseAttentionLikeOpIterationDims(op, foldedIterationDims,
                                              rewriter);
}

void populateFoldReshapeOpsByExpansionPatterns(
    RewritePatternSet &patterns,
    const linalg::ControlFusionFn &controlFoldingReshapes) {
  patterns.add<FoldWithProducerReshapeByExpansion<AttentionOp>,
               FoldWithConsumerReshapeByExpansion<AttentionOp>,
               FoldWithProducerReshapeByExpansion<OnlineAttentionOp>,
               FoldWithConsumerReshapeByExpansion<OnlineAttentionOp>,
               FoldWithProducerReshapeByExpansion<ScatterOp>,
               FoldWithConsumerReshapeByExpansion<ScatterOp>,
               FoldWithProducerReshapeByExpansion<GatherOp>,
               FoldWithConsumerReshapeByExpansion<GatherOp>>(
      patterns.getContext(), controlFoldingReshapes);
}

SmallVector<unsigned> defaultControlDropUnitDims(Operation *op) {
  auto fusionOp = cast<LinalgFusionOpInterface>(op);
  return llvm::to_vector(llvm::seq<unsigned>(0, fusionOp.getNumLoops()));
}

static void
replaceUnitDimIndexOps(Block *body,
                       const llvm::SmallDenseSet<unsigned> &unitDims,
                       RewriterBase &rewriter) {
  for (IndexOp indexOp : llvm::make_early_inc_range(body->getOps<IndexOp>())) {
    OpBuilder::InsertionGuard guard(rewriter);
    rewriter.setInsertionPoint(indexOp);
    if (unitDims.contains(indexOp.getDim())) {
      rewriter.replaceOpWithNewOp<arith::ConstantIndexOp>(indexOp, 0);
      continue;
    }

    unsigned droppedDimsBefore = llvm::count_if(
        unitDims, [&](unsigned dim) { return dim < indexOp.getDim(); });
    if (droppedDimsBefore != 0) {
      rewriter.replaceOpWithNewOp<IndexOp>(indexOp, indexOp.getDim() -
                                                        droppedDimsBefore);
    }
  }
}

template <typename AttentionOpType>
struct DropAttentionLikeUnitDims final : OpRewritePattern<AttentionOpType> {
  DropAttentionLikeUnitDims(MLIRContext *context,
                            linalg::ControlDropUnitDims options,
                            PatternBenefit benefit = 1)
      : OpRewritePattern<AttentionOpType>(context, benefit),
        options(std::move(options)) {}

  LogicalResult matchAndRewrite(AttentionOpType attentionOp,
                                PatternRewriter &rewriter) const override {
    linalg::DroppedUnitDimsBuilder builder =
        [](Location loc, OpBuilder &b, IndexingMapOpInterface op,
           ArrayRef<Value> newOperands, ArrayRef<AffineMap> newIndexingMaps,
           const llvm::SmallDenseSet<unsigned> &droppedDims)
        -> IndexingMapOpInterface {
      auto attentionOp = cast<AttentionOpType>(op);
      auto resultTypes = llvm::map_to_vector(
          newOperands.take_back(attentionOp.getNumDpsInits()),
          [](Value v) { return v.getType(); });

      IRMapping mapping;
      for (auto [oldOperand, newOperand] :
           llvm::zip(attentionOp->getOperands(), newOperands)) {
        mapping.map(oldOperand, newOperand);
      }
      auto *cloned = b.clone(*attentionOp, mapping);
      auto newOp = cast<AttentionOpType>(cloned);
      for (auto [result, type] : llvm::zip(newOp->getResults(), resultTypes)) {
        result.setType(type);
      }
      newOp.setIndexingMapsAttr(b.getAffineMapArrayAttr(newIndexingMaps));
      IRRewriter rewriter(b);
      replaceUnitDimIndexOps(newOp.getBody(), droppedDims, rewriter);
      return newOp;
    };
    FailureOr<linalg::DropUnitDimsResult> result = linalg::dropUnitDims(
        rewriter, cast<IndexingMapOpInterface>(attentionOp.getOperation()),
        builder, options);
    if (failed(result)) {
      return failure();
    }

    rewriter.replaceOp(attentionOp, result->replacements);
    return success();
  }

private:
  linalg::ControlDropUnitDims options;
};

void populateFoldUnitExtentDimsPatterns(
    RewritePatternSet &patterns, const linalg::ControlDropUnitDims &options) {
  patterns.add<DropScatterUnitIndexDepth>(patterns.getContext());
  patterns
      .add<DropGatherUnitDims, DropScatterUnitDims,
           DropAttentionLikeUnitDims<AttentionOp>,
           DropAttentionLikeUnitDims<OnlineAttentionOp>, DropMapStoreUnitDims>(
          patterns.getContext(), options);
}

} // namespace mlir::iree_compiler::IREE::LinalgExt
