// Copyright 2022 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/LinalgExtDialect.h"
#include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.h"
#include "iree/compiler/Dialect/LinalgExt/Transforms/Passes.h"
#include "iree/compiler/Dialect/LinalgExt/Transforms/Transforms.h"
#include "iree/compiler/Dialect/LinalgExt/Utils/Utils.h"
#include "llvm/ADT/STLExtras.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/AffineExpr.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"

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

#define GEN_PASS_DEF_TOPKSPLITREDUCTIONPASS
#include "iree/compiler/Dialect/LinalgExt/Transforms/Passes.h.inc"

namespace {

// Marker used as attribute the depth of the split reduction transformations.
const StringLiteral kSplitReductionDepthMarker = "__split_reduction_depth__";

SmallVector<int64_t> getExpandedShape(ArrayRef<int64_t> shape,
                                      int64_t splitReductionRatio,
                                      int64_t splitDimParallel) {
  SmallVector<int64_t> ans;
  ans.reserve(shape.size() + 1);
  ans.assign(shape.begin(), shape.end());
  ans[splitDimParallel] = splitReductionRatio;
  ans.insert(std::next(ans.begin(), splitDimParallel + 1),
             shape[splitDimParallel] / splitReductionRatio);

  return ans;
}

SmallVector<int64_t> getCollapsedShape(ArrayRef<int64_t> shape,
                                       int64_t splitReductionRatio, int64_t k,
                                       int64_t targetDim) {
  SmallVector<int64_t> ans(shape);
  ans[targetDim] = k * splitReductionRatio;
  return ans;
}

SmallVector<ReassociationIndices>
getReassociationIndices(int64_t rank, int64_t splitDimParallel) {
  SmallVector<ReassociationIndices> reassociationIndices;
  for (int i = 0; i < rank; ++i) {
    if (i < splitDimParallel) {
      reassociationIndices.push_back({i});
    } else if (i == splitDimParallel) {
      reassociationIndices.push_back({i, i + 1});
    } else if (i > splitDimParallel) {
      reassociationIndices.push_back({i + 1});
    }
  }
  return reassociationIndices;
}

LogicalResult shouldParallelTopk(iree_compiler::IREE::LinalgExt::TopkOp topkOp,
                                 RewriterBase &rewriter, int64_t kDimOrig,
                                 int64_t splitReductionRatio,
                                 int64_t splitReductionDepth) {
  // Determine if we should split the reduction. Requires aligned static shapes
  // and no input indices.
  auto valuesOriginalType = topkOp.getInputType();
  if (valuesOriginalType.isDynamicDim(kDimOrig)) {
    return rewriter.notifyMatchFailure(topkOp,
                                       "cannot split dynamic dimension");
  }
  if (topkOp.getIndices() && splitReductionDepth == 0) {
    return rewriter.notifyMatchFailure(
        topkOp, "input indices aren't supported for first split");
  }
  if (splitReductionRatio <= 1) {
    return rewriter.notifyMatchFailure(topkOp, "reduction ratio <= 1");
  }
  if (valuesOriginalType.getDimSize(kDimOrig) % splitReductionRatio != 0) {
    return rewriter.notifyMatchFailure(
        topkOp,
        "reduction dimension must be perfectly aligned to (divisible by) the "
        "split ratio");
  }
  return success();
}

// Creates the first phase of the topk split reduction by reshaping the input
// into parallel computations then feeding them into a topk op.
iree_compiler::IREE::LinalgExt::TopkOp
computeParallelTopk(Location loc, RewriterBase &rewriter,
                    iree_compiler::IREE::LinalgExt::TopkOp topkOp,
                    ArrayRef<ReassociationIndices> reassociationIndices,
                    int64_t splitReductionRatio, int64_t splitDimParallel,
                    int64_t kDimParallel, int64_t kSize) {
  Value valuesOrig = topkOp.getValues();
  auto valuesOriginalType = cast<ShapedType>(valuesOrig.getType());
  Type valueElementType = valuesOriginalType.getElementType();
  Type indicesElementType =
      cast<ShapedType>(topkOp.getResultTypes()[1]).getElementType();

  SmallVector<int64_t> expandedShape = getExpandedShape(
      valuesOriginalType.getShape(), splitReductionRatio, splitDimParallel);
  auto valuesExpandedType =
      RankedTensorType::get(expandedShape, valueElementType);

  // Expand input values shape for parallel processing
  Value valuesExpanded = tensor::ExpandShapeOp::create(
      rewriter, loc, valuesExpandedType, valuesOrig, reassociationIndices);

  // Expand input indices shape for parallel processing if they exist
  Value indicesExpanded;
  if (Value inputIndices = topkOp.getIndices()) {
    Type indicesExpandedType =
        RankedTensorType::get(expandedShape, indicesElementType);
    indicesExpanded = tensor::ExpandShapeOp::create(
        rewriter, loc, indicesExpandedType, inputIndices, reassociationIndices);
  }

  // Define the expanded output types
  SmallVector<int64_t> expandedResultShape = expandedShape;
  expandedResultShape[kDimParallel] = kSize;
  auto outputValuesExpandedType =
      RankedTensorType::get(expandedResultShape, valueElementType);
  auto outputIndicesExpandedType =
      RankedTensorType::get(expandedResultShape, indicesElementType);

  // Initialize the expanded output values
  SmallVector<Value> dynSizes;
  for (auto i : llvm::seq<int64_t>(0, valuesExpandedType.getRank())) {
    if (valuesExpandedType.isDynamicDim(i)) {
      dynSizes.push_back(
          tensor::DimOp::create(rewriter, loc, valuesExpanded, i));
    }
  }
  Value emptyTensorOutputValues = tensor::EmptyOp::create(
      rewriter, loc, outputValuesExpandedType.getShape(), valueElementType,
      dynSizes);
  Value emptyTensorOutputIndices = tensor::EmptyOp::create(
      rewriter, loc, outputIndicesExpandedType.getShape(), indicesElementType,
      dynSizes);

  // Initialize indices to positive infinity and values to negative infinity
  // for a top (maxk) comparison.
  TypedAttr negInfAttr;
  if (auto intType = dyn_cast<IntegerType>(valueElementType)) {
    negInfAttr = rewriter.getIntegerAttr(
        intType, APInt::getSignedMinValue(intType.getWidth()));
  } else {
    auto negApFloat =
        APFloat::getInf(cast<FloatType>(valueElementType).getFloatSemantics(),
                        /*Negative=*/true);
    negInfAttr = rewriter.getFloatAttr(valueElementType, negApFloat);
  }
  Value negInf = arith::ConstantOp::create(rewriter, loc, negInfAttr);
  TypedAttr posInfAttr =
      rewriter.getIntegerAttr(indicesElementType, APInt::getSignedMaxValue(32));
  Value posInf = arith::ConstantOp::create(rewriter, loc, posInfAttr);
  Value negInfTensor =
      linalg::FillOp::create(rewriter, loc, negInf, emptyTensorOutputValues)
          .result();
  Value posInfTensor =
      linalg::FillOp::create(rewriter, loc, posInf, emptyTensorOutputIndices)
          .result();

  SmallVector<Type> parallelTopkResultTypes = {outputValuesExpandedType,
                                               outputIndicesExpandedType};
  SmallVector<Value> parallelTopkIns = {valuesExpanded};
  if (indicesExpanded) {
    parallelTopkIns.push_back(indicesExpanded);
  }
  SmallVector<Value> parallelTopkOuts = {negInfTensor, posInfTensor};

  // Parallel topk
  auto parallelTopkOp = iree_compiler::IREE::LinalgExt::TopkOp::create(
      rewriter, loc, parallelTopkResultTypes, parallelTopkIns[0],
      /*indices=*/parallelTopkIns.size() > 1 ? parallelTopkIns[1] : Value(),
      parallelTopkOuts[0], parallelTopkOuts[1], kDimParallel);
  rewriter.cloneRegionBefore(topkOp.getRegion(), parallelTopkOp.getRegion(),
                             parallelTopkOp.getRegion().end());

  return parallelTopkOp;
}

// Update the output indices from the parallel TopK with the correct offsets.
// Each parallel computation uses implicit indices (starting from 0) during
// selection, but the values are part of the large input space split into M =
// splitReductionFn() ways. The following linalg.generic adds the appropriate
// offset to reflect to values original position. "Updated pos" = "initial
// pos" + "splitDimParallel size * "splitDimParallel index"
Value offsetParallelIndices(Location loc, RewriterBase &rewriter,
                            Value parallelIndices, int64_t kDimParallelSize,
                            int64_t splitDimParallel) {
  auto parallelIndicesType = cast<ShapedType>(parallelIndices.getType());
  size_t parallelIndicesRank = parallelIndicesType.getRank();
  AffineMap mapIdentity = rewriter.getMultiDimIdentityMap(parallelIndicesRank);
  SmallVector<AffineMap> indexingMaps = {mapIdentity};
  SmallVector<utils::IteratorType> iterators(parallelIndicesRank,
                                             utils::IteratorType::parallel);
  Value mSplitVal = arith::ConstantIntOp::create(
      rewriter, loc, parallelIndicesType.getElementType(), kDimParallelSize);
  return linalg::GenericOp::create(
             rewriter, loc,
             /*resultType=*/parallelIndicesType,
             /*inputs=*/ValueRange{},
             /*outputs=*/ValueRange{parallelIndices}, indexingMaps, iterators,
             [&](OpBuilder &b, Location loc, ValueRange args) {
               Value splitIndex =
                   linalg::IndexOp::create(b, loc, splitDimParallel);
               Value splitIndexInt = arith::IndexCastOp::create(
                   b, loc, parallelIndicesType.getElementType(), splitIndex);
               Value mOffset =
                   arith::MulIOp::create(b, loc, mSplitVal, splitIndexInt);
               Value updatedParallelIndex =
                   arith::AddIOp::create(b, loc, mOffset, args[0]);
               linalg::YieldOp::create(b, loc, updatedParallelIndex);
             })
      .getResult(0);
}

// Creates the second phase of the topk split reduction by collapsing output
// from parallel topk and computing the final combined result.
TopkOp computeReductionTopk(Location loc, RewriterBase &rewriter, TopkOp topkOp,
                            TopkOp parallelTopkOp, Value updatedParallelIndices,
                            ArrayRef<ReassociationIndices> reassociationIndices,
                            int64_t splitReductionRatio, int64_t kDimOrig,
                            int64_t kSize) {
  Value valuesOrig = topkOp.getValues();
  auto valuesOriginalType = cast<ShapedType>(valuesOrig.getType());
  Type valueElementType = valuesOriginalType.getElementType();
  Type indicesElementType =
      cast<ShapedType>(topkOp.getResultTypes()[1]).getElementType();

  // Define the collapsed input shapes
  SmallVector<int64_t> collapsedShape = getCollapsedShape(
      valuesOriginalType.getShape(), splitReductionRatio, kSize, kDimOrig);
  auto valuesCollapsedType =
      RankedTensorType::get(collapsedShape, valueElementType);
  auto indicesCollapsedType =
      RankedTensorType::get(collapsedShape, indicesElementType);

  // Collapse collapse parallel output for the input of final reduction
  Value valuesCollapsed = tensor::CollapseShapeOp::create(
      rewriter, loc, valuesCollapsedType, parallelTopkOp.getResults()[0],
      reassociationIndices);
  Value indicesCollapsed = tensor::CollapseShapeOp::create(
      rewriter, loc, indicesCollapsedType, updatedParallelIndices,
      reassociationIndices);

  // Combined final topk
  auto reductionTopkOp = iree_compiler::IREE::LinalgExt::TopkOp::create(
      rewriter, loc, topkOp->getResultTypes(), valuesCollapsed,
      indicesCollapsed, topkOp.getOutputValues(), topkOp.getOutputIndices(),
      kDimOrig);
  rewriter.cloneRegionBefore(topkOp.getRegion(), reductionTopkOp.getRegion(),
                             reductionTopkOp.getRegion().end());
  return reductionTopkOp;
}

int64_t getSplitReductionDepth(TopkOp topkOp) {
  auto attr =
      topkOp->template getAttrOfType<IntegerAttr>(kSplitReductionDepthMarker);
  if (attr) {
    return attr.getInt();
  } else {
    return 0;
  }
}

void setSplitReductionDepth(TopkOp topkOp, RewriterBase &rewriter,
                            int64_t depth) {
  topkOp->setAttr(kSplitReductionDepthMarker,
                  rewriter.getI64IntegerAttr(depth));
}
} // namespace

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

