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

#include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtInterfaces.h"
#include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.h"
#include "iree/compiler/Dialect/LinalgExt/Utils/IndexingUtils.h"
#include "iree/compiler/Dialect/LinalgExt/Utils/Utils.h"
#include "llvm/ADT/TypeSwitch.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Affine/Utils.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Linalg/Utils/Utils.h"
#include "mlir/Dialect/Math/IR/Math.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/IR/IRMapping.h"
#include "mlir/IR/OpDefinition.h"

#define DEBUG_TYPE "linalg-ext-tiling"

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

//===----------------------------------------------------------------------===//
// Utils.
//===----------------------------------------------------------------------===//

/// Returns the size and offset scaled by some scale factor, and clamped to a
/// dimSize for the dimension. `(offset + size) * scale` will be clamped to the
/// `dimSize`.
static std::pair<OpFoldResult, OpFoldResult>
getScaledSizeAndOffset(OpBuilder &builder, Location loc, OpFoldResult size,
                       OpFoldResult offset, OpFoldResult dimSize,
                       int64_t offsetScale, int64_t sizeScale) {
  AffineExpr dim0, dim1, dim2;
  auto ctx = builder.getContext();
  bindDims(ctx, dim0, dim1, dim2);
  auto imageOffset = affine::makeComposedFoldedAffineApply(
      builder, loc, {dim0 * offsetScale}, offset);
  auto dimSizeValue = getValueOrCreateConstantIndexOp(builder, loc, dimSize);
  AffineMap sizeMap =
      AffineMap::get(3, 0, {dim0 - dim1, dim2 * sizeScale}, ctx);
  auto imageSize = affine::makeComposedFoldedAffineMin(
      builder, loc, sizeMap, {dimSizeValue, imageOffset, size});
  return std::make_pair(imageSize, imageOffset);
}

/// If the input has a fully static shape, return the static sizes. Otherwise,
/// attempt to reify the shape of the input from its defining op. Input dims
/// are store into `reifiedInputDims`.
static LogicalResult
getStaticOrReifiedInputDims(OpBuilder &builder, Location loc, Value input,
                            ReifiedRankedShapedTypeDims &reifiedInputDims) {
  if (auto reifyOp = input.getDefiningOp<ReifyRankedShapedTypeOpInterface>()) {
    return reifyOp.reifyResultShapes(builder, reifiedInputDims);
  }
  auto inputType = cast<ShapedType>(input.getType());
  if (!inputType.hasStaticShape()) {
    return failure();
  }
  reifiedInputDims.push_back(
      getAsIndexOpFoldResult(builder.getContext(), inputType.getShape()));
  return success();
}

/// Returns the rank of the Value's type, and 0 if it is not a ShapedType.
static int64_t getRank(Value v) {
  auto type = dyn_cast<ShapedType>(v.getType());
  if (type) {
    return type.getRank();
  }
  return 0;
}

static Value normalizeMaskValue(OpBuilder &builder, Location loc, Value mask) {
  auto intType = dyn_cast<IntegerType>(mask.getType());
  assert(intType && "expected integer mask type");
  if (intType.getWidth() == 1) {
    return mask;
  }
  return arith::TruncIOp::create(builder, loc, builder.getI1Type(), mask);
}

/// Method similar to `LinalgOp`s that concatenates shapes of all operands.
static SmallVector<OpFoldResult>
createFlatListOfOperandDims(OpBuilder &b, Location loc, Operation *op) {
  SmallVector<OpFoldResult> res;
  for (OpOperand &opOperand : op->getOpOperands()) {
    for (auto dim : llvm::seq(getRank(opOperand.get()))) {
      res.push_back(linalg::createFoldedDimOp(b, loc, opOperand.get(), dim));
    }
  }
  return res;
}

/// Permutes the offset and size arrays by the result indexes of the provided
/// affine map.
static SmallVector<Range> getPermutedRange(AffineMap permutation,
                                           ArrayRef<OpFoldResult> offsets,
                                           ArrayRef<OpFoldResult> sizes) {
  auto one = IntegerAttr::get(IndexType::get(permutation.getContext()), 1);
  assert(permutation.isProjectedPermutation() &&
         "Affine map should be a projected permutation");
  SmallVector<Range> output;
  for (AffineExpr dimExpr : permutation.getResults()) {
    int dim = cast<AffineDimExpr>(dimExpr).getPosition();
    output.push_back(Range{offsets[dim], sizes[dim], one});
  }
  return output;
}

//===----------------------------------------------------------------------===//
// ScatterOp
//===----------------------------------------------------------------------===//

SmallVector<utils::IteratorType> ScatterOp::getLoopIteratorTypes() {
  SmallVector<utils::IteratorType> iteratorTypes(getUpdateType().getRank(),
                                                 utils::IteratorType::parallel);
  if (!getUniqueIndices()) {
    int64_t batchRank = getBatchRank();
    for (auto i : llvm::seq<int64_t>(0, batchRank)) {
      iteratorTypes[i] = utils::IteratorType::reduction;
    }
  }
  return iteratorTypes;
}

SmallVector<Range> ScatterOp::getIterationDomain(OpBuilder &builder) {
  Location loc = getLoc();
  OpFoldResult zero = builder.getIndexAttr(0);
  OpFoldResult one = builder.getIndexAttr(1);
  SmallVector<Range> ranges;
  for (auto dim : llvm::seq<int64_t>(0, getUpdateType().getRank())) {
    OpFoldResult ub = getDim(builder, loc, getUpdates(), dim);
    ranges.push_back(Range{zero, ub, one});
  }
  return ranges;
}

FailureOr<TilingResult>
ScatterOp::getTiledImplementation(OpBuilder &builder,
                                  ArrayRef<OpFoldResult> offsets,
                                  ArrayRef<OpFoldResult> sizes) {
  assert(offsets.size() >= 1 && sizes.size() >= 1);
  Location loc = getLoc();
  auto zeroAttr = builder.getI64IntegerAttr(0);
  auto oneAttr = builder.getI64IntegerAttr(1);
  SmallVector<Operation *> slices;

  // Slice of the updates.
  auto updateRank = getUpdateType().getRank();
  SmallVector<OpFoldResult> updateStrides(updateRank, oneAttr);
  Operation *updateSlice =
      getSlice(builder, loc, getUpdates(), offsets, sizes, updateStrides);
  Value tiledUpdate = updateSlice->getResult(0);
  slices.push_back(updateSlice);

  // Slice of indices.
  auto indicesRank = getIndicesType().getRank();
  SmallVector<OpFoldResult> indicesOffsets(offsets.take_front(getBatchRank()));
  SmallVector<OpFoldResult> indicesSizes(sizes.take_front(getBatchRank()));
  if (getBatchRank() != getIndicesType().getRank()) {
    indicesOffsets.push_back(zeroAttr);
    indicesSizes.push_back(builder.getIndexAttr(getIndexDepth()));
  }
  SmallVector<OpFoldResult> indicesStrides(indicesRank, oneAttr);

  Operation *indicesSlice = getSlice(builder, loc, getIndices(), indicesOffsets,
                                     indicesSizes, indicesStrides);
  Value tiledIndices = indicesSlice->getResult(0);
  slices.push_back(indicesSlice);

  Value tiledMask;
  if (Value mask = getMask()) {
    std::optional<ShapedType> maskType = getMaskType();
    if (maskType->getRank() == 0) {
      tiledMask = mask;
    } else {
      SmallVector<OpFoldResult> maskOffsets(offsets.take_front(getBatchRank()));
      SmallVector<OpFoldResult> maskSizes(sizes.take_front(getBatchRank()));
      SmallVector<OpFoldResult> maskStrides(maskType->getRank(), oneAttr);
      Operation *maskSlice =
          getSlice(builder, loc, mask, maskOffsets, maskSizes, maskStrides);
      tiledMask = maskSlice->getResult(0);
      slices.push_back(maskSlice);
    }
  }

  // Slice of the original.
  SmallVector<OpFoldResult> originalOffsets, originalSizes;
  if (failed(getResultTilePosition(builder, 0, offsets, sizes, originalOffsets,
                                   originalSizes))) {
    return {};
  }
  auto originalRank = getOriginalType().getRank();
  SmallVector<OpFoldResult> originalStrides(originalRank, oneAttr);
  Operation *originalSlice =
      getSlice(builder, loc, getOriginal(), originalOffsets, originalSizes,
               originalStrides);
  Value tiledOriginal = originalSlice->getResult(0);
  slices.push_back(originalSlice);

  SmallVector<Type> resultTypes;
  if (getNumResults()) {
    resultTypes.push_back(tiledOriginal.getType());
  }
  SmallVector<Value> tiledOperands = {tiledUpdate, tiledIndices};
  if (tiledMask) {
    tiledOperands.push_back(tiledMask);
  }
  tiledOperands.push_back(tiledOriginal);
  Operation *tiledScatterOp =
      mlir::clone(builder, getOperation(), resultTypes, tiledOperands);
  return TilingResult{{tiledScatterOp},
                      SmallVector<Value>(tiledScatterOp->getResults()),
                      slices};
}

LogicalResult ScatterOp::getResultTilePosition(
    OpBuilder &builder, unsigned resultNumber, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, SmallVector<OpFoldResult> &resultOffsets,
    SmallVector<OpFoldResult> &resultSizes) {
  auto zeroAttr = builder.getI64IntegerAttr(0);
  // Slice of the original.
  auto originalRank = getOriginalType().getRank();
  resultOffsets.resize(originalRank, zeroAttr);
  resultSizes.resize(originalRank);

  auto updateRank = getUpdateType().getRank();
  Location loc = getLoc();
  for (auto dim : llvm::seq<int64_t>(0, originalRank - getUpdateSliceRank())) {
    resultSizes[dim] = getDim(builder, loc, getOriginal(), dim);
  }
  for (auto dim :
       llvm::seq<int64_t>(originalRank - getUpdateSliceRank(), originalRank)) {
    resultOffsets[dim] = offsets[dim - (originalRank - updateRank)];
    resultSizes[dim] = sizes[dim - (originalRank - updateRank)];
  }
  return success();
}

/// Method to return the position of the result tile computed by the tiled
/// operation.
LogicalResult ScatterOp::getIterationDomainTileFromOperandTiles(
    OpBuilder &b, ArrayRef<unsigned> operandNumbers,
    ArrayRef<SmallVector<OpFoldResult>> allOffsets,
    ArrayRef<SmallVector<OpFoldResult>> allSizes,
    SmallVectorImpl<OpFoldResult> &iterDomainOffsets,
    SmallVectorImpl<OpFoldResult> &iterDomainSizes) {
  // Fusion with producers is not possible in general if `unique_indices` is not
  // true as reductions along the scattered indices are not tilable in parallel.
  if (!getUniqueIndices()) {
    return failure();
  }
  // TODO: Support fusion along the index operand. For the index operand, the
  // offset + size must be the full size for the inner most dim.
  if (operandNumbers.size() != 1 ||
      getUpdatesMutable().getOperandNumber() != operandNumbers.front()) {
    return failure();
  }
  ArrayRef<OpFoldResult> offsets(allOffsets[0]);
  ArrayRef<OpFoldResult> sizes(allSizes[0]);

  // The iteration domain is defined in terms of the |input|, so simply
  // use the given offsets/sizes.
  iterDomainOffsets.assign(offsets.begin(), offsets.end());
  iterDomainSizes.assign(sizes.begin(), sizes.end());
  return success();
}

/// Method to generate the tiled implementation of an operation from the tile
/// of the operand.
FailureOr<TilingResult> ScatterOp::getTiledImplementationFromOperandTiles(
    OpBuilder &b, ArrayRef<unsigned> operandNumbers,
    ArrayRef<SmallVector<OpFoldResult>> allOffsets,
    ArrayRef<SmallVector<OpFoldResult>> allSizes) {
  SmallVector<OpFoldResult> mappedOffsets, mappedSizes;
  if (failed(getIterationDomainTileFromOperandTiles(
          b, operandNumbers, allOffsets, allSizes, mappedOffsets,
          mappedSizes))) {
    return failure();
  }
  return getTiledImplementation(b, mappedOffsets, mappedSizes);
}

LogicalResult ScatterOp::generateScalarImplementation(OpBuilder &b,
                                                      Location loc,
                                                      ValueRange ivs) {
  auto indexDepth = getIndexDepth();
  Value update = memref::LoadOp::create(b, loc, getUpdates(), ivs);
  SmallVector<Value> starts;
  SmallVector<Value> loadIndices;
  append_range(loadIndices, ivs.take_front(getBatchRank()));

  // Populate with empty values.
  auto originalTy = getOriginalType();
  starts.resize(originalTy.getRank(), Value());
  auto updateIvs = ivs.drop_front(getBatchRank());

  int64_t offset = starts.size() - updateIvs.size();
  for (auto [idx, iv] : llvm::enumerate(updateIvs)) {
    starts[idx + offset] = iv;
  }

  ArrayRef<int64_t> dimMap = getDimensionMap();

  if (getIndicesType().getRank() > getBatchRank()) {
    loadIndices.push_back(Value());
  }
  for (auto i : llvm::seq<unsigned>(0, indexDepth)) {
    if (getIndicesType().getRank() > getBatchRank()) {
      loadIndices.back() = arith::ConstantIndexOp::create(b, loc, i);
    }
    Value idx = memref::LoadOp::create(b, loc, getIndices(), loadIndices);
    Value ret = arith::IndexCastOp::create(b, loc, b.getIndexType(), idx);

    auto dim = dimMap[i];

    if (starts[dim]) {
      ret = arith::AddIOp::create(b, loc, ret, starts[dim]);
    }
    starts[dim] = ret;
  }

  auto emitScatterUpdate = [&](OpBuilder &builder, Location nestedLoc) {
    Value init =
        memref::LoadOp::create(builder, nestedLoc, getOriginal(), starts);

    IRMapping bvm;
    Block &block = getRegion().front();
    bvm.map(block.getArgument(0), update);
    bvm.map(block.getArgument(1), init);
    for (auto &blockOp : block.without_terminator()) {
      builder.clone(blockOp, bvm);
    }
    // The last op is linalg_ext.yield op. Store the operand to destination.
    memref::StoreOp::create(
        builder, nestedLoc,
        bvm.lookupOrDefault(block.getTerminator()->getOperand(0)),
        getOriginal(), starts);
  };

  if (Value mask = getMask()) {
    SmallVector<Value> maskIndices(ivs.take_front(getBatchRank()));
    Value maskValue = memref::LoadOp::create(b, loc, mask, maskIndices);
    maskValue = normalizeMaskValue(b, loc, maskValue);
    scf::IfOp::create(b, loc, maskValue,
                      [&](OpBuilder &thenBuilder, Location thenLoc) {
                        emitScatterUpdate(thenBuilder, thenLoc);
                        scf::YieldOp::create(thenBuilder, thenLoc);
                      });
    return success();
  }

  emitScatterUpdate(b, loc);
  return success();
}

//===----------------------------------------------------------------------===//
// GatherOp
//===----------------------------------------------------------------------===//

SmallVector<utils::IteratorType> GatherOp::getLoopIteratorTypes() {
  return SmallVector<utils::IteratorType>(getOutputType().getRank(),
                                          utils::IteratorType::parallel);
}

SmallVector<Range> GatherOp::getIterationDomain(OpBuilder &builder) {
  Location loc = getLoc();
  OpFoldResult zero = builder.getIndexAttr(0);
  OpFoldResult one = builder.getIndexAttr(1);
  SmallVector<Range> ranges;
  for (auto dim : llvm::seq<int64_t>(0, getOutputType().getRank())) {
    OpFoldResult ub = getDim(builder, loc, getOutput(), dim);
    ranges.push_back(Range{zero, ub, one});
  }
  return ranges;
}

FailureOr<TilingResult>
GatherOp::getTiledImplementation(OpBuilder &builder,
                                 ArrayRef<OpFoldResult> offsets,
                                 ArrayRef<OpFoldResult> sizes) {
  assert(offsets.size() >= 1 && sizes.size() >= 1);
  Location loc = getLoc();
  auto zeroAttr = builder.getI64IntegerAttr(0);
  auto oneAttr = builder.getI64IntegerAttr(1);
  SmallVector<Operation *> slices;

  // Slice of the result.
  auto resultRank = getOutputType().getRank();
  SmallVector<OpFoldResult> resultStrides(resultRank, oneAttr);
  Operation *resultSlice =
      getSlice(builder, loc, getOutput(), offsets, sizes, resultStrides);
  Value tiledResult = resultSlice->getResult(0);

  // Slice of indices.
  auto indicesRank = getIndicesType().getRank();
  SmallVector<OpFoldResult> indicesOffsets(offsets.take_front(getBatchRank()));
  SmallVector<OpFoldResult> indicesSizes(sizes.take_front(getBatchRank()));
  if (getBatchRank() != getIndicesType().getRank()) {
    indicesOffsets.push_back(zeroAttr);
    indicesSizes.push_back(builder.getIndexAttr(getIndexDepth()));
  }
  SmallVector<OpFoldResult> indicesStrides(indicesRank, oneAttr);

  Operation *indicesSlice = getSlice(builder, loc, getIndices(), indicesOffsets,
                                     indicesSizes, indicesStrides);
  Value tiledIndices = indicesSlice->getResult(0);

  Value tiledMask;
  if (Value mask = getMask()) {
    std::optional<ShapedType> maskType = getMaskType();
    if (maskType->getRank() == 0) {
      tiledMask = mask;
    } else {
      SmallVector<OpFoldResult> maskOffsets(offsets.take_front(getBatchRank()));
      SmallVector<OpFoldResult> maskSizes(sizes.take_front(getBatchRank()));
      SmallVector<OpFoldResult> maskStrides(maskType->getRank(), oneAttr);
      Operation *maskSlice =
          getSlice(builder, loc, mask, maskOffsets, maskSizes, maskStrides);
      tiledMask = maskSlice->getResult(0);
      slices.push_back(maskSlice);
    }
  }

  // Slice of the source.
  auto sourceRank = getSourceType().getRank();
  auto indexDepth = getIndexDepth();

  // The first `indexDepth` dims are not tiled
  SmallVector<OpFoldResult> sourceOffsets, sourceSizes;
  for (auto dim : llvm::seq<int64_t>(0, indexDepth)) {
    sourceOffsets.push_back(zeroAttr);
    sourceSizes.push_back(getDim(builder, loc, getSource(), dim));
  }
  llvm::append_range(sourceOffsets,
                     offsets.slice(getBatchRank(), sourceRank - indexDepth));
  llvm::append_range(sourceSizes,
                     sizes.slice(getBatchRank(), sourceRank - indexDepth));
  SmallVector<OpFoldResult> sourceStrides(sourceRank, oneAttr);
  Operation *sourceSlice = getSlice(builder, loc, getSource(), sourceOffsets,
                                    sourceSizes, sourceStrides);
  Value tiledSource = sourceSlice->getResult(0);

  slices.push_back(sourceSlice);
  slices.push_back(indicesSlice);
  slices.push_back(resultSlice);

  SmallVector<Type> resultTypes;
  if (getNumResults()) {
    resultTypes.push_back(tiledResult.getType());
  }
  SmallVector<Value> tiledOperands = {tiledSource, tiledIndices};
  if (tiledMask) {
    tiledOperands.push_back(tiledMask);
  }
  tiledOperands.push_back(tiledResult);
  Operation *tiledGatherOp =
      mlir::clone(builder, getOperation(), resultTypes, tiledOperands);
  return TilingResult{
      {tiledGatherOp}, SmallVector<Value>(tiledGatherOp->getResults()), slices};
}

LogicalResult GatherOp::getResultTilePosition(
    OpBuilder &builder, unsigned resultNumber, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, SmallVector<OpFoldResult> &resultOffsets,
    SmallVector<OpFoldResult> &resultSizes) {
  resultOffsets.assign(offsets.begin(), offsets.end());
  resultSizes.assign(sizes.begin(), sizes.end());
  return success();
}

FailureOr<TilingResult>
GatherOp::generateResultTileValue(OpBuilder &builder, unsigned resultNumber,
                                  ArrayRef<OpFoldResult> offsets,
                                  ArrayRef<OpFoldResult> sizes) {
  return getTiledImplementation(builder, offsets, sizes);
}

LogicalResult GatherOp::generateScalarImplementation(OpBuilder &b, Location loc,
                                                     ValueRange ivs) {
  auto indexDepth = getIndexDepth();
  SmallVector<Value> loadIndices(ivs.take_front(getBatchRank()));

  // Populate with empty values.
  auto sourceTy = getSourceType();
  auto resultIvs = ivs.drop_front(getBatchRank());
  SmallVector<Value> starts(sourceTy.getRank() - resultIvs.size(), Value());
  llvm::append_range(starts, resultIvs);

  // The innermost dim of `indices` having an innermost dim for each index.
  bool hasIndexDim = getIndicesType().getRank() > getBatchRank();
  if (hasIndexDim) {
    loadIndices.push_back(Value());
  }

  // Populate `starts` by loading indices from `indices`
  ArrayRef<int64_t> dimMap = getDimensionMap();
  for (int64_t i = 0; i < indexDepth; ++i) {
    if (hasIndexDim) {
      loadIndices.back() = arith::ConstantIndexOp::create(b, loc, i);
    }
    Value idx = memref::LoadOp::create(b, loc, getIndices(), loadIndices);
    Value ret = arith::IndexCastOp::create(b, loc, b.getIndexType(), idx);
    auto dim = dimMap[i];
    if (starts[dim]) {
      ret = arith::AddIOp::create(b, loc, ret, starts[dim]);
    }
    starts[dim] = ret;
  }

  auto emitGatherStore = [&](OpBuilder &builder, Location nestedLoc) {
    Value init =
        memref::LoadOp::create(builder, nestedLoc, getSource(), starts);
    // The last op is linalg_ext.yield op. Store the operand to destination.
    memref::StoreOp::create(builder, nestedLoc, init, getOutput(), ivs);
  };

  if (Value mask = getMask()) {
    SmallVector<Value> maskIndices(ivs.take_front(getBatchRank()));
    Value maskValue = memref::LoadOp::create(b, loc, mask, maskIndices);
    maskValue = normalizeMaskValue(b, loc, maskValue);
    scf::IfOp::create(b, loc, maskValue,
                      [&](OpBuilder &thenBuilder, Location thenLoc) {
                        emitGatherStore(thenBuilder, thenLoc);
                        scf::YieldOp::create(thenBuilder, thenLoc);
                      });
    return success();
  }

  emitGatherStore(b, loc);
  return success();
}

//===----------------------------------------------------------------------===//
// MapLoadOp
//===----------------------------------------------------------------------===//

SmallVector<utils::IteratorType> MapLoadOp::getLoopIteratorTypes() {
  SmallVector<utils::IteratorType> iteratorTypes(getOutputRank(),
                                                 utils::IteratorType::parallel);
  return iteratorTypes;
}

SmallVector<Range> MapLoadOp::getIterationDomain(OpBuilder &builder) {
  Location loc = getLoc();
  OpFoldResult zero = builder.getIndexAttr(0);
  OpFoldResult one = builder.getIndexAttr(1);
  SmallVector<Range> ranges;
  for (auto dim : llvm::seq<int64_t>(0, getOutputRank())) {
    OpFoldResult ub = getDim(builder, loc, getOutput(), dim);
    ranges.push_back(Range{zero, ub, one});
  }
  return ranges;
}