namespace {
struct TopkSplitReductionPass final
    : impl::TopkSplitReductionPassBase<TopkSplitReductionPass> {
  using Base::Base;

  void getDependentDialects(DialectRegistry &registry) const override {
    registry
        .insert<linalg::LinalgDialect, arith::ArithDialect, scf::SCFDialect>();
  }

  void runOnOperation() override {
    if (splitRatios.empty()) {
      return;
    }

    TopkSplitReductionControlFn splitReductionFn =
        [&](int64_t splitReductionDepth) -> int64_t {
      SmallVector<int64_t> reductionRatios(splitRatios.begin(),
                                           splitRatios.end());
      if (splitReductionDepth >= reductionRatios.size()) {
        return -1;
      } else {
        return reductionRatios[splitReductionDepth];
      }
    };

    IRRewriter rewriter(&getContext());
    mlir::FunctionOpInterface funcOp = getOperation();
    SmallVector<LinalgExt::TopkOp> topkCandidates;
    funcOp->walk([&](LinalgExt::TopkOp op) { topkCandidates.push_back(op); });
    for (auto op : topkCandidates) {
      (void)splitReduction(rewriter, op, splitReductionFn);
    }
  }
};
} // namespace

LogicalResult
splitReduction(RewriterBase &rewriter, LinalgExt::TopkOp topkOp,
               const TopkSplitReductionControlFn &splitReductionFn) {
  Location loc = topkOp.getLoc();
  OpBuilder::InsertionGuard guard(rewriter);
  rewriter.setInsertionPoint(topkOp);
  // Original reduction dimension used for the final combined reduction
  int64_t kDimOrig = topkOp.getDimension();
  // For parallel topk: the dimension that we compute parallel reductions
  int64_t splitDimParallel = kDimOrig;
  // For parallel topk: the dimension that we reduce
  int64_t kDimParallel = kDimOrig + 1;
  int64_t kSize =
      cast<ShapedType>(topkOp.getResult(0).getType()).getDimSize(kDimOrig);
  int64_t splitReductionDepth = getSplitReductionDepth(topkOp);
  int64_t splitReductionRatio = splitReductionFn(splitReductionDepth);
  SmallVector<ReassociationIndices> reassociationIndices =
      getReassociationIndices(topkOp.getInputRank(), splitDimParallel);

  // Determine if should compute parallel topk
  LogicalResult shouldParallelTopkResult = shouldParallelTopk(
      topkOp, rewriter, kDimOrig, splitReductionRatio, splitReductionDepth);
  if (shouldParallelTopkResult.failed()) {
    return shouldParallelTopkResult;
  }

  // Topk parallel reduction
  TopkOp parallelTopkOp = computeParallelTopk(
      loc, rewriter, topkOp, reassociationIndices, splitReductionRatio,
      splitDimParallel, kDimParallel, kSize);

  // Update parallel indices to correct offsets if input indices weren't
  // provided. If input indices were provided, no offsetting is needed as
  // original original indices are already known.
  Value updatedParallelIndices = parallelTopkOp.getResult(1);
  if (!topkOp.getIndices()) {
    Value parallelIndices = parallelTopkOp.getResult(1);
    SmallVector<int64_t> expandedShape = getExpandedShape(
        cast<ShapedType>(topkOp.getValues().getType()).getShape(),
        splitReductionRatio, splitDimParallel);
    int64_t kDimParallelSize = expandedShape[kDimParallel];
    updatedParallelIndices = offsetParallelIndices(
        loc, rewriter, parallelIndices, kDimParallelSize, splitDimParallel);
  }

  // Topk final reduction
  TopkOp reductionTopkOp = computeReductionTopk(
      loc, rewriter, topkOp, parallelTopkOp, updatedParallelIndices,
      reassociationIndices, splitReductionRatio, kDimOrig, kSize);

  // Replace and update result
  rewriter.replaceOp(topkOp, reductionTopkOp.getResults());
  setSplitReductionDepth(reductionTopkOp, rewriter, splitReductionDepth + 1);

  // Recursively apply split reduction until reaching the target depth.
  if (failed(splitReduction(rewriter, reductionTopkOp, splitReductionFn))) {
    reductionTopkOp->removeAttr(LinalgExt::kSplitReductionDepthMarker);
  }

  return success();
}

static Value expandValue(OpBuilder &builder, Location loc, Value value,
                         RankedTensorType expandedType) {
  RankedTensorType originalType = cast<RankedTensorType>(value.getType());
  if (originalType == expandedType) {
    return value;
  }
  auto reassociation =
      getReassociationIndicesForReshape(originalType, expandedType);
  assert(reassociation && "failed to infer reassociation indices from types");
  return tensor::ExpandShapeOp::create(builder, loc, expandedType, value,
                                       *reassociation);
}

// Returns an expanded input AffineMap by splitting the given reduction
// dimension into [parallel, reduction] or [reduction, parallel], inserting a
// new dimension at `insertSplitDim`.
static AffineMap getExpandedInputIndexingMap(AffineMap oldMap,
                                             unsigned reductionDim,
                                             unsigned insertSplitDim,
                                             bool innerParallel,
                                             MLIRContext *ctx) {
  SmallVector<AffineExpr> exprs;
  for (unsigned idx = 0; idx < oldMap.getNumResults(); ++idx) {
    unsigned dim = oldMap.getDimPosition(idx);
    if (dim != reductionDim) {
      unsigned shifted = (dim < insertSplitDim) ? dim : dim + 1;
      exprs.push_back(getAffineDimExpr(shifted, ctx));
      continue;
    }
    // Expand the reduction dimension into [reduction, parallel] or [parallel,
    // reduction].
    if (innerParallel) {
      exprs.push_back(
          getAffineDimExpr(dim < insertSplitDim ? dim : dim + 1, ctx));
      exprs.push_back(getAffineDimExpr(insertSplitDim, ctx));
    } else {
      exprs.push_back(getAffineDimExpr(insertSplitDim, ctx));
      exprs.push_back(
          getAffineDimExpr(dim < insertSplitDim ? dim : dim + 1, ctx));
    }
  }

  return AffineMap::get(oldMap.getNumDims() + 1, oldMap.getNumSymbols(), exprs,
                        ctx);
}