FailureOr<TilingResult>
MapLoadOp::getTiledImplementation(OpBuilder &builder,
                                  ArrayRef<OpFoldResult> offsets,
                                  ArrayRef<OpFoldResult> sizes) {
  Location loc = getLoc();

  // Get a slice of the output (the dest/init operand).
  SmallVector<OpFoldResult> outputStrides(getOutputRank(),
                                          builder.getI64IntegerAttr(1));
  Operation *outputSlice =
      getSlice(builder, loc, getOutput(), offsets, sizes, outputStrides);

  // Clone the operation with the full source but sliced output, and then
  // compose the tiling offsets with the index transformation of the
  // map_load op, because the space of the transformation output indices
  // is now local to the new output tile.
  Value tiledOutput = outputSlice->getResult(0);
  SmallVector<Type> resultTypes;
  if (getNumResults()) {
    resultTypes.push_back(tiledOutput.getType());
  }
  Operation *tiledOp = mlir::clone(builder, getOperation(), resultTypes,
                                   {getSource(), tiledOutput});
  auto tiledMapLoadOp = cast<MapLoadOp>(tiledOp);
  auto indexTransformBuilder =
      [&](ArrayRef<BlockArgument> outputIndices) -> SmallVector<Value> {
    SmallVector<OpFoldResult> offsetIndices;
    auto addMap = AffineMap::get(
        2, 0, {builder.getAffineDimExpr(0) + builder.getAffineDimExpr(1)});
    for (auto [outIdx, offset] : llvm::zip_equal(outputIndices, offsets)) {
      offsetIndices.push_back(affine::makeComposedFoldedAffineApply(
          builder, loc, addMap, {OpFoldResult(outIdx), offset}));
    }
    return getValueOrCreateConstantIndexOp(builder, loc, offsetIndices);
  };
  tiledMapLoadOp.insertTransformationAtStart(builder, indexTransformBuilder,
                                             offsets.size());
  return TilingResult{{tiledOp}, {tiledOp->getResults()}, {outputSlice}};
}

LogicalResult MapLoadOp::getResultTilePosition(
    OpBuilder &builder, unsigned resultNumber, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, SmallVector<OpFoldResult> &resultOffsets,
    SmallVector<OpFoldResult> &resultSizes) {
  resultOffsets.assign(offsets.begin(), offsets.end());
  resultSizes.assign(sizes.begin(), sizes.end());
  return success();
}

FailureOr<TilingResult>
MapLoadOp::generateResultTileValue(OpBuilder &builder, unsigned resultNumber,
                                   ArrayRef<OpFoldResult> offsets,
                                   ArrayRef<OpFoldResult> sizes) {
  return getTiledImplementation(builder, offsets, sizes);
}

LogicalResult MapLoadOp::getIterationDomainTileFromOperandTiles(
    OpBuilder &b, ArrayRef<unsigned> operandNumbers,
    ArrayRef<SmallVector<OpFoldResult>> allOffsets,
    ArrayRef<SmallVector<OpFoldResult>> allSizes,
    SmallVectorImpl<OpFoldResult> &iterDomainOffsets,
    SmallVectorImpl<OpFoldResult> &iterDomainSizes) {
  if (operandNumbers.size() != 1 ||
      operandNumbers.front() != getOutputMutable().getOperandNumber()) {
    return failure();
  }
  ArrayRef<OpFoldResult> offsets(allOffsets[0]);
  ArrayRef<OpFoldResult> sizes(allSizes[0]);

  // The iteration domain is defined in terms of the `output`, so simply
  // use the given offsets/sizes.
  iterDomainOffsets.assign(offsets.begin(), offsets.end());
  iterDomainSizes.assign(sizes.begin(), sizes.end());
  return success();
}

FailureOr<TilingResult> MapLoadOp::getTiledImplementationFromOperandTiles(
    OpBuilder &b, ArrayRef<unsigned> operandNumbers,
    ArrayRef<SmallVector<OpFoldResult>> allOffsets,
    ArrayRef<SmallVector<OpFoldResult>> allSizes) {
  SmallVector<OpFoldResult> mappedOffsets, mappedSizes;
  if (failed(getIterationDomainTileFromOperandTiles(
          b, operandNumbers, allOffsets, allSizes, mappedOffsets,
          mappedSizes))) {
    return failure();
  }
  return getTiledImplementation(b, mappedOffsets, mappedSizes);
}

/// The body of the transformation_region is inlined, and the yielded indices
/// are used to read values from the source and write to the output. Bounds
/// checking is performed on the source indices, and the padding value is used
/// if the indices are out of bounds.
LogicalResult MapLoadOp::generateScalarImplementation(OpBuilder &b,
                                                      Location loc,
                                                      ValueRange ivs) {
  // The scalar implementation is currently only implemented for buffer
  // semantics.
  if (!hasPureBufferSemantics()) {
    return failure();
  }

  auto bodyBuilder = [&](OpBuilder nestedBuilder, Location nestedLoc,
                         ArrayRef<Value> yieldedValues) {
    // The last yielded Value is the padding, the rest are source indices.
    Value paddingValue = yieldedValues.back();
    ArrayRef<Value> loadIndices = yieldedValues.drop_back();

    // Check bounds for each source dimension. Start with true so that
    // for 0-D sources, inBounds is always true.
    Value inBounds = nestedBuilder.createOrFold<arith::ConstantIntOp>(
        nestedLoc, /*value=*/1, /*width=*/1);
    Value zero =
        nestedBuilder.createOrFold<arith::ConstantIndexOp>(nestedLoc, 0);
    for (auto [dim, idx] : llvm::enumerate(loadIndices)) {
      Value dimSize =
          memref::DimOp::create(nestedBuilder, nestedLoc, getSource(), dim);

      // Check: idx >= 0
      Value geZero = arith::CmpIOp::create(
          nestedBuilder, nestedLoc, arith::CmpIPredicate::sge, idx, zero);
      // Check: idx < dimSize
      Value ltDim = arith::CmpIOp::create(
          nestedBuilder, nestedLoc, arith::CmpIPredicate::slt, idx, dimSize);
      // Combine: idx >= 0 && idx < dimSize
      Value dimInBounds =
          arith::AndIOp::create(nestedBuilder, nestedLoc, geZero, ltDim);

      inBounds = arith::AndIOp::create(nestedBuilder, nestedLoc, inBounds,
                                       dimInBounds);
    }

    // Create if-else: if in bounds, load from source; else use padding.
    // The if yields the value to store.
    auto ifOp = scf::IfOp::create(nestedBuilder, nestedLoc,
                                  TypeRange{paddingValue.getType()}, inBounds,
                                  /*addThenBlock=*/true, /*addElseBlock=*/true);
    {
      auto thenBuilder = ifOp.getThenBodyBuilder();
      Value loaded = memref::LoadOp::create(thenBuilder, nestedLoc, getSource(),
                                            loadIndices);
      scf::YieldOp::create(thenBuilder, nestedLoc, loaded);
    }
    {
      auto elseBuilder = ifOp.getElseBodyBuilder();
      scf::YieldOp::create(elseBuilder, nestedLoc, paddingValue);
    }
    memref::StoreOp::create(nestedBuilder, nestedLoc, ifOp.getResult(0),
                            getOutput(), ivs);
  };
  inlineMapLoadBody(b, loc, ivs, bodyBuilder);
  return success();
}

//===----------------------------------------------------------------------===//
// MapStoreOp
//===----------------------------------------------------------------------===//

SmallVector<utils::IteratorType> MapStoreOp::getLoopIteratorTypes() {
  SmallVector<utils::IteratorType> iteratorTypes(getInputRank(),
                                                 utils::IteratorType::parallel);
  return iteratorTypes;
}

SmallVector<Range> MapStoreOp::getIterationDomain(OpBuilder &builder) {
  Location loc = getLoc();
  OpFoldResult zero = builder.getIndexAttr(0);
  OpFoldResult one = builder.getIndexAttr(1);
  SmallVector<Range> ranges;
  for (auto dim : llvm::seq<int64_t>(0, getInputRank())) {
    OpFoldResult ub = getDim(builder, loc, getInput(), dim);
    ranges.push_back(Range{zero, ub, one});
  }
  return ranges;
}

FailureOr<TilingResult>
MapStoreOp::getTiledImplementation(OpBuilder &builder,
                                   ArrayRef<OpFoldResult> offsets,
                                   ArrayRef<OpFoldResult> sizes) {
  Location loc = getLoc();

  // Get a slice of the input.
  SmallVector<OpFoldResult> inputStrides(getInputRank(),
                                         builder.getI64IntegerAttr(1));
  Operation *inputSlice =
      getSlice(builder, loc, getInput(), offsets, sizes, inputStrides);

  // Clone the operation with the slice of the input, and then compose the
  // tiling offsets with the index transformation of the map_store op,
  // because the space of the transformation source indices is now local to
  // the new input tile.
  Operation *tiledOp = mlir::clone(builder, getOperation(), getResultTypes(),
                                   {inputSlice->getResult(0), getOutput()});
  auto tiledMapStoreOp = cast<MapStoreOp>(tiledOp);
  auto indexTransformBuilder =
      [&](ArrayRef<BlockArgument> srcIndices) -> SmallVector<Value> {
    SmallVector<OpFoldResult> offsetIndices;
    auto addMap = AffineMap::get(
        2, 0, {builder.getAffineDimExpr(0) + builder.getAffineDimExpr(1)});
    for (auto [srcIdx, offset] : llvm::zip_equal(srcIndices, offsets)) {
      offsetIndices.push_back(affine::makeComposedFoldedAffineApply(
          builder, loc, addMap, {OpFoldResult(srcIdx), offset}));
    }
    return getValueOrCreateConstantIndexOp(builder, loc, offsetIndices);
  };
  tiledMapStoreOp.insertTransformationAtStart(builder, indexTransformBuilder,
                                              offsets.size());
  return TilingResult{{tiledOp}, {tiledOp->getResults()}, {inputSlice}};
}

LogicalResult MapStoreOp::getResultTilePosition(
    OpBuilder &builder, unsigned resultNumber, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, SmallVector<OpFoldResult> &resultOffsets,
    SmallVector<OpFoldResult> &resultSizes) {
  auto zeroAttr = builder.getI64IntegerAttr(0);
  resultOffsets = SmallVector<OpFoldResult>(getOutputRank(), zeroAttr);
  for (auto dim : llvm::seq<int64_t>(getOutputRank())) {
    resultSizes.push_back(getDim(builder, getLoc(), getOutput(), dim));
  }
  return success();
}

LogicalResult MapStoreOp::getIterationDomainTileFromOperandTiles(
    OpBuilder &b, ArrayRef<unsigned> operandNumbers,
    ArrayRef<SmallVector<OpFoldResult>> allOffsets,
    ArrayRef<SmallVector<OpFoldResult>> allSizes,
    SmallVectorImpl<OpFoldResult> &iterDomainOffsets,
    SmallVectorImpl<OpFoldResult> &iterDomainSizes) {
  if (operandNumbers.size() != 1 ||
      operandNumbers.front() != getInputMutable().getOperandNumber()) {
    return failure();
  }
  ArrayRef<OpFoldResult> offsets(allOffsets[0]);
  ArrayRef<OpFoldResult> sizes(allSizes[0]);

  // The iteration domain is defined in terms of the `input`, so simply
  // use the given offsets/sizes.
  iterDomainOffsets.assign(offsets.begin(), offsets.end());
  iterDomainSizes.assign(sizes.begin(), sizes.end());
  return success();
}

FailureOr<TilingResult> MapStoreOp::getTiledImplementationFromOperandTiles(
    OpBuilder &b, ArrayRef<unsigned> operandNumbers,
    ArrayRef<SmallVector<OpFoldResult>> allOffsets,
    ArrayRef<SmallVector<OpFoldResult>> allSizes) {
  SmallVector<OpFoldResult> mappedOffsets, mappedSizes;
  if (failed(getIterationDomainTileFromOperandTiles(
          b, operandNumbers, allOffsets, allSizes, mappedOffsets,
          mappedSizes))) {
    return failure();
  }
  return getTiledImplementation(b, mappedOffsets, mappedSizes);
}

/// The body of the transformation_region is inlined, and the yielded indices
/// are used to write input values to the output. The reads and writes are
/// wrapped in an scf.if, conditioned on the yielded mask value of the
/// transformation body.
LogicalResult MapStoreOp::generateScalarImplementation(OpBuilder &b,
                                                       Location loc,
                                                       ValueRange ivs) {
  // The scalar implementation is currently only implemented for buffer
  // semantics, because we need to conditionally write values based on the
  // mask. Vectorized map_store ops should be decomposed, not tiled to loops,
  // so vector types are not allowed.
  if (!hasPureBufferSemantics() || isa<VectorType>(getInputType())) {
    return failure();
  }
  auto bodyBuilder = [&](OpBuilder nestedBuilder, Location nestedLoc,
                         ArrayRef<Value> yieldedValues) {
    // The last yielded Value is the mask, so use it as the if condition instead
    // of the store indices.
    Value ifCond = yieldedValues.back();
    ArrayRef<Value> storeIndices = yieldedValues.drop_back();
    auto thenBuilder = [&](OpBuilder &ifBuilder, Location ifLoc) {
      SmallVector<OpFoldResult> ivsOfr(ivs);
      Value input = memref::LoadOp::create(ifBuilder, ifLoc, getInput(), ivs);
      memref::StoreOp::create(ifBuilder, ifLoc, input, getOutput(),
                              storeIndices);
      scf::YieldOp::create(ifBuilder, ifLoc);
    };
    scf::IfOp::create(nestedBuilder, nestedLoc, ifCond, thenBuilder);
  };
  inlineMapStoreBody(b, loc, ivs, bodyBuilder);
  return success();
}

//===----------------------------------------------------------------------===//
// SortOp
//===----------------------------------------------------------------------===//

SmallVector<utils::IteratorType> SortOp::getLoopIteratorTypes() {
  // All loops except the dimension to sort along are parallel.
  SmallVector<utils::IteratorType> iteratorTypes(getOperandRank(),
                                                 utils::IteratorType::parallel);
  iteratorTypes[getDimension()] = utils::IteratorType::reduction;
  return iteratorTypes;
}

SmallVector<Range> SortOp::getIterationDomain(OpBuilder &builder) {
  int64_t operandRank = getOperandRank();
  SmallVector<Range> loopBounds(operandRank);
  Location loc = getLoc();
  OpFoldResult zero = builder.getIndexAttr(0);
  OpFoldResult one = builder.getIndexAttr(1);
  Value source = getOperand(0);
  for (auto dim : llvm::seq<int64_t>(0, operandRank)) {
    loopBounds[dim].offset = zero;
    loopBounds[dim].size = getDim(builder, loc, source, dim);
    loopBounds[dim].stride = one;
  }
  return loopBounds;
}

FailureOr<TilingResult>
SortOp::getTiledImplementation(OpBuilder &builder,
                               ArrayRef<OpFoldResult> offsets,
                               ArrayRef<OpFoldResult> sizes) {
  int64_t rank = getOperandRank();
  assert(offsets.size() == static_cast<size_t>(rank) &&
         sizes.size() == static_cast<size_t>(rank));
  auto oneAttr = builder.getI64IntegerAttr(1);
  SmallVector<OpFoldResult> strides(rank, oneAttr);
  SmallVector<Operation *> slices;
  SmallVector<Value> tiledOperands(getOutputs().size());
  for (auto [idx, output] : llvm::enumerate(getOutputs())) {
    Operation *slice =
        getSlice(builder, getLoc(), output, offsets, sizes, strides);
    tiledOperands[idx] = slice->getResult(0);
    slices.push_back(slice);
  }
  SmallVector<Type, 4> resultTypes;
  if (getNumResults()) {
    resultTypes = llvm::map_to_vector<4>(tiledOperands,
                                         [&](Value v) { return v.getType(); });
  }
  Operation *tiledSortOp =
      mlir::clone(builder, getOperation(), resultTypes, tiledOperands);
  return TilingResult{
      {tiledSortOp}, SmallVector<Value>{tiledSortOp->getResults()}, slices};
}

LogicalResult SortOp::getResultTilePosition(
    OpBuilder &builder, unsigned resultNumber, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, SmallVector<OpFoldResult> &resultOffsets,
    SmallVector<OpFoldResult> &resultSizes) {
  resultOffsets = llvm::to_vector(offsets);
  resultSizes = llvm::to_vector(sizes);
  return success();
}

/// Emit one bubble-sort sweep: compare adjacent elements along sortDim and
/// conditionally swap. compareOperands are fed to the comparator region (as
/// interleaved pairs); carryOperands are swapped in lockstep but not compared.
/// Convention: comparator returning true means keep current order (no swap).
static void emitBubbleSortSweep(OpBuilder &b, Location loc, Value zero,
                                Value one, Value ub, int64_t sortDim,
                                ValueRange ivs, ValueRange compareOperands,
                                ValueRange carryOperands,
                                Region &comparatorRegion) {
  Block &cmpBlock = comparatorRegion.front();
  scf::ForOp::create(
      b, loc, zero, ub, one, /*iterArgs=*/ValueRange{},
      [&](OpBuilder &b, Location loc, Value j, ValueRange) {
        // Compare adjacent elements at positions j (lhs) and j+1 (rhs).
        Value jPlusOne = arith::AddIOp::create(b, loc, j, one);
        SmallVector<Value> lhsIndices(ivs);
        lhsIndices[sortDim] = j;
        SmallVector<Value> rhsIndices(ivs);
        rhsIndices[sortDim] = jPlusOne;

        SmallVector<Value> cmpArgs;
        SmallVector<Value> lhsVals, rhsVals;
        for (Value operand : compareOperands) {
          Value lhs = memref::LoadOp::create(b, loc, operand, lhsIndices);
          Value rhs = memref::LoadOp::create(b, loc, operand, rhsIndices);
          cmpArgs.push_back(lhs);
          cmpArgs.push_back(rhs);
          lhsVals.push_back(lhs);
          rhsVals.push_back(rhs);
        }

        for (Value operand : carryOperands) {
          lhsVals.push_back(
              memref::LoadOp::create(b, loc, operand, lhsIndices));
          rhsVals.push_back(
              memref::LoadOp::create(b, loc, operand, rhsIndices));
        }

        IRMapping sortMap;
        sortMap.map(cmpBlock.getArguments(), cmpArgs);
        for (auto &blockOp : cmpBlock.without_terminator()) {
          b.clone(blockOp, sortMap);
        }
        Value keepOrder =
            sortMap.lookup(cmpBlock.getTerminator()->getOperand(0));

        SmallVector<Value> allOps;
        allOps.append(compareOperands.begin(), compareOperands.end());
        allOps.append(carryOperands.begin(), carryOperands.end());
        for (auto [operand, lhs, rhs] :
             llvm::zip_equal(allOps, lhsVals, rhsVals)) {
          Value newLhs = arith::SelectOp::create(b, loc, keepOrder, lhs, rhs);
          Value newRhs = arith::SelectOp::create(b, loc, keepOrder, rhs, lhs);
          memref::StoreOp::create(b, loc, newLhs, operand, lhsIndices);
          memref::StoreOp::create(b, loc, newRhs, operand, rhsIndices);
        }
        scf::YieldOp::create(b, loc, /*operands=*/ValueRange{});
      });
}

LogicalResult SortOp::generateScalarImplementation(OpBuilder &b, Location loc,
                                                   ValueRange ivs) {
  auto sortDim = getDimension();
  Value zero = arith::ConstantIndexOp::create(b, loc, 0);
  Value one = arith::ConstantIndexOp::create(b, loc, 1);
  Value ub;
  if (getOperandType(0).isDynamicDim(sortDim)) {
    ub = memref::DimOp::create(b, loc, getOperand(0), sortDim);
  } else {
    ub = arith::ConstantIndexOp::create(b, loc,
                                        getOperandType(0).getDimSize(sortDim));
  }
  ub = arith::SubIOp::create(b, loc, ub, one);
  SmallVector<Value> compareOperands(getDpsInits().begin(),
                                     getDpsInits().end());
  emitBubbleSortSweep(b, loc, zero, one, ub, sortDim, ivs, compareOperands,
                      /*carryOperands=*/{}, getRegion());
  return success();
}

//===----------------------------------------------------------------------===//
// FftOp
//===----------------------------------------------------------------------===//

SmallVector<utils::IteratorType> FftOp::getLoopIteratorTypes() {
  // There are `rank-1` outer loops. The fft itself has one loop for each
  // stage, which handles the merge step -- taking two half size tensors and
  // merge them into one tensor.
  SmallVector<utils::IteratorType> iteratorTypes(getOperandRank(),
                                                 utils::IteratorType::parallel);
  return iteratorTypes;
}

SmallVector<Range> FftOp::getIterationDomain(OpBuilder &builder) {
  SmallVector<Range> res;
  Location loc = getLoc();
  Value zero = arith::ConstantIndexOp::create(builder, loc, 0);
  Value one = arith::ConstantIndexOp::create(builder, loc, 1);
  for (auto [idx, val] : llvm::enumerate(getOperandShape().drop_back())) {
    OpFoldResult size;
    if (ShapedType::isDynamic(val)) {
      size = getDimValue(builder, loc, getReal(), idx);
    } else {
      size = builder.getIndexAttr(val);
    }
    res.emplace_back(Range{/*offset=*/zero, size, /*stride=*/one});
  }

  OpFoldResult size = getDim(builder, loc, getReal(), getOperandRank() - 1);
  Value stride = arith::ShLIOp::create(builder, loc, one, getStage());
  res.emplace_back(Range{/*offset=*/zero, size, /*stride=*/stride});
  return res;
}

void FftOp::generateScalarImplWithoutCoeffBuf(OpBuilder &b, Location loc,
                                              ArrayRef<Value> operands,
                                              Value wholeSize) {
  auto rank = getOperandRank();
  SmallVector<AffineMap> maps(operands.size(), b.getMultiDimIdentityMap(rank));

  auto f32Type = b.getF32Type();
  auto indexToF32 = [](OpBuilder &builder, Location loc, Value v) -> Value {
    v = arith::IndexCastOp::create(builder, loc, builder.getI32Type(), v);
    return arith::SIToFPOp::create(builder, loc, builder.getF32Type(), v);
  };

  // We will need exp(-2 * PI * j / m * I), compute "-2 * PI / m" for imag part
  // first.
  Value coeff = arith::ConstantFloatOp::create(
      b, loc, f32Type, llvm::APFloat(static_cast<float>(-2 * acos(-1))));
  coeff = arith::DivFOp::create(b, loc, coeff, indexToF32(b, loc, wholeSize));

  linalg::GenericOp::create(
      b, loc, TypeRange{}, ValueRange{}, operands, maps, getLoopIteratorTypes(),
      [&](OpBuilder &b, Location loc, ValueRange args) {
        Value lhsReal = args[0];
        Value lhsImag = args[1];
        Value rhsReal = args[2];
        Value rhsImag = args[3];

        // Compute "-2 * PI / m * j"
        Value w = arith::MulFOp::create(
            b, loc, coeff,
            indexToF32(b, loc, linalg::IndexOp::create(b, loc, rank - 1)));
        Value wReal = math::CosOp::create(b, loc, w);
        Value wImag = math::SinOp::create(b, loc, w);

        // t = w * a[k + j + mh];
        // ->  (x + yi)(u + vi) = (xu - yv) + (xv + yu)i
        Value xu = arith::MulFOp::create(b, loc, wReal, rhsReal);
        Value yv = arith::MulFOp::create(b, loc, wImag, rhsImag);
        Value xv = arith::MulFOp::create(b, loc, wReal, rhsImag);
        Value yu = arith::MulFOp::create(b, loc, wImag, rhsReal);
        Value tReal = arith::SubFOp::create(b, loc, xu, yv);
        Value tImag = arith::AddFOp::create(b, loc, xv, yu);

        // cplx u = a[k + j];
        // a[k + j] = u + t;
        // a[k + j + mh] = u - t;
        Value r1 = arith::AddFOp::create(b, loc, lhsReal, tReal);
        Value r2 = arith::AddFOp::create(b, loc, lhsImag, tImag);
        Value r3 = arith::SubFOp::create(b, loc, lhsReal, tReal);
        Value r4 = arith::SubFOp::create(b, loc, lhsImag, tImag);
        linalg::YieldOp::create(b, loc, ValueRange{r1, r2, r3, r4});
      });
}