// Returns an expanded output AffineMap with a dimension inserted at
// `insertSplitDim`.
static AffineMap getExpandedOutputIndexingMap(AffineMap oldMap,
                                              unsigned insertSplitDim,
                                              MLIRContext *ctx) {
  SmallVector<AffineExpr> exprs;
  for (unsigned idx = 0; idx <= oldMap.getNumResults(); ++idx) {
    if (idx == insertSplitDim) {
      exprs.push_back(getAffineDimExpr(insertSplitDim, ctx));
    }
    if (idx < oldMap.getNumResults()) {
      unsigned dim = oldMap.getDimPosition(idx);
      unsigned shifted = (dim < insertSplitDim) ? dim : dim + 1;
      exprs.push_back(getAffineDimExpr(shifted, ctx));
    }
  }

  return AffineMap::get(oldMap.getNumDims() + 1, oldMap.getNumSymbols(), exprs,
                        ctx);
}

static Value getSplitReductionInit(OpBuilder &builder, Location loc,
                                   Value origInit, Attribute identityAttr,
                                   int64_t insertDimSize,
                                   unsigned insertSplitIndex) {
  Type elemType = getElementTypeOrSelf(origInit.getType());
  SmallVector<OpFoldResult> shape =
      tensor::getMixedSizes(builder, loc, origInit);
  shape.insert(shape.begin() + insertSplitIndex,
               builder.getIndexAttr(insertDimSize));
  Value identityVal = arith::ConstantOp::create(builder, loc, elemType,
                                                cast<TypedAttr>(identityAttr));
  Value empty = tensor::EmptyOp::create(builder, loc, shape, elemType);
  return linalg::FillOp::create(builder, loc, identityVal, empty).getResult(0);
}