void FftOp::generateScalarImplWithCoeffBuf(OpBuilder &b, Location loc,
                                           ArrayRef<Value> operands) {
  auto rank = getOperandRank();
  SmallVector<AffineMap> maps;
  // The size of coefficient buffer is expected to match `2^(stage-1)`, which
  // equals to the last dim of operands.
  maps.append(
      2, AffineMap::get(rank, 0, b.getAffineDimExpr(rank - 1), b.getContext()));
  maps.append(operands.size(), b.getMultiDimIdentityMap(rank));

  linalg::GenericOp::create(
      b, loc, TypeRange{}, ValueRange{getRealCoeff(), getImagCoeff()}, operands,
      maps, getLoopIteratorTypes(),
      [&](OpBuilder &b, Location loc, ValueRange args) {
        Value wReal = args[0];
        Value wImag = args[1];
        Value lhsReal = args[2];
        Value lhsImag = args[3];
        Value rhsReal = args[4];
        Value rhsImag = args[5];

        // t = w * a[k + j + mh];
        // ->  (x + yi)(u + vi) = (xu - yv) + (xv + yu)i
        Value xu = arith::MulFOp::create(b, loc, wReal, rhsReal);
        Value yv = arith::MulFOp::create(b, loc, wImag, rhsImag);
        Value xv = arith::MulFOp::create(b, loc, wReal, rhsImag);
        Value yu = arith::MulFOp::create(b, loc, wImag, rhsReal);
        Value tReal = arith::SubFOp::create(b, loc, xu, yv);
        Value tImag = arith::AddFOp::create(b, loc, xv, yu);

        // cplx u = a[k + j];
        // a[k + j] = u + t;
        // a[k + j + mh] = u - t;
        Value r1 = arith::AddFOp::create(b, loc, lhsReal, tReal);
        Value r2 = arith::AddFOp::create(b, loc, lhsImag, tImag);
        Value r3 = arith::SubFOp::create(b, loc, lhsReal, tReal);
        Value r4 = arith::SubFOp::create(b, loc, lhsImag, tImag);
        linalg::YieldOp::create(b, loc, ValueRange{r1, r2, r3, r4});
      });
}

// Generates FFT stage scalar implementation. This follows Cooley–Tukey FFT
// algorithm. The pseudo reference code is:
//   let s <- stage of linalg_ext.fft
//   int m = 1 << s;
//   int mh = m >> 1;
//   for (int k = 0; k < n; k += m) {
//     for (int j = 0; j < mh; ++j) {
//       cplx w = exp(-2 * PI * j / m * I);
//       cplx t = w * a[k + j + mh];
//       cplx u = a[k + j];
//       a[k + j] = u + t;
//       a[k + j + mh] = u - t;
//     }
//   }
LogicalResult FftOp::generateScalarImplementation(OpBuilder &b, Location loc,
                                                  ValueRange ivs) {
  Value real = getReal();
  Value imag = getImag();
  Value stage = getStage();
  Value one = arith::ConstantIndexOp::create(b, loc, 1);
  Value wholeSize = arith::ShLIOp::create(b, loc, one, stage);
  Value halfSize = arith::ShRSIOp::create(b, loc, wholeSize, one);

  auto rank = getOperandRank();
  SmallVector<Value> operands;
  SmallVector<OpFoldResult> lhsIvs(ivs);
  SmallVector<OpFoldResult> ones(rank, b.getIndexAttr(1));
  SmallVector<OpFoldResult> sizes(rank, b.getIndexAttr(1));
  sizes.back() = halfSize;
  operands.push_back(
      memref::SubViewOp::create(b, loc, real, lhsIvs, sizes, ones));
  operands.push_back(
      memref::SubViewOp::create(b, loc, imag, lhsIvs, sizes, ones));

  SmallVector<OpFoldResult> rhsIvs(ivs);
  rhsIvs.back() =
      arith::AddIOp::create(b, loc, ivs.back(), halfSize).getResult();
  operands.push_back(
      memref::SubViewOp::create(b, loc, real, rhsIvs, sizes, ones));
  operands.push_back(
      memref::SubViewOp::create(b, loc, imag, rhsIvs, sizes, ones));

  if (hasCoeff()) {
    generateScalarImplWithCoeffBuf(b, loc, operands);
  } else {
    generateScalarImplWithoutCoeffBuf(b, loc, operands, wholeSize);
  }

  return success();
}

FailureOr<TilingResult>
FftOp::getTiledImplementation(OpBuilder &builder,
                              ArrayRef<OpFoldResult> offsets,
                              ArrayRef<OpFoldResult> sizes) {
  int64_t rank = getOperandRank();
  SmallVector<OpFoldResult> strides(rank, builder.getI64IntegerAttr(1));
  SmallVector<Operation *> slices;
  SmallVector<Value> tiledOperands(3);
  tiledOperands[0] = getStage();
  tiledOperands[1] = getRealCoeff();
  tiledOperands[2] = getImagCoeff();
  SmallVector<Type, 4> resultTypes;

  for (Value out : getDpsInits()) {
    Operation *slice =
        getSlice(builder, getLoc(), out, offsets, sizes, strides);
    tiledOperands.push_back(slice->getResult(0));
    slices.push_back(slice);
    if (hasPureTensorSemantics()) {
      resultTypes.push_back(tiledOperands.back().getType());
    }
  }
  Operation *tiledFftOp =
      mlir::clone(builder, getOperation(), resultTypes, tiledOperands);
  return TilingResult{
      {tiledFftOp}, SmallVector<Value>(tiledFftOp->getResults()), slices};
}

LogicalResult FftOp::getResultTilePosition(
    OpBuilder &builder, unsigned resultNumber, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, SmallVector<OpFoldResult> &resultOffsets,
    SmallVector<OpFoldResult> &resultSizes) {
  resultOffsets.assign(offsets.begin(), offsets.end());
  resultSizes.assign(sizes.begin(), sizes.end());
  return success();
}

//===----------------------------------------------------------------------===//
// ScanOp
//===----------------------------------------------------------------------===//

SmallVector<Range> ScanOp::getIterationDomain(OpBuilder &builder) {
  int64_t operandRank = getOperandRank();
  SmallVector<Range> loopBounds(operandRank);
  Location loc = getLoc();
  OpFoldResult zero = builder.getIndexAttr(0);
  OpFoldResult one = builder.getIndexAttr(1);
  Value source = getInput();
  for (auto dim : llvm::seq<int64_t>(0, operandRank)) {
    loopBounds[dim].offset = zero;
    loopBounds[dim].size = getDim(builder, loc, source, dim);
    loopBounds[dim].stride = one;
  }
  return loopBounds;
}

SmallVector<utils::IteratorType> ScanOp::getLoopIteratorTypes() {
  SmallVector<utils::IteratorType> iteratorTypes(getOperandRank(),
                                                 utils::IteratorType::parallel);
  iteratorTypes[getDimension()] = utils::IteratorType::reduction;
  return iteratorTypes;
}

// Generates naive scalar implementation of scan for a given operator f.
// For inclusive,
//     output[0] = input[0]
//     output[i] = f(output[i-1], input[i])
//
// For exclusive,
//     output[0] = 0
//     output[i] = f(output[i-1], input[i-1])

LogicalResult ScanOp::generateScalarImplementation(OpBuilder &b, Location loc,
                                                   ValueRange ivs) {
  SmallVector<Value> indices, scanBlkArgs;
  indices.append(ivs.begin(), ivs.end());
  Value zero = arith::ConstantIndexOp::create(b, loc, 0);
  Value one = arith::ConstantIndexOp::create(b, loc, 1);
  auto scanDim = getDimension();
  auto cond = arith::CmpIOp::create(b, loc, arith::CmpIPredicate::eq,
                                    indices[scanDim], zero);
  bool isInclusive = getInclusive();
  SmallVector<Value> accIndices;
  for (int i = 0; i < indices.size(); i++) {
    if (i != scanDim) {
      accIndices.push_back(indices[i]);
    }
  }

  auto scfIf = scf::IfOp::create(
      b, loc, cond,
      [&](OpBuilder &b, Location loc) {
        if (isInclusive) {
          auto value = memref::LoadOp::create(b, loc, getInput(), indices);
          memref::StoreOp::create(b, loc, value, getOutput(), indices);
        } else {
          auto value =
              memref::LoadOp::create(b, loc, getAccumulator(), accIndices);
          memref::StoreOp::create(b, loc, value, getOutput(), indices);
        }
        scf::YieldOp::create(b, loc);
      },
      [&](OpBuilder &b, Location loc) {
        SmallVector<Value> indices(ivs);
        Value iv = indices[scanDim];
        Value ivMinusOne = arith::SubIOp::create(b, loc, iv, one);
        indices[scanDim] = ivMinusOne;
        scanBlkArgs.push_back(
            memref::LoadOp::create(b, loc, getOutput(), indices));
        Value i0;
        if (!isInclusive) {
          i0 = memref::LoadOp::create(b, loc, getInput(), indices);
        }
        indices[scanDim] = iv;
        if (isInclusive) {
          i0 = memref::LoadOp::create(b, loc, getInput(), indices);
        }
        scanBlkArgs.push_back(i0);
      });

  auto &srcBlock = getRegion().front();
  Region &region = scfIf.getElseRegion();
  IRMapping bvm;
  {
    OpBuilder::InsertionGuard guard(b);
    auto &block = region.front();
    b.setInsertionPointToEnd(&block);
    for (auto it : llvm::zip_equal(srcBlock.getArguments(), scanBlkArgs)) {
      bvm.map(std::get<0>(it), std::get<1>(it));
    }
    for (auto &blockOp : srcBlock.without_terminator()) {
      b.clone(blockOp, bvm);
    }
    memref::StoreOp::create(
        b, loc, bvm.lookupOrDefault(srcBlock.getTerminator()->getOperand(0)),
        getOutput(), indices);
    memref::StoreOp::create(
        b, loc, bvm.lookupOrDefault(srcBlock.getTerminator()->getOperand(0)),
        getAccumulator(), accIndices);
    scf::YieldOp::create(b, loc);
  }
  return success();
}

FailureOr<TilingResult>
ScanOp::getTiledImplementation(OpBuilder &builder,
                               ArrayRef<OpFoldResult> offsets,
                               ArrayRef<OpFoldResult> sizes) {
  int64_t rank = getOperandRank();
  assert(offsets.size() == static_cast<size_t>(rank) &&
         sizes.size() == static_cast<size_t>(rank));
  auto oneAttr = builder.getI64IntegerAttr(1);
  SmallVector<OpFoldResult> strides(rank, oneAttr);
  SmallVector<Value> tiledOperands;
  SmallVector<Operation *> slices;

  // Input
  {
    Operation *inputSlice =
        getSlice(builder, getLoc(), getInput(), offsets, sizes, strides);
    tiledOperands.emplace_back(inputSlice->getResult(0));
    slices.push_back(inputSlice);
  }

  // Output 0
  {
    Operation *output0Slice =
        getSlice(builder, getLoc(), getOutput(), offsets, sizes, strides);
    tiledOperands.emplace_back(output0Slice->getResult(0));
    slices.push_back(output0Slice);
  }

  if (rank > 1) {
    SmallVector<OpFoldResult> accumOffsets, accumSizes;
    if (failed(getResultTilePosition(builder, 1, offsets, sizes, accumOffsets,
                                     accumSizes))) {
      return {};
    }
    SmallVector<OpFoldResult> accumStrides(rank - 1, oneAttr);
    Operation *output1Slice = getSlice(builder, getLoc(), getAccumulator(),
                                       accumOffsets, accumSizes, accumStrides);
    tiledOperands.emplace_back(output1Slice->getResult(0));
    slices.push_back(output1Slice);
  } else {
    tiledOperands.emplace_back(getAccumulator());
  }

  SmallVector<Type, 4> resultTypes;
  if (hasPureTensorSemantics()) {
    resultTypes.push_back(tiledOperands[1].getType());
    resultTypes.push_back(tiledOperands[2].getType());
  }

  Operation *tiledScanOp =
      mlir::clone(builder, getOperation(), resultTypes, tiledOperands);
  return TilingResult{
      {tiledScanOp}, SmallVector<Value>(tiledScanOp->getResults()), slices};
}

LogicalResult ScanOp::getResultTilePosition(
    OpBuilder &builder, unsigned resultNumber, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, SmallVector<OpFoldResult> &resultOffsets,
    SmallVector<OpFoldResult> &resultSizes) {
  if (resultNumber == 0) {
    resultOffsets.assign(offsets.begin(), offsets.end());
    resultSizes.assign(sizes.begin(), sizes.end());
    return success();
  }
  if (resultNumber == 1) {
    int64_t rank = getOperandRank();
    if (rank > 1) {
      for (auto i : llvm::seq<int64_t>(0, rank)) {
        if (i == getDimension()) {
          continue;
        }
        resultOffsets.push_back(offsets[i]);
        resultSizes.push_back(sizes[i]);
      }
    }
    return success();
  }
  return failure();
}

//===----------------------------------------------------------------------===//
// TopkOp
//===----------------------------------------------------------------------===//

SmallVector<Range> TopkOp::getIterationDomain(OpBuilder &builder) {
  int64_t operandRank = getInputRank();
  SmallVector<Range> loopBounds(operandRank);
  Location loc = getLoc();
  OpFoldResult zero = builder.getIndexAttr(0);
  OpFoldResult one = builder.getIndexAttr(1);
  Value source = getValues();
  for (auto [idx, val] : llvm::enumerate(getInputType().getShape())) {
    loopBounds[idx].offset = zero;
    loopBounds[idx].size = getDim(builder, loc, source, idx);
    loopBounds[idx].stride = one;
  }
  return loopBounds;
}

SmallVector<utils::IteratorType> TopkOp::getLoopIteratorTypes() {
  SmallVector<utils::IteratorType> iteratorTypes(getInputRank(),
                                                 utils::IteratorType::parallel);
  iteratorTypes[getDimension()] = utils::IteratorType::reduction;
  return iteratorTypes;
}

LogicalResult TopkOp::generateScalarImplementation(OpBuilder &b, Location loc,
                                                   ValueRange ivs) {
  int64_t kDim = getDimension();
  Value zero = arith::ConstantIndexOp::create(b, loc, 0);
  Value one = arith::ConstantIndexOp::create(b, loc, 1);
  Value initialValue = memref::LoadOp::create(b, loc, getValues(), ivs);

  // If the indices tensor is not provided, the value index is derived from the
  // loop induction variables.
  Value initialIndex;
  if (getIndices()) {
    initialIndex = memref::LoadOp::create(b, loc, getIndices(), ivs);
  } else {
    Value rawInitialIndex = ivs[kDim];
    initialIndex =
        arith::IndexCastOp::create(b, loc, b.getI32Type(), rawInitialIndex);
  }

  // Compute K (ub) from the selected dim of the output
  Value ub = memref::DimOp::create(b, loc, getOutputValues(), getDimension());

  // Inner K loop functions:
  //   Load current K value and index
  //   Compare N/K using inserted block compare
  //   Check if N == K using strict weak ordering, select which index came first
  //   Select new K value from N/K comparison
  //   Select new K index from N/K comparison or which index came first
  //   Store new k value and index
  //   Yield loop carry values after K selection
  Value kValue, kIndex;
  auto scfFor = scf::ForOp::create(
      b, loc, zero, ub, one, ValueRange{initialValue, initialIndex},
      [&](OpBuilder &b, Location loc, Value iv, ValueRange loopCarryValues) {
        SmallVector<Value> indices(ivs);
        indices[kDim] = iv;
        kValue = memref::LoadOp::create(b, loc, getOutputValues(), indices);
        kIndex = memref::LoadOp::create(b, loc, getOutputIndices(), indices);
      });

  SmallVector<Value> indices(ivs);
  indices[kDim] = scfFor.getInductionVar();
  auto loopCarryValues = scfFor.getRegionIterArgs();

  // Retrieve region as black box comparison function f(x,y). Plug into op.
  auto &srcBlock = getRegion().front();
  IRMapping bvmF; // f(x,y)
  IRMapping bvmR; // f(y,x)
  {
    // Save previous insertion point. Continue within loop body.
    OpBuilder::InsertionGuard guard(b);
    b.setInsertionPointToEnd(&scfFor.getRegion().front());
    SmallVector<Value> forwardValues = {loopCarryValues[0], kValue};
    SmallVector<Value> reverseValues = {kValue, loopCarryValues[0]};
    bvmF.map(srcBlock.getArguments(), forwardValues);
    bvmR.map(srcBlock.getArguments(), reverseValues);
    for (auto &blockOp : srcBlock.without_terminator()) {
      b.clone(blockOp, bvmF);
      b.clone(blockOp, bvmR);
    }
    Value forwardCmpRes = bvmF.lookup(srcBlock.getTerminator()->getOperand(0));
    Value reverseCmpRes = bvmR.lookup(srcBlock.getTerminator()->getOperand(0));

    // Check value equality using strictly weak ordering from the region:
    //   f(x,y) --> forwardCmpRes
    //   f(y,x) --> reverseCmpRes
    //   if forwardCmpRes == reverseCmpRes then select which came first
    Value cmpValuesEqual = arith::CmpIOp::create(
        b, loc, arith::CmpIPredicate::eq, forwardCmpRes, reverseCmpRes);
    Value cmpFirstIndex = arith::CmpIOp::create(
        b, loc, arith::CmpIPredicate::slt, loopCarryValues[1], kIndex);
    Value combinedCmpEqRes =
        arith::AndIOp::create(b, loc, cmpValuesEqual, cmpFirstIndex);
    // True if N > K or N came before K
    Value indexCmpRes =
        arith::OrIOp::create(b, loc, forwardCmpRes, combinedCmpEqRes);
    // Select results for K based on comparisons
    Value resultKValue = arith::SelectOp::create(b, loc, forwardCmpRes,
                                                 loopCarryValues[0], kValue);
    Value resultKIndex = arith::SelectOp::create(b, loc, indexCmpRes,
                                                 loopCarryValues[1], kIndex);
    memref::StoreOp::create(b, loc, resultKValue, getOutputValues(), indices);
    memref::StoreOp::create(b, loc, resultKIndex, getOutputIndices(), indices);
    // Select loop carry, opposite of K results
    Value resultCarryValue = arith::SelectOp::create(
        b, loc, forwardCmpRes, kValue, loopCarryValues[0]);
    Value resultCarryIndex = arith::SelectOp::create(
        b, loc, indexCmpRes, kIndex, loopCarryValues[1]);
    scf::YieldOp::create(b, loc,
                         ValueRange{resultCarryValue, resultCarryIndex});
  }
  return success();
}

FailureOr<TilingResult>
TopkOp::getTiledImplementation(OpBuilder &builder,
                               ArrayRef<OpFoldResult> offsets,
                               ArrayRef<OpFoldResult> sizes) {
  int64_t rank = getInputRank();
  assert(offsets.size() == static_cast<size_t>(rank) &&
         sizes.size() == static_cast<size_t>(rank));
  SmallVector<OpFoldResult> strides(rank, builder.getI64IntegerAttr(1));
  Location loc = getLoc();

  SmallVector<OpFoldResult> outputOffsets, outputSizes;
  if (failed(getResultTilePosition(builder, 0, offsets, sizes, outputOffsets,
                                   outputSizes))) {
    return {};
  }

  SmallVector<Value> tiledOperands;
  SmallVector<Operation *> slices;

  // Values
  {
    Operation *valuesSlice =
        getSlice(builder, loc, getValues(), offsets, sizes, strides);
    tiledOperands.emplace_back(valuesSlice->getResult(0));
    slices.push_back(valuesSlice);
  }

  if (getIndices()) {
    Operation *indicesSlice =
        getSlice(builder, loc, getIndices(), offsets, sizes, strides);
    tiledOperands.emplace_back(indicesSlice->getResult(0));
    slices.push_back(indicesSlice);
  }

  // Replace the tile size for the K dimension to use the output size instead of
  // the input size.
  Value kSize =
      getDimValue(builder, getLoc(), getOutputValues(), getDimension());
  outputSizes[getDimension()] = getAsOpFoldResult(kSize);

  // Output values
  {
    Operation *outputValuesSlice = getSlice(builder, loc, getOutputValues(),
                                            offsets, outputSizes, strides);
    tiledOperands.emplace_back(outputValuesSlice->getResult(0));
    slices.push_back(outputValuesSlice);
  }

  // Output indices
  {
    Operation *outputIndicesSlice = getSlice(builder, loc, getOutputIndices(),
                                             offsets, outputSizes, strides);
    tiledOperands.emplace_back(outputIndicesSlice->getResult(0));
    slices.push_back(outputIndicesSlice);
  }

  SmallVector<Type, 2> resultTypes;
  if (hasPureTensorSemantics()) {
    resultTypes.push_back(tiledOperands[tiledOperands.size() - 2].getType());
    resultTypes.push_back(tiledOperands[tiledOperands.size() - 1].getType());
  }

  Operation *tiledTopkOp =
      mlir::clone(builder, getOperation(), resultTypes, tiledOperands);
  return TilingResult{
      {tiledTopkOp}, SmallVector<Value>(tiledTopkOp->getResults()), slices};
}

LogicalResult TopkOp::getResultTilePosition(
    OpBuilder &builder, unsigned resultNumber, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, SmallVector<OpFoldResult> &resultOffsets,
    SmallVector<OpFoldResult> &resultSizes) {
  resultOffsets.assign(offsets.begin(), offsets.end());
  resultSizes.assign(sizes.begin(), sizes.end());
  Value kSize = getDimValue(builder, getLoc(), getDpsInits()[resultNumber],
                            getDimension());
  resultSizes[getDimension()] = getAsOpFoldResult(kSize);
  return success();
}

/// Emit a full bubble sort over K output positions using the shared sweep
/// helper. The outer double loop provides the O(k^2) passes needed for a
/// complete sort.
static void emitBubbleSort(OpBuilder &b, Location loc, Value ub, Value zero,
                           Value one, int64_t kDim, ValueRange ivs,
                           Value outputValues, Value outputIndices,
                           Region &comparatorRegion) {
  Value kMinus1 = arith::SubIOp::create(b, loc, ub, one);
  SmallVector<Value> carryOps;
  if (outputIndices) {
    carryOps.push_back(outputIndices);
  }
  scf::ForOp::create(
      b, loc, zero, kMinus1, one, /*iterArgs=*/ValueRange{},
      [&](OpBuilder &b, Location loc, Value i, ValueRange) {
        Value innerUb = arith::SubIOp::create(b, loc, kMinus1, i);
        emitBubbleSortSweep(b, loc, zero, one, innerUb, kDim, ivs,
                            /*compareOperands=*/{outputValues}, carryOps,
                            comparatorRegion);
        scf::YieldOp::create(b, loc, /*operands=*/ValueRange{});
      });
}

//===----------------------------------------------------------------------===//
// TopkV2Op
//===----------------------------------------------------------------------===//

SmallVector<Range> TopkV2Op::getIterationDomain(OpBuilder &builder) {
  unsigned rank = getInputType().getRank();
  SmallVector<Range> loopBounds(rank);
  Location loc = getLoc();
  OpFoldResult zero = builder.getIndexAttr(0);
  OpFoldResult one = builder.getIndexAttr(1);
  Value source = getValues();
  for (unsigned idx = 0; idx < rank; ++idx) {
    loopBounds[idx].offset = zero;
    loopBounds[idx].size = getDim(builder, loc, source, idx);
    loopBounds[idx].stride = one;
  }
  return loopBounds;
}

SmallVector<utils::IteratorType> TopkV2Op::getLoopIteratorTypes() {
  SmallVector<utils::IteratorType> iteratorTypes(getInputRank(),
                                                 utils::IteratorType::parallel);
  iteratorTypes[getDimension()] = utils::IteratorType::reduction;
  return iteratorTypes;
}

LogicalResult TopkV2Op::generateScalarImplementation(OpBuilder &b, Location loc,
                                                     ValueRange ivs) {
  int64_t kDim = getDimension();
  Value zero = arith::ConstantIndexOp::create(b, loc, 0);
  Value one = arith::ConstantIndexOp::create(b, loc, 1);
  Value inputValues = getValues();
  Value outputValues = getOutputValues();
  Value outputIndices = getOutputIndices();
  Value initialValue = memref::LoadOp::create(b, loc, inputValues, ivs);

  // Compute K (ub) from the selected dim of the output.
  Value ub = memref::DimOp::create(b, loc, outputValues, kDim);

  SmallVector<Value> initValues = {initialValue};
  if (outputIndices) {
    Value initialIndex;
    Value inputIndices = getInputIndices();
    if (inputIndices) {
      initialIndex = memref::LoadOp::create(b, loc, inputIndices, ivs);
    } else {
      Type indexElemTy =
          cast<ShapedType>(outputIndices.getType()).getElementType();
      initialIndex = arith::IndexCastOp::create(b, loc, indexElemTy, ivs[kDim]);
    }
    initValues.push_back(initialIndex);
  }

  Block &srcBlock = getRegion().front();
  scf::ForOp::create(
      b, loc, zero, ub, one, initValues,
      [&](OpBuilder &b, Location loc, Value iv, ValueRange loopCarryValues) {
        SmallVector<Value> indices(ivs);
        indices[kDim] = iv;
        Value kValue = memref::LoadOp::create(b, loc, outputValues, indices);
        Value kIndex;
        if (outputIndices) {
          kIndex = memref::LoadOp::create(b, loc, outputIndices, indices);
        }

        // Clone the comparator region with forward mapping f(x,y).
        IRMapping bvmF;
        SmallVector<Value> forwardValues = {loopCarryValues[0], kValue};
        bvmF.map(srcBlock.getArguments(), forwardValues);
        for (auto &blockOp : srcBlock.without_terminator()) {
          b.clone(blockOp, bvmF);
        }
        Value forwardCmpRes =
            bvmF.lookup(srcBlock.getTerminator()->getOperand(0));

        Value resultKValue = arith::SelectOp::create(
            b, loc, forwardCmpRes, loopCarryValues[0], kValue);
        memref::StoreOp::create(b, loc, resultKValue, outputValues, indices);
        Value resultCarryValue = arith::SelectOp::create(
            b, loc, forwardCmpRes, kValue, loopCarryValues[0]);

        SmallVector<Value> yieldValues = {resultCarryValue};

        // With output indices, also clone the reverse mapping f(y,x) for
        // strict weak ordering tie-breaking.
        if (outputIndices) {
          IRMapping bvmR;
          SmallVector<Value> reverseValues = {kValue, loopCarryValues[0]};
          bvmR.map(srcBlock.getArguments(), reverseValues);
          for (auto &blockOp : srcBlock.without_terminator()) {
            b.clone(blockOp, bvmR);
          }
          Value reverseCmpRes =
              bvmR.lookup(srcBlock.getTerminator()->getOperand(0));

          // Strict weak ordering tie-breaking:
          //   f(x,y) --> forwardCmpRes.
          //   f(y,x) --> reverseCmpRes.
          //   If forwardCmpRes == reverseCmpRes, select which came first.
          Value cmpValuesEqual = arith::CmpIOp::create(
              b, loc, arith::CmpIPredicate::eq, forwardCmpRes, reverseCmpRes);
          Value cmpFirstIndex = arith::CmpIOp::create(
              b, loc, arith::CmpIPredicate::slt, loopCarryValues[1], kIndex);
          Value combinedCmpEqRes =
              arith::AndIOp::create(b, loc, cmpValuesEqual, cmpFirstIndex);
          Value indexCmpRes =
              arith::OrIOp::create(b, loc, forwardCmpRes, combinedCmpEqRes);
          Value resultKIndex = arith::SelectOp::create(
              b, loc, indexCmpRes, loopCarryValues[1], kIndex);
          memref::StoreOp::create(b, loc, resultKIndex, outputIndices, indices);
          Value resultCarryIndex = arith::SelectOp::create(
              b, loc, indexCmpRes, kIndex, loopCarryValues[1]);
          yieldValues.push_back(resultCarryIndex);
        }

        scf::YieldOp::create(b, loc, yieldValues);
      });

  // When is_sorted is true, add a bubble sort pass after each insertion to
  // ensure the output buffer is sorted regardless of its initial state. The
  // insertion-sort-style K loop above maintains sorted order only if the
  // output buffer was already sorted; this pass fixes any initial disorder.
  if (getIsSorted()) {
    emitBubbleSort(b, loc, ub, zero, one, kDim, ivs, outputValues,
                   outputIndices, getRegion());
  }

  return success();
}

FailureOr<TilingResult>
TopkV2Op::getTiledImplementation(OpBuilder &builder,
                                 ArrayRef<OpFoldResult> offsets,
                                 ArrayRef<OpFoldResult> sizes) {
  int64_t rank = getInputRank();
  assert(offsets.size() == static_cast<size_t>(rank) &&
         sizes.size() == static_cast<size_t>(rank));
  SmallVector<OpFoldResult> strides(rank, builder.getI64IntegerAttr(1));
  Location loc = getLoc();
  int64_t kDim = getDimension();
  Value outputValues = getOutputValues();
  Value outputIndices = getOutputIndices();

  SmallVector<OpFoldResult> outputOffsets, outputSizes;
  if (failed(getResultTilePosition(builder, 0, offsets, sizes, outputOffsets,
                                   outputSizes))) {
    return {};
  }

  SmallVector<Value> tiledOperands;
  SmallVector<Operation *> slices;

  // Input values.
  Operation *valuesSlice =
      getSlice(builder, loc, getValues(), offsets, sizes, strides);
  tiledOperands.emplace_back(valuesSlice->getResult(0));
  slices.push_back(valuesSlice);

  // Optional input indices.
  Value inputIndices = getInputIndices();
  if (inputIndices) {
    Operation *indicesSlice =
        getSlice(builder, loc, inputIndices, offsets, sizes, strides);
    tiledOperands.emplace_back(indicesSlice->getResult(0));
    slices.push_back(indicesSlice);
  }

  // Replace the tile size for the K dimension to use the output size instead of
  // the input size.
  Value kSize = getDimValue(builder, loc, outputValues, kDim);
  outputSizes[kDim] = getAsOpFoldResult(kSize);

  // Output values.
  Operation *outputValuesSlice =
      getSlice(builder, loc, outputValues, outputOffsets, outputSizes, strides);
  tiledOperands.emplace_back(outputValuesSlice->getResult(0));
  slices.push_back(outputValuesSlice);

  // Optional output indices.
  Operation *outputIndicesSlice = nullptr;
  if (outputIndices) {
    outputIndicesSlice = getSlice(builder, loc, outputIndices, outputOffsets,
                                  outputSizes, strides);
    tiledOperands.emplace_back(outputIndicesSlice->getResult(0));
    slices.push_back(outputIndicesSlice);
  }

  SmallVector<Type, 2> resultTypes;
  if (hasPureTensorSemantics()) {
    resultTypes.push_back(outputValuesSlice->getResult(0).getType());
    if (outputIndicesSlice) {
      resultTypes.push_back(outputIndicesSlice->getResult(0).getType());
    }
  }

  Operation *tiledTopkV2Op =
      mlir::clone(builder, getOperation(), resultTypes, tiledOperands);
  return TilingResult{
      {tiledTopkV2Op}, SmallVector<Value>(tiledTopkV2Op->getResults()), slices};
}

LogicalResult TopkV2Op::getResultTilePosition(
    OpBuilder &builder, unsigned resultNumber, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, SmallVector<OpFoldResult> &resultOffsets,
    SmallVector<OpFoldResult> &resultSizes) {
  resultOffsets.assign(offsets.begin(), offsets.end());
  resultSizes.assign(sizes.begin(), sizes.end());
  Value kSize = getDimValue(builder, getLoc(), getDpsInits()[resultNumber],
                            getDimension());
  resultSizes[getDimension()] = getAsOpFoldResult(kSize);
  return success();
}

//===----------------------------------------------------------------------===//
// ArgCompareOp
//===----------------------------------------------------------------------===//

SmallVector<Range> ArgCompareOp::getIterationDomain(OpBuilder &builder) {
  Location loc = getLoc();
  OpFoldResult zero = builder.getIndexAttr(0);
  OpFoldResult one = builder.getIndexAttr(1);
  int64_t rank = getInputRank();
  SmallVector<Range> ranges;
  for (int64_t dim = 0; dim < rank; ++dim) {
    OpFoldResult ub = getDim(builder, loc, getInputValue(), dim);
    ranges.push_back(Range{zero, ub, one});
  }
  return ranges;
}

SmallVector<utils::IteratorType> ArgCompareOp::getLoopIteratorTypes() {
  SmallVector<utils::IteratorType> iteratorTypes(getInputRank(),
                                                 utils::IteratorType::parallel);
  iteratorTypes[getDimension()] = utils::IteratorType::reduction;
  return iteratorTypes;
}

FailureOr<TilingResult>
ArgCompareOp::getTiledImplementation(OpBuilder &builder,
                                     ArrayRef<OpFoldResult> offsets,
                                     ArrayRef<OpFoldResult> sizes) {
  Location loc = getLoc();
  int64_t rank = getInputRank();
  assert(offsets.size() == static_cast<size_t>(rank) &&
         "Unexpected offsets size");
  assert(sizes.size() == static_cast<size_t>(rank) && "Unexpected sizes size");

  SmallVector<Operation *> slices;
  SmallVector<Value> tiledOperands;

  SmallVector<OpFoldResult> strides(rank, builder.getIndexAttr(1));
  Operation *inputSlice =
      getSlice(builder, loc, getInputValue(), offsets, sizes, strides);
  tiledOperands.push_back(inputSlice->getResult(0));
  slices.push_back(inputSlice);

  if (hasExplicitIndexInput()) {
    Operation *inputIndexSlice =
        getSlice(builder, loc, getInputIndex(), offsets, sizes, strides);
    tiledOperands.push_back(inputIndexSlice->getResult(0));
    slices.push_back(inputIndexSlice);
  }

  SmallVector<OpFoldResult> outputOffsets, outputSizes;
  if (failed(getResultTilePosition(builder, 0, offsets, sizes, outputOffsets,
                                   outputSizes))) {
    return emitOpError("failed to compute output tile position");
  }

  SmallVector<OpFoldResult> outputStrides(outputOffsets.size(),
                                          builder.getIndexAttr(1));
  Operation *outputValSlice =
      getSlice(builder, loc, getOutputValue(), outputOffsets, outputSizes,
               outputStrides);
  tiledOperands.push_back(outputValSlice->getResult(0));
  slices.push_back(outputValSlice);

  Operation *outputIdxSlice =
      getSlice(builder, loc, getOutputIndex(), outputOffsets, outputSizes,
               outputStrides);
  tiledOperands.push_back(outputIdxSlice->getResult(0));
  slices.push_back(outputIdxSlice);

  if (getIndexBase()) {
    tiledOperands.push_back(getIndexBase());
  }

  SmallVector<Type> resultTypes;
  if (hasPureTensorSemantics()) {
    resultTypes.push_back(outputValSlice->getResult(0).getType());
    resultTypes.push_back(outputIdxSlice->getResult(0).getType());
  }

  Operation *tiledArgmaxOp =
      mlir::clone(builder, getOperation(), resultTypes, tiledOperands);

  return TilingResult{
      {tiledArgmaxOp}, SmallVector<Value>(tiledArgmaxOp->getResults()), slices};
}

LogicalResult ArgCompareOp::getResultTilePosition(
    OpBuilder &builder, unsigned resultNumber, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, SmallVector<OpFoldResult> &resultOffsets,
    SmallVector<OpFoldResult> &resultSizes) {
  int64_t dim = getDimension();
  int64_t inputRank = getInputRank();

  resultOffsets.clear();
  resultSizes.clear();

  for (int64_t i = 0; i < inputRank; ++i) {
    if (i == dim) {
      continue;
    }
    resultOffsets.push_back(offsets[i]);
    resultSizes.push_back(sizes[i]);
  }

  return success();
}

LogicalResult ArgCompareOp::getIterationDomainTileFromResultTile(
    OpBuilder &builder, unsigned resultNumber, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes,
    SmallVectorImpl<OpFoldResult> &iterDomainOffsets,
    SmallVectorImpl<OpFoldResult> &iterDomainSizes) {
  // Result tiles are in the output coordinate space (reduction dim removed).
  // Re-insert the reduction dimension as the full range to get input-rank
  // coordinates for the iteration domain.
  int64_t reductionDim = getDimension();
  OpFoldResult zero = builder.getIndexAttr(0);
  OpFoldResult reductionDimSize =
      getDim(builder, getLoc(), getInputValue(), reductionDim);
  int64_t inputRank = offsets.size() + 1;
  iterDomainOffsets.reserve(inputRank);
  iterDomainSizes.reserve(inputRank);
  int64_t outIdx = 0;
  for (int64_t i = 0; i < inputRank; ++i) {
    if (i == reductionDim) {
      iterDomainOffsets.push_back(zero);
      iterDomainSizes.push_back(reductionDimSize);
    } else {
      iterDomainOffsets.push_back(offsets[outIdx]);
      iterDomainSizes.push_back(sizes[outIdx]);
      ++outIdx;
    }
  }
  return success();
}

FailureOr<TilingResult>
ArgCompareOp::generateResultTileValue(OpBuilder &builder, unsigned resultNumber,
                                      ArrayRef<OpFoldResult> offsets,
                                      ArrayRef<OpFoldResult> sizes) {
  SmallVector<OpFoldResult> mappedOffsets, mappedSizes;
  if (failed(getIterationDomainTileFromResultTile(
          builder, resultNumber, offsets, sizes, mappedOffsets, mappedSizes))) {
    return failure();
  }
  FailureOr<TilingResult> tilingResult =
      getTiledImplementation(builder, mappedOffsets, mappedSizes);
  if (failed(tilingResult)) {
    return failure();
  }
  return TilingResult{
      tilingResult->tiledOps,
      SmallVector<Value>{tilingResult->tiledValues[resultNumber]},
      tilingResult->generatedSlices};
}

LogicalResult ArgCompareOp::generateScalarImplementation(OpBuilder &b,
                                                         Location loc,
                                                         ValueRange ivs) {
  uint64_t reductionDim = getDimension();
  SmallVector<Value> parallelIndices;
  for (size_t i = 0, rank = ivs.size(); i < rank; ++i) {
    if (i == reductionDim) {
      continue;
    }
    parallelIndices.push_back(ivs[i]);
  }

  Value bestValueSoFar =
      memref::LoadOp::create(b, loc, getOutputValue(), parallelIndices);
  Value bestIndexSoFar =
      memref::LoadOp::create(b, loc, getOutputIndex(), parallelIndices);

  Value candidateValue = memref::LoadOp::create(b, loc, getInputValue(), ivs);

  auto &srcBlock = getRegion().front();
  IRMapping regionMap;
  regionMap.map(srcBlock.getArgument(0), candidateValue);
  regionMap.map(srcBlock.getArgument(1), bestValueSoFar);
  for (Operation &op : srcBlock.without_terminator()) {
    b.clone(op, regionMap);
  }

  Value cmpResult = regionMap.lookup(srcBlock.getTerminator()->getOperand(0));
  Value selectedValue = arith::SelectOp::create(b, loc, cmpResult,
                                                candidateValue, bestValueSoFar);

  Value candidateIndex;
  if (hasExplicitIndexInput()) {
    // Explicit-index mode: load from input.
    // The verifier ensures input and output index types match, so no cast
    // needed.
    candidateIndex = memref::LoadOp::create(b, loc, getInputIndex(), ivs);
  } else {
    // Implicit-index mode: compute from induction variable.
    Value indexOffset = ivs[reductionDim];
    if (getIndexBase()) {
      indexOffset = arith::AddIOp::create(b, loc, getIndexBase(), indexOffset);
    }
    candidateIndex = indexOffset;

    // Cast if induction variable type doesn't match output index type.
    if (candidateIndex.getType() != bestIndexSoFar.getType()) {
      candidateIndex = arith::IndexCastOp::create(
          b, loc, bestIndexSoFar.getType(), candidateIndex);
    }
  }

  Value selectedIndex = arith::SelectOp::create(b, loc, cmpResult,
                                                candidateIndex, bestIndexSoFar);
  memref::StoreOp::create(b, loc, selectedValue, getOutputValue(),
                          parallelIndices);
  memref::StoreOp::create(b, loc, selectedIndex, getOutputIndex(),
                          parallelIndices);
  return success();
}

FailureOr<SmallVector<Value>>
ArgCompareOp::generateInitialTensorForPartialReduction(
    OpBuilder &b, Location loc, ArrayRef<OpFoldResult> sizes,
    const llvm::SetVector<unsigned> &reductionDims) {
  // Get the original init tensors.
  Value valueInit = getOutputValue();
  Value indexInit = getOutputIndex();

  // Create tensors with the partial result shape.
  Type valueElTy = getElementTypeOrSelf(valueInit.getType());
  Type indexElTy = getElementTypeOrSelf(indexInit.getType());
  SmallVector<OpFoldResult> partialResultShape(sizes);
  Value emptyValueTensor =
      tensor::EmptyOp::create(b, loc, partialResultShape, valueElTy);
  Value emptyIndexTensor =
      tensor::EmptyOp::create(b, loc, partialResultShape, indexElTy);

  // Broadcast init values to partial result shape for slicing inside
  // scf.forall. Each tile in the parallel loop will extract slices from
  // these broadcasted tensors to get initialized values for the ArgCompareOp.
  // Example: tensor<64xf32> -> tensor<64x32xf32> for 32 reduction tiles along
  // dim 1.
  auto broadcastDims = llvm::to_vector_of<int64_t>(reductionDims);
  Operation *valueBroadcastOp = linalg::BroadcastOp::create(
      b, loc, valueInit, emptyValueTensor, broadcastDims);
  Operation *indexBroadcastOp = linalg::BroadcastOp::create(
      b, loc, indexInit, emptyIndexTensor, broadcastDims);

  return SmallVector<Value>{valueBroadcastOp->getResult(0),
                            indexBroadcastOp->getResult(0)};
}

// Helper function to create a tiled operation for the OuterReduction strategy.
//
// OuterReduction uses scf.for with iter_args to accumulate results across
// iterations. Within each iteration, we perform element-wise comparisons (not
// reductions), so we always use linalg.generic. The only difference between
// explicit and implicit-index modes is whether indices come from the input
// tensor or are computed from the tile offset.
static std::pair<Operation *, SmallVector<Operation *>>
createTiledOpForOuterReduction(ArgCompareOp op, OpBuilder &b, Location loc,
                               ValueRange init, ArrayRef<OpFoldResult> offsets,
                               SmallVector<Value> &tiledOperands,
                               int64_t reductionDim, int64_t rank,
                               bool isExplicitIndexMode) {
  SmallVector<Operation *> slices;
  Value initValResult = init[0];
  Value initIdxResult = init[1];

  SmallVector<Type, 2> resultTypes;
  if (op.hasPureTensorSemantics()) {
    resultTypes = {initValResult.getType(), initIdxResult.getType()};
  }

  // For OuterReduction, we always use linalg.generic to perform element-wise
  // comparison, regardless of index mode. The difference is just how we obtain
  // the indices: from input (explicit) or computed (implicit).
  AffineMap identityMap = b.getMultiDimIdentityMap(rank);

  // Base setup: input_value, accum_value, accum_index.
  SmallVector<AffineMap> indexingMaps = {identityMap, identityMap, identityMap};
  SmallVector<Value> inputs = {tiledOperands[0]};

  if (isExplicitIndexMode) {
    // Explicit-index mode: Also read indices from the sliced input_index
    // tensor.
    indexingMaps.push_back(identityMap);
    inputs.push_back(tiledOperands[1]);
  }

  SmallVector<utils::IteratorType> iterators(rank,
                                             utils::IteratorType::parallel);

  auto genericOp = linalg::GenericOp::create(
      b, loc, TypeRange{initValResult.getType(), initIdxResult.getType()},
      inputs, ValueRange{initValResult, initIdxResult}, indexingMaps, iterators,
      [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange args) {
        Value newVal, accVal, accIdx, newIdx;

        newVal = args[0];
        if (isExplicitIndexMode) {
          // Args: [input_value, input_index, accum_value, accum_index].
          newIdx = args[1];
          accVal = args[2];
          accIdx = args[3];
        } else {
          // Args: [input_value, accum_value, accum_index].
          accVal = args[1];
          accIdx = args[2];

          // Compute the global index for this element.
          Value localIdx =
              linalg::IndexOp::create(nestedBuilder, nestedLoc, reductionDim);
          Value offsetValue = getValueOrCreateConstantIndexOp(
              nestedBuilder, nestedLoc, offsets[reductionDim]);
          Value globalIdx = arith::AddIOp::create(nestedBuilder, nestedLoc,
                                                  offsetValue, localIdx);

          if (Value indexBase = op.getIndexBase()) {
            globalIdx = arith::AddIOp::create(nestedBuilder, nestedLoc,
                                              indexBase, globalIdx);
          }

          // Cast to match the output index tensor's element type.
          Type idxElemType = accIdx.getType();
          newIdx = globalIdx;
          if (!isa<IndexType>(idxElemType)) {
            newIdx = arith::IndexCastOp::create(nestedBuilder, nestedLoc,
                                                idxElemType, globalIdx);
          }
        }

        // Clone the user's comparison logic to decide whether to update.
        Block &originalBlock = op.getRegion().front();
        IRMapping mapper;
        mapper.map(originalBlock.getArgument(0), newVal);
        mapper.map(originalBlock.getArgument(1), accVal);
        for (Operation &innerOp : originalBlock) {
          if (innerOp.hasTrait<OpTrait::IsTerminator>()) {
            break;
          }
          nestedBuilder.clone(innerOp, mapper);
        }

        Value predicate =
            mapper.lookup(originalBlock.getTerminator()->getOperand(0));
        Value selectedVal = arith::SelectOp::create(nestedBuilder, nestedLoc,
                                                    predicate, newVal, accVal);
        Value selectedIdx = arith::SelectOp::create(nestedBuilder, nestedLoc,
                                                    predicate, newIdx, accIdx);
        linalg::YieldOp::create(nestedBuilder, nestedLoc,
                                ValueRange{selectedVal, selectedIdx});
      });

  return {genericOp, slices};
}