static void setPreservedAttrs(Operation *op, DictionaryAttr attrs) {
  if (attrs) {
    op->setAttrs(attrs);
  }
}

struct ArgmaxCombinerResults {
  Value maxValue;
  Value selectedIndex;
};

static ArgmaxCombinerResults buildArgmaxCombiner(OpBuilder &b, Location loc,
                                                 const ArgmaxCombinerInfo &info,
                                                 Value val, Value idx,
                                                 Value outVal, Value outIdx) {
  // Split reduction creates new reductions over local and global indices, so
  // rebuild the recognized argmax semantics instead of cloning the original
  // combiner ops. Preserve the matched combiner attributes explicitly so
  // frontend flags such as arith fast-math follow the regenerated ops.
  if (info.kind == ArgmaxKind::Canonical) {
    auto maxOp = arith::MaximumFOp::create(b, loc, val, outVal);
    setPreservedAttrs(maxOp.getOperation(), info.maximumAttrs);
    auto cmpOp =
        arith::CmpFOp::create(b, loc, arith::CmpFPredicate::OGT, val, outVal);
    setPreservedAttrs(cmpOp.getOperation(), info.greaterThanCmpAttrs);
    auto selectOp = arith::SelectOp::create(b, loc, cmpOp, idx, outIdx);
    setPreservedAttrs(selectOp.getOperation(), info.indexSelectAttrs);
    return {/*maxValue=*/maxOp.getResult(),
            /*selectedIndex=*/selectOp.getResult()};
  }

  auto greater =
      arith::CmpFOp::create(b, loc, arith::CmpFPredicate::OGT, val, outVal);
  setPreservedAttrs(greater.getOperation(), info.greaterThanCmpAttrs);
  auto isNan =
      arith::CmpFOp::create(b, loc, arith::CmpFPredicate::UNE, val, val);
  setPreservedAttrs(isNan.getOperation(), info.isNanCmpAttrs);
  Value valueCond = arith::OrIOp::create(b, loc, greater, isNan);
  auto equal =
      arith::CmpFOp::create(b, loc, arith::CmpFPredicate::OEQ, val, outVal);
  setPreservedAttrs(equal.getOperation(), info.equalCmpAttrs);
  Value lowerIndex =
      arith::CmpIOp::create(b, loc, arith::CmpIPredicate::slt, idx, outIdx);
  Value tieCond = arith::AndIOp::create(b, loc, equal, lowerIndex);
  Value indexCond = arith::OrIOp::create(b, loc, valueCond, tieCond);
  auto valueSelect = arith::SelectOp::create(b, loc, valueCond, val, outVal);
  setPreservedAttrs(valueSelect.getOperation(), info.valueSelectAttrs);
  auto indexSelect = arith::SelectOp::create(b, loc, indexCond, idx, outIdx);
  setPreservedAttrs(indexSelect.getOperation(), info.indexSelectAttrs);
  return {/*maxValue=*/valueSelect.getResult(),
          /*selectedIndex=*/indexSelect.getResult()};
}