// Helper function to create a tiled operation for the OuterParallel strategy.
//
// OuterParallel uses scf.forall with parallel_insert_slice where each parallel
// iteration handles an independent chunk of the reduction dimension. Each chunk
// performs a complete reduction and writes its result to a specific slot in the
// output. Unlike OuterReduction, we can use ArgCompareOp for both index modes
// because each chunk has a fixed offset that can be baked into index_base.
static std::pair<Operation *, SmallVector<Operation *>>
createTiledOpForOuterParallel(ArgCompareOp op, OpBuilder &b, Location loc,
                              ValueRange init, ArrayRef<OpFoldResult> offsets,
                              ArrayRef<OpFoldResult> sizes,
                              ArrayRef<OpFoldResult> splitReductionIvs,
                              SmallVector<Value> &tiledOperands,
                              int64_t reductionDim, int64_t rank,
                              bool isExplicitIndexMode) {
  SmallVector<Operation *> slices;

  // Extract a slice of the output for this parallel chunk.
  SmallVector<OpFoldResult> initOffsets, initSizes, initStrides;
  for (int64_t i = 0; i < rank; ++i) {
    if (i == reductionDim) {
      initOffsets.push_back(splitReductionIvs[0]);
      initSizes.push_back(b.getIndexAttr(1));
    } else {
      initOffsets.push_back(offsets[i]);
      initSizes.push_back(sizes[i]);
    }
    initStrides.push_back(b.getIndexAttr(1));
  }

  SmallVector<int64_t> resultValShape, resultIdxShape;
  for (int64_t i = 0; i < rank; ++i) {
    if (i == reductionDim) {
      continue;
    }
    int64_t size =
        getConstantIntValue(initSizes[i]).value_or(ShapedType::kDynamic);
    resultValShape.push_back(size);
    resultIdxShape.push_back(size);
  }

  auto sliceValResultType = RankedTensorType::get(
      resultValShape, op.getOutputValueType().getElementType());
  auto sliceIdxResultType =
      RankedTensorType::get(resultIdxShape, op.getOutputIndexElementType());

  auto initValSlice = tensor::ExtractSliceOp::create(
      b, loc, sliceValResultType, init[0], initOffsets, initSizes, initStrides);
  slices.push_back(initValSlice);

  auto initIdxSlice = tensor::ExtractSliceOp::create(
      b, loc, sliceIdxResultType, init[1], initOffsets, initSizes, initStrides);
  slices.push_back(initIdxSlice);

  Value initValResult = initValSlice.getResult();
  Value initIdxResult = initIdxSlice.getResult();

  tiledOperands.push_back(initValResult);
  tiledOperands.push_back(initIdxResult);

  SmallVector<Type, 2> resultTypes;
  if (op.hasPureTensorSemantics()) {
    resultTypes = {initValResult.getType(), initIdxResult.getType()};
  }

  ArgCompareOp tiledArgmaxOp;
  if (isExplicitIndexMode) {
    // Indices are provided as input, no index_base needed.
    tiledArgmaxOp = ArgCompareOp::create(
        b, loc, resultTypes,
        /*input_value=*/tiledOperands[0], /*input_index=*/tiledOperands[1],
        /*output_value=*/tiledOperands[2], /*output_index=*/tiledOperands[3],
        /*index_base=*/Value(), /*dimension=*/reductionDim);
  } else {
    // Compute index_base for this chunk so that generated indices correspond
    // to positions in the original tensor. We use tileIndex * tileSize to get
    // the starting offset of this chunk in the reduction dimension.
    Value tileIndex =
        getValueOrCreateConstantIndexOp(b, loc, splitReductionIvs[0]);
    Value tileSize =
        getValueOrCreateConstantIndexOp(b, loc, sizes[reductionDim]);
    Value newIndexBase = arith::MulIOp::create(b, loc, tileIndex, tileSize);

    if (Value globalIndexBase = op.getIndexBase()) {
      newIndexBase =
          arith::AddIOp::create(b, loc, globalIndexBase, newIndexBase);
    }

    tiledArgmaxOp = ArgCompareOp::create(
        b, loc, resultTypes,
        /*input_value=*/tiledOperands[0], /*input_index=*/Value(),
        /*output_value=*/tiledOperands[1], /*output_index=*/tiledOperands[2],
        /*index_base=*/newIndexBase, /*dimension=*/reductionDim);
  }

  return {tiledArgmaxOp, slices};
}

FailureOr<TilingResult> ArgCompareOp::tileToPartialReduction(
    OpBuilder &b, Location loc, ReductionTilingStrategy strategy,
    ValueRange init, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes,
    const llvm::SetVector<unsigned> &reductionDims,
    ArrayRef<OpFoldResult> splitReductionIvs) {
  assert(
      (strategy == ReductionTilingStrategy::PartialReductionOuterParallel ||
       strategy == ReductionTilingStrategy::PartialReductionOuterReduction) &&
      "Unexpected partial reduction tiling strategy");

  OpBuilder::InsertionGuard guard(b);

  int64_t rank = getInputRank();
  int64_t reductionDim = getDimension();

  if (reductionDims.size() != 1 || reductionDims[0] != reductionDim) {
    return failure();
  }

  assert(offsets.size() == static_cast<size_t>(rank) &&
         "Unexpected offsets size");
  assert(sizes.size() == static_cast<size_t>(rank) && "Unexpected sizes size");

  SmallVector<Operation *> slices;
  SmallVector<Value> tiledOperands;

  SmallVector<OpFoldResult> strides(rank, b.getIndexAttr(1));
  Operation *inputSlice =
      getSlice(b, loc, getInputValue(), offsets, sizes, strides);
  tiledOperands.push_back(inputSlice->getResult(0));
  slices.push_back(inputSlice);

  bool isExplicitIndexMode = hasExplicitIndexInput();
  if (isExplicitIndexMode) {
    Operation *inputIndexSlice =
        getSlice(b, loc, getInputIndex(), offsets, sizes, strides);
    tiledOperands.push_back(inputIndexSlice->getResult(0));
    slices.push_back(inputIndexSlice);
  }

  // Dispatch to appropriate helper based on strategy.
  // The two strategies handle partial results differently:
  // - OuterReduction: Uses scf.for iter_args to accumulate across iterations.
  // - OuterParallel: Uses scf.forall with parallel_insert_slice for parallel
  // chunks.
  Operation *tiledOp;
  SmallVector<Operation *> strategySlices;
  if (strategy == ReductionTilingStrategy::PartialReductionOuterReduction) {
    auto [op, extraSlices] = createTiledOpForOuterReduction(
        *this, b, loc, init, offsets, tiledOperands, reductionDim, rank,
        isExplicitIndexMode);
    tiledOp = op;
    strategySlices = std::move(extraSlices);
  } else {
    auto [op, extraSlices] = createTiledOpForOuterParallel(
        *this, b, loc, init, offsets, sizes, splitReductionIvs, tiledOperands,
        reductionDim, rank, isExplicitIndexMode);
    tiledOp = op;
    strategySlices = std::move(extraSlices);
  }

  slices.insert(slices.end(), strategySlices.begin(), strategySlices.end());

  // Clone the user's comparison region into the tiled operation. This is only
  // needed for OuterParallel which creates ArgCompareOp; OuterReduction uses
  // linalg.generic which already has its region fully defined.
  if (strategy == ReductionTilingStrategy::PartialReductionOuterParallel) {
    auto argCompareOp = cast<ArgCompareOp>(tiledOp);
    Region &targetRegion = argCompareOp.getRegion();
    Region &sourceRegion = getRegion();
    IRMapping mapper;
    sourceRegion.cloneInto(&targetRegion, mapper);
  }

  return TilingResult{
      {tiledOp}, SmallVector<Value>(tiledOp->getResults()), slices};
}

FailureOr<MergeResult>
ArgCompareOp::mergeReductions(OpBuilder &b, Location loc,
                              ValueRange partialReduce,
                              const llvm::SetVector<unsigned> &reductionDims) {
  int64_t reductionDim = reductionDims.front();

  // Create arg_compare in explicit-index mode to merge the partial results.
  // Explicit-index mode: 2 inputs (value + index), no index_base.
  ArgCompareOp mergeOp = ArgCompareOp::create(
      b, loc, hasPureTensorSemantics() ? TypeRange(getDpsInits()) : TypeRange{},
      /*input_value=*/partialReduce[0], /*input_index=*/partialReduce[1],
      /*output_value=*/getDpsInits()[0], /*output_index=*/getDpsInits()[1],
      /*index_base=*/Value(),
      /*dimension=*/reductionDim);

  // Clone the comparator region from the original operation.
  Region &targetRegion = mergeOp->getRegion(0);
  Region &sourceRegion = getRegion();
  IRMapping mapper;
  sourceRegion.cloneInto(&targetRegion, mapper);

  return MergeResult{{mergeOp.getOperation()},
                     {mergeOp.getResult(0), mergeOp.getResult(1)}};
}

LogicalResult ArgCompareOp::getPartialResultTilePosition(
    OpBuilder &b, unsigned resultNumber, ReductionTilingStrategy tilingStrategy,
    ArrayRef<OpFoldResult> offsets, ArrayRef<OpFoldResult> sizes,
    const llvm::SetVector<unsigned> &reductionDims,
    ArrayRef<OpFoldResult> splitReductionIvs,
    SmallVector<OpFoldResult> &resultOffsets,
    SmallVector<OpFoldResult> &resultSizes) {
  int64_t reductionDim = getDimension();
  resultOffsets.clear();
  resultSizes.clear();

  bool isOuterReduction =
      tilingStrategy == ReductionTilingStrategy::PartialReductionOuterReduction;

  for (auto [i, size, offset] : llvm::enumerate(sizes, offsets)) {
    if (i == reductionDim) {
      // OuterReduction accumulates in place, always writing to offset 0.
      // OuterParallel writes each chunk to a separate slot indexed by chunkIdx.
      resultOffsets.push_back(isOuterReduction ? b.getIndexAttr(0)
                                               : splitReductionIvs[0]);
      resultSizes.push_back(isOuterReduction ? size : b.getIndexAttr(1));
    } else {
      resultOffsets.push_back(offset);
      resultSizes.push_back(size);
    }
  }

  return success();
}

//===----------------------------------------------------------------------===//
// ExpReductionOp
//===----------------------------------------------------------------------===//

SmallVector<utils::IteratorType> ExpReductionOp::getLoopIteratorTypes() {
  return llvm::to_vector(getIteratorTypes()
                             .getAsValueRange<IREE::LinalgExt::IteratorTypeAttr,
                                              utils::IteratorType>());
}

SmallVector<Range> ExpReductionOp::getIterationDomain(OpBuilder &b) {
  Location loc = getLoc();
  OpFoldResult zero = b.getIndexAttr(0);
  OpFoldResult one = b.getIndexAttr(1);

  SmallVector<OpFoldResult> allShapesSizes =
      createFlatListOfOperandDims(b, loc, getOperation());
  AffineMap map = getShapesToLoopsMap();
  return llvm::map_to_vector(map.getResults(), [&](AffineExpr loopExpr) {
    OpFoldResult ofr =
        affine::makeComposedFoldedAffineApply(b, loc, loopExpr, allShapesSizes);
    return Range{zero, ofr, one};
  });
}

FailureOr<TilingResult>
ExpReductionOp::getTiledImplementation(OpBuilder &b,
                                       ArrayRef<OpFoldResult> offsets,
                                       ArrayRef<OpFoldResult> sizes) {
  Location loc = getLoc();
  auto indexingMapOp = cast<IndexingMapOpInterface>(getOperation());
  SmallVector<Value> tiledOperands;
  SmallVector<Operation *> generatedSlices;
  for (OpOperand &opOperand : getOperation()->getOpOperands()) {
    AffineMap map = indexingMapOp.getMatchingIndexingMap(&opOperand);
    SmallVector<Range> slice = getPermutedRange(map, offsets, sizes);
    Operation *sliceOp = getSlice(b, loc, opOperand.get(), slice);
    tiledOperands.emplace_back(sliceOp->getResult(0));
    generatedSlices.push_back(sliceOp);
  }

  SmallVector<Type, 4> resultTensorTypes;
  if (getNumResults()) {
    resultTensorTypes = llvm::map_to_vector<4>(
        getDpsInitsMutable(), [&generatedSlices](OpOperand &opOperand) {
          return generatedSlices[opOperand.getOperandNumber()]
              ->getResultTypes()[0];
        });
  }

  Operation *tiledOp = mlir::clone(b, *this, resultTensorTypes, tiledOperands);
  return TilingResult{
      {tiledOp}, SmallVector<Value>(tiledOp->getResults()), generatedSlices};
}

LogicalResult ExpReductionOp::getResultTilePosition(
    OpBuilder &b, unsigned resultNumber, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, SmallVector<OpFoldResult> &resultOffsets,
    SmallVector<OpFoldResult> &resultSizes) {
  auto indexingMapOp = cast<IndexingMapOpInterface>(getOperation());
  OpOperand *outOperand = getDpsInitOperand(resultNumber);
  AffineMap indexingMap = indexingMapOp.getMatchingIndexingMap(outOperand);
  SmallVector<Range> range = getPermutedRange(indexingMap, offsets, sizes);
  resultOffsets.resize(range.size());
  resultSizes.resize(range.size());
  for (auto [index, r] : llvm::enumerate(range)) {
    resultOffsets[index] = r.offset;
    resultSizes[index] = r.size;
  }
  return success();
}

FailureOr<TilingResult>
ExpReductionOp::generateResultTileValue(OpBuilder &b, unsigned resultNumber,
                                        ArrayRef<OpFoldResult> offsets,
                                        ArrayRef<OpFoldResult> sizes) {
  SmallVector<OpFoldResult> mappedOffsets, mappedSizes;
  if (failed(getIterationDomainTileFromResultTile(
          b, resultNumber, offsets, sizes, mappedOffsets, mappedSizes))) {
    return failure();
  }
  FailureOr<TilingResult> tilingResult =
      getTiledImplementation(b, mappedOffsets, mappedSizes);
  if (failed(tilingResult)) {
    return failure();
  }
  if (tilingResult->tiledOps.size() != 1) {
    return emitOpError("failed to generate tiled implementation");
  }
  return TilingResult{
      tilingResult->tiledOps,
      SmallVector<Value>{tilingResult->tiledValues[resultNumber]},
      tilingResult->generatedSlices};
}

//===----------------------------------------------------------------------===//
// Im2colOp
//===----------------------------------------------------------------------===//

SmallVector<Range> Im2colOp::getIterationDomain(OpBuilder &builder) {
  Location loc = getLoc();
  OpFoldResult zero = builder.getIndexAttr(0);
  OpFoldResult one = builder.getIndexAttr(1);
  Value dest = getOutput();
  SmallVector<Range> loopBounds(getOutputRank());
  for (int dim = 0; dim < getOutputRank(); ++dim) {
    loopBounds[dim].offset = zero;
    loopBounds[dim].size = getDim(builder, loc, dest, dim);
    loopBounds[dim].stride = one;
  }
  return loopBounds;
}

SmallVector<utils::IteratorType> Im2colOp::getLoopIteratorTypes() {
  SmallVector<utils::IteratorType> iteratorTypes(getOutputRank(),
                                                 utils::IteratorType::parallel);
  return iteratorTypes;
}

FailureOr<TilingResult>
Im2colOp::getTiledImplementation(OpBuilder &builder,
                                 ArrayRef<OpFoldResult> offsets,
                                 ArrayRef<OpFoldResult> sizes) {
  Location loc = getLoc();
  OpFoldResult one = builder.getIndexAttr(1);

  Value inputValue = getInput();

  SmallVector<OpFoldResult> outputStrides(getOutputRank(), one);
  Operation *outputSlice =
      getSlice(builder, loc, getOutput(), offsets, sizes, outputStrides);

  // Adjust offsets by adding the tiling offsets. The offsets are in canonical
  // [Batch, M, K] order, and output_perm[actual] = canonical, so we use
  // output_perm directly to map actual tensor dims to canonical positions.
  SmallVector<OpFoldResult> newOffsets(getMixedOffsets());
  ArrayRef<int64_t> outPerm = getOutputPerm();
  for (int64_t actual = 0; actual < getOutputRank(); ++actual) {
    int64_t canonical = outPerm[actual];
    newOffsets[canonical] =
        addOfrs(builder, loc, offsets[actual], newOffsets[canonical]);
  }

  // Compute tile-local output padding amounts.
  // new_pad_low[d] = max(output_pad_low[d] - tileOff[d], 0)
  // new_pad_high[d] = max(tileOff[d] + tileSize[d] - (origDim[d] -
  //                       output_pad_high[d]), 0)
  SmallVector<OpFoldResult> existingOutPadLow = getMixedOutputPadLow();
  SmallVector<OpFoldResult> existingOutPadHigh = getMixedOutputPadHigh();
  SmallVector<OpFoldResult> newOutPadLow, newOutPadHigh;
  if (!existingOutPadLow.empty()) {
    MLIRContext *ctx = builder.getContext();
    AffineExpr d0 = getAffineDimExpr(0, ctx);
    AffineExpr d1 = getAffineDimExpr(1, ctx);
    AffineMap posMaxMap =
        AffineMap::get(1, 0, {d0, getAffineConstantExpr(0, ctx)}, ctx);
    AffineMap subMap = AffineMap::get(2, 0, d0 - d1, ctx);

    SmallVector<OpFoldResult> origDims =
        tensor::getMixedSizes(builder, loc, getOutput());
    OpFoldResult zero = builder.getIndexAttr(0);

    for (int64_t d = 0; d < getOutputRank(); ++d) {
      // new_pad_low = max(pad_low - tileOff, 0)
      // When pad_low is statically zero, skip the computation and keep zero.
      OpFoldResult tilePadLow;
      if (isZeroInteger(existingOutPadLow[d])) {
        tilePadLow = zero;
      } else {
        OpFoldResult lowDiff = affine::makeComposedFoldedAffineApply(
            builder, loc, subMap, {existingOutPadLow[d], offsets[d]});
        tilePadLow = affine::makeComposedFoldedAffineMax(builder, loc,
                                                         posMaxMap, {lowDiff});
      }

      // validEnd = origDim - pad_high; overrun = tileOff + tileSize - validEnd
      // When pad_high is statically zero, skip the computation and keep zero.
      OpFoldResult tilePadHigh;
      if (isZeroInteger(existingOutPadHigh[d])) {
        tilePadHigh = zero;
      } else {
        OpFoldResult validEnd = affine::makeComposedFoldedAffineApply(
            builder, loc, subMap, {origDims[d], existingOutPadHigh[d]});
        OpFoldResult tileEnd = affine::makeComposedFoldedAffineApply(
            builder, loc, AffineMap::get(2, 0, d0 + d1, ctx),
            {offsets[d], sizes[d]});
        OpFoldResult highDiff = affine::makeComposedFoldedAffineApply(
            builder, loc, subMap, {tileEnd, validEnd});
        tilePadHigh = affine::makeComposedFoldedAffineMax(
            builder, loc, posMaxMap, {highDiff});
      }

      newOutPadLow.push_back(tilePadLow);
      newOutPadHigh.push_back(tilePadHigh);
    }
  }

  // Create the tiled op. The input is not sliced (batch offset is in the op).
  SmallVector<Value> operands = {inputValue, outputSlice->getResult(0)};
  // Copy all metadata operands from the untiled operation.
  operands.append(getOperation()->getOperands().begin() + 2,
                  getOperation()->getOperands().end());
  Im2colOp tiledOp =
      mlir::clone(builder, *this, outputSlice->getResultTypes(), operands);
  // Set the new offsets, since they have changed with tiling.
  // output_sizes remain unchanged by tiling (key design property).
  tiledOp.setMixedOffsets(newOffsets);
  // Set tile-local output padding amounts.
  if (!newOutPadLow.empty()) {
    tiledOp.setMixedOutputPadLow(newOutPadLow);
    tiledOp.setMixedOutputPadHigh(newOutPadHigh);
  }

  return TilingResult{
      {tiledOp}, SmallVector<Value>(tiledOp->getResults()), {outputSlice}};
}

FailureOr<TilingResult>
Im2colOp::generateResultTileValue(OpBuilder &builder, unsigned resultNumber,
                                  ArrayRef<OpFoldResult> offsets,
                                  ArrayRef<OpFoldResult> sizes) {
  return getTiledImplementation(builder, offsets, sizes);
}

LogicalResult Im2colOp::getResultTilePosition(
    OpBuilder &builder, unsigned resultNumber, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, SmallVector<OpFoldResult> &resultOffsets,
    SmallVector<OpFoldResult> &resultSizes) {
  resultOffsets = SmallVector<OpFoldResult>(offsets);
  resultSizes = SmallVector<OpFoldResult>(sizes);
  return success();
}

//===----------------------------------------------------------------------===//
// WinogradInputTransformOp
//===----------------------------------------------------------------------===//

SmallVector<Range>
WinogradInputTransformOp::getIterationDomain(OpBuilder &builder) {
  Location loc = getLoc();
  OpFoldResult zero = builder.getIndexAttr(0);
  OpFoldResult one = builder.getIndexAttr(1);
  Value dest = getOutput();
  SmallVector<Range> loopBounds(getIterationDomainRank());
  int count = 0;
  for (auto dim :
       llvm::seq<int64_t>(getImageDimensions().size(), getOutputRank())) {
    loopBounds[count].offset = zero;
    loopBounds[count].size = getDim(builder, loc, dest, dim);
    loopBounds[count].stride = one;
    count++;
  }
  return loopBounds;
}

SmallVector<utils::IteratorType>
WinogradInputTransformOp::getLoopIteratorTypes() {
  SmallVector<utils::IteratorType> iteratorTypes(getIterationDomainRank(),
                                                 utils::IteratorType::parallel);
  return iteratorTypes;
}

FailureOr<TilingResult>
WinogradInputTransformOp::getTiledImplementation(OpBuilder &builder,
                                                 ArrayRef<OpFoldResult> offsets,
                                                 ArrayRef<OpFoldResult> sizes) {
  Location loc = getLoc();
  auto one = builder.getIndexAttr(1);
  auto zero = builder.getIndexAttr(0);
  const int cDim = getChannelDim();

  assert(offsets.size() == 4);
  SmallVector<OpFoldResult> inputOffsets(getInputRank(), zero);
  SmallVector<OpFoldResult> outputOffsets(getOutputRank(), zero);
  const auto hDim = getImageDimensions()[0];
  const auto wDim = getImageDimensions()[1];
  outputOffsets[2] = inputOffsets[0] = offsets[0];
  outputOffsets[3] = offsets[1];
  outputOffsets[4] = offsets[2];
  outputOffsets[5] = inputOffsets[cDim] = offsets[3];

  SmallVector<OpFoldResult> inputStrides(getInputRank(), one);
  SmallVector<OpFoldResult> outputStrides(getOutputRank(), one);
  ReifiedRankedShapedTypeDims reifiedResultShapes, reifiedInputShapes;
  if (failed(reifyResultShapes(builder, reifiedResultShapes))) {
    return failure();
  }
  SmallVector<OpFoldResult> outputSizes = reifiedResultShapes[0];
  if (failed(getStaticOrReifiedInputDims(builder, loc, getInput(),
                                         reifiedInputShapes))) {
    return failure();
  }
  SmallVector<OpFoldResult> inputSizes = reifiedInputShapes[0];

  assert(sizes.size() == 4);
  outputSizes[2] = inputSizes[0] = sizes[0];
  outputSizes[3] = sizes[1];
  outputSizes[4] = sizes[2];
  outputSizes[5] = inputSizes[cDim] = sizes[3];

  auto hSizeAndOffset = getScaledSizeAndOffset(
      builder, loc, sizes[1], offsets[1], inputSizes[hDim], getOutputTileSize(),
      getInputTileSize());
  auto wSizeAndOffset = getScaledSizeAndOffset(
      builder, loc, sizes[2], offsets[2], inputSizes[wDim], getOutputTileSize(),
      getInputTileSize());

  inputSizes[hDim] = hSizeAndOffset.first;
  inputSizes[wDim] = wSizeAndOffset.first;
  inputOffsets[hDim] = hSizeAndOffset.second;
  inputOffsets[wDim] = wSizeAndOffset.second;

  SmallVector<Value> tiledOperands;
  SmallVector<Operation *> slices;
  // Input
  {
    Operation *inputSlice = getSlice(builder, loc, getInput(), inputOffsets,
                                     inputSizes, inputStrides);
    tiledOperands.emplace_back(inputSlice->getResult(0));
    slices.push_back(inputSlice);
  }

  // Output
  {
    Operation *outputSlice = getSlice(builder, loc, getOutput(), outputOffsets,
                                      outputSizes, outputStrides);
    tiledOperands.emplace_back(outputSlice->getResult(0));
    slices.push_back(outputSlice);
  }

  SmallVector<Type, 4> resultTypes;
  if (hasPureTensorSemantics()) {
    resultTypes.push_back(tiledOperands[1].getType());
  }

  Operation *tiledOp =
      mlir::clone(builder, getOperation(), resultTypes, tiledOperands);

  return TilingResult{
      {tiledOp}, SmallVector<Value>(tiledOp->getResults()), slices};
}

LogicalResult WinogradInputTransformOp::getResultTilePosition(
    OpBuilder &builder, unsigned resultNumber, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, SmallVector<OpFoldResult> &resultOffsets,
    SmallVector<OpFoldResult> &resultSizes) {
  if (resultNumber == 0) {
    auto resultShape = cast<ShapedType>(getOutput().getType()).getShape();
    resultSizes = getAsOpFoldResult(builder.getIndexArrayAttr(resultShape));
    resultOffsets =
        SmallVector<OpFoldResult>(getOutputRank(), builder.getIndexAttr(0));
    resultOffsets[2] = offsets[0];
    resultOffsets[3] = offsets[1];
    resultOffsets[4] = offsets[2];
    resultOffsets[5] = offsets[3];
    resultSizes[2] = sizes[0];
    resultSizes[3] = sizes[1];
    resultSizes[4] = sizes[2];
    resultSizes[5] = sizes[3];
    return success();
  }
  return failure();
}

//===----------------------------------------------------------------------===//
// WinogradFilterTransformOp
//===----------------------------------------------------------------------===//

SmallVector<Range>
WinogradFilterTransformOp::getIterationDomain(OpBuilder &builder) {
  Location loc = getLoc();
  OpFoldResult zero = builder.getIndexAttr(0);
  OpFoldResult one = builder.getIndexAttr(1);
  Value source = getOutput();
  int64_t numKernelDims = getKernelDimensions().size();
  auto outRank = getOutputRank();
  SmallVector<Range> loopBounds(outRank - numKernelDims);
  for (auto dim : llvm::seq<int64_t>(numKernelDims, outRank)) {
    int64_t loopDim = dim - numKernelDims;
    loopBounds[loopDim].offset = zero;
    loopBounds[loopDim].size = getDim(builder, loc, source, dim);
    loopBounds[loopDim].stride = one;
  }
  return loopBounds;
}

SmallVector<utils::IteratorType>
WinogradFilterTransformOp::getLoopIteratorTypes() {
  SmallVector<utils::IteratorType> iteratorTypes(getIterationDomainRank(),
                                                 utils::IteratorType::parallel);
  return iteratorTypes;
}

FailureOr<TilingResult> WinogradFilterTransformOp::getTiledImplementation(
    OpBuilder &builder, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes) {
  Location loc = getLoc();
  OpFoldResult one = builder.getIndexAttr(1);
  OpFoldResult zero = builder.getIndexAttr(0);
  const int cDim = getChannelDim();
  const int fDim = getFilterDim();

  assert(offsets.size() == 2);
  SmallVector<OpFoldResult> inputOffsets(getInputRank(), zero);
  SmallVector<OpFoldResult> outputOffsets(getOutputRank(), zero);
  outputOffsets[2] = inputOffsets[cDim] = offsets[0];
  outputOffsets[3] = inputOffsets[fDim] = offsets[1];

  SmallVector<OpFoldResult> inputStrides(getInputRank(), one);
  SmallVector<OpFoldResult> outputStrides(getOutputRank(), one);

  assert(sizes.size() == 2);
  ArrayRef<int64_t> inputShape = getInputType().getShape();
  ArrayRef<int64_t> outputShape = getOutputType().getShape();
  SmallVector<OpFoldResult> inputSizes =
      getAsOpFoldResult(builder.getIndexArrayAttr(inputShape));
  SmallVector<OpFoldResult> outputSizes =
      getAsOpFoldResult(builder.getIndexArrayAttr(outputShape));
  outputSizes[2] = inputSizes[cDim] = sizes[0];
  outputSizes[3] = inputSizes[fDim] = sizes[1];

  SmallVector<Value> tiledOperands;
  SmallVector<Operation *> slices;

  // Input
  {
    Operation *inputSlice = getSlice(builder, loc, getInput(), inputOffsets,
                                     inputSizes, inputStrides);
    tiledOperands.emplace_back(inputSlice->getResult(0));
    slices.push_back(inputSlice);
  }

  // Output
  {
    Operation *outputSlice = getSlice(builder, loc, getOutput(), outputOffsets,
                                      outputSizes, outputStrides);
    tiledOperands.emplace_back(outputSlice->getResult(0));
    slices.push_back(outputSlice);
  }

  SmallVector<Type> resultTypes;
  if (hasPureTensorSemantics()) {
    resultTypes.push_back(tiledOperands[1].getType());
  }

  Operation *tiledOp =
      mlir::clone(builder, getOperation(), resultTypes, tiledOperands);

  return TilingResult{
      {tiledOp}, SmallVector<Value>(tiledOp->getResults()), slices};
}

LogicalResult WinogradFilterTransformOp::getResultTilePosition(
    OpBuilder &builder, unsigned resultNumber, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, SmallVector<OpFoldResult> &resultOffsets,
    SmallVector<OpFoldResult> &resultSizes) {
  if (resultNumber != 0) {
    return failure();
  }
  ArrayRef<int64_t> resultShape = getOutputType().getShape();
  resultSizes = getAsOpFoldResult(builder.getIndexArrayAttr(resultShape));
  resultOffsets =
      SmallVector<OpFoldResult>(getOutputRank(), builder.getIndexAttr(0));
  resultOffsets[2] = offsets[0];
  resultOffsets[3] = offsets[1];
  resultSizes[2] = sizes[0];
  resultSizes[3] = sizes[1];
  return success();
}

//===----------------------------------------------------------------------===//
// WinogradOutputTransformOp
//===----------------------------------------------------------------------===//

SmallVector<Range>
WinogradOutputTransformOp::getIterationDomain(OpBuilder &builder) {
  Location loc = getLoc();
  OpFoldResult zero = builder.getIndexAttr(0);
  OpFoldResult one = builder.getIndexAttr(1);
  Value source = getInput();
  SmallVector<Range> loopBounds(getIterationDomainRank());
  int count = 0;
  for (auto dim :
       llvm::seq<int64_t>(getImageDimensions().size(), getInputRank())) {
    loopBounds[count].offset = zero;
    loopBounds[count].size = getDim(builder, loc, source, dim);
    loopBounds[count].stride = one;
    count++;
  }
  return loopBounds;
}

SmallVector<utils::IteratorType>
WinogradOutputTransformOp::getLoopIteratorTypes() {
  SmallVector<utils::IteratorType> iteratorTypes(getIterationDomainRank(),
                                                 utils::IteratorType::parallel);
  return iteratorTypes;
}

FailureOr<TilingResult> WinogradOutputTransformOp::getTiledImplementation(
    OpBuilder &builder, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes) {
  Location loc = getLoc();
  auto one = builder.getIndexAttr(1);
  auto zero = builder.getIndexAttr(0);
  const int cDim = getChannelDim();

  assert(offsets.size() == 4);
  const auto hDim = getImageDimensions()[0];
  const auto wDim = getImageDimensions()[1];
  SmallVector<OpFoldResult> inputOffsets(getInputRank(), zero);
  SmallVector<OpFoldResult> outputOffsets(getOutputRank(), zero);

  inputOffsets[2] = outputOffsets[0] = offsets[0];
  inputOffsets[3] = offsets[1];
  inputOffsets[4] = offsets[2];
  inputOffsets[5] = outputOffsets[cDim] = offsets[3];

  SmallVector<OpFoldResult> inputStrides(getInputRank(), one);
  SmallVector<OpFoldResult> outputStrides(getOutputRank(), one);

  ReifiedRankedShapedTypeDims reifiedResultShapes, reifiedInputShapes;
  if (failed(reifyResultShapes(builder, reifiedResultShapes))) {
    return failure();
  }
  SmallVector<OpFoldResult> outputSizes = reifiedResultShapes[0];
  if (failed(getStaticOrReifiedInputDims(builder, loc, getInput(),
                                         reifiedInputShapes))) {
    return failure();
  }
  SmallVector<OpFoldResult> inputSizes = reifiedInputShapes[0];

  inputSizes[2] = outputSizes[0] = sizes[0];
  inputSizes[5] = outputSizes[cDim] = sizes[3];

  assert(sizes.size() == 4);
  inputSizes[2] = outputSizes[0] = sizes[0];
  inputSizes[3] = sizes[1];
  inputSizes[4] = sizes[2];
  inputSizes[5] = outputSizes[cDim] = sizes[3];

  auto hSizeAndOffset = getScaledSizeAndOffset(
      builder, loc, sizes[1], offsets[1], outputSizes[hDim],
      getOutputTileSize(), getOutputTileSize());
  auto wSizeAndOffset = getScaledSizeAndOffset(
      builder, loc, sizes[2], offsets[2], outputSizes[wDim],
      getOutputTileSize(), getOutputTileSize());

  outputSizes[hDim] = hSizeAndOffset.first;
  outputSizes[wDim] = wSizeAndOffset.first;
  outputOffsets[hDim] = hSizeAndOffset.second;
  outputOffsets[wDim] = wSizeAndOffset.second;

  Operation *outputSlice = getSlice(builder, loc, getOutput(), outputOffsets,
                                    outputSizes, outputStrides);
  // The image dims of the winograd.output_transform result will always be a
  // multiple of the static output_tile_size, so insert a tensor.cast op to
  // maintain more static information in the IR.
  auto outSliceType = cast<ShapedType>(outputSlice->getResultTypes().front());
  SmallVector<int64_t> staticOutShape(outSliceType.getShape());
  auto constSizeH = getConstantIntValue(sizes[1]);
  if (constSizeH.has_value()) {
    staticOutShape[hDim] = constSizeH.value() * getOutputTileSize();
  }
  auto constSizeW = getConstantIntValue(sizes[2]);
  if (constSizeW.has_value()) {
    staticOutShape[wDim] = constSizeW.value() * getOutputTileSize();
  }
  Value staticOutputSlice = castValue(builder, loc, outputSlice->getResult(0),
                                      outSliceType.clone(staticOutShape));

  SmallVector<Value> tiledOperands;
  auto inputSlice = getSlice(builder, loc, getInput(), inputOffsets, inputSizes,
                             inputStrides);
  tiledOperands.emplace_back(inputSlice->getResult(0));
  tiledOperands.emplace_back(staticOutputSlice);

  SmallVector<Type, 4> resultTypes;
  if (hasPureTensorSemantics()) {
    resultTypes.push_back(tiledOperands[1].getType());
  }

  Operation *tiledOp =
      mlir::clone(builder, getOperation(), resultTypes, tiledOperands);

  SmallVector<Value> results(tiledOp->getResults());
  if (!results.empty()) {
    results.front() = castValue(builder, loc, results.front(), outSliceType);
  }
  return TilingResult{{tiledOp}, results, {inputSlice, outputSlice}};
}

LogicalResult WinogradOutputTransformOp::getResultTilePosition(
    OpBuilder &builder, unsigned resultNumber, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, SmallVector<OpFoldResult> &resultOffsets,
    SmallVector<OpFoldResult> &resultSizes) {
  if (resultNumber == 0) {
    auto resultShape = cast<ShapedType>(getOutput().getType()).getShape();
    resultSizes = getAsOpFoldResult(builder.getIndexArrayAttr(resultShape));
    resultOffsets =
        SmallVector<OpFoldResult>(getOutputRank(), builder.getIndexAttr(0));
    const int cDim = getChannelDim();
    const auto hDim = getImageDimensions()[0];
    const auto wDim = getImageDimensions()[1];
    auto loc = getLoc();
    resultOffsets[0] = offsets[0];
    resultOffsets[cDim] = offsets[3];
    resultSizes[0] = sizes[0];
    resultSizes[cDim] = sizes[3];
    SmallVector<SmallVector<OpFoldResult>> reifiedResultShapes;
    if (failed(reifyResultShapes(builder, reifiedResultShapes))) {
      return failure();
    }
    auto hSizeAndOffset = getScaledSizeAndOffset(
        builder, loc, sizes[1], offsets[1], reifiedResultShapes[0][hDim],
        getOutputTileSize(), getOutputTileSize());
    auto wSizeAndOffset = getScaledSizeAndOffset(
        builder, loc, sizes[2], offsets[2], reifiedResultShapes[0][wDim],
        getOutputTileSize(), getOutputTileSize());

    resultSizes[hDim] = hSizeAndOffset.first;
    resultSizes[wDim] = wSizeAndOffset.first;
    resultOffsets[hDim] = hSizeAndOffset.second;
    resultOffsets[wDim] = wSizeAndOffset.second;
    return success();
  }
  return failure();
}

//===----------------------------------------------------------------------===//
// Attention Helpers
//===----------------------------------------------------------------------===//

static SmallVector<Range>
getAttentionIterationDomain(Location loc, OpBuilder &b, int64_t domainRank,
                            ArrayRef<Value> values,
                            ArrayRef<AffineMap> indexingMaps) {
  SmallVector<Range> loopBounds(domainRank);
  OpFoldResult zero = b.getIndexAttr(0);
  OpFoldResult one = b.getIndexAttr(1);

  for (Range &bound : loopBounds) {
    bound.offset = zero;
    bound.stride = one;
  }

  SmallVector<bool> dimsFound(domainRank, false);
  auto fillSizes = [&](Value val, AffineMap indexingMap) {
    for (auto [idx, dimExpr] : llvm::enumerate(indexingMap.getResults())) {
      auto dim = cast<AffineDimExpr>(dimExpr);
      int64_t pos = dim.getPosition();
      if (dimsFound[pos]) {
        continue;
      }
      dimsFound[pos] = true;
      loopBounds[pos].size = getDim(b, loc, val, idx);
    }
  };

  for (auto [val, indexingMap] : llvm::zip_equal(values, indexingMaps)) {
    fillSizes(val, indexingMap);
  }

  return loopBounds;
}

static SmallVector<utils::IteratorType>
getAttentionIteratorTypes(int64_t domainRank, AffineMap qMap, AffineMap kMap,
                          AffineMap vMap, AffineMap oMap) {
  FailureOr<AttentionOpDetail> maybeOpInfo =
      AttentionOpDetail::get(qMap, kMap, vMap, oMap);
  assert(succeeded(maybeOpInfo) && "Failed to infer attention op details");
  AttentionOpDetail opInfo = maybeOpInfo.value();

  // All dimensions other than k1 and k2 are parallel.
  SmallVector<utils::IteratorType> iteratorTypes(domainRank,
                                                 utils::IteratorType::parallel);

  for (auto dim :
       llvm::concat<const int64_t>(opInfo.getK1Dims(), opInfo.getK2Dims())) {
    iteratorTypes[dim] = utils::IteratorType::reduction;
  }

  return iteratorTypes;
}

static Operation *getPermutedSlice(OpBuilder &b, Location loc, Value val,
                                   AffineMap permutation,
                                   ArrayRef<OpFoldResult> offsets,
                                   ArrayRef<OpFoldResult> sizes) {
  SmallVector<Range> slice = getPermutedRange(permutation, offsets, sizes);
  Operation *querySliceOp = getSlice(b, loc, val, slice);
  return querySliceOp;
}

//===----------------------------------------------------------------------===//
// AttentionOp
//===----------------------------------------------------------------------===//

SmallVector<Range> AttentionOp::getIterationDomain(OpBuilder &b) {
  // Attention shape can be determined from Q, K, V alone.
  SmallVector<Value> shapedValues = {getQuery(), getKey(), getValue()};
  SmallVector<AffineMap> indexingMaps = {getQueryMap(), getKeyMap(),
                                         getValueMap()};
  return getAttentionIterationDomain(getLoc(), b, getIterationDomainRank(),
                                     shapedValues, indexingMaps);
}

SmallVector<utils::IteratorType> AttentionOp::getLoopIteratorTypes() {
  return getAttentionIteratorTypes(getIterationDomainRank(), getQueryMap(),
                                   getKeyMap(), getValueMap(), getOutputMap());
}

FailureOr<TilingResult>
AttentionOp::getTiledImplementation(OpBuilder &builder,
                                    ArrayRef<OpFoldResult> offsets,
                                    ArrayRef<OpFoldResult> sizes) {
  assert(offsets.size() == getIterationDomainRank());
  assert(sizes.size() == getIterationDomainRank());

  // TODO: Add support for linalg_ext.index operations in the region.
  // Currently, tiling will break if index operations are present.
  if (!getBody()->getOps<IREE::LinalgExt::IndexOp>().empty()) {
    return failure();
  }

  Location loc = getLoc();

  SmallVector<Range> querySlice =
      getPermutedRange(getQueryMap(), offsets, sizes);
  SmallVector<Range> keySlice = getPermutedRange(getKeyMap(), offsets, sizes);
  SmallVector<Range> valueSlice =
      getPermutedRange(getValueMap(), offsets, sizes);
  SmallVector<Range> outputSlice =
      getPermutedRange(getOutputMap(), offsets, sizes);

  Value scale = getScale();

  SmallVector<Value> tiledOperands;
  SmallVector<Operation *> slices;

  // Query
  {
    Operation *querySliceOp = getSlice(builder, loc, getQuery(), querySlice);
    tiledOperands.emplace_back(querySliceOp->getResult(0));
    slices.push_back(querySliceOp);
  }

  // Key
  {
    Operation *keySliceOp = getSlice(builder, loc, getKey(), keySlice);
    tiledOperands.emplace_back(keySliceOp->getResult(0));
    slices.push_back(keySliceOp);
  }

  // Value
  {
    Operation *valueSliceOp = getSlice(builder, loc, getValue(), valueSlice);
    tiledOperands.emplace_back(valueSliceOp->getResult(0));
    slices.push_back(valueSliceOp);
  }

  // Scale
  tiledOperands.emplace_back(scale);

  // Mask
  Value attnMask = getMask();
  if (attnMask) {
    SmallVector<Range> maskSlice =
        getPermutedRange(*getMaskMap(), offsets, sizes);
    Operation *maskSliceOp = getSlice(builder, loc, attnMask, maskSlice);
    tiledOperands.emplace_back(maskSliceOp->getResult(0));
    slices.push_back(maskSliceOp);
  }

  // Output
  {
    Operation *outputSliceOp = getSlice(builder, loc, getOutput(), outputSlice);
    tiledOperands.emplace_back(outputSliceOp->getResult(0));
    slices.push_back(outputSliceOp);
  }

  SmallVector<Type> resultTypes;
  if (hasPureTensorSemantics()) {
    int64_t baseIdx = attnMask ? 5 : 4;
    resultTypes.push_back(tiledOperands[baseIdx].getType());
  }

  Operation *tiledOp =
      mlir::clone(builder, getOperation(), resultTypes, tiledOperands);

  return TilingResult{
      {tiledOp}, SmallVector<Value>(tiledOp->getResults()), slices};
}

LogicalResult AttentionOp::getResultTilePosition(
    OpBuilder &builder, unsigned resultNumber, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, SmallVector<OpFoldResult> &resultOffsets,
    SmallVector<OpFoldResult> &resultSizes) {
  resultOffsets.clear();
  resultSizes.clear();

  AffineMap resultIndexingMap;
  switch (resultNumber) {
  case 0:
    resultIndexingMap = getOutputMap();
    break;
  default:
    return failure();
  }

  for (AffineExpr dimExpr : resultIndexingMap.getResults()) {
    int dim = cast<AffineDimExpr>(dimExpr).getPosition();
    resultOffsets.push_back(offsets[dim]);
    resultSizes.push_back(sizes[dim]);
  }
  return success();
}

FailureOr<TilingResult>
AttentionOp::generateResultTileValue(OpBuilder &builder, unsigned resultNumber,
                                     ArrayRef<OpFoldResult> offsets,
                                     ArrayRef<OpFoldResult> sizes) {
  // Input offsets and sizes here are from the POV of the outputMap. We need to
  // normalize these offsets and size for it to be useful.

  // Initialize normalized offsets with 0s and normalized sizes with original
  // size.
  SmallVector<Range> iterationDomain(getIterationDomain(builder));
  SmallVector<OpFoldResult> normalizedSizes =
      llvm::map_to_vector(iterationDomain, [](Range x) { return x.size; });
  SmallVector<OpFoldResult> normalizedOffsets(getIterationDomainRank(),
                                              builder.getIndexAttr(0));
  ArrayRef<AffineExpr> outputDims = getOutputMap().getResults();
  for (int i = 0; i < outputDims.size(); i++) {
    int dim = cast<AffineDimExpr>(outputDims[i]).getPosition();
    normalizedOffsets[dim] = offsets[i];
    normalizedSizes[dim] = sizes[i];
  }
  return getTiledImplementation(builder, normalizedOffsets, normalizedSizes);
}

//===----------------------------------------------------------------------===//
// OnlineAttentionOp
//===----------------------------------------------------------------------===//