FailureOr<linalg::SplitReductionResult>
splitArgmaxReduction(RewriterBase &rewriter, linalg::GenericOp genericOp,
                     linalg::ControlSplitReductionFn controlSplitReductionFn) {
  std::optional<ArgmaxCombinerInfo> argmaxInfo =
      getArgmaxCombinerInfo(genericOp);
  if (!argmaxInfo) {
    return rewriter.notifyMatchFailure(
        genericOp, "operation is not a recognized argmax pattern");
  }
  ArgmaxCombinerInfo info = *argmaxInfo;

  OpBuilder::InsertionGuard guard(rewriter);
  rewriter.setInsertionPoint(genericOp);
  Location loc = genericOp->getLoc();

  linalg::SplitReductionOptions control = controlSplitReductionFn(genericOp);
  if (control.innerParallel) {
    // TODO(Bangtian): Currently bail out when innerParallel is true.
    // Much of the logic is already in place based on the upstream
    // implementation, but full support requires early detection and additional
    // handling before performing the split. Will revisit when support becomes
    // necessary.
    return rewriter.notifyMatchFailure(
        genericOp, "innerParallel split is not supported yet");
  }

  int64_t ratio = control.ratio;
  unsigned insertSplitIndex = control.index;
  unsigned insertSplitDimension = control.index;
  if (ratio <= 1) {
    return rewriter.notifyMatchFailure(
        genericOp, "split ratio needs to be greater than 1");
  }

  SmallVector<unsigned> dims;
  genericOp.getReductionDims(dims);
  unsigned reductionDim = dims[0];
  if (control.innerParallel) {
    insertSplitDimension = reductionDim + 1;
  }

  SmallVector<int64_t> loopRanges = genericOp.getStaticLoopRanges();
  int64_t reductionDimSize = loopRanges[reductionDim];
  if (ShapedType::isDynamic(reductionDimSize) ||
      reductionDimSize % ratio != 0) {
    return rewriter.notifyMatchFailure(
        genericOp, "Reduction dimension not divisible by split ratio");
  }

  if (insertSplitIndex >
      genericOp.getShape(genericOp.getDpsInitOperand(0)).size()) {
    return rewriter.notifyMatchFailure(genericOp,
                                       "Insert dimension position too large "
                                       "compared to intermediate tensor size");
  }

  std::optional<TypedAttr> identity;
  if (info.kind == ArgmaxKind::Canonical) {
    identity = arith::getNeutralElement(info.maximumOp);
    if (!identity) {
      return rewriter.notifyMatchFailure(
          genericOp, "unknown identity value for the reduction");
    }
  } else {
    auto valueType =
        dyn_cast<FloatType>(genericOp.getRegionOutputArgs()[0].getType());
    if (!valueType) {
      return rewriter.notifyMatchFailure(
          genericOp, "expected floating point value output for argmax");
    }
    identity =
        FloatAttr::get(valueType, APFloat::getInf(valueType.getFloatSemantics(),
                                                  /*Negative=*/true));
  }

  SmallVector<Value> newInputs;
  for (OpOperand *operand : genericOp.getDpsInputOperands()) {
    AffineMap map = genericOp.getMatchingIndexingMap(operand);
    ArrayRef<int64_t> oldShape = genericOp.getShape(operand);
    Type elementType = getElementTypeOrSelf(operand->get().getType());
    SmallVector<int64_t> expandedShape;
    for (auto idx : llvm::seq<unsigned>(0, map.getNumResults())) {
      unsigned dim = map.getDimPosition(idx);
      if (dim != reductionDim) {
        expandedShape.push_back(oldShape[idx]);
        continue;
      }
      int64_t orig = oldShape[idx];
      int64_t outer = orig / ratio;
      if (control.innerParallel) {
        expandedShape.push_back(ratio); // reduction.
        expandedShape.push_back(outer); // parallel.
      } else {
        expandedShape.push_back(outer); // parallel.
        expandedShape.push_back(ratio); // reduction.
      }
    }

    auto expandedType = RankedTensorType::get(expandedShape, elementType);
    Value expanded = expandValue(rewriter, loc, operand->get(), expandedType);
    newInputs.push_back(expanded);
  }

  // The total number of output elements along this new dimension is
  // reductionDimSize / ratio.
  int64_t outputDimSize = reductionDimSize / ratio;
  OpOperand *valueInit = genericOp.getDpsInitOperand(0);
  OpOperand *indexInit = genericOp.getDpsInitOperand(1);
  // Value identity.
  Value identityValue =
      getSplitReductionInit(rewriter, loc, valueInit->get(), *identity,
                            outputDimSize, insertSplitIndex);
  // Index identity.
  Type indexElemType = genericOp.getRegionOutputArgs()[1].getType();
  Value identityIndex = getSplitReductionInit(
      rewriter, loc, indexInit->get(), rewriter.getZeroAttr(indexElemType),
      outputDimSize, insertSplitIndex);

  SmallVector<utils::IteratorType> newIteratorTypes =
      genericOp.getIteratorTypesArray();
  newIteratorTypes.insert(newIteratorTypes.begin() + insertSplitDimension,
                          utils::IteratorType::parallel);

  SmallVector<AffineMap> newMaps;
  unsigned numInputs = genericOp.getNumDpsInputs();
  for (OpOperand &operand : genericOp->getOpOperands()) {
    AffineMap map = genericOp.getMatchingIndexingMap(&operand);
    if (operand.getOperandNumber() < numInputs) {
      newMaps.push_back(getExpandedInputIndexingMap(
          map, reductionDim, insertSplitDimension, control.innerParallel,
          rewriter.getContext()));
    } else {
      newMaps.push_back(getExpandedOutputIndexingMap(map, insertSplitDimension,
                                                     rewriter.getContext()));
    }
  }

  // Step 1: Create an argmax that performs a partial reduction over the split
  // (tile) dimension. The result yields the local maximum values and their
  // corresponding local indices within each tile. These local indices will be
  // adjusted to global indices in step 2.
  auto partialArgmax = linalg::GenericOp::create(
      rewriter, loc,
      TypeRange{identityValue.getType(), identityIndex.getType()}, newInputs,
      ValueRange{identityValue, identityIndex}, newMaps, newIteratorTypes,
      [reductionDim, info](OpBuilder &b, Location loc, ValueRange args) {
        Value in = args[0];
        Value outVal = args[1];
        Value outIdx = args[2];
        Value reductionIdx = linalg::IndexOp::create(b, loc, reductionDim + 1);
        if (outIdx.getType() != reductionIdx.getType()) {
          reductionIdx = arith::IndexCastOp::create(b, loc, outIdx.getType(),
                                                    reductionIdx);
        }
        Value inCast = in;
        Type inType = in.getType();
        Type outType = outVal.getType();
        unsigned inBitWidth = inType.getIntOrFloatBitWidth();
        unsigned outBitWidth = outType.getIntOrFloatBitWidth();

        if (outBitWidth > inBitWidth) {
          inCast = arith::ExtFOp::create(b, loc, outVal.getType(), in);
        } else if (outBitWidth < inBitWidth) {
          inCast = arith::TruncFOp::create(b, loc, outVal.getType(), in);
        }

        ArgmaxCombinerResults results = buildArgmaxCombiner(
            b, loc, info, inCast, reductionIdx, outVal, outIdx);
        linalg::YieldOp::create(
            b, loc, ValueRange{results.maxValue, results.selectedIndex});
      });

  // Step 2: Final reduction that computes global indices and selects the one
  // corresponding to the maximum value.
  Value tileSize = arith::ConstantIndexOp::create(rewriter, loc, ratio);
  unsigned intermRank =
      cast<RankedTensorType>(identityValue.getType()).getRank();
  AffineMap valueMap = rewriter.getMultiDimIdentityMap(intermRank);
  AffineMap indexMap = valueMap;

  SmallVector<AffineExpr> resultExprs;
  SmallVector<utils::IteratorType> reductionIteratorTypes(
      intermRank, utils::IteratorType::parallel);
  reductionIteratorTypes[insertSplitIndex] = utils::IteratorType::reduction;
  for (auto i : llvm::seq<unsigned>(0, intermRank)) {
    if (i != insertSplitIndex) {
      resultExprs.push_back(rewriter.getAffineDimExpr(i));
    }
  }

  AffineMap outputMap =
      AffineMap::get(intermRank, 0, resultExprs, rewriter.getContext());
  SmallVector<AffineMap> finalReductionMaps = {valueMap, indexMap, outputMap,
                                               outputMap};
  auto finalReduction = linalg::GenericOp::create(
      rewriter, loc, genericOp.getResultTypes(),
      ValueRange{partialArgmax.getResult(0), partialArgmax.getResult(1)},
      genericOp.getDpsInits(), finalReductionMaps, reductionIteratorTypes,
      [tileSize, insertSplitDimension, info](OpBuilder &b, Location loc,
                                             ValueRange inputs) {
        Value val = inputs[0];
        Value local = inputs[1];
        Value outVal = inputs[2];
        Value outIdx = inputs[3];
        Value outer = linalg::IndexOp::create(b, loc, insertSplitDimension);
        Value offset = arith::MulIOp::create(b, loc, outer, tileSize);
        if (offset.getType() != local.getType()) {
          offset = arith::IndexCastOp::create(b, loc, local.getType(), offset);
        }
        // gidx = outer * ratio + local.
        Value gidx = arith::AddIOp::create(b, loc, offset, local);
        ArgmaxCombinerResults results =
            buildArgmaxCombiner(b, loc, info, val, gidx, outVal, outIdx);
        linalg::YieldOp::create(
            b, loc, ValueRange{results.maxValue, results.selectedIndex});
      });

  rewriter.replaceOp(genericOp, finalReduction.getResults());
  // Init or alloc and fillOp are not applicable for argmax op; set to nullptr.
  return linalg::SplitReductionResult{
      /*initOrAlloc=*/nullptr, /*fillOp=*/nullptr,
      cast<linalg::LinalgOp>(partialArgmax.getOperation()), finalReduction};
}

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