SmallVector<Range> OnlineAttentionOp::getIterationDomain(OpBuilder &b) {
  // Attention shape can be determined from Q, K, V alone.
  SmallVector<Value> shapedValues = {getQuery(), getKey(), getValue()};
  SmallVector<AffineMap> indexingMaps = {getQueryMap(), getKeyMap(),
                                         getValueMap()};
  return getAttentionIterationDomain(getLoc(), b, getIterationDomainRank(),
                                     shapedValues, indexingMaps);
}

SmallVector<utils::IteratorType> OnlineAttentionOp::getLoopIteratorTypes() {
  return getAttentionIteratorTypes(getIterationDomainRank(), getQueryMap(),
                                   getKeyMap(), getValueMap(), getOutputMap());
}

FailureOr<TilingResult>
OnlineAttentionOp::getTiledImplementation(OpBuilder &builder,
                                          ArrayRef<OpFoldResult> offsets,
                                          ArrayRef<OpFoldResult> sizes) {
  assert(offsets.size() == getIterationDomainRank());
  assert(sizes.size() == getIterationDomainRank());

  Location loc = getLoc();

  SmallVector<Range> querySlice =
      getPermutedRange(getQueryMap(), offsets, sizes);
  SmallVector<Range> keySlice = getPermutedRange(getKeyMap(), offsets, sizes);
  SmallVector<Range> valueSlice =
      getPermutedRange(getValueMap(), offsets, sizes);
  std::optional<SmallVector<Range>> maskSlice;
  if (auto maskMap = getMaskMap()) {
    maskSlice = getPermutedRange(*maskMap, offsets, sizes);
  }

  SmallVector<Range> outputSlice =
      getPermutedRange(getOutputMap(), offsets, sizes);
  SmallVector<Range> maxSlice = getPermutedRange(getMaxMap(), offsets, sizes);
  SmallVector<Range> sumSlice = getPermutedRange(getSumMap(), offsets, sizes);

  Value scale = getScale();

  SmallVector<Value> tiledOperands;
  SmallVector<Operation *> slices;
  /// Query
  {
    Operation *querySliceOp = getSlice(builder, loc, getQuery(), querySlice);
    tiledOperands.emplace_back(querySliceOp->getResult(0));
    slices.push_back(querySliceOp);
  }

  /// Key
  {
    Operation *keySliceOp = getSlice(builder, loc, getKey(), keySlice);
    tiledOperands.emplace_back(keySliceOp->getResult(0));
    slices.push_back(keySliceOp);
  }

  /// Value
  {
    Operation *valueSliceOp = getSlice(builder, loc, getValue(), valueSlice);
    tiledOperands.emplace_back(valueSliceOp->getResult(0));
    slices.push_back(valueSliceOp);
  }

  tiledOperands.emplace_back(scale);

  // Mask
  Value attnMask = getMask();
  if (attnMask) {
    SmallVector<Range> maskSlice =
        getPermutedRange(*getMaskMap(), offsets, sizes);
    Operation *maskSliceOp = getSlice(builder, loc, attnMask, maskSlice);
    tiledOperands.emplace_back(maskSliceOp->getResult(0));
    slices.push_back(maskSliceOp);
  }

  /// Output
  {
    Operation *outputSliceOp = getSlice(builder, loc, getOutput(), outputSlice);
    tiledOperands.emplace_back(outputSliceOp->getResult(0));
    slices.push_back(outputSliceOp);
  }

  /// Max
  {
    Operation *maxSliceOp = getSlice(builder, loc, getMax(), maxSlice);
    tiledOperands.emplace_back(maxSliceOp->getResult(0));
    slices.push_back(maxSliceOp);
  }

  /// Sum
  {
    Operation *sumSliceOp = getSlice(builder, loc, getSum(), sumSlice);
    tiledOperands.emplace_back(sumSliceOp->getResult(0));
    slices.push_back(sumSliceOp);
  }

  SmallVector<Type> resultTypes;
  resultTypes.push_back(tiledOperands[tiledOperands.size() - 3].getType());
  resultTypes.push_back(tiledOperands[tiledOperands.size() - 2].getType());
  resultTypes.push_back(tiledOperands[tiledOperands.size() - 1].getType());

  Operation *tiledOp =
      mlir::clone(builder, getOperation(), resultTypes, tiledOperands);

  return TilingResult{
      {tiledOp}, SmallVector<Value>(tiledOp->getResults()), slices};
}

LogicalResult OnlineAttentionOp::getResultTilePosition(
    OpBuilder &builder, unsigned resultNumber, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, SmallVector<OpFoldResult> &resultOffsets,
    SmallVector<OpFoldResult> &resultSizes) {
  resultOffsets.clear();
  resultSizes.clear();

  AffineMap resultIndexingMap;
  switch (resultNumber) {
  case 0:
    resultIndexingMap = getOutputMap();
    break;
  case 1:
    resultIndexingMap = getMaxMap();
    break;
  case 2:
    resultIndexingMap = getSumMap();
    break;
  default:
    return failure();
  }

  for (AffineExpr dimExpr : resultIndexingMap.getResults()) {
    int dim = cast<AffineDimExpr>(dimExpr).getPosition();
    resultOffsets.push_back(offsets[dim]);
    resultSizes.push_back(sizes[dim]);
  }
  return success();
}

static AffineMap getPartialResultMap(AffineMap map, AttentionOpDetail &opInfo) {
  // Append K2 dimensions at start.
  for (auto [idx, dim] : llvm::enumerate(opInfo.getK2Dims())) {
    map = map.insertResult(getAffineDimExpr(dim, map.getContext()), idx);
  }
  return map;
}

FailureOr<SmallVector<Value>>
OnlineAttentionOp::generateInitialTensorForPartialReduction(
    OpBuilder &b, Location loc, ArrayRef<OpFoldResult> sizes,
    const llvm::SetVector<unsigned> &reductionDims) {
  FailureOr<AttentionOpDetail> maybeOpInfo = AttentionOpDetail::get(
      getQueryMap(), getKeyMap(), getValueMap(), getOutputMap());
  if (failed(maybeOpInfo)) {
    return emitOpError("failed to verify op's indexing maps");
  }
  AttentionOpDetail &opInfo = maybeOpInfo.value();

  SmallVector<OpFoldResult> shape = llvm::map_to_vector(
      getIterationDomain(b), [](Range x) { return x.size; });

  SmallVector<OpFoldResult> tiledShape;
  for (auto [tileSize, dimSize] : llvm::zip_equal(sizes, shape)) {
    if (isZeroInteger(tileSize)) {
      tiledShape.push_back(dimSize);
    } else {
      tiledShape.push_back(tileSize);
    }
  }

  SmallVector<OpFoldResult> accSize = applyPermutationMap<OpFoldResult>(
      getPartialResultMap(getOutputMap(), opInfo), tiledShape);
  SmallVector<OpFoldResult> maxSize = applyPermutationMap<OpFoldResult>(
      getPartialResultMap(getMaxMap(), opInfo), tiledShape);
  SmallVector<OpFoldResult> sumSize = applyPermutationMap<OpFoldResult>(
      getPartialResultMap(getSumMap(), opInfo), tiledShape);

  Type accElTy = getElementTypeOrSelf(getOutput().getType());
  Type maxElTy = getElementTypeOrSelf(getMax().getType());
  Type sumElTy = getElementTypeOrSelf(getSum().getType());

  Value partialAcc = tensor::EmptyOp::create(b, loc, accSize, accElTy);
  Value partialMax = tensor::EmptyOp::create(b, loc, maxSize, maxElTy);
  Value partialSum = tensor::EmptyOp::create(b, loc, sumSize, sumElTy);

  Value accInit = arith::getIdentityValue(arith::AtomicRMWKind::addf, accElTy,
                                          b, loc, /*useOnlyFiniteValue=*/true);
  Value maxInit =
      arith::getIdentityValue(arith::AtomicRMWKind::maximumf, maxElTy, b, loc,
                              /*useOnlyFiniteValue=*/true);
  Value sumInit =
      arith::getIdentityValue(arith::AtomicRMWKind::addf, sumElTy, b, loc);

  Value accFill =
      linalg::FillOp::create(b, loc, ValueRange{accInit}, partialAcc)
          .getResult(0);
  Value maxFill =
      linalg::FillOp::create(b, loc, ValueRange{maxInit}, partialMax)
          .getResult(0);
  Value sumFill =
      linalg::FillOp::create(b, loc, ValueRange{sumInit}, partialSum)
          .getResult(0);

  return SmallVector<Value>{accFill, maxFill, sumFill};
}

FailureOr<TilingResult> OnlineAttentionOp::tileToPartialReduction(
    OpBuilder &b, Location loc, ReductionTilingStrategy strategy,
    ValueRange init, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes,
    const llvm::SetVector<unsigned> &reductionDims,
    ArrayRef<OpFoldResult> splitReductionIvs) {
  FailureOr<AttentionOpDetail> maybeOpInfo = AttentionOpDetail::get(
      getQueryMap(), getKeyMap(), getValueMap(), getOutputMap());
  if (failed(maybeOpInfo)) {
    return emitOpError("failed to verify op's indexing maps");
  }
  AttentionOpDetail &opInfo = maybeOpInfo.value();

  // Extend result maps, keeping everything else the same.
  AffineMap partialAccMap = getPartialResultMap(getOutputMap(), opInfo);
  AffineMap partialMaxMap = getPartialResultMap(getMaxMap(), opInfo);
  AffineMap partialSumMap = getPartialResultMap(getSumMap(), opInfo);

  SmallVector<AffineMap> indexingMaps = getIndexingMapsArray();
  indexingMaps[getNumDpsInputs()] = partialAccMap;
  indexingMaps[getNumDpsInputs() + 1] = partialMaxMap;
  indexingMaps[getNumDpsInputs() + 2] = partialSumMap;

  SmallVector<Value> tiledOperands;
  SmallVector<Operation *> slices;

  auto appendSlice = [&](Value val, AffineMap map,
                         ArrayRef<OpFoldResult> offsets) -> LogicalResult {
    Operation *sliceOp = getPermutedSlice(b, loc, val, map, offsets, sizes);
    if (!sliceOp) {
      return emitOpError("failed to get slice");
    }
    tiledOperands.emplace_back(sliceOp->getResult(0));
    slices.push_back(sliceOp);
    return success();
  };

  if (failed(appendSlice(getQuery(), getQueryMap(), offsets))) {
    return failure();
  }
  if (failed(appendSlice(getKey(), getKeyMap(), offsets))) {
    return failure();
  }
  if (failed(appendSlice(getValue(), getValueMap(), offsets))) {
    return failure();
  }

  tiledOperands.emplace_back(getScale());

  if (Value mask = getMask()) {
    if (failed(appendSlice(mask, *getMaskMap(), offsets))) {
      return failure();
    }
  }

  // For results, we set offset of the iterated reduction dims to 0.
  SmallVector<OpFoldResult> initOffsets(offsets);
  for (int dim : opInfo.getK2Dims()) {
    initOffsets[dim] = b.getIndexAttr(0);
  }

  if (failed(appendSlice(init[0], partialAccMap, initOffsets))) {
    return failure();
  }
  if (failed(appendSlice(init[1], partialMaxMap, initOffsets))) {
    return failure();
  }
  if (failed(appendSlice(init[2], partialSumMap, initOffsets))) {
    return failure();
  }

  // Get the initial values.
  ValueRange slicedInits = ArrayRef(tiledOperands).take_back(3);

  auto tiledOp = cast<OnlineAttentionOp>(
      mlir::clone(b, getOperation(), slicedInits.getTypes(), tiledOperands));
  tiledOp.setIndexingMapsAttr(b.getAffineMapArrayAttr(indexingMaps));

  return TilingResult{
      {tiledOp}, SmallVector<Value>(tiledOp->getResults()), slices};
}

template <typename CombinerOp>
static linalg::ReduceOp reduceOnK2(OnlineAttentionOp attn, AffineMap partialMap,
                                   AttentionOpDetail &opInfo, OpBuilder &b,
                                   Location loc, Value partialResult,
                                   Value init) {
  // linalg.reduce's iteration space is the result's iteration space (and
  // not the operations iteration space). To account for this, permute the
  // reduction dimensions based on the partial result map.
  SmallVector<int64_t> partialReductionDims;
  for (auto [resultNum, dimExpr] : llvm::enumerate(partialMap.getResults())) {
    unsigned dim = cast<AffineDimExpr>(dimExpr).getPosition();
    if (llvm::find(opInfo.getK2Dims(), dim) != opInfo.getK2Dims().end()) {
      partialReductionDims.push_back(resultNum);
    }
  }

  return linalg::ReduceOp::create(
      b, loc, partialResult, init, partialReductionDims,
      [&](OpBuilder &b, Location loc, ValueRange inputs) {
        Value reduced = CombinerOp::create(b, loc, inputs[0], inputs[1]);
        linalg::YieldOp::create(b, loc, reduced);
      });
};

template <typename T>
static Value elementwiseValueInPlace(OpBuilder &builder, Location loc,
                                     AffineMap inputMap, AffineMap scaleMap,
                                     Value value, Value scale) {
  SmallVector<AffineMap> compressedMaps =
      compressUnusedDims(SmallVector<AffineMap>{inputMap, scaleMap});
  inputMap = compressedMaps[0];
  scaleMap = compressedMaps[1];

  SmallVector<utils::IteratorType> iteratorTypes(inputMap.getNumDims(),
                                                 utils::IteratorType::parallel);

  auto genericOp = linalg::GenericOp::create(
      builder, loc, value.getType(), scale, value,
      SmallVector<AffineMap>{scaleMap, inputMap}, iteratorTypes,
      [&](OpBuilder &b, Location loc, ValueRange args) {
        // Convert scale to the same datatype as input.
        Value scale = convertScalarToDtype(b, loc, args[0], args[1].getType(),
                                           /*isUnsignedCast=*/false);
        Value result = T::create(b, loc, scale, args[1]);
        linalg::YieldOp::create(b, loc, result);
      });
  return genericOp.getResult(0);
}

// Compute output = exp2(output - input)
static Value computeSubAndExp2(OpBuilder &builder, Location loc,
                               AffineMap inputMap, AffineMap outputMap,
                               Value input, Value output) {
  SmallVector<AffineMap> compressedMaps =
      compressUnusedDims(SmallVector<AffineMap>{inputMap, outputMap});
  inputMap = compressedMaps[0];
  outputMap = compressedMaps[1];

  SmallVector<utils::IteratorType> iteratorTypes(inputMap.getNumDims(),
                                                 utils::IteratorType::parallel);
  auto genericOp = linalg::GenericOp::create(
      builder, loc, output.getType(), input, output,
      SmallVector<AffineMap>{inputMap, outputMap}, iteratorTypes,
      [&](OpBuilder &b, Location loc, ValueRange args) {
        // Convert input to the same datatype as output.
        Value in = convertScalarToDtype(b, loc, args[0], args[1].getType(),
                                        /*isUnsignedCast=*/false);
        Value diff = arith::SubFOp::create(b, loc, args[1], in);
        Value weight = math::Exp2Op::create(b, loc, diff);
        linalg::YieldOp::create(b, loc, weight);
      });
  return genericOp.getResult(0);
}

FailureOr<MergeResult> OnlineAttentionOp::mergeReductions(
    OpBuilder &b, Location loc, ValueRange partialReduce,
    const llvm::SetVector<unsigned> &reductionDims) {
  FailureOr<AttentionOpDetail> maybeOpInfo = AttentionOpDetail::get(
      getQueryMap(), getKeyMap(), getValueMap(), getOutputMap());
  if (failed(maybeOpInfo)) {
    return emitOpError("failed to verify op's indexing maps");
  }
  AttentionOpDetail &opInfo = maybeOpInfo.value();

  AffineMap partialAccMap = getPartialResultMap(getOutputMap(), opInfo);
  AffineMap partialMaxMap = getPartialResultMap(getMaxMap(), opInfo);
  AffineMap partialSumMap = getPartialResultMap(getSumMap(), opInfo);

  // newMax = max(maxInit, rowMax(partialMax))
  linalg::ReduceOp reducedMax = reduceOnK2<arith::MaximumFOp>(
      *this, partialMaxMap, opInfo, b, loc, partialReduce[1], getMax());

  // norm = exp2(partialMax - newMax)
  Value norm = computeSubAndExp2(b, loc, getMaxMap(), partialMaxMap,
                                 reducedMax.getResult(0), partialReduce[1]);

  // normSum = norm * partialSum
  Value normSum = elementwiseValueInPlace<arith::MulFOp>(
      b, loc, partialSumMap, partialMaxMap, partialReduce[2], norm);

  // newSum = sumInit + rowSum(partialSum)
  linalg::ReduceOp reducedSum = reduceOnK2<arith::AddFOp>(
      *this, partialSumMap, opInfo, b, loc, normSum, getSum());

  // normAcc = norm * partialAcc
  Value normAcc = elementwiseValueInPlace<arith::MulFOp>(
      b, loc, partialAccMap, partialMaxMap, partialReduce[0], norm);

  // newAcc = accInit + rowMax(partialAcc)
  linalg::ReduceOp reducedAcc = reduceOnK2<arith::AddFOp>(
      *this, partialAccMap, opInfo, b, loc, normAcc, getOutput());

  return MergeResult{{reducedAcc, reducedMax, reducedSum},
                     {reducedAcc.getResult(0), reducedMax.getResult(0),
                      reducedSum.getResult(0)}};
}

LogicalResult OnlineAttentionOp::getPartialResultTilePosition(
    OpBuilder &b, unsigned resultNumber, ReductionTilingStrategy tilingStrategy,
    ArrayRef<OpFoldResult> offsets, ArrayRef<OpFoldResult> sizes,
    const llvm::SetVector<unsigned> &reductionDims,
    ArrayRef<OpFoldResult> splitReductionIvs,
    SmallVector<OpFoldResult> &resultOffsets,
    SmallVector<OpFoldResult> &resultSizes) {

  FailureOr<AttentionOpDetail> maybeOpInfo = AttentionOpDetail::get(
      getQueryMap(), getKeyMap(), getValueMap(), getOutputMap());
  if (failed(maybeOpInfo)) {
    return emitOpError("failed to verify op's indexing maps");
  }
  AttentionOpDetail &opInfo = maybeOpInfo.value();

  resultOffsets.clear();
  resultSizes.clear();

  AffineMap resultIndexingMap;
  switch (resultNumber) {
  case 0:
    resultIndexingMap = getOutputMap();
    break;
  case 1:
    resultIndexingMap = getMaxMap();
    break;
  case 2:
    resultIndexingMap = getSumMap();
    break;
  default:
    return failure();
  }

  AffineMap partialMap = getPartialResultMap(resultIndexingMap, opInfo);

  for (AffineExpr dimExpr : partialMap.getResults()) {
    int dim = cast<AffineDimExpr>(dimExpr).getPosition();
    resultSizes.push_back(sizes[dim]);

    if (llvm::find(opInfo.getK2Dims(), dim) != opInfo.getK2Dims().end()) {
      // Reduction dims are reduced, and are always outputted in the same
      // place. So use offset 0 for them.
      resultOffsets.push_back(b.getIndexAttr(0));
    } else {
      resultOffsets.push_back(offsets[dim]);
    }
  }
  return success();
}

//===---------------------------------------------------------------------===//
// CustomOp
//===---------------------------------------------------------------------===//

/// These methods copied/modified from `TilingInterface` implementation of
/// `getIterationDomain` of `LinalgOp`s.

SmallVector<utils::IteratorType> CustomOp::getLoopIteratorTypes() {
  return llvm::map_to_vector(getIteratorTypes(), [](Attribute attr) {
    return cast<IREE::LinalgExt::IteratorTypeAttr>(attr).getValue();
  });
}

SmallVector<Range> CustomOp::getIterationDomainForDimensions(
    OpBuilder &builder, ArrayRef<unsigned> dims, ArrayRef<unsigned> symbols) {
  CustomOp customOp = *this;
  SmallVector<AffineMap> maps = customOp.getIndexingMapsArray();
  if (maps.empty()) {
    return SmallVector<Range>{};
  }

  Location loc = getLoc();

  // 1. Create a flat list of all the operand shapes (similar to Linalg)
  SmallVector<OpFoldResult> allShapesSizes =
      createFlatListOfOperandDims(builder, loc, customOp);

  // 2a. Next we need to get a map from shapes to loop. Since `custom_op`
  //     has indexing maps that have symbols, to make this work correctly
  //     compute new maps that replaces the symbols with "new" dims.
  unsigned numDims = getNumLoops();
  unsigned numSymbols = getNumNonLoopDimensions();
  MLIRContext *context = getContext();
  SmallVector<AffineMap> modifiedMaps =
      convertDimsToSymbols(context, maps, numDims, numSymbols);

  // 2b. Concat the affine maps.
  AffineMap concatMap =
      inversePermutation(concatAffineMaps(modifiedMaps, context));
  // TODO: Ideally we should bail if the map is invalid, i.e. we abort from
  // applying the transformation. We could add this to the verifier as well, but
  // it is unclear if this makes the op invalid. Revisit after more experience
  // of how this operation is used.
  assert(concatMap && "failure in inverting indexing maps");
  SmallVector<Range> ranges;

  OpFoldResult zero = builder.getIndexAttr(0), one = builder.getIndexAttr(1);
  auto getRange = [&](AffineExpr expr) {
    OpFoldResult ofr = affine::makeComposedFoldedAffineApply(builder, loc, expr,
                                                             allShapesSizes);
    return Range{zero, ofr, one};
  };
  ranges = llvm::map_to_vector(
      dims, [&](unsigned dim) { return getRange(concatMap.getResult(dim)); });
  ranges.append(llvm::map_to_vector(symbols, [&](unsigned symbol) {
    return getRange(concatMap.getResult(symbol + numSymbols));
  }));
  return ranges;
}

SmallVector<Range> CustomOp::getIterationDomain(OpBuilder &builder) {
  auto dims = llvm::to_vector(llvm::seq<unsigned>(0, getNumLoops()));
  return getIterationDomainForDimensions(builder, dims,
                                         /*symbols=*/ArrayRef<unsigned>{});
}

/// During tiling, the tiling implementation generates offsets and sizes
/// to use for the dimensions of the operation that correspond to loops.
/// The dimensions of the operation corresponding to symbols isnt tiled.
/// This method extends the offsets and sizes for a tile with the
/// offsets (which are zero) and sizes (which are same as the untiled sizes) for
/// the dimension represented by symbols.
static std::pair<SmallVector<OpFoldResult>, SmallVector<OpFoldResult>>
appendOffsetAndSizeForSymbolDimensions(OpBuilder &builder, CustomOp customOp,
                                       ArrayRef<OpFoldResult> offsets,
                                       ArrayRef<OpFoldResult> sizes) {
  unsigned numSymbols = customOp.getNumNonLoopDimensions();
  if (numSymbols == 0) {
    return {llvm::to_vector(offsets), llvm::to_vector(sizes)};
  }

  auto appendedOffsets = llvm::to_vector(offsets);
  appendedOffsets.append(numSymbols, builder.getIndexAttr(0));

  auto appendedSizes = llvm::to_vector(sizes);
  auto symbols = llvm::to_vector(llvm::seq<unsigned>(0, numSymbols));
  auto symbolRanges = customOp.getIterationDomainForDimensions(
      builder, /*dims=*/ArrayRef<unsigned>{}, symbols);
  auto symbolRangesUb =
      llvm::map_range(symbolRanges, [](Range r) { return r.size; });
  appendedSizes.append(symbolRangesUb.begin(), symbolRangesUb.end());
  return {appendedOffsets, appendedSizes};
}

/// This method is adapted from the method `linalg::computeAllSliceParameters`,
/// adapted to work with map that have symbols, and empty maps.
static SmallVector<std::optional<linalg::SliceParameters>>
computeCustomOpAllSliceParameters(OpBuilder &builder, Location loc,
                                  CustomOp customOp, ValueRange valuesToTile,
                                  ArrayRef<OpFoldResult> ivs,
                                  SmallVector<OpFoldResult> tileSizes) {
  assert(ivs.size() == static_cast<size_t>(llvm::count_if(
                           llvm::make_range(tileSizes.begin(), tileSizes.end()),
                           [](OpFoldResult v) { return !isZeroInteger(v); })) &&
         "expected as many ivs as non-zero sizes");
  unsigned numDims = customOp.getNumLoops();
  unsigned numSymbols = customOp.getNumNonLoopDimensions();

  // Construct (potentially temporary) mins and maxes on which to apply maps
  // that define tile subshapes.
  SmallVector<OpFoldResult> lbs =
      linalg::computeTileOffsets(builder, loc, ivs, tileSizes);
  SmallVector<OpFoldResult> sizeBounds = tileSizes;

  std::tie(lbs, sizeBounds) = appendOffsetAndSizeForSymbolDimensions(
      builder, customOp, lbs, sizeBounds);

  tileSizes.append(numSymbols, builder.getIndexAttr(0));
  SmallVector<OpFoldResult> subShapeSizes =
      linalg::computeTileSizes(builder, loc, tileSizes, sizeBounds);

  SmallVector<AffineExpr> symbolReplacements;
  if (numSymbols != 0) {
    symbolReplacements =
        getDimExprsForSymbols(builder.getContext(), numDims, numSymbols);
  }

  assert(static_cast<int64_t>(valuesToTile.size()) <=
             customOp->getNumOperands() &&
         "more value to tile than operands.");
  SmallVector<std::optional<linalg::SliceParameters>> allSliceParams;
  allSliceParams.reserve(valuesToTile.size());
  for (auto [opOperand, val] :
       llvm::zip(customOp->getOpOperands(), valuesToTile)) {
    Value shapedOp = val;
    LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: for operand " << shapedOp);
    AffineMap map = customOp.getMatchingIndexingMap(&opOperand);
    if (numSymbols != 0) {
      map = convertDimsToSymbols(map, numDims, numSymbols, symbolReplacements);
    }

    // If the map is empty, we dont tile this operand.
    Type operandType = opOperand.get().getType();
    if (map.isEmpty() || !isa<ShapedType>(operandType)) {
      allSliceParams.push_back(std::nullopt);
      LLVM_DEBUG(llvm::dbgs()
                 << ": not tiled: use shape: " << operandType << "\n");
      continue;
    }

    allSliceParams.push_back(linalg::computeSliceParameters(
        builder, loc, shapedOp, tileSizes, map, lbs, sizeBounds, subShapeSizes,
        /*omitPartialTileCheck=*/true));
  }

  return allSliceParams;
}

/// This method is same as `materializeTiledShape` method defined in
/// `mlir/Dialect/Linalg/Utils/Utils.[h|cpp]`.
static Operation *
materializeTiledShape(OpBuilder &builder, Location loc, Value valueToTile,
                      const linalg::SliceParameters &sliceParams) {
  auto shapedType = dyn_cast<ShapedType>(valueToTile.getType());
  auto *sliceOp = TypeSwitch<ShapedType, Operation *>(shapedType)
                      .Case([&](RankedTensorType) {
                        return tensor::ExtractSliceOp::create(
                            builder, loc, valueToTile, sliceParams.offsets,
                            sliceParams.sizes, sliceParams.strides);
                      })
                      .Default([](ShapedType) -> Operation * {
                        llvm_unreachable("Unexpected shaped type");
                      });
  return sliceOp;
}

/// This method is adapted from the method `linalg::makeTiledShapes`,
/// adapted to work with map that have symbols, and empty maps.
static SmallVector<Value>
makeCustomOpTiledShapes(OpBuilder &builder, Location loc, CustomOp customOp,
                        ValueRange valuesToTile, ArrayRef<OpFoldResult> ivs,
                        ArrayRef<OpFoldResult> tileSizes) {
  SmallVector<std::optional<linalg::SliceParameters>> allSliceParameter =
      computeCustomOpAllSliceParameters(builder, loc, customOp, valuesToTile,
                                        ivs, llvm::to_vector(tileSizes));
  SmallVector<Value> tiledShapes;
  for (auto [valueToTile, sliceParams] :
       llvm::zip_equal(valuesToTile, allSliceParameter)) {
    tiledShapes.push_back(
        sliceParams.has_value()
            ? materializeTiledShape(builder, loc, valueToTile, *sliceParams)
                  ->getResult(0)
            : valueToTile);
  }
  return tiledShapes;
}

static void offsetCustomOpIndices(OpBuilder &b, CustomOp customOp,
                                  ArrayRef<OpFoldResult> offsets) {
  IRRewriter rewriter(b);
  for (auto indexOp : customOp.getBody()->getOps<IREE::LinalgExt::IndexOp>()) {
    if (indexOp.getDim() >= offsets.size() || !offsets[indexOp.getDim()]) {
      continue;
    }
    OpBuilder::InsertionGuard guard(b);
    rewriter.setInsertionPointAfter(indexOp);
    AffineExpr index, offset;
    bindDims(b.getContext(), index, offset);
    OpFoldResult applied = affine::makeComposedFoldedAffineApply(
        rewriter, indexOp.getLoc(), index + offset,
        {getAsOpFoldResult(indexOp.getResult()), offsets[indexOp.getDim()]});
    Value materialized =
        getValueOrCreateConstantIndexOp(b, indexOp.getLoc(), applied);
    rewriter.replaceUsesWithIf(indexOp, materialized, [&](OpOperand &use) {
      return use.getOwner() != materialized.getDefiningOp();
    });
  }
}

FailureOr<TilingResult>
CustomOp::getTiledImplementation(OpBuilder &builder,
                                 ArrayRef<OpFoldResult> offsets,
                                 ArrayRef<OpFoldResult> sizes) {
  CustomOp customOp = *this;
  // Leave the `sizeBounds` value empty. That is only needed when the `sizes`
  // specified could lead to out of bounds accesses.
  Location loc = getLoc();

  SmallVector<Value> valuesToTile = getOperation()->getOperands();
  SmallVector<Value> tiledOperands = makeCustomOpTiledShapes(
      builder, loc, *this, valuesToTile, offsets, sizes);
  SmallVector<Operation *> generatedSlices = llvm::map_to_vector(
      llvm::make_filter_range(
          tiledOperands,
          [](Value v) -> bool {
            return isa_and_nonnull<tensor::ExtractSliceOp, memref::SubViewOp>(
                v.getDefiningOp());
          }),
      [](Value v) -> Operation * { return v.getDefiningOp(); });

  SmallVector<Type> resultTensorTypes =
      llvm::map_to_vector(getDpsInitsMutable(), [&](OpOperand &opOperand) {
        return tiledOperands[opOperand.getOperandNumber()].getType();
      });

  Operation *tiledOp =
      mlir::clone(builder, customOp, resultTensorTypes, tiledOperands);
  offsetCustomOpIndices(builder, cast<CustomOp>(tiledOp), offsets);

  return TilingResult{
      {tiledOp}, SmallVector<Value>(tiledOp->getResults()), generatedSlices};
}

/// Methods copied/modified from `TilingInterface` implementation of
/// `getTiledImplementation` of `LinalgOp`s.

LogicalResult CustomOp::getResultTilePosition(
    OpBuilder &builder, unsigned resultNumber, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, SmallVector<OpFoldResult> &resultOffsets,
    SmallVector<OpFoldResult> &resultSizes) {
  CustomOp customOp = *this;
  Location loc = getLoc();

  AffineExpr d0 = builder.getAffineDimExpr(0);
  SmallVector<OpFoldResult> appendedOffsets, appendedSizes;
  std::tie(appendedOffsets, appendedSizes) =
      appendOffsetAndSizeForSymbolDimensions(builder, customOp, offsets, sizes);

  SmallVector<OpFoldResult> subShapeSizes =
      llvm::map_to_vector(appendedSizes, [&](OpFoldResult ofr) {
        return affine::makeComposedFoldedAffineApply(builder, loc, d0 - 1, ofr);
      });

  OpOperand *outOperand = customOp.getDpsInitOperand(resultNumber);
  linalg::SliceParameters sliceParams = linalg::computeSliceParameters(
      builder, loc, outOperand->get(), sizes,
      customOp.getMatchingIndexingMap(outOperand), appendedOffsets,
      /*ubs*/ {}, subShapeSizes, /*omitPartialTileCheck=*/true);
  resultOffsets = sliceParams.offsets;
  resultSizes = sliceParams.sizes;
  return success();
}

//===---------------------------------------------------------------------===//
// tensor::ConcatOp
//===---------------------------------------------------------------------===//

// TilingInterface overloads with InnerTileAlignment parameter.
// Only `linalg.pack/unpack` operations utilize this parameter, therefore,
// these forward to the overloads without the parameter.
FailureOr<TilingResult> ArgCompareOp::getTiledImplementation(
    OpBuilder &builder, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, ArrayRef<mlir::InnerTileAlignment>) {
  return getTiledImplementation(builder, offsets, sizes);
}

FailureOr<TilingResult> AttentionOp::getTiledImplementation(
    OpBuilder &builder, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, ArrayRef<mlir::InnerTileAlignment>) {
  return getTiledImplementation(builder, offsets, sizes);
}

FailureOr<TilingResult> CustomOp::getTiledImplementation(
    OpBuilder &builder, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, ArrayRef<mlir::InnerTileAlignment>) {
  return getTiledImplementation(builder, offsets, sizes);
}

FailureOr<TilingResult> ExpReductionOp::getTiledImplementation(
    OpBuilder &builder, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, ArrayRef<mlir::InnerTileAlignment>) {
  return getTiledImplementation(builder, offsets, sizes);
}

FailureOr<TilingResult> FftOp::getTiledImplementation(
    OpBuilder &builder, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, ArrayRef<mlir::InnerTileAlignment>) {
  return getTiledImplementation(builder, offsets, sizes);
}

FailureOr<TilingResult> GatherOp::getTiledImplementation(
    OpBuilder &builder, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, ArrayRef<mlir::InnerTileAlignment>) {
  return getTiledImplementation(builder, offsets, sizes);
}

FailureOr<TilingResult> MapLoadOp::getTiledImplementation(
    OpBuilder &builder, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, ArrayRef<mlir::InnerTileAlignment>) {
  return getTiledImplementation(builder, offsets, sizes);
}

FailureOr<TilingResult> MapStoreOp::getTiledImplementation(
    OpBuilder &builder, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, ArrayRef<mlir::InnerTileAlignment>) {
  return getTiledImplementation(builder, offsets, sizes);
}

FailureOr<TilingResult> OnlineAttentionOp::getTiledImplementation(
    OpBuilder &builder, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, ArrayRef<mlir::InnerTileAlignment>) {
  return getTiledImplementation(builder, offsets, sizes);
}

FailureOr<TilingResult> ScanOp::getTiledImplementation(
    OpBuilder &builder, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, ArrayRef<mlir::InnerTileAlignment>) {
  return getTiledImplementation(builder, offsets, sizes);
}

FailureOr<TilingResult> ScatterOp::getTiledImplementation(
    OpBuilder &builder, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, ArrayRef<mlir::InnerTileAlignment>) {
  return getTiledImplementation(builder, offsets, sizes);
}

FailureOr<TilingResult> SortOp::getTiledImplementation(
    OpBuilder &builder, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, ArrayRef<mlir::InnerTileAlignment>) {
  return getTiledImplementation(builder, offsets, sizes);
}

FailureOr<TilingResult> TopkOp::getTiledImplementation(
    OpBuilder &builder, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, ArrayRef<mlir::InnerTileAlignment>) {
  return getTiledImplementation(builder, offsets, sizes);
}

FailureOr<TilingResult> WinogradInputTransformOp::getTiledImplementation(
    OpBuilder &builder, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, ArrayRef<mlir::InnerTileAlignment>) {
  return getTiledImplementation(builder, offsets, sizes);
}

FailureOr<TilingResult> ArgCompareOp::generateResultTileValue(
    OpBuilder &builder, unsigned resultNumber, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, ArrayRef<mlir::InnerTileAlignment>) {
  return generateResultTileValue(builder, resultNumber, offsets, sizes);
}

FailureOr<TilingResult> AttentionOp::generateResultTileValue(
    OpBuilder &builder, unsigned resultNumber, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, ArrayRef<mlir::InnerTileAlignment>) {
  return generateResultTileValue(builder, resultNumber, offsets, sizes);
}

FailureOr<TilingResult> ExpReductionOp::generateResultTileValue(
    OpBuilder &builder, unsigned resultNumber, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, ArrayRef<mlir::InnerTileAlignment>) {
  return generateResultTileValue(builder, resultNumber, offsets, sizes);
}

FailureOr<TilingResult> GatherOp::generateResultTileValue(
    OpBuilder &builder, unsigned resultNumber, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, ArrayRef<mlir::InnerTileAlignment>) {
  return generateResultTileValue(builder, resultNumber, offsets, sizes);
}

FailureOr<TilingResult> MapLoadOp::generateResultTileValue(
    OpBuilder &builder, unsigned resultNumber, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, ArrayRef<mlir::InnerTileAlignment>) {
  return generateResultTileValue(builder, resultNumber, offsets, sizes);
}

FailureOr<TilingResult> Im2colOp::getTiledImplementation(
    OpBuilder &builder, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, ArrayRef<mlir::InnerTileAlignment>) {
  return getTiledImplementation(builder, offsets, sizes);
}

FailureOr<TilingResult> TopkV2Op::getTiledImplementation(
    OpBuilder &builder, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, ArrayRef<mlir::InnerTileAlignment>) {
  return getTiledImplementation(builder, offsets, sizes);
}

FailureOr<TilingResult> WinogradFilterTransformOp::getTiledImplementation(
    OpBuilder &builder, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, ArrayRef<mlir::InnerTileAlignment>) {
  return getTiledImplementation(builder, offsets, sizes);
}

FailureOr<TilingResult> WinogradOutputTransformOp::getTiledImplementation(
    OpBuilder &builder, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, ArrayRef<mlir::InnerTileAlignment>) {
  return getTiledImplementation(builder, offsets, sizes);
}

FailureOr<TilingResult> Im2colOp::generateResultTileValue(
    OpBuilder &builder, unsigned resultNumber, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, ArrayRef<mlir::InnerTileAlignment>) {
  return generateResultTileValue(builder, resultNumber, offsets, sizes);
}

FailureOr<TilingResult> MapLoadOp::getTiledImplementationFromOperandTiles(
    OpBuilder &b, ArrayRef<unsigned> operandNumbers,
    ArrayRef<SmallVector<OpFoldResult>> allOffsets,
    ArrayRef<SmallVector<OpFoldResult>> allSizes,
    ArrayRef<mlir::InnerTileAlignment>) {
  return getTiledImplementationFromOperandTiles(b, operandNumbers, allOffsets,
                                                allSizes);
}

LogicalResult MapLoadOp::getIterationDomainTileFromOperandTiles(
    OpBuilder &b, ArrayRef<unsigned> operandNumbers,
    ArrayRef<SmallVector<OpFoldResult>> allOffsets,
    ArrayRef<SmallVector<OpFoldResult>> allSizes,
    SmallVectorImpl<OpFoldResult> &iterDomainOffsets,
    SmallVectorImpl<OpFoldResult> &iterDomainSizes,
    ArrayRef<mlir::InnerTileAlignment>) {
  return getIterationDomainTileFromOperandTiles(b, operandNumbers, allOffsets,
                                                allSizes, iterDomainOffsets,
                                                iterDomainSizes);
}

FailureOr<TilingResult> MapStoreOp::getTiledImplementationFromOperandTiles(
    OpBuilder &b, ArrayRef<unsigned> operandNumbers,
    ArrayRef<SmallVector<OpFoldResult>> allOffsets,
    ArrayRef<SmallVector<OpFoldResult>> allSizes,
    ArrayRef<mlir::InnerTileAlignment>) {
  return getTiledImplementationFromOperandTiles(b, operandNumbers, allOffsets,
                                                allSizes);
}

LogicalResult MapStoreOp::getIterationDomainTileFromOperandTiles(
    OpBuilder &b, ArrayRef<unsigned> operandNumbers,
    ArrayRef<SmallVector<OpFoldResult>> allOffsets,
    ArrayRef<SmallVector<OpFoldResult>> allSizes,
    SmallVectorImpl<OpFoldResult> &iterDomainOffsets,
    SmallVectorImpl<OpFoldResult> &iterDomainSizes,
    ArrayRef<mlir::InnerTileAlignment>) {
  return getIterationDomainTileFromOperandTiles(b, operandNumbers, allOffsets,
                                                allSizes, iterDomainOffsets,
                                                iterDomainSizes);
}

FailureOr<TilingResult> ScatterOp::getTiledImplementationFromOperandTiles(
    OpBuilder &b, ArrayRef<unsigned> operandNumbers,
    ArrayRef<SmallVector<OpFoldResult>> allOffsets,
    ArrayRef<SmallVector<OpFoldResult>> allSizes,
    ArrayRef<mlir::InnerTileAlignment>) {
  return getTiledImplementationFromOperandTiles(b, operandNumbers, allOffsets,
                                                allSizes);
}

LogicalResult ScatterOp::getIterationDomainTileFromOperandTiles(
    OpBuilder &b, ArrayRef<unsigned> operandNumbers,
    ArrayRef<SmallVector<OpFoldResult>> allOffsets,
    ArrayRef<SmallVector<OpFoldResult>> allSizes,
    SmallVectorImpl<OpFoldResult> &iterDomainOffsets,
    SmallVectorImpl<OpFoldResult> &iterDomainSizes,
    ArrayRef<mlir::InnerTileAlignment>) {
  return getIterationDomainTileFromOperandTiles(b, operandNumbers, allOffsets,
                                                allSizes, iterDomainOffsets,
                                                iterDomainSizes);
}

namespace {
struct ConcatOpTilingExternalModel
    : TilingInterface::ExternalModel<ConcatOpTilingExternalModel,
                                     tensor::ConcatOp> {
  SmallVector<utils::IteratorType> getLoopIteratorTypes(Operation *op) const {
    auto concatOp = cast<tensor::ConcatOp>(op);
    SmallVector<utils::IteratorType> iteratorTypes(
        concatOp.getResultType().getRank(), utils::IteratorType::parallel);
    iteratorTypes[concatOp.getDim()] = utils::IteratorType::reduction;
    return iteratorTypes;
  }
  SmallVector<Range> getIterationDomain(Operation *op,
                                        OpBuilder &builder) const;
  FailureOr<TilingResult>
  getTiledImplementation(Operation *op, OpBuilder &builder,
                         ArrayRef<OpFoldResult> offsets,
                         ArrayRef<OpFoldResult> sizes) const;

  FailureOr<TilingResult>
  generateResultTileValue(Operation *op, OpBuilder &builder,
                          unsigned resultNumber, ArrayRef<OpFoldResult> offsets,
                          ArrayRef<OpFoldResult> sizes) const {
    return getTiledImplementation(op, builder, offsets, sizes);
  }

  FailureOr<TilingResult> getTiledImplementation(
      Operation *op, OpBuilder &builder, ArrayRef<OpFoldResult> offsets,
      ArrayRef<OpFoldResult> sizes, ArrayRef<mlir::InnerTileAlignment>) const {
    return getTiledImplementation(op, builder, offsets, sizes);
  }

  FailureOr<TilingResult>
  generateResultTileValue(Operation *op, OpBuilder &builder,
                          unsigned resultNumber, ArrayRef<OpFoldResult> offsets,
                          ArrayRef<OpFoldResult> sizes,
                          ArrayRef<mlir::InnerTileAlignment>) const {
    return generateResultTileValue(op, builder, resultNumber, offsets, sizes);
  }

  LogicalResult
  getResultTilePosition(Operation *op, OpBuilder &builder,
                        unsigned resultNumber, ArrayRef<OpFoldResult> offsets,
                        ArrayRef<OpFoldResult> sizes,
                        SmallVector<OpFoldResult> &resultOffsets,
                        SmallVector<OpFoldResult> &resultSizes) const {
    resultOffsets.assign(offsets.begin(), offsets.end());
    resultSizes.assign(sizes.begin(), sizes.end());
    return success();
  }
};
} // namespace

SmallVector<Range>
ConcatOpTilingExternalModel::getIterationDomain(Operation *op,
                                                OpBuilder &builder) const {
  auto concatOp = cast<tensor::ConcatOp>(op);
  OpFoldResult zero = builder.getIndexAttr(0);
  OpFoldResult one = builder.getIndexAttr(1);

  ReifiedRankedShapedTypeDims reifiedDims;
  LogicalResult result = concatOp.reifyResultShapes(builder, reifiedDims);
  (void)result;
  assert(succeeded(result));

  SmallVector<Range> ranges;
  for (auto dim : llvm::seq<int64_t>(0, concatOp.getResultType().getRank())) {
    ranges.push_back(Range{zero, reifiedDims[0][dim], one});
    continue;
  }

  return ranges;
}

FailureOr<TilingResult> ConcatOpTilingExternalModel::getTiledImplementation(
    Operation *op, OpBuilder &builder, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes) const {
  auto concatOp = cast<tensor::ConcatOp>(op);
  Location loc = concatOp.getLoc();
  uint64_t dim = concatOp.getDim();
  RankedTensorType resultType = concatOp.getResultType();

  // Don't support tiling along the concated dimension.
  std::optional<int64_t> concatDimSize = getConstantIntValue(sizes[dim]);
  if (!concatDimSize || concatDimSize.value() != resultType.getDimSize(dim)) {
    return failure();
  }

  // Collect slices of each operand.
  SmallVector<OpFoldResult> strides(resultType.getRank(),
                                    builder.getI64IntegerAttr(1));
  SmallVector<OpFoldResult> operandSizes(sizes);
  operandSizes[dim] = getDim(builder, loc, concatOp.getOperand(0), dim);
  SmallVector<Operation *> slices =
      llvm::map_to_vector(concatOp.getOperands(), [&](Value operand) {
        return getSlice(builder, loc, operand, offsets, operandSizes, strides);
      });
  SmallVector<Value> sliceValues = llvm::map_to_vector(
      slices, [](Operation *slice) -> Value { return slice->getResult(0); });

  // Create tiled `tensor.concat`.
  SmallVector<int64_t> resultStaticShape;
  SmallVector<Value> resultDynamicShape;
  dispatchIndexOpFoldResults(sizes, resultDynamicShape, resultStaticShape);
  auto newResultType = resultType.clone(resultStaticShape);
  Operation *tiledConcatOp =
      mlir::clone(builder, concatOp.getOperation(), TypeRange{newResultType},
                  ValueRange(sliceValues));
  return TilingResult{
      {tiledConcatOp}, SmallVector<Value>(tiledConcatOp->getResults()), slices};
}

void registerConcatOpTilingInterfaceExternalModel(MLIRContext *ctx) {
  tensor::ConcatOp::attachInterface<ConcatOpTilingExternalModel>(*ctx);
}

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