Repository: iree-org/iree
Pull request #20717: [Codegen] split-k on argmax op
Author: login: bangtianliu
Created: 2025-05-03T00:11:20Z | Merged: 2025-05-15T17:46:18Z

This PR adds support for split reduction on the argmax operation.

The main question is: why couldn't upstream `mlir::linalg::splitReduction` be reused? The answer will also implicitly show what this PR is about.
The upstream mlir::linalg::splitReduction implementation is designed for single-result reductions (e.g., sum, max) and assumes a specific combiner behavior. In contrast, argmax is a two-result reduction, where both the value and the index must be tracked across the split. 

Supporting this would require significant changes compared to the upstream utility, The relevant changes made by this PR include:
- Handle two outputs in the reduction body and yield
- Preserve the value-index pairing semantics throughout the split and recombine phases
- Match a more complex combiner structure (maximumf + cmpf+ select) and use all three combiner ops to generate the corresponding split and combine logic.

Introduce two configurable options to control when split reduction should apply:
- `splitArgmaxReductionThreshold`: the minimum size of the reduction dimension required to trigger splitting (default: 128 * 1024).
- `splitArgmaxTileSize`: the size of the reduction tile in the split operation (default: 128).

Also add tests to verify that: 1) Split is correctly applied when the threshold is met; 2) Split is not applied when the reduction size is below the threshold.

Other changes from this PR:

- Format the utility function isArgmaxOp, move it from iree/compiler/src/iree/compiler/Codegen/Utils/Utils.cpp to compiler/src/iree/compiler/Dialect/LinalgExt/Utils/Utils.cpp
- Add support for dynamic dimensions in argmax, allowing parallel dimensions to be dynamic.

Issue: https://github.com/iree-org/iree/issues/20650


Diff:
diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/GPULowerToUKernels.cpp b/compiler/src/iree/compiler/Codegen/Common/GPU/GPULowerToUKernels.cpp
index 6dfb75a5060a..e9162f8d10e8 100644
--- a/compiler/src/iree/compiler/Codegen/Common/GPU/GPULowerToUKernels.cpp
+++ b/compiler/src/iree/compiler/Codegen/Common/GPU/GPULowerToUKernels.cpp
@@ -12,6 +12,7 @@
 #include "iree/compiler/Codegen/Dialect/GPU/IR/GPULoweringConfigUtils.h"
 #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUOps.h"
 #include "iree/compiler/Codegen/Utils/Utils.h"
+#include "iree/compiler/Dialect/LinalgExt/Utils/Utils.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
@@ -69,7 +70,7 @@ struct LowerArgmaxToUKernelPattern : OpRewritePattern<linalg::GenericOp> {
 
   LogicalResult matchAndRewrite(linalg::GenericOp op,
                                 PatternRewriter &rewriter) const override {
-    if (failed(isArgmaxOp(op))) {
+    if (!IREE::LinalgExt::isArgmaxOp(op)) {
       return failure();
     }
     FailureOr<IREE::Codegen::UKernelOpInterface> ukernelOp =
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/Utils/LLVMGPUSelectUKernels.cpp b/compiler/src/iree/compiler/Codegen/LLVMGPU/Utils/LLVMGPUSelectUKernels.cpp
index 4e1952dd104b..4954510b0a03 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/Utils/LLVMGPUSelectUKernels.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/Utils/LLVMGPUSelectUKernels.cpp
@@ -8,6 +8,7 @@
 #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUOps.h"
 #include "iree/compiler/Codegen/Utils/GPUUtils.h"
 #include "iree/compiler/Codegen/Utils/Utils.h"
+#include "iree/compiler/Dialect/LinalgExt/Utils/Utils.h"
 #include "iree/compiler/Utils/EmbeddedDataDirectory.h"
 #include "llvm/ExecutionEngine/ExecutionEngine.h"
 #include "llvm/ExecutionEngine/GenericValue.h"
@@ -52,7 +53,7 @@ getUKernelNameAndSuffixForMultiMma(IREE::GPU::MultiMmaOp op) {
 // Returns ukernel name and suffix for any op. Empty name = no ukernel.
 static UKernelNameAndSuffix getUKernelNameAndSuffix(Operation *op) {
   if (auto genericOp = dyn_cast<linalg::GenericOp>(op)) {
-    if (succeeded(isArgmaxOp(genericOp))) {
+    if (IREE::LinalgExt::isArgmaxOp(genericOp)) {
       return getUKernelNameAndSuffixForArgmax(genericOp);
     }
   } else if (auto multiMmaOp = dyn_cast<IREE::GPU::MultiMmaOp>(op)) {
diff --git a/compiler/src/iree/compiler/Codegen/Utils/Utils.cpp b/compiler/src/iree/compiler/Codegen/Utils/Utils.cpp
index 6a1bbc5802d0..13dfe14e5f8a 100644
--- a/compiler/src/iree/compiler/Codegen/Utils/Utils.cpp
+++ b/compiler/src/iree/compiler/Codegen/Utils/Utils.cpp
@@ -1317,104 +1317,6 @@ Operation *dropEncodingAndCloneOp(OpBuilder &builder, Operation *op,
       operands);
 }
 
-LogicalResult isArgmaxOp(linalg::GenericOp genericOp) {
-  // Check for 2 results(value, index), and 1 input
-  if (genericOp.getNumDpsInits() != 2) {
-    return failure();
-  }
-  if (genericOp.getNumDpsInputs() != 1) {
-    return failure();
-  }
-
-  // If max value is being used, it is not a pure argmax.
-  if (!genericOp.getResults()[0].use_empty()) {
-    return failure();
-  }
-
-  // Check that the rank is at least 3 and all loops are parallel
-  unsigned numLoops = genericOp.getNumLoops();
-  unsigned numParallelLoops = genericOp.getNumParallelLoops();
-
-  // Argmax will require 1D reduction.
-  if (numParallelLoops != (numLoops - 1)) {
-    return failure();
-  }
-  // TODO: Add better affine map checks.
-  auto indexing_maps = genericOp.getIndexingMapsArray();
-  if (!indexing_maps[0].isIdentity())
-    return failure();
-
-  // Check that initial value is negative Infinite.
-  // TODO: Move this check to ukernel once we implement
-  //       variant to handle non neg-Inf initial value.
-  Value initVal = genericOp.getDpsInitOperand(0)->get();
-  auto fillOp = initVal.getDefiningOp<linalg::FillOp>();
-  if (!fillOp)
-    return failure();
-  Value fillVal = fillOp.getDpsInputOperand(0)->get();
-  if (!matchPattern(fillVal, m_NegInfFloat()))
-    return failure();
-
-  // Work back from linalg.yield and check body of genericOp.
-  // The genericOp should yield the result of an arith.select,
-  // preceded by an arith.cmpf, arith.maximumf, and arith.extui
-  auto yieldOp = cast<linalg::YieldOp>(genericOp.getBody()->getTerminator());
-  Value producerOutput;
-  Operation *producer;
-
-  // Producer of linalg.yield 1st arg is arith.maximumf
-  {
-    producerOutput = yieldOp->getOperand(0);
-    producer = producerOutput.getDefiningOp();
-    if (!producer || producer->getNumOperands() == 0) {
-      return failure();
-    }
-    if (!matchPattern(producer, m_Op<arith::MaximumFOp>())) {
-      return failure();
-    }
-  }
-
-  // Producer of linalg.yield op 2nd arg is arith.select
-  // TODO: Add check that select is selecting between linalg.index and index of
-  // current max.
-  {
-    producerOutput = yieldOp->getOperand(1);
-    producer = producerOutput.getDefiningOp();
-    if (!producer || producer->getNumOperands() == 0) {
-      return failure();
-    }
-    if (!matchPattern(producer, m_Op<arith::SelectOp>())) {
-      return failure();
-    }
-  }
-
-  // Producer of arith.select op is arith.cmpf
-  {
-    producerOutput = producer->getOperand(0);
-    producer = producerOutput.getDefiningOp();
-    if (!producer || producer->getNumOperands() == 0) {
-      return failure();
-    }
-    auto producerCmpFOp = dyn_cast<arith::CmpFOp>(producer);
-    if (!producerCmpFOp) {
-      return failure();
-    }
-    if (producerCmpFOp.getPredicate() != arith::CmpFPredicate::OGT) {
-      return failure();
-    }
-
-    // Check that in and out of cmpf are loop variables.
-    // Currently first operand is disabled because it may be mixed type
-    // which would lead it to be extf(%arg0).
-    // TODO: Add better mixed type support check.
-    if (producer->getOperand(1) != genericOp.getBody()->getArgument(1)) {
-      return failure();
-    }
-  }
-
-  return success();
-}
-
 //===---------------------------------------------------------------------===//
 // Replace Memref users (transitively)
 //===---------------------------------------------------------------------===//
diff --git a/compiler/src/iree/compiler/Codegen/Utils/Utils.h b/compiler/src/iree/compiler/Codegen/Utils/Utils.h
index 57bb673571f5..5ff621debcee 100644
--- a/compiler/src/iree/compiler/Codegen/Utils/Utils.h
+++ b/compiler/src/iree/compiler/Codegen/Utils/Utils.h
@@ -239,9 +239,6 @@ Operation *dropEncodingAndCloneOp(OpBuilder &builder, Operation *op,
                                   ValueRange convertedInputOperands,
                                   ValueRange convertedOutputOperands);
 
-/// Check if a linalg.generic is representing an argmax operation.
-LogicalResult isArgmaxOp(linalg::GenericOp genericOp);
-
 /// Replace the uses of memref value `origValue` with the given
 /// `replacementValue`. Some uses of the memref value might require changes to
 /// the operation itself. Create new operations which can carry the change, and
diff --git a/compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/SplitReduction.cpp b/compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/SplitReduction.cpp
index 83695f61e10a..fe25e65a457a 100644
--- a/compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/SplitReduction.cpp
+++ b/compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/SplitReduction.cpp
@@ -8,6 +8,7 @@
 #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"
@@ -389,4 +390,331 @@ splitReduction(RewriterBase &rewriter, LinalgExt::TopkOp topkOp,
   return success();
 }
 
+struct ArgmaxCombinerOps {
+  Operation *maxOp = nullptr;    // arith.maximumf
+  Operation *selectOp = nullptr; // arith.select
+  Operation *cmpOp = nullptr;    // arith.cmpf
+};
+
+// Matches the combiner pattern in a linalg.generic argmax-style reduction:
+// Example MLIR:
+// %4:2 = linalg.generic {
+//     indexing_maps = [...],
+//     iterator_types = ["parallel", "reduction"]
+//   } ins(%arg0 : tensor<?x128xbf16>) outs(%1, %3 : tensor<?xbf16>,
+//   tensor<?xi64>) {
+// ^bb0(%in: bf16, %out: bf16, %out_0: i64):
+//   %5 = linalg.index 1 : index
+//   %6 = arith.index_cast %5 : index to i64
+//   %7 = arith.maximumf %in, %out : bf16
+//   %8 = arith.cmpf ogt, %in, %out : bf16
+//   %9 = arith.select %8, %6, %out_0 : i64
+//   linalg.yield %7, %9 : bf16, i64
+// } -> (tensor<?xbf16>, tensor<?xi64>)
+//
+// This function extracts the `arith.maximumf`, `arith.cmpf`, and `arith.select`
+// operations from the body to facilitate transformations such as split
+// reduction.
+static ArgmaxCombinerOps collectArgmaxCombinerOps(linalg::GenericOp genericOp) {
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  auto yieldOp = cast<linalg::YieldOp>(genericOp.getBody()->getTerminator());
+
+  // Extract max value producer: arith.maximumf.
+  Value maxResult = yieldOp.getOperand(0);
+  auto maxOp = cast<arith::MaximumFOp>(maxResult.getDefiningOp());
+
+  // Extract index result producer: arith.select.
+  Value indexResult = yieldOp.getOperand(1);
+  auto selectOp = cast<arith::SelectOp>(indexResult.getDefiningOp());
+
+  // Extract the condition of the select, expected to be arith.cmpf with
+  // predicate OGT.
+  auto cmpOp = cast<arith::CmpFOp>(selectOp.getCondition().getDefiningOp());
+
+  ArgmaxCombinerOps ops;
+  ops.maxOp = maxOp;
+  ops.selectOp = selectOp;
+  ops.cmpOp = cmpOp;
+  return ops;
+}
+
+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 builder.create<tensor::ExpandShapeOp>(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 = builder.create<arith::ConstantOp>(
+      loc, elemType, cast<TypedAttr>(identityAttr));
+  Value empty = builder.create<tensor::EmptyOp>(loc, shape, elemType);
+  return builder.create<linalg::FillOp>(loc, identityVal, empty).getResult(0);
+}
+
+FailureOr<linalg::SplitReductionResult>
+splitArgmaxReduction(RewriterBase &rewriter, linalg::GenericOp genericOp,
+                     linalg::ControlSplitReductionFn controlSplitReductionFn) {
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  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, 4> 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");
+  }
+
+  ArgmaxCombinerOps combinerOps = collectArgmaxCombinerOps(genericOp);
+  Operation *reductionOp = combinerOps.maxOp;
+  std::optional<TypedAttr> identity = arith::getNeutralElement(reductionOp);
+  if (!identity.has_value()) {
+    return rewriter.notifyMatchFailure(
+        genericOp, "Unknown identity value for the reduction");
+  }
+
+  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()));
+    }
+  }
+
+  // Create partial linalg.generic op with global index computation.
+  Value tileSize = rewriter.create<arith::ConstantIndexOp>(loc, ratio);
+  auto partialOp = rewriter.create<linalg::GenericOp>(
+      loc, TypeRange{identityValue.getType(), identityIndex.getType()},
+      newInputs, ValueRange{identityValue, identityIndex}, newMaps,
+      newIteratorTypes);
+
+  rewriter.inlineRegionBefore(genericOp.getRegion(), partialOp.getRegion(),
+                              partialOp.getRegion().begin());
+
+  Block &body = partialOp.getRegion().front();
+  rewriter.setInsertionPointToStart(&body);
+
+  unsigned innerIdxDim = reductionDim + 1;
+  unsigned outerIdxDim = insertSplitDimension;
+
+  // Compute global index (gidx) for reduction when the original reduction
+  // dimension is split into [outerIdx, innerIdx] using `ratio`. This is used to
+  // correctly compute the global index for comparisons and index selection.
+  Value outerIdx = rewriter.create<linalg::IndexOp>(loc, outerIdxDim);
+  Value innerIdx = rewriter.create<linalg::IndexOp>(loc, innerIdxDim);
+  Value offset = rewriter.create<arith::MulIOp>(loc, outerIdx, tileSize);
+  Value gidx = rewriter.create<arith::AddIOp>(loc, offset, innerIdx);
+
+  auto selectOp = dyn_cast<arith::SelectOp>(combinerOps.selectOp);
+  Value oldIdx = selectOp.getTrueValue();
+  Value newIdx = gidx;
+  if (oldIdx.getType() != gidx.getType()) {
+    newIdx = rewriter.create<arith::IndexCastOp>(loc, oldIdx.getType(), gidx);
+  }
+  selectOp.setOperand(1, newIdx);
+  rewriter.setInsertionPointAfter(partialOp);
+
+  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};
+
+  // Create block for final reduction region.
+  auto finalReduction = rewriter.create<linalg::GenericOp>(
+      loc, genericOp.getResultTypes(),
+      ValueRange{partialOp.getResult(0), partialOp.getResult(1)},
+      genericOp.getDpsInits(), finalReductionMaps, reductionIteratorTypes,
+      [combinerOps](OpBuilder &b, Location loc, ValueRange inputs) {
+        Operation *clonedMax = b.clone(*combinerOps.maxOp);
+        clonedMax->setOperands({inputs[0], inputs[2]});
+        Operation *clonedCmp = b.clone(*combinerOps.cmpOp);
+        clonedCmp->setOperands({inputs[0], inputs[2]});
+        Operation *clonedSel = b.clone(*combinerOps.selectOp);
+        clonedSel->setOperands({clonedCmp->getResult(0), inputs[1], inputs[3]});
+        b.create<linalg::YieldOp>(
+            loc, ValueRange{clonedMax->getResult(0), clonedSel->getResult(0)});
+      });
+
+  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>(partialOp.getOperation()), finalReduction};
+}
+
 } // namespace mlir::iree_compiler::IREE::LinalgExt
diff --git a/compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/Transforms.h b/compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/Transforms.h
index added441d812..47dd42a8599a 100644
--- a/compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/Transforms.h
+++ b/compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/Transforms.h
@@ -118,4 +118,59 @@ FailureOr<std::pair<Value, Value>> rewriteFft(Operation *op, Value operand,
                                               int64_t fftLength,
                                               PatternRewriter &rewriter);
 
+/// Apply transformation to split a linalg.generic argmax reduction
+/// into a two-stage reduction using an additional parallel dimension.
+/// The transformation first computes a partial argmax over tiles (parallel),
+/// then reduces those results into a final result (reduction).
+///
+/// This pattern is specialized for reductions that yield both the maximum
+/// value and its index, using the combination of `arith.maximumf`,
+/// `arith.cmpf`, and `arith.select` ops. It assumes a known structure of the
+/// region and injects index computations to track global indices.
+///
+/// Returns the resulting partial and final linalg.generic ops, or failure
+/// if the pattern does not match or cannot be split.
+///
+/// Example: original argmax op reducing over dim=512
+/// %4:2 = linalg.generic {
+///   indexing_maps = [...],
+///   iterator_types = ["parallel", "reduction"]
+/// } ins(%arg0 : tensor<?x512xbf16>)
+///   outs(%out_val, %out_idx : tensor<?xbf16>, tensor<?xi64>) {
+/// ^bb0(%in: bf16, %out: bf16, %out_0: i64):
+///   %idx = linalg.index 1 : index
+///   %cast = arith.index_cast %idx : index to i64
+///   %max = arith.maximumf %in, %out : bf16
+///   %cmp = arith.cmpf ogt, %in, %out : bf16
+///   %sel = arith.select %cmp, %cast, %out_0 : i64
+///   linalg.yield %max, %sel : bf16, i64
+/// } -> (tensor<?xbf16>, tensor<?xi64>)
+///
+/// To: splitting K=512 into 4 x 128 + final argmax over the tile dimension
+///     (dim=1 of ?x4)
+///
+/// %expanded = tensor.expand_shape %arg0 [[0], [1, 2]] : tensor<?x512xbf16>
+///     into tensor<?x4x128xbf16>
+/// %init_val = linalg.fill ins(%cst : bf16) outs(%empty : tensor<?x4xbf16>)
+/// %init_idx = linalg.fill ins(%zero : i64) outs(%empty : tensor<?x4xi64>)
+/// %partial:2 = linalg.generic {
+///   indexing_maps = [...],
+///   iterator_types = ["parallel", "reduction"]
+/// } ins(%expanded : tensor<?x4x128xbf16>)
+///   outs(%init_val, %init_idx : tensor<?x4xbf16>, tensor<?x4xi64>) {
+///   // compute global index: outer_idx * 128 + inner_idx
+///   ...
+/// }
+/// %final:2 = linalg.generic {
+///   indexing_maps = [...],
+///   iterator_types = ["reduction"]
+/// } ins(%partial#0, %partial#1)
+///   outs(%out_val, %out_idx : tensor<?xbf16>, tensor<?xi64>) {
+///   // same combiner: maximumf, cmpf, select
+///   ...
+/// }
+FailureOr<linalg::SplitReductionResult>
+splitArgmaxReduction(RewriterBase &rewriter, linalg::GenericOp genericOp,
+                     linalg::ControlSplitReductionFn controlSplitReductionFn);
+
 }; // namespace mlir::iree_compiler::IREE::LinalgExt
diff --git a/compiler/src/iree/compiler/Dialect/LinalgExt/Utils/Utils.cpp b/compiler/src/iree/compiler/Dialect/LinalgExt/Utils/Utils.cpp
index 255cab4a2479..c7d1dcc69c1f 100644
--- a/compiler/src/iree/compiler/Dialect/LinalgExt/Utils/Utils.cpp
+++ b/compiler/src/iree/compiler/Dialect/LinalgExt/Utils/Utils.cpp
@@ -826,4 +826,98 @@ bool isaHorizontallyFusedContraction(Operation *op) {
   return true;
 }
 
+bool isArgmaxOp(linalg::GenericOp genericOp) {
+  // Check for 2 results(value, index), and 1 input
+  if (genericOp.getNumDpsInits() != 2) {
+    return false;
+  }
+
+  if (genericOp.getNumDpsInputs() != 1) {
+    return false;
+  }
+
+  // If max value is being used, it is not a pure argmax.
+  if (!genericOp.getResults()[0].use_empty()) {
+    return false;
+  }
+
+  // Argmax will require 1D reduction.
+  if (genericOp.getNumReductionLoops() != 1) {
+    return false;
+  }
+
+  // TODO: Add better affine map checks.
+  auto indexing_maps = genericOp.getIndexingMapsArray();
+  if (!indexing_maps[0].isIdentity())
+    return false;
+
+  // Check that initial value is negative Infinite.
+  // TODO: Move this check to ukernel once we implement
+  //       variant to handle non neg-Inf initial value.
+  Value initVal = genericOp.getDpsInitOperand(0)->get();
+  auto fillOp = initVal.getDefiningOp<linalg::FillOp>();
+  if (!fillOp)
+    return false;
+  Value fillVal = fillOp.getDpsInputOperand(0)->get();
+  if (!matchPattern(fillVal, m_NegInfFloat()))
+    return false;
+
+  // Work back from linalg.yield and check body of genericOp.
+  // The genericOp should yield the result of an arith.select,
+  // preceded by an arith.cmpf, arith.maximumf, and arith.extui
+  auto yieldOp = cast<linalg::YieldOp>(genericOp.getBody()->getTerminator());
+  Value producerOutput;
+  Operation *producer;
+
+  // Producer of linalg.yield 1st arg is arith.maximumf
+  {
+    producerOutput = yieldOp->getOperand(0);
+    producer = producerOutput.getDefiningOp();
+    if (!producer || producer->getNumOperands() == 0) {
+      return false;
+    }
+    if (!matchPattern(producer, m_Op<arith::MaximumFOp>())) {
+      return false;
+    }
+  }
+
+  // Producer of linalg.yield op 2nd arg is arith.select
+  // TODO: Add check that select is selecting between linalg.index and index of
+  // current max.
+  {
+    producerOutput = yieldOp->getOperand(1);
+    producer = producerOutput.getDefiningOp();
+    if (!producer || producer->getNumOperands() == 0) {
+      return false;
+    }
+    if (!matchPattern(producer, m_Op<arith::SelectOp>())) {
+      return false;
+    }
+  }
+
+  // Producer of arith.select op is arith.cmpf
+  {
+    producerOutput = producer->getOperand(0);
+    producer = producerOutput.getDefiningOp();
+    if (!producer || producer->getNumOperands() == 0) {
+      return false;
+    }
+    auto producerCmpFOp = dyn_cast<arith::CmpFOp>(producer);
+    if (!producerCmpFOp ||
+        producerCmpFOp.getPredicate() != arith::CmpFPredicate::OGT) {
+      return false;
+    }
+
+    // Check that in and out of cmpf are loop variables.
+    // Currently first operand is disabled because it may be mixed type
+    // which would lead it to be extf(%arg0).
+    // TODO: Add better mixed type support check.
+    if (producer->getOperand(1) != genericOp.getBody()->getArgument(1)) {
+      return false;
+    }
+  }
+
+  return true;
+}
+
 } // namespace mlir::iree_compiler::IREE::LinalgExt
diff --git a/compiler/src/iree/compiler/Dialect/LinalgExt/Utils/Utils.h b/compiler/src/iree/compiler/Dialect/LinalgExt/Utils/Utils.h
index 9bc1d58c1b86..4b111bbb1087 100644
--- a/compiler/src/iree/compiler/Dialect/LinalgExt/Utils/Utils.h
+++ b/compiler/src/iree/compiler/Dialect/LinalgExt/Utils/Utils.h
@@ -224,5 +224,8 @@ bool isGatherlikeOp(Operation *op);
 /// different RHS.
 bool isaHorizontallyFusedContraction(Operation *op);
 
+/// Check if a linalg.generic is representing an argmax operation.
+bool isArgmaxOp(linalg::GenericOp genericOp);
+
 } // namespace mlir::iree_compiler::IREE::LinalgExt
 #endif // IREE_COMPILER_DIALECT_LINALGEXT_UTILS_UTILS_H_
diff --git a/compiler/src/iree/compiler/DispatchCreation/Passes.td b/compiler/src/iree/compiler/DispatchCreation/Passes.td
index 7c7fc1edecbb..3e18911d0ef3 100644
--- a/compiler/src/iree/compiler/DispatchCreation/Passes.td
+++ b/compiler/src/iree/compiler/DispatchCreation/Passes.td
@@ -164,6 +164,7 @@ def SplitReductionPass :
   let summary = "Split reduction dimension to increase parallelism.";
   let dependentDialects = [
     "mlir::linalg::LinalgDialect",
+    "IREE::LinalgExt::IREELinalgExtDialect",
   ];
 }
 
diff --git a/compiler/src/iree/compiler/DispatchCreation/SplitReduction.cpp b/compiler/src/iree/compiler/DispatchCreation/SplitReduction.cpp
index 87eeb80b92f3..efdb6287e99f 100644
--- a/compiler/src/iree/compiler/DispatchCreation/SplitReduction.cpp
+++ b/compiler/src/iree/compiler/DispatchCreation/SplitReduction.cpp
@@ -11,7 +11,9 @@
 //===----------------------------------------------------------------------===//
 
 #include "iree/compiler/Dialect/LinalgExt/Transforms/Transforms.h"
+#include "iree/compiler/Dialect/LinalgExt/Utils/Utils.h"
 #include "iree/compiler/DispatchCreation/Passes.h"
+#include "mlir/Analysis/SliceAnalysis.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
@@ -24,8 +26,20 @@ namespace mlir::iree_compiler::DispatchCreation {
 
 // TODO(thomasraoux): Move to attributes.
 static llvm::cl::opt<int64_t>
-    splitReductionRatio("iree-dispatch-creation-split-matmul-reduction",
-                        llvm::cl::desc("split ratio"), llvm::cl::init(1));
+    splitMatmulReductionRatio("iree-dispatch-creation-split-matmul-reduction",
+                              llvm::cl::desc("split ratio"), llvm::cl::init(1));
+
+static llvm::cl::opt<int64_t> splitArgmaxReductionThreshold(
+    "iree-dispatch-creation-split-argmax-threshold",
+    llvm::cl::desc(
+        "Minimum size of the reduction dimension to trigger argmax split"),
+    llvm::cl::init(128 * 1024));
+
+static llvm::cl::opt<int64_t> splitArgmaxTileSize(
+    "iree-dispatch-creation-split-argmax-tile-size",
+    llvm::cl::desc("Tile size of the reduction dimension after splitting "
+                   "(i.e., the chunk size)"),
+    llvm::cl::init(128));
 
 static llvm::cl::list<int64_t> topkSplitReductionRatio(
     "iree-dispatch-creation-topk-split-reduction",
@@ -63,28 +77,73 @@ namespace {
 struct SplitReductionPass final
     : public impl::SplitReductionPassBase<SplitReductionPass> {
   void runOnOperation() override {
-    if (splitReductionRatio.getValue() <= 1 &&
-        topkSplitReductionRatio.empty()) {
+    if (splitMatmulReductionRatio.getValue() <= 1 &&
+        topkSplitReductionRatio.empty() &&
+        (splitArgmaxTileSize.getValue() <= 1 ||
+         splitArgmaxReductionThreshold.getValue() <= 1)) {
       return;
     }
 
     MLIRContext *context = &getContext();
     auto funcOp = getOperation();
 
+    SmallVector<linalg::MatmulOp> matmulCandidates;
+    SmallVector<IREE::LinalgExt::TopkOp> topkCandidates;
+    SmallVector<linalg::GenericOp> argmaxCandidates;
+
+    IRRewriter rewriter(context);
+    funcOp->walk([&](Operation *op) {
+      TypeSwitch<Operation *>(op)
+          .Case<linalg::MatmulOp>([&](auto matmulOp) {
+            if (splitMatmulReductionRatio > 1) {
+              matmulCandidates.push_back(matmulOp);
+            }
+          })
+          .Case<IREE::LinalgExt::TopkOp>([&](auto topkOp) {
+            if (!topkSplitReductionRatio.empty()) {
+              topkCandidates.push_back(topkOp);
+            }
+          })
+          .Case<linalg::GenericOp>([&](auto genericOp) {
+            if (splitArgmaxTileSize > 1 &&
+                IREE::LinalgExt::isArgmaxOp(genericOp)) {
+              // Due to isArgmaxOp, we support exactly one reduction dimension.
+              SmallVector<unsigned> dims;
+              genericOp.getReductionDims(dims);
+              unsigned reductionDim = dims[0];
+              SmallVector<int64_t, 4> loopRanges =
+                  genericOp.getStaticLoopRanges();
+              int64_t reductionDimSize = loopRanges[reductionDim];
+              if (reductionDimSize >= splitArgmaxReductionThreshold) {
+                argmaxCandidates.push_back(genericOp);
+              }
+            }
+          });
+    });
+
+    // Split matmul ops.
     auto matmulSplitReductionControlFn =
         [&](linalg::LinalgOp op) -> linalg::SplitReductionOptions {
       // For matmul make the new parallel dimension first so that it looks
       // like a batch_matmul and can follow the same codegen.
-      return {int64_t(splitReductionRatio), 0, /*innerParallel=*/false};
+      return {splitMatmulReductionRatio, 0, /*innerParallel=*/false};
     };
-
-    SmallVector<linalg::MatmulOp> matmulCandidates;
-    IRRewriter rewriter(context);
-    funcOp->walk([&](linalg::MatmulOp op) { matmulCandidates.push_back(op); });
     for (auto op : matmulCandidates) {
       (void)splitReductionOnMatmul(rewriter, op, matmulSplitReductionControlFn);
     }
 
+    // Split argmax ops.
+    auto argmaxSplitReductionControlFn =
+        [&](linalg::LinalgOp op) -> linalg::SplitReductionOptions {
+      return {splitArgmaxTileSize, op.getNumLoops() - 1,
+              /*innerParallel=*/false};
+    };
+    for (auto op : argmaxCandidates) {
+      (void)IREE::LinalgExt::splitArgmaxReduction(
+          rewriter, op, argmaxSplitReductionControlFn);
+    }
+
+    // Split topk ops.
     IREE::LinalgExt::TopkSplitReductionControlFn topkSplitReductionControlFn =
         [&](int64_t splitReductionDepth) -> int64_t {
       SmallVector<int64_t> reductionRatios(topkSplitReductionRatio.begin(),
@@ -95,10 +154,6 @@ struct SplitReductionPass final
         return reductionRatios[splitReductionDepth];
       }
     };
-
-    SmallVector<IREE::LinalgExt::TopkOp> topkCandidates;
-    funcOp->walk(
-        [&](IREE::LinalgExt::TopkOp op) { topkCandidates.push_back(op); });
     for (auto op : topkCandidates) {
       (void)splitReduction(rewriter, op, topkSplitReductionControlFn);
     }
diff --git a/compiler/src/iree/compiler/DispatchCreation/test/split_reduction.mlir b/compiler/src/iree/compiler/DispatchCreation/test/split_reduction.mlir
index ad07f39a0576..484fb4a8c55c 100644
--- a/compiler/src/iree/compiler/DispatchCreation/test/split_reduction.mlir
+++ b/compiler/src/iree/compiler/DispatchCreation/test/split_reduction.mlir
@@ -27,3 +27,110 @@ util.func public @matmul_with_encoding(%arg0: tensor<100x200xf32, #encoding>, %a
 }
 // CHECK-LABEL: util.func public @matmul_with_encoding(
 // CHECK:         linalg.matmul
+
+// -----
+
+util.func public @argmax(%arg0: tensor<?x131072xbf16>, %arg1: index) -> tensor<?xi64> {
+  %cst = arith.constant 0xFF80 : bf16
+  %c0_i64 = arith.constant 0 : i64
+  %0 = tensor.empty(%arg1) : tensor<?xbf16>
+  %1 = linalg.fill ins(%cst : bf16) outs(%0 : tensor<?xbf16>) -> tensor<?xbf16>
+  %2 = tensor.empty(%arg1) : tensor<?xi64>
+  %3 = linalg.fill ins(%c0_i64 : i64) outs(%2 : tensor<?xi64>) -> tensor<?xi64>
+  %4:2 = linalg.generic {
+      indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
+                       affine_map<(d0, d1) -> (d0)>,
+                       affine_map<(d0, d1) -> (d0)>],
+      iterator_types = ["parallel", "reduction"]}
+      ins(%arg0 : tensor<?x131072xbf16>) outs(%1, %3 : tensor<?xbf16>, tensor<?xi64>) {
+  ^bb0(%in: bf16, %out: bf16, %out_0: i64):
+    %5 = linalg.index 1 : index
+    %6 = arith.index_cast %5 : index to i64
+    %7 = arith.maximumf %in, %out : bf16
+    %8 = arith.cmpf ogt, %in, %out : bf16
+    %9 = arith.select %8, %6, %out_0 : i64
+    linalg.yield %7, %9 : bf16, i64
+  } -> (tensor<?xbf16>, tensor<?xi64>)
+  util.return %4#1 : tensor<?xi64>
+}
+
+// CHECK: #[[$MAP0:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+// CHECK: #[[$MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+// CHECK: #[[$MAP2:.+]] = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK: #[[$MAP3:.+]] = affine_map<(d0, d1) -> (d0)>
+
+// CHECK-LABEL:   util.func public @argmax
+// Check identity value preparation
+// CHECK: %[[CST:.+]] = arith.constant 0xFF80 : bf16
+// CHECK: %[[ZERO:.+]] = arith.constant 0 : i64
+// CHECK: %[[FINALVAL_EMPTY:.+]] = tensor.empty(%[[ARG1:.+]]) : tensor<?xbf16>
+// CHECK: %[[FINALVAL:.+]] = linalg.fill ins(%[[CST]] : bf16) outs(%[[FINALVAL_EMPTY]] : tensor<?xbf16>) -> tensor<?xbf16>
+// CHECK: %[[FINALIDX_EMPTY:.+]] = tensor.empty(%[[ARG1]]) : tensor<?xi64>
+// CHECK: %[[FINALIDX:.+]] = linalg.fill ins(%[[ZERO]] : i64) outs(%[[FINALIDX_EMPTY]] : tensor<?xi64>) -> tensor<?xi64>
+
+// Check partial reduction.
+// CHECK: %[[EXPAND:.+]] = tensor.expand_shape %arg0 {{\[}}[0], [1, 2]] output_shape [%{{.+}}, 1024, 128] : tensor<?x131072xbf16> into tensor<?x1024x128xbf16>
+// CHECK: %[[INITVAL:.+]] = tensor.empty(%{{.+}}) : tensor<?x1024xbf16>
+// CHECK: %[[FILLVAL:.+]] = linalg.fill ins(%{{.+}} : bf16) outs(%[[INITVAL]] : tensor<?x1024xbf16>) -> tensor<?x1024xbf16>
+// CHECK: %[[INITIDX:.+]] = tensor.empty(%{{.+}}) : tensor<?x1024xi64>
+// CHECK: %[[FILLIDX:.+]] = linalg.fill ins(%{{.+}} : i64) outs(%[[INITIDX]] : tensor<?x1024xi64>) -> tensor<?x1024xi64>
+
+// CHECK: %[[PARTIAL:.+]]:2 = linalg.generic
+// CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP1]]]
+// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]
+// CHECK-SAME: ins(%[[EXPAND]] : tensor<?x1024x128xbf16>)
+// CHECK-SAME: outs(%[[FILLVAL]], %[[FILLIDX]] : tensor<?x1024xbf16>, tensor<?x1024xi64>)
+// CHECK: ^bb0(%[[VAL:.+]]: bf16, %[[ACC:.+]]: bf16, %[[IDX:.+]]: i64)
+// CHECK: %[[OUTER:.+]] = linalg.index 1 : index
+// CHECK: %[[INNER:.+]] = linalg.index 2 : index
+// CHECK: %[[OFFSET:.+]] = arith.muli %[[OUTER]], %{{.+}} : index
+// CHECK: %[[GIDX:.+]] = arith.addi %[[OFFSET]], %[[INNER]] : index
+// CHECK: %[[CAST:.+]] = arith.index_cast %[[GIDX]] : index to i64
+// CHECK: %[[MAX:.+]] = arith.maximumf %[[VAL]], %[[ACC]] : bf16
+// CHECK: %[[CMP:.+]] = arith.cmpf ogt, %[[VAL]], %[[ACC]] : bf16
+// CHECK: %[[SEL:.+]] = arith.select %[[CMP]], %[[CAST]], %[[IDX]] : i64
+// CHECK: linalg.yield %[[MAX]], %[[SEL]] : bf16, i64
+
+// Check Final reduction.
+// CHECK: %[[FINAL:.+]]:2 = linalg.generic
+// CHECK-SAME: indexing_maps = [#[[$MAP2]], #[[$MAP2]], #[[$MAP3]], #[[$MAP3]]]
+// CHECK-SAME: iterator_types = ["parallel", "reduction"]
+// CHECK-SAME: ins(%[[PARTIAL]]#0, %[[PARTIAL]]#1 : tensor<?x1024xbf16>, tensor<?x1024xi64>)
+// CHECK-SAME: outs(%[[FINALVAL]], %[[FINALIDX]] : tensor<?xbf16>, tensor<?xi64>)
+// CHECK: ^bb0(%[[V1:.+]]: bf16, %[[I1:.+]]: i64, %[[V2:.+]]: bf16, %[[I2:.+]]: i64)
+// CHECK: %[[MAX2:.+]] = arith.maximumf %[[V1]], %[[V2]] : bf16
+// CHECK: %[[CMP2:.+]] = arith.cmpf ogt, %[[V1]], %[[V2]] : bf16
+// CHECK: %[[SEL2:.+]] = arith.select %[[CMP2]], %[[I1]], %[[I2]] : i64
+// CHECK: linalg.yield %[[MAX2]], %[[SEL2]] : bf16, i64
+
+// Check final return.
+// CHECK: util.return %[[FINAL]]#1 : tensor<?xi64>
+
+// -----
+
+util.func public @argmax_no_split(%arg0: tensor<?x512xbf16>, %arg1: index) -> tensor<?xi64> {
+  %cst = arith.constant 0xFF80 : bf16
+  %c0_i64 = arith.constant 0 : i64
+  %0 = tensor.empty(%arg1) : tensor<?xbf16>
+  %1 = linalg.fill ins(%cst : bf16) outs(%0 : tensor<?xbf16>) -> tensor<?xbf16>
+  %2 = tensor.empty(%arg1) : tensor<?xi64>
+  %3 = linalg.fill ins(%c0_i64 : i64) outs(%2 : tensor<?xi64>) -> tensor<?xi64>
+  %4:2 = linalg.generic {
+      indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
+                       affine_map<(d0, d1) -> (d0)>,
+                       affine_map<(d0, d1) -> (d0)>],
+      iterator_types = ["parallel", "reduction"]}
+      ins(%arg0 : tensor<?x512xbf16>) outs(%1, %3 : tensor<?xbf16>, tensor<?xi64>) {
+  ^bb0(%in: bf16, %out: bf16, %out_0: i64):
+    %5 = linalg.index 1 : index
+    %6 = arith.index_cast %5 : index to i64
+    %7 = arith.maximumf %in, %out : bf16
+    %8 = arith.cmpf ogt, %in, %out : bf16
+    %9 = arith.select %8, %6, %out_0 : i64
+    linalg.yield %7, %9 : bf16, i64
+  } -> (tensor<?xbf16>, tensor<?xi64>)
+  util.return %4#1 : tensor<?xi64>
+}
+
+// CHECK-LABEL:   util.func public @argmax_no_split
+// CHECK-NOT:     tensor.expand_shape


Reviews:
user: login: MaheshRavishankar
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-05T18:26:49Z
user: login: pashu123
body: NITs. I'll take a look once again.
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-05T18:46:56Z
user: login: bangtianliu
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-05T20:22:03Z
user: login: bangtianliu
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-06T14:54:37Z
user: login: MaheshRavishankar
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-06T18:03:00Z
user: login: hanhanW
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-06T18:40:43Z
user: login: hanhanW
body: I want to review the template part if you don't mind wait me a bit. I can do the review later today or tomorrow. No needs to wait for me if others approve the change. I can send post review comments or raise a PR myself after you land the PR.
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-06T18:44:01Z
user: login: pashu123
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-06T18:46:46Z
user: login: pashu123
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T16:30:57Z
user: login: hanhanW
body: I did not do a detailed review, but here are a few high-level comments so far:

1. We already have separation in the main function, i.e., `SplitReductionPass::runOnOperation`. Can we avoid the template? Template is good for some cases, but the downside is hard to debug and maintain. It is also hard to trace code. I'd prefer to avoid it as much as possible. This falls to the case that we can avoid template. I don't see a need to use template here. Naming the method as splitArgmaxReduction looks much easier to me. https://google.github.io/styleguide/cppguide.html#Template_metaprogramming 
2. Can you elaborate why the upstream method can't be reused? I can see some reasons when I browse the code, but it is easier if you can list down why and what the difference is in PR description. It is also helpful for reviewers and future contributors, as they would track back to the commit.
3. There are too many dup code between your implementation and upstream splitReduction. Can you think a way to refactor the code? I can help if you need some guidance.
4. I think it is better to move the `splitArgmaxReduction()` to `LinalgExt/Transforms/Transforms.[h|cpp]`. Because (a) I feel that we will upstream it in the future, (b) this makes the pass file cleaner, (c) it belongs to transformations about LinalgExt (linalg extention), IMO, (d) it can be reused by other passes when needed. RE(d): we do have a good example of reusing splitReduction in CPU codegen to perform innermost reduction. I think it is better to expose the implementation outside this pass in the first place.
5. Similar to (3), please add more details to the PR description. The PR is not trivial, we should have informative PR description. Also, it'd be a plus if you mention that there is a refactoring about `isArgmaxOp` in the PR description. https://google.github.io/eng-practices/review/developer/cl-descriptions.html

The topK variant locates at wrong files, I'll help fix it in a separate PR.
state: CHANGES_REQUESTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T17:56:05Z
user: login: hanhanW
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T17:59:32Z
user: login: hanhanW
body: I left some nits, and a big comment about how we generate expanded operands. Please take a look, thanks.

EDIT: I can't share the link before I submit the comment, here is the link to the big comment: https://github.com/iree-org/iree/pull/20717#discussion_r2082458893
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-09T20:49:10Z
user: login: bangtianliu
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-09T21:09:06Z
user: login: hanhanW
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-09T21:11:45Z
user: login: MaheshRavishankar
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-10T17:32:15Z
user: login: MaheshRavishankar
state: CHANGES_REQUESTED
author_association: COLLABORATOR
submitted_at: 2025-05-10T17:47:25Z
user: login: hanhanW
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-13T00:22:48Z
user: login: bangtianliu
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-13T02:40:14Z
user: login: bangtianliu
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-13T03:19:46Z
user: login: hanhanW
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-13T23:04:31Z
user: login: hanhanW
body: It looks much better to me, just some style nits and the earlier outs operand implementation improvement. Thanks for addressing my comments!
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-14T04:05:21Z
user: login: hanhanW
body: LGTM % some nits. (I missed some in previous reviews.)
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-14T23:48:53Z
user: login: pashu123
body: LGTM!
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-15T15:53:42Z
user: login: MaheshRavishankar
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-05-15T17:36:42Z

Review comments:
diff_hunk: @@ -826,4 +826,102 @@ bool isaHorizontallyFusedContraction(Operation *op) {
   return true;
 }
 
+bool isArgmaxOp(linalg::GenericOp genericOp) {
+  // Check for 2 results(value, index), and 1 input
+  if (genericOp.getNumDpsInits() != 2) {
+    return false;
+  }
+  if (genericOp.getNumDpsInputs() != 1) {
+    return false;
+  }
+
+  // If max value is being used, it is not a pure argmax.
+  if (!genericOp.getResults()[0].use_empty()) {
+    return false;
+  }
+
+  // Check that the rank is at least 3 and all loops are parallel
+  unsigned numLoops = genericOp.getNumLoops();
+  unsigned numParallelLoops = genericOp.getNumParallelLoops();
+
+  // Argmax will require 1D reduction.
+  if (numParallelLoops != (numLoops - 1)) {
path: compiler/src/iree/compiler/Dialect/LinalgExt/Utils/Utils.cpp
user: login: pashu123
body: `op.getNumReductionLoops() != 1`
created_at: 2025-05-05T18:18:47Z
author_association: CONTRIBUTOR
diff_hunk: @@ -826,4 +826,102 @@ bool isaHorizontallyFusedContraction(Operation *op) {
   return true;
 }
 
+bool isArgmaxOp(linalg::GenericOp genericOp) {
+  // Check for 2 results(value, index), and 1 input
+  if (genericOp.getNumDpsInits() != 2) {
+    return false;
+  }
+  if (genericOp.getNumDpsInputs() != 1) {
+    return false;
+  }
+
+  // If max value is being used, it is not a pure argmax.
+  if (!genericOp.getResults()[0].use_empty()) {
+    return false;
+  }
+
+  // Check that the rank is at least 3 and all loops are parallel
+  unsigned numLoops = genericOp.getNumLoops();
+  unsigned numParallelLoops = genericOp.getNumParallelLoops();
+
+  // Argmax will require 1D reduction.
+  if (numParallelLoops != (numLoops - 1)) {
+    return false;
+  }
+  // TODO: Add better affine map checks.
+  auto indexing_maps = genericOp.getIndexingMapsArray();
+  if (!indexing_maps[0].isIdentity())
+    return false;
+
+  // Check that initial value is negative Infinite.
+  // TODO: Move this check to ukernel once we implement
+  //       variant to handle non neg-Inf initial value.
+  Value initVal = genericOp.getDpsInitOperand(0)->get();
+  auto fillOp = initVal.getDefiningOp<linalg::FillOp>();
+  if (!fillOp)
+    return false;
+  Value fillVal = fillOp.getDpsInputOperand(0)->get();
+  if (!matchPattern(fillVal, m_NegInfFloat()))
+    return false;
+
path: compiler/src/iree/compiler/Dialect/LinalgExt/Utils/Utils.cpp
user: login: pashu123
body: Why do we need this check?
created_at: 2025-05-05T18:22:36Z
author_association: CONTRIBUTOR
diff_hunk: @@ -826,4 +826,102 @@ bool isaHorizontallyFusedContraction(Operation *op) {
   return true;
 }
 
+bool isArgmaxOp(linalg::GenericOp genericOp) {
path: compiler/src/iree/compiler/Dialect/LinalgExt/Utils/Utils.cpp
user: login: pashu123
body: I suggest using `LogicalResult` with better error messages so that if there's a case we want to support, people will know exactly where to add the support.
created_at: 2025-05-05T18:25:38Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,71 @@
+// RUN: iree-opt --pass-pipeline='builtin.module(util.func(iree-dispatch-creation-split-reduction-ops,cse, canonicalize))' --iree-dispatch-creation-split-argmax-reduction=4 %s | FileCheck %s
+
+util.func public @argmax(%arg0: tensor<?x131072xbf16>, %arg1: index) -> tensor<?xi64> {
+  %cst = arith.constant 0xFF80 : bf16
+  %c0_i64 = arith.constant 0 : i64
+  %0 = tensor.empty(%arg1) : tensor<?xbf16>
+  %1 = linalg.fill ins(%cst : bf16) outs(%0 : tensor<?xbf16>) -> tensor<?xbf16>
+  %2 = tensor.empty(%arg1) : tensor<?xi64>
+  %3 = linalg.fill ins(%c0_i64 : i64) outs(%2 : tensor<?xi64>) -> tensor<?xi64>
+  %4:2 = linalg.generic {
+      indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
+                       affine_map<(d0, d1) -> (d0)>,
+                       affine_map<(d0, d1) -> (d0)>],
+      iterator_types = ["parallel", "reduction"]}
+      ins(%arg0 : tensor<?x131072xbf16>) outs(%1, %3 : tensor<?xbf16>, tensor<?xi64>) {
+  ^bb0(%in: bf16, %out: bf16, %out_0: i64):
+    %5 = linalg.index 1 : index
+    %6 = arith.index_cast %5 : index to i64
+    %7 = arith.maximumf %in, %out : bf16
+    %8 = arith.cmpf ogt, %in, %out : bf16
+    %9 = arith.select %8, %6, %out_0 : i64
+    linalg.yield %7, %9 : bf16, i64
+  } -> (tensor<?xbf16>, tensor<?xi64>)
+  util.return %4#1 : tensor<?xi64>
+}
+
+// CHECK-LABEL:   util.func public @argmax
+
+// Check identity value preparation
+// CHECK: %[[CST:.+]] = arith.constant 0xFF80 : bf16
+// CHECK: %[[ZERO:.+]] = arith.constant 0 : i64
+// CHECK: %[[FINALVAL_EMPTY:.+]] = tensor.empty(%[[ARG1:.+]]) : tensor<?xbf16>
path: compiler/src/iree/compiler/DispatchCreation/test/split_argmax_reduction.mlir
user: login: MaheshRavishankar
body: I am assuming you didnt need to create these operations, i.e they are the same as what existed before.
created_at: 2025-05-05T18:26:30Z
author_association: COLLABORATOR
diff_hunk: @@ -24,14 +25,344 @@ namespace mlir::iree_compiler::DispatchCreation {
 
 // TODO(thomasraoux): Move to attributes.
 static llvm::cl::opt<int64_t>
-    splitReductionRatio("iree-dispatch-creation-split-matmul-reduction",
-                        llvm::cl::desc("split ratio"), llvm::cl::init(1));
+    splitMatmulReductionRatio("iree-dispatch-creation-split-matmul-reduction",
+                              llvm::cl::desc("split ratio"), llvm::cl::init(1));
+
+static llvm::cl::opt<int64_t> splitArgmaxReductionRatio(
+    "iree-dispatch-creation-split-argmax-reduction",
+    llvm::cl::desc("Ratio to split argmax. Set to 0 or 1 to disable"),
+    llvm::cl::init(1));
 
 static llvm::cl::list<int64_t> topkSplitReductionRatio(
     "iree-dispatch-creation-topk-split-reduction",
     llvm::cl::desc("comma separated list of split ratios"),
     llvm::cl::CommaSeparated);
 
+template <typename LinalgOpTy>
+static FailureOr<linalg::SplitReductionResult>
+splitReductionImpl(RewriterBase &rewriter, LinalgOpTy op,
+                   linalg::ControlSplitReductionFn controlSplitReductionFn) {
+  return linalg::splitReduction(rewriter, op, controlSplitReductionFn);
+}
+
+// Matches the combiner pattern in a linalg.generic argmax-style reduction:
+// Example MLIR:
+// %4:2 = linalg.generic {
+//     indexing_maps = [...],
+//     iterator_types = ["parallel", "reduction"]
+//   } ins(%arg0 : tensor<?x128xbf16>) outs(%1, %3 : tensor<?xbf16>,
+//   tensor<?xi64>) {
+// ^bb0(%in: bf16, %out: bf16, %out_0: i64):
+//   %5 = linalg.index 1 : index
+//   %6 = arith.index_cast %5 : index to i64
+//   %7 = arith.maximumf %in, %out : bf16
+//   %8 = arith.cmpf ogt, %in, %out : bf16
+//   %9 = arith.select %8, %6, %out_0 : i64
+//   linalg.yield %7, %9 : bf16, i64
+// } -> (tensor<?xbf16>, tensor<?xi64>)
+//
+// This function extracts the `arith.maximumf`, `arith.cmpf`, and `arith.select`
+// operations from the body to facilitate transformations such as split
+// reduction.
+static LogicalResult
+collectArgmaxCombinerOps(linalg::GenericOp genericOp,
+                         SmallVector<Operation *, 4> &combinerOps) {
+  assert(IREE::LinalgExt::isArgmaxOp(genericOp) &&
+         "expected operation to be an argmax op");
+
+  auto yieldOp = cast<linalg::YieldOp>(genericOp.getBody()->getTerminator());
+
+  // Extract max value producer: arith.maximumf.
+  Value maxResult = yieldOp.getOperand(0);
+  auto maxOp = dyn_cast<arith::MaximumFOp>(maxResult.getDefiningOp());
+  if (!maxOp)
path: compiler/src/iree/compiler/DispatchCreation/SplitReduction.cpp
user: login: pashu123
body: This has already been checked. It is with the `isArgmaxOp` check.
created_at: 2025-05-05T18:29:13Z
author_association: CONTRIBUTOR
diff_hunk: @@ -24,14 +25,344 @@ namespace mlir::iree_compiler::DispatchCreation {
 
 // TODO(thomasraoux): Move to attributes.
 static llvm::cl::opt<int64_t>
-    splitReductionRatio("iree-dispatch-creation-split-matmul-reduction",
-                        llvm::cl::desc("split ratio"), llvm::cl::init(1));
+    splitMatmulReductionRatio("iree-dispatch-creation-split-matmul-reduction",
+                              llvm::cl::desc("split ratio"), llvm::cl::init(1));
+
+static llvm::cl::opt<int64_t> splitArgmaxReductionRatio(
+    "iree-dispatch-creation-split-argmax-reduction",
+    llvm::cl::desc("Ratio to split argmax. Set to 0 or 1 to disable"),
+    llvm::cl::init(1));
 
 static llvm::cl::list<int64_t> topkSplitReductionRatio(
     "iree-dispatch-creation-topk-split-reduction",
     llvm::cl::desc("comma separated list of split ratios"),
     llvm::cl::CommaSeparated);
 
+template <typename LinalgOpTy>
+static FailureOr<linalg::SplitReductionResult>
+splitReductionImpl(RewriterBase &rewriter, LinalgOpTy op,
+                   linalg::ControlSplitReductionFn controlSplitReductionFn) {
+  return linalg::splitReduction(rewriter, op, controlSplitReductionFn);
+}
+
+// Matches the combiner pattern in a linalg.generic argmax-style reduction:
+// Example MLIR:
+// %4:2 = linalg.generic {
+//     indexing_maps = [...],
+//     iterator_types = ["parallel", "reduction"]
+//   } ins(%arg0 : tensor<?x128xbf16>) outs(%1, %3 : tensor<?xbf16>,
+//   tensor<?xi64>) {
+// ^bb0(%in: bf16, %out: bf16, %out_0: i64):
+//   %5 = linalg.index 1 : index
+//   %6 = arith.index_cast %5 : index to i64
+//   %7 = arith.maximumf %in, %out : bf16
+//   %8 = arith.cmpf ogt, %in, %out : bf16
+//   %9 = arith.select %8, %6, %out_0 : i64
+//   linalg.yield %7, %9 : bf16, i64
+// } -> (tensor<?xbf16>, tensor<?xi64>)
+//
+// This function extracts the `arith.maximumf`, `arith.cmpf`, and `arith.select`
+// operations from the body to facilitate transformations such as split
+// reduction.
+static LogicalResult
+collectArgmaxCombinerOps(linalg::GenericOp genericOp,
+                         SmallVector<Operation *, 4> &combinerOps) {
+  assert(IREE::LinalgExt::isArgmaxOp(genericOp) &&
+         "expected operation to be an argmax op");
+
+  auto yieldOp = cast<linalg::YieldOp>(genericOp.getBody()->getTerminator());
+
+  // Extract max value producer: arith.maximumf.
+  Value maxResult = yieldOp.getOperand(0);
+  auto maxOp = dyn_cast<arith::MaximumFOp>(maxResult.getDefiningOp());
+  if (!maxOp)
+    return failure();
+
+  // Extract index result producer: arith.select.
+  Value indexResult = yieldOp.getOperand(1);
+  auto selectOp = dyn_cast<arith::SelectOp>(indexResult.getDefiningOp());
+  if (!selectOp)
path: compiler/src/iree/compiler/DispatchCreation/SplitReduction.cpp
user: login: pashu123
body: same.
created_at: 2025-05-05T18:29:20Z
author_association: CONTRIBUTOR
diff_hunk: @@ -24,14 +25,344 @@ namespace mlir::iree_compiler::DispatchCreation {
 
 // TODO(thomasraoux): Move to attributes.
 static llvm::cl::opt<int64_t>
-    splitReductionRatio("iree-dispatch-creation-split-matmul-reduction",
-                        llvm::cl::desc("split ratio"), llvm::cl::init(1));
+    splitMatmulReductionRatio("iree-dispatch-creation-split-matmul-reduction",
+                              llvm::cl::desc("split ratio"), llvm::cl::init(1));
+
+static llvm::cl::opt<int64_t> splitArgmaxReductionRatio(
+    "iree-dispatch-creation-split-argmax-reduction",
+    llvm::cl::desc("Ratio to split argmax. Set to 0 or 1 to disable"),
+    llvm::cl::init(1));
 
 static llvm::cl::list<int64_t> topkSplitReductionRatio(
     "iree-dispatch-creation-topk-split-reduction",
     llvm::cl::desc("comma separated list of split ratios"),
     llvm::cl::CommaSeparated);
 
+template <typename LinalgOpTy>
+static FailureOr<linalg::SplitReductionResult>
+splitReductionImpl(RewriterBase &rewriter, LinalgOpTy op,
+                   linalg::ControlSplitReductionFn controlSplitReductionFn) {
+  return linalg::splitReduction(rewriter, op, controlSplitReductionFn);
+}
+
+// Matches the combiner pattern in a linalg.generic argmax-style reduction:
+// Example MLIR:
+// %4:2 = linalg.generic {
+//     indexing_maps = [...],
+//     iterator_types = ["parallel", "reduction"]
+//   } ins(%arg0 : tensor<?x128xbf16>) outs(%1, %3 : tensor<?xbf16>,
+//   tensor<?xi64>) {
+// ^bb0(%in: bf16, %out: bf16, %out_0: i64):
+//   %5 = linalg.index 1 : index
+//   %6 = arith.index_cast %5 : index to i64
+//   %7 = arith.maximumf %in, %out : bf16
+//   %8 = arith.cmpf ogt, %in, %out : bf16
+//   %9 = arith.select %8, %6, %out_0 : i64
+//   linalg.yield %7, %9 : bf16, i64
+// } -> (tensor<?xbf16>, tensor<?xi64>)
+//
+// This function extracts the `arith.maximumf`, `arith.cmpf`, and `arith.select`
+// operations from the body to facilitate transformations such as split
+// reduction.
+static LogicalResult
+collectArgmaxCombinerOps(linalg::GenericOp genericOp,
+                         SmallVector<Operation *, 4> &combinerOps) {
path: compiler/src/iree/compiler/DispatchCreation/SplitReduction.cpp
user: login: pashu123
body: You can pass `MutableArrayRef<Operation*> combinerOps`.
created_at: 2025-05-05T18:33:39Z
author_association: CONTRIBUTOR
diff_hunk: @@ -24,14 +25,344 @@ namespace mlir::iree_compiler::DispatchCreation {
 
 // TODO(thomasraoux): Move to attributes.
 static llvm::cl::opt<int64_t>
-    splitReductionRatio("iree-dispatch-creation-split-matmul-reduction",
-                        llvm::cl::desc("split ratio"), llvm::cl::init(1));
+    splitMatmulReductionRatio("iree-dispatch-creation-split-matmul-reduction",
+                              llvm::cl::desc("split ratio"), llvm::cl::init(1));
+
+static llvm::cl::opt<int64_t> splitArgmaxReductionRatio(
+    "iree-dispatch-creation-split-argmax-reduction",
+    llvm::cl::desc("Ratio to split argmax. Set to 0 or 1 to disable"),
+    llvm::cl::init(1));
 
 static llvm::cl::list<int64_t> topkSplitReductionRatio(
     "iree-dispatch-creation-topk-split-reduction",
     llvm::cl::desc("comma separated list of split ratios"),
     llvm::cl::CommaSeparated);
 
+template <typename LinalgOpTy>
+static FailureOr<linalg::SplitReductionResult>
+splitReductionImpl(RewriterBase &rewriter, LinalgOpTy op,
+                   linalg::ControlSplitReductionFn controlSplitReductionFn) {
+  return linalg::splitReduction(rewriter, op, controlSplitReductionFn);
+}
+
+// Matches the combiner pattern in a linalg.generic argmax-style reduction:
+// Example MLIR:
+// %4:2 = linalg.generic {
+//     indexing_maps = [...],
+//     iterator_types = ["parallel", "reduction"]
+//   } ins(%arg0 : tensor<?x128xbf16>) outs(%1, %3 : tensor<?xbf16>,
+//   tensor<?xi64>) {
+// ^bb0(%in: bf16, %out: bf16, %out_0: i64):
+//   %5 = linalg.index 1 : index
+//   %6 = arith.index_cast %5 : index to i64
+//   %7 = arith.maximumf %in, %out : bf16
+//   %8 = arith.cmpf ogt, %in, %out : bf16
+//   %9 = arith.select %8, %6, %out_0 : i64
+//   linalg.yield %7, %9 : bf16, i64
+// } -> (tensor<?xbf16>, tensor<?xi64>)
+//
+// This function extracts the `arith.maximumf`, `arith.cmpf`, and `arith.select`
+// operations from the body to facilitate transformations such as split
+// reduction.
+static LogicalResult
+collectArgmaxCombinerOps(linalg::GenericOp genericOp,
+                         SmallVector<Operation *, 4> &combinerOps) {
+  assert(IREE::LinalgExt::isArgmaxOp(genericOp) &&
+         "expected operation to be an argmax op");
+
+  auto yieldOp = cast<linalg::YieldOp>(genericOp.getBody()->getTerminator());
+
+  // Extract max value producer: arith.maximumf.
+  Value maxResult = yieldOp.getOperand(0);
+  auto maxOp = dyn_cast<arith::MaximumFOp>(maxResult.getDefiningOp());
+  if (!maxOp)
+    return failure();
+
+  // Extract index result producer: arith.select.
+  Value indexResult = yieldOp.getOperand(1);
+  auto selectOp = dyn_cast<arith::SelectOp>(indexResult.getDefiningOp());
+  if (!selectOp)
+    return failure();
+
+  // Extract the condition of the select, expected to be arith.cmpf with
+  // predicate OGT.
+  auto cmpOp = dyn_cast<arith::CmpFOp>(selectOp.getCondition().getDefiningOp());
+  if (!cmpOp || cmpOp.getPredicate() != arith::CmpFPredicate::OGT)
+    return failure();
+
+  combinerOps.push_back(maxOp);
+  combinerOps.push_back(selectOp);
+  combinerOps.push_back(cmpOp);
+
+  return success();
+}
+
+template <>
+FailureOr<linalg::SplitReductionResult> splitReductionImpl<linalg::GenericOp>(
+    RewriterBase &rewriter, linalg::GenericOp genericOp,
+    linalg::ControlSplitReductionFn controlSplitReductionFn) {
+  assert(IREE::LinalgExt::isArgmaxOp(genericOp) &&
+         "expected operation to be an argmax op");
+
+  OpBuilder::InsertionGuard guard(rewriter);
+  rewriter.setInsertionPoint(genericOp);
+  Location loc = genericOp->getLoc();
+
+  linalg::SplitReductionOptions control = controlSplitReductionFn(genericOp);
+  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);
+
+  if (dims.size() != 1) {
+    return rewriter.notifyMatchFailure(genericOp,
+                                       "needs a single reduction dimension");
+  }
path: compiler/src/iree/compiler/DispatchCreation/SplitReduction.cpp
user: login: pashu123
body: I don't think you need this check again. This has already been checked by the isArgMaxOp helper.
created_at: 2025-05-05T18:36:49Z
author_association: CONTRIBUTOR
diff_hunk: @@ -24,14 +25,344 @@ namespace mlir::iree_compiler::DispatchCreation {
 
 // TODO(thomasraoux): Move to attributes.
 static llvm::cl::opt<int64_t>
-    splitReductionRatio("iree-dispatch-creation-split-matmul-reduction",
-                        llvm::cl::desc("split ratio"), llvm::cl::init(1));
+    splitMatmulReductionRatio("iree-dispatch-creation-split-matmul-reduction",
+                              llvm::cl::desc("split ratio"), llvm::cl::init(1));
+
+static llvm::cl::opt<int64_t> splitArgmaxReductionRatio(
+    "iree-dispatch-creation-split-argmax-reduction",
+    llvm::cl::desc("Ratio to split argmax. Set to 0 or 1 to disable"),
+    llvm::cl::init(1));
 
 static llvm::cl::list<int64_t> topkSplitReductionRatio(
     "iree-dispatch-creation-topk-split-reduction",
     llvm::cl::desc("comma separated list of split ratios"),
     llvm::cl::CommaSeparated);
 
+template <typename LinalgOpTy>
+static FailureOr<linalg::SplitReductionResult>
+splitReductionImpl(RewriterBase &rewriter, LinalgOpTy op,
+                   linalg::ControlSplitReductionFn controlSplitReductionFn) {
+  return linalg::splitReduction(rewriter, op, controlSplitReductionFn);
+}
+
+// Matches the combiner pattern in a linalg.generic argmax-style reduction:
+// Example MLIR:
+// %4:2 = linalg.generic {
+//     indexing_maps = [...],
+//     iterator_types = ["parallel", "reduction"]
+//   } ins(%arg0 : tensor<?x128xbf16>) outs(%1, %3 : tensor<?xbf16>,
+//   tensor<?xi64>) {
+// ^bb0(%in: bf16, %out: bf16, %out_0: i64):
+//   %5 = linalg.index 1 : index
+//   %6 = arith.index_cast %5 : index to i64
+//   %7 = arith.maximumf %in, %out : bf16
+//   %8 = arith.cmpf ogt, %in, %out : bf16
+//   %9 = arith.select %8, %6, %out_0 : i64
+//   linalg.yield %7, %9 : bf16, i64
+// } -> (tensor<?xbf16>, tensor<?xi64>)
+//
+// This function extracts the `arith.maximumf`, `arith.cmpf`, and `arith.select`
+// operations from the body to facilitate transformations such as split
+// reduction.
+static LogicalResult
+collectArgmaxCombinerOps(linalg::GenericOp genericOp,
+                         SmallVector<Operation *, 4> &combinerOps) {
+  assert(IREE::LinalgExt::isArgmaxOp(genericOp) &&
+         "expected operation to be an argmax op");
+
+  auto yieldOp = cast<linalg::YieldOp>(genericOp.getBody()->getTerminator());
+
+  // Extract max value producer: arith.maximumf.
+  Value maxResult = yieldOp.getOperand(0);
+  auto maxOp = dyn_cast<arith::MaximumFOp>(maxResult.getDefiningOp());
+  if (!maxOp)
+    return failure();
+
+  // Extract index result producer: arith.select.
+  Value indexResult = yieldOp.getOperand(1);
+  auto selectOp = dyn_cast<arith::SelectOp>(indexResult.getDefiningOp());
+  if (!selectOp)
+    return failure();
+
+  // Extract the condition of the select, expected to be arith.cmpf with
+  // predicate OGT.
+  auto cmpOp = dyn_cast<arith::CmpFOp>(selectOp.getCondition().getDefiningOp());
+  if (!cmpOp || cmpOp.getPredicate() != arith::CmpFPredicate::OGT)
+    return failure();
+
+  combinerOps.push_back(maxOp);
+  combinerOps.push_back(selectOp);
+  combinerOps.push_back(cmpOp);
+
+  return success();
+}
+
+template <>
+FailureOr<linalg::SplitReductionResult> splitReductionImpl<linalg::GenericOp>(
+    RewriterBase &rewriter, linalg::GenericOp genericOp,
+    linalg::ControlSplitReductionFn controlSplitReductionFn) {
+  assert(IREE::LinalgExt::isArgmaxOp(genericOp) &&
+         "expected operation to be an argmax op");
+
+  OpBuilder::InsertionGuard guard(rewriter);
+  rewriter.setInsertionPoint(genericOp);
+  Location loc = genericOp->getLoc();
+
+  linalg::SplitReductionOptions control = controlSplitReductionFn(genericOp);
+  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);
+
+  if (dims.size() != 1) {
+    return rewriter.notifyMatchFailure(genericOp,
+                                       "needs a single reduction dimension");
+  }
+  unsigned reductionDim = dims[0];
+  if (control.innerParallel) {
+    insertSplitDimension = reductionDim + 1;
+  }
+
+  SmallVector<int64_t, 4> loopRanges = genericOp.getStaticLoopRanges();
+  int64_t reductionDimSize = loopRanges[reductionDim];
+  if (reductionDimSize == ShapedType::kDynamic ||
+      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");
+  }
+
+  SmallVector<Operation *, 4> combinerOps;
+  if (failed(collectArgmaxCombinerOps(genericOp, combinerOps)))
+    return rewriter.notifyMatchFailure(genericOp, "invalid combiner");
path: compiler/src/iree/compiler/DispatchCreation/SplitReduction.cpp
user: login: pashu123
body: NIT: "invalid combiner for argmax"
created_at: 2025-05-05T18:40:31Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,71 @@
+// RUN: iree-opt --pass-pipeline='builtin.module(util.func(iree-dispatch-creation-split-reduction-ops,cse, canonicalize))' --iree-dispatch-creation-split-argmax-reduction=4 %s | FileCheck %s
+
+util.func public @argmax(%arg0: tensor<?x131072xbf16>, %arg1: index) -> tensor<?xi64> {
+  %cst = arith.constant 0xFF80 : bf16
+  %c0_i64 = arith.constant 0 : i64
+  %0 = tensor.empty(%arg1) : tensor<?xbf16>
+  %1 = linalg.fill ins(%cst : bf16) outs(%0 : tensor<?xbf16>) -> tensor<?xbf16>
+  %2 = tensor.empty(%arg1) : tensor<?xi64>
+  %3 = linalg.fill ins(%c0_i64 : i64) outs(%2 : tensor<?xi64>) -> tensor<?xi64>
+  %4:2 = linalg.generic {
+      indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
+                       affine_map<(d0, d1) -> (d0)>,
+                       affine_map<(d0, d1) -> (d0)>],
+      iterator_types = ["parallel", "reduction"]}
+      ins(%arg0 : tensor<?x131072xbf16>) outs(%1, %3 : tensor<?xbf16>, tensor<?xi64>) {
+  ^bb0(%in: bf16, %out: bf16, %out_0: i64):
+    %5 = linalg.index 1 : index
+    %6 = arith.index_cast %5 : index to i64
+    %7 = arith.maximumf %in, %out : bf16
+    %8 = arith.cmpf ogt, %in, %out : bf16
+    %9 = arith.select %8, %6, %out_0 : i64
+    linalg.yield %7, %9 : bf16, i64
+  } -> (tensor<?xbf16>, tensor<?xi64>)
+  util.return %4#1 : tensor<?xi64>
+}
+
+// CHECK-LABEL:   util.func public @argmax
+
+// Check identity value preparation
+// CHECK: %[[CST:.+]] = arith.constant 0xFF80 : bf16
+// CHECK: %[[ZERO:.+]] = arith.constant 0 : i64
+// CHECK: %[[FINALVAL_EMPTY:.+]] = tensor.empty(%[[ARG1:.+]]) : tensor<?xbf16>
path: compiler/src/iree/compiler/DispatchCreation/test/split_argmax_reduction.mlir
user: login: bangtianliu
body: These are identity (or initial) value and index used for reduction.
created_at: 2025-05-05T20:22:03Z
author_association: CONTRIBUTOR
diff_hunk: @@ -826,4 +826,102 @@ bool isaHorizontallyFusedContraction(Operation *op) {
   return true;
 }
 
+bool isArgmaxOp(linalg::GenericOp genericOp) {
+  // Check for 2 results(value, index), and 1 input
+  if (genericOp.getNumDpsInits() != 2) {
+    return false;
+  }
+  if (genericOp.getNumDpsInputs() != 1) {
+    return false;
+  }
+
+  // If max value is being used, it is not a pure argmax.
+  if (!genericOp.getResults()[0].use_empty()) {
+    return false;
+  }
+
+  // Check that the rank is at least 3 and all loops are parallel
+  unsigned numLoops = genericOp.getNumLoops();
+  unsigned numParallelLoops = genericOp.getNumParallelLoops();
+
+  // Argmax will require 1D reduction.
+  if (numParallelLoops != (numLoops - 1)) {
+    return false;
+  }
+  // TODO: Add better affine map checks.
+  auto indexing_maps = genericOp.getIndexingMapsArray();
+  if (!indexing_maps[0].isIdentity())
+    return false;
+
+  // Check that initial value is negative Infinite.
+  // TODO: Move this check to ukernel once we implement
+  //       variant to handle non neg-Inf initial value.
+  Value initVal = genericOp.getDpsInitOperand(0)->get();
+  auto fillOp = initVal.getDefiningOp<linalg::FillOp>();
+  if (!fillOp)
+    return false;
+  Value fillVal = fillOp.getDpsInputOperand(0)->get();
+  if (!matchPattern(fillVal, m_NegInfFloat()))
+    return false;
+
path: compiler/src/iree/compiler/Dialect/LinalgExt/Utils/Utils.cpp
user: login: bangtianliu
body: I didn't originally write this logic for this function. You can refer to the comment here for the reasoning: https://github.com/iree-org/iree/blob/main/compiler/src/iree/compiler/Codegen/Utils/Utils.cpp#L1347.

I just moved the code from iree/compiler/src/iree/compiler/Codegen/Utils/Utils.cpp to compiler/src/iree/compiler/Dialect/LinalgExt/Utils/Utils.cpp, following Mahesh’s branch: https://github.com/MaheshRavishankar/iree/commit/2e87623675840153525d21248d41beb57d516f82.
created_at: 2025-05-06T14:54:37Z
author_association: CONTRIBUTOR
diff_hunk: @@ -826,4 +826,102 @@ bool isaHorizontallyFusedContraction(Operation *op) {
   return true;
 }
 
+bool isArgmaxOp(linalg::GenericOp genericOp) {
path: compiler/src/iree/compiler/Dialect/LinalgExt/Utils/Utils.cpp
user: login: MaheshRavishankar
body: I think that would be inconsistent with other methods here.
created_at: 2025-05-06T18:03:00Z
author_association: COLLABORATOR
diff_hunk: @@ -826,4 +826,102 @@ bool isaHorizontallyFusedContraction(Operation *op) {
   return true;
 }
 
+bool isArgmaxOp(linalg::GenericOp genericOp) {
path: compiler/src/iree/compiler/Dialect/LinalgExt/Utils/Utils.cpp
user: login: hanhanW
body: > I suggest using LogicalResult with better error messages so that if there's a case we want to support, people will know exactly where to add the support.

I'm curious about how it works in your mind? I.e, how do we propagate the error message to listener (or whatever stuff that I'm not familiar with)?

My feeling is that people write code like below, and I don't connect how returning `LogicalResult` connects to the code properly. The "better error message" is dropped in this kind of implementation.

```cpp
if (failed(isArgmaxOp(op)) {
  return rewriter.noitfyMatchFailure(op, "not a argmax op");
}
```

If this is the most common usage, I'm -1 returning LogicalResult.

(I personally don't like returning LogicalResult when possible. Boolean is simple and covers the needs for many cases, IMO. I use LogicalResult when I see the values or for style consistency.)
created_at: 2025-05-06T18:40:43Z
author_association: CONTRIBUTOR
diff_hunk: @@ -826,4 +826,102 @@ bool isaHorizontallyFusedContraction(Operation *op) {
   return true;
 }
 
+bool isArgmaxOp(linalg::GenericOp genericOp) {
path: compiler/src/iree/compiler/Dialect/LinalgExt/Utils/Utils.cpp
user: login: pashu123
body: The isArgmaxOp function is too big and contains many if-else's, just returning the boolean doesn't pinpoint the exact location inside the isArgmaxOp where the matcher isn't working.
created_at: 2025-05-06T18:46:46Z
author_association: CONTRIBUTOR
diff_hunk: @@ -24,14 +25,355 @@ namespace mlir::iree_compiler::DispatchCreation {
 
 // TODO(thomasraoux): Move to attributes.
 static llvm::cl::opt<int64_t>
-    splitReductionRatio("iree-dispatch-creation-split-matmul-reduction",
-                        llvm::cl::desc("split ratio"), llvm::cl::init(1));
+    splitMatmulReductionRatio("iree-dispatch-creation-split-matmul-reduction",
+                              llvm::cl::desc("split ratio"), llvm::cl::init(1));
+
+static llvm::cl::opt<bool> enableDefaultArgmaxSplitPattern(
+    "iree-dispatch-creation-enable-default-argmax-split-pattern",
+    llvm::cl::desc("Enable default argmax split-k for known reduction patterns "
+                   "with a reduction dimension of 128."),
+    llvm::cl::init(true));
+
+// Controls the tile size used when applying split-k to argmax reductions.
+// This value defines how many elements along the reduction dimension are
+// processed per tile (e.g., a value of 128 means each tile reduces over 128
+// elements).
+static llvm::cl::opt<int64_t> splitArgmaxReductionRatio(
+    "iree-dispatch-creation-split-argmax-reduction",
+    llvm::cl::desc("Ratio to split argmax. Set to 0 or 1 to disable"),
+    llvm::cl::init(1));
 
 static llvm::cl::list<int64_t> topkSplitReductionRatio(
     "iree-dispatch-creation-topk-split-reduction",
     llvm::cl::desc("comma separated list of split ratios"),
     llvm::cl::CommaSeparated);
 
+template <typename LinalgOpTy>
+static FailureOr<linalg::SplitReductionResult>
+splitReductionImpl(RewriterBase &rewriter, LinalgOpTy op,
+                   linalg::ControlSplitReductionFn controlSplitReductionFn) {
+  return linalg::splitReduction(rewriter, op, controlSplitReductionFn);
+}
+
+// Matches the combiner pattern in a linalg.generic argmax-style reduction:
+// Example MLIR:
+// %4:2 = linalg.generic {
+//     indexing_maps = [...],
+//     iterator_types = ["parallel", "reduction"]
+//   } ins(%arg0 : tensor<?x128xbf16>) outs(%1, %3 : tensor<?xbf16>,
+//   tensor<?xi64>) {
+// ^bb0(%in: bf16, %out: bf16, %out_0: i64):
+//   %5 = linalg.index 1 : index
+//   %6 = arith.index_cast %5 : index to i64
+//   %7 = arith.maximumf %in, %out : bf16
+//   %8 = arith.cmpf ogt, %in, %out : bf16
+//   %9 = arith.select %8, %6, %out_0 : i64
+//   linalg.yield %7, %9 : bf16, i64
+// } -> (tensor<?xbf16>, tensor<?xi64>)
+//
+// This function extracts the `arith.maximumf`, `arith.cmpf`, and `arith.select`
+// operations from the body to facilitate transformations such as split
+// reduction.
+static LogicalResult
+collectArgmaxCombinerOps(linalg::GenericOp genericOp,
+                         MutableArrayRef<Operation *> combinerOps) {
+  if (combinerOps.size() < 3) {
+    return genericOp->emitError(
+        "combinerOps must have space for exactly 3 elements");
+  }
+
+  assert(IREE::LinalgExt::isArgmaxOp(genericOp) &&
+         "expected operation to be an argmax op");
+
+  auto yieldOp = cast<linalg::YieldOp>(genericOp.getBody()->getTerminator());
+
+  // Extract max value producer: arith.maximumf.
+  Value maxResult = yieldOp.getOperand(0);
+  auto maxOp = dyn_cast<arith::MaximumFOp>(maxResult.getDefiningOp());
+
+  // Extract index result producer: arith.select.
+  Value indexResult = yieldOp.getOperand(1);
+  auto selectOp = dyn_cast<arith::SelectOp>(indexResult.getDefiningOp());
+
+  // Extract the condition of the select, expected to be arith.cmpf with
+  // predicate OGT.
+  auto cmpOp = dyn_cast<arith::CmpFOp>(selectOp.getCondition().getDefiningOp());
+
+  combinerOps[0] = maxOp;
+  combinerOps[1] = selectOp;
+  combinerOps[2] = cmpOp;
path: compiler/src/iree/compiler/DispatchCreation/SplitReduction.cpp
user: login: pashu123
body: I would rather create a struct combinerOps with Operation* max, ..., so you can directly query the maxOp and don't care about the extra bookkeeping.
created_at: 2025-05-07T16:26:54Z
author_association: CONTRIBUTOR
diff_hunk: @@ -24,14 +25,355 @@ namespace mlir::iree_compiler::DispatchCreation {
 
 // TODO(thomasraoux): Move to attributes.
 static llvm::cl::opt<int64_t>
-    splitReductionRatio("iree-dispatch-creation-split-matmul-reduction",
-                        llvm::cl::desc("split ratio"), llvm::cl::init(1));
+    splitMatmulReductionRatio("iree-dispatch-creation-split-matmul-reduction",
+                              llvm::cl::desc("split ratio"), llvm::cl::init(1));
+
+static llvm::cl::opt<bool> enableDefaultArgmaxSplitPattern(
+    "iree-dispatch-creation-enable-default-argmax-split-pattern",
+    llvm::cl::desc("Enable default argmax split-k for known reduction patterns "
+                   "with a reduction dimension of 128."),
+    llvm::cl::init(true));
path: compiler/src/iree/compiler/DispatchCreation/SplitReduction.cpp
user: login: hanhanW
body: What is default argmax split-k?
created_at: 2025-05-07T17:25:07Z
author_association: CONTRIBUTOR
diff_hunk: @@ -24,14 +25,355 @@ namespace mlir::iree_compiler::DispatchCreation {
 
 // TODO(thomasraoux): Move to attributes.
 static llvm::cl::opt<int64_t>
-    splitReductionRatio("iree-dispatch-creation-split-matmul-reduction",
-                        llvm::cl::desc("split ratio"), llvm::cl::init(1));
+    splitMatmulReductionRatio("iree-dispatch-creation-split-matmul-reduction",
+                              llvm::cl::desc("split ratio"), llvm::cl::init(1));
+
+static llvm::cl::opt<bool> enableDefaultArgmaxSplitPattern(
+    "iree-dispatch-creation-enable-default-argmax-split-pattern",
+    llvm::cl::desc("Enable default argmax split-k for known reduction patterns "
+                   "with a reduction dimension of 128."),
+    llvm::cl::init(true));
+
+// Controls the tile size used when applying split-k to argmax reductions.
+// This value defines how many elements along the reduction dimension are
+// processed per tile (e.g., a value of 128 means each tile reduces over 128
+// elements).
+static llvm::cl::opt<int64_t> splitArgmaxReductionRatio(
+    "iree-dispatch-creation-split-argmax-reduction",
+    llvm::cl::desc("Ratio to split argmax. Set to 0 or 1 to disable"),
+    llvm::cl::init(1));
 
 static llvm::cl::list<int64_t> topkSplitReductionRatio(
     "iree-dispatch-creation-topk-split-reduction",
     llvm::cl::desc("comma separated list of split ratios"),
     llvm::cl::CommaSeparated);
 
+template <typename LinalgOpTy>
+static FailureOr<linalg::SplitReductionResult>
+splitReductionImpl(RewriterBase &rewriter, LinalgOpTy op,
+                   linalg::ControlSplitReductionFn controlSplitReductionFn) {
+  return linalg::splitReduction(rewriter, op, controlSplitReductionFn);
+}
+
+// Matches the combiner pattern in a linalg.generic argmax-style reduction:
+// Example MLIR:
+// %4:2 = linalg.generic {
+//     indexing_maps = [...],
+//     iterator_types = ["parallel", "reduction"]
+//   } ins(%arg0 : tensor<?x128xbf16>) outs(%1, %3 : tensor<?xbf16>,
+//   tensor<?xi64>) {
+// ^bb0(%in: bf16, %out: bf16, %out_0: i64):
+//   %5 = linalg.index 1 : index
+//   %6 = arith.index_cast %5 : index to i64
+//   %7 = arith.maximumf %in, %out : bf16
+//   %8 = arith.cmpf ogt, %in, %out : bf16
+//   %9 = arith.select %8, %6, %out_0 : i64
+//   linalg.yield %7, %9 : bf16, i64
+// } -> (tensor<?xbf16>, tensor<?xi64>)
+//
+// This function extracts the `arith.maximumf`, `arith.cmpf`, and `arith.select`
+// operations from the body to facilitate transformations such as split
+// reduction.
+static LogicalResult
+collectArgmaxCombinerOps(linalg::GenericOp genericOp,
+                         MutableArrayRef<Operation *> combinerOps) {
+  if (combinerOps.size() < 3) {
+    return genericOp->emitError(
+        "combinerOps must have space for exactly 3 elements");
+  }
+
+  assert(IREE::LinalgExt::isArgmaxOp(genericOp) &&
+         "expected operation to be an argmax op");
+
+  auto yieldOp = cast<linalg::YieldOp>(genericOp.getBody()->getTerminator());
+
+  // Extract max value producer: arith.maximumf.
+  Value maxResult = yieldOp.getOperand(0);
+  auto maxOp = dyn_cast<arith::MaximumFOp>(maxResult.getDefiningOp());
+
+  // Extract index result producer: arith.select.
+  Value indexResult = yieldOp.getOperand(1);
+  auto selectOp = dyn_cast<arith::SelectOp>(indexResult.getDefiningOp());
+
+  // Extract the condition of the select, expected to be arith.cmpf with
+  // predicate OGT.
+  auto cmpOp = dyn_cast<arith::CmpFOp>(selectOp.getCondition().getDefiningOp());
+
+  combinerOps[0] = maxOp;
+  combinerOps[1] = selectOp;
+  combinerOps[2] = cmpOp;
+
+  return success();
+}
+
+template <>
+FailureOr<linalg::SplitReductionResult> splitReductionImpl<linalg::GenericOp>(
+    RewriterBase &rewriter, linalg::GenericOp genericOp,
+    linalg::ControlSplitReductionFn controlSplitReductionFn) {
+  assert(IREE::LinalgExt::isArgmaxOp(genericOp) &&
+         "expected operation to be an argmax op");
+
+  OpBuilder::InsertionGuard guard(rewriter);
+  rewriter.setInsertionPoint(genericOp);
+  Location loc = genericOp->getLoc();
+
+  linalg::SplitReductionOptions control = controlSplitReductionFn(genericOp);
+  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, 4> loopRanges = genericOp.getStaticLoopRanges();
+  int64_t reductionDimSize = loopRanges[reductionDim];
+
+  //  The total number of output elements along this new dimension is
+  //  reductionDimSize / ratio.
+  int64_t output_dimsize = reductionDimSize / ratio;
path: compiler/src/iree/compiler/DispatchCreation/SplitReduction.cpp
user: login: hanhanW
body: We follow camelCase style.

```suggestion
  int64_t outputDimSize = reductionDimSize / ratio;
```
created_at: 2025-05-07T17:34:34Z
author_association: CONTRIBUTOR
diff_hunk: @@ -24,14 +25,355 @@ namespace mlir::iree_compiler::DispatchCreation {
 
 // TODO(thomasraoux): Move to attributes.
 static llvm::cl::opt<int64_t>
-    splitReductionRatio("iree-dispatch-creation-split-matmul-reduction",
-                        llvm::cl::desc("split ratio"), llvm::cl::init(1));
+    splitMatmulReductionRatio("iree-dispatch-creation-split-matmul-reduction",
+                              llvm::cl::desc("split ratio"), llvm::cl::init(1));
+
+static llvm::cl::opt<bool> enableDefaultArgmaxSplitPattern(
+    "iree-dispatch-creation-enable-default-argmax-split-pattern",
+    llvm::cl::desc("Enable default argmax split-k for known reduction patterns "
+                   "with a reduction dimension of 128."),
+    llvm::cl::init(true));
+
+// Controls the tile size used when applying split-k to argmax reductions.
+// This value defines how many elements along the reduction dimension are
+// processed per tile (e.g., a value of 128 means each tile reduces over 128
+// elements).
+static llvm::cl::opt<int64_t> splitArgmaxReductionRatio(
+    "iree-dispatch-creation-split-argmax-reduction",
+    llvm::cl::desc("Ratio to split argmax. Set to 0 or 1 to disable"),
+    llvm::cl::init(1));
 
 static llvm::cl::list<int64_t> topkSplitReductionRatio(
     "iree-dispatch-creation-topk-split-reduction",
     llvm::cl::desc("comma separated list of split ratios"),
     llvm::cl::CommaSeparated);
 
+template <typename LinalgOpTy>
+static FailureOr<linalg::SplitReductionResult>
+splitReductionImpl(RewriterBase &rewriter, LinalgOpTy op,
+                   linalg::ControlSplitReductionFn controlSplitReductionFn) {
+  return linalg::splitReduction(rewriter, op, controlSplitReductionFn);
+}
+
+// Matches the combiner pattern in a linalg.generic argmax-style reduction:
+// Example MLIR:
+// %4:2 = linalg.generic {
+//     indexing_maps = [...],
+//     iterator_types = ["parallel", "reduction"]
+//   } ins(%arg0 : tensor<?x128xbf16>) outs(%1, %3 : tensor<?xbf16>,
+//   tensor<?xi64>) {
+// ^bb0(%in: bf16, %out: bf16, %out_0: i64):
+//   %5 = linalg.index 1 : index
+//   %6 = arith.index_cast %5 : index to i64
+//   %7 = arith.maximumf %in, %out : bf16
+//   %8 = arith.cmpf ogt, %in, %out : bf16
+//   %9 = arith.select %8, %6, %out_0 : i64
+//   linalg.yield %7, %9 : bf16, i64
+// } -> (tensor<?xbf16>, tensor<?xi64>)
+//
+// This function extracts the `arith.maximumf`, `arith.cmpf`, and `arith.select`
+// operations from the body to facilitate transformations such as split
+// reduction.
+static LogicalResult
+collectArgmaxCombinerOps(linalg::GenericOp genericOp,
+                         MutableArrayRef<Operation *> combinerOps) {
+  if (combinerOps.size() < 3) {
+    return genericOp->emitError(
+        "combinerOps must have space for exactly 3 elements");
+  }
+
+  assert(IREE::LinalgExt::isArgmaxOp(genericOp) &&
+         "expected operation to be an argmax op");
+
+  auto yieldOp = cast<linalg::YieldOp>(genericOp.getBody()->getTerminator());
+
+  // Extract max value producer: arith.maximumf.
+  Value maxResult = yieldOp.getOperand(0);
+  auto maxOp = dyn_cast<arith::MaximumFOp>(maxResult.getDefiningOp());
+
+  // Extract index result producer: arith.select.
+  Value indexResult = yieldOp.getOperand(1);
+  auto selectOp = dyn_cast<arith::SelectOp>(indexResult.getDefiningOp());
+
+  // Extract the condition of the select, expected to be arith.cmpf with
+  // predicate OGT.
+  auto cmpOp = dyn_cast<arith::CmpFOp>(selectOp.getCondition().getDefiningOp());
+
+  combinerOps[0] = maxOp;
+  combinerOps[1] = selectOp;
+  combinerOps[2] = cmpOp;
+
+  return success();
+}
+
+template <>
+FailureOr<linalg::SplitReductionResult> splitReductionImpl<linalg::GenericOp>(
+    RewriterBase &rewriter, linalg::GenericOp genericOp,
+    linalg::ControlSplitReductionFn controlSplitReductionFn) {
+  assert(IREE::LinalgExt::isArgmaxOp(genericOp) &&
+         "expected operation to be an argmax op");
+
+  OpBuilder::InsertionGuard guard(rewriter);
+  rewriter.setInsertionPoint(genericOp);
+  Location loc = genericOp->getLoc();
+
+  linalg::SplitReductionOptions control = controlSplitReductionFn(genericOp);
+  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, 4> loopRanges = genericOp.getStaticLoopRanges();
+  int64_t reductionDimSize = loopRanges[reductionDim];
+
+  //  The total number of output elements along this new dimension is
+  //  reductionDimSize / ratio.
+  int64_t output_dimsize = reductionDimSize / ratio;
+
+  if (reductionDimSize == ShapedType::kDynamic ||
+      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");
+  }
+
+  Operation *ops[3] = {nullptr, nullptr, nullptr};
+  MutableArrayRef<Operation *> combinerOps(ops);
+  if (failed(collectArgmaxCombinerOps(genericOp, combinerOps)))
+    return rewriter.notifyMatchFailure(genericOp,
+                                       "invalid combiner for argmax");
+
+  Operation *reductionOp = combinerOps[0];
+  std::optional<TypedAttr> identity = arith::getNeutralElement(reductionOp);
+  if (!identity.has_value())
+    return rewriter.notifyMatchFailure(
+        genericOp, "Unknown identity value for the reduction");
+
+  SmallVector<Value> newInputs;
+  SmallVector<AffineMap> newMaps;
+  // Calculate the new shapes and indexing maps of the input operands.
+  for (OpOperand *operand : genericOp.getDpsInputOperands()) {
+    AffineMap map = genericOp.getMatchingIndexingMap(operand);
+    SmallVector<int64_t> newShape;
+    SmallVector<AffineExpr> exprs;
+    SmallVector<ReassociationIndices> reassociation;
+    unsigned index = 0;
+    for (unsigned idx : llvm::seq<unsigned>(0, map.getNumResults())) {
+      unsigned dim = map.getDimPosition(idx);
+      if (reductionDim == dim) {
+        if (control.innerParallel) {
+          newShape.push_back(ratio); // reduce
+          newShape.push_back(genericOp.getShape(operand)[idx] /
+                             ratio); // parallel (insert)
+          exprs.push_back(rewriter.getAffineDimExpr(
+              dim < insertSplitDimension ? dim : dim + 1));
+          exprs.push_back(rewriter.getAffineDimExpr(insertSplitDimension));
+        } else {
+          newShape.push_back(genericOp.getShape(operand)[idx] /
+                             ratio); // parallel (insert)
+          newShape.push_back(ratio); // reduce
+          exprs.push_back(rewriter.getAffineDimExpr(insertSplitDimension));
+          exprs.push_back(rewriter.getAffineDimExpr(
+              dim < insertSplitDimension ? dim : dim + 1));
+        }
+        reassociation.push_back({index++, index++});
+        continue;
+      }
+      newShape.push_back(genericOp.getShape(operand)[idx]);
+      exprs.push_back(rewriter.getAffineDimExpr(
+          dim < insertSplitDimension ? dim : dim + 1));
+      reassociation.push_back({index++});
+    }
+    newMaps.push_back(
+        AffineMap::get(map.getNumDims() + 1, 0, exprs, genericOp.getContext()));
+    // If the shape is unchanged the input doesn't change.
+    if (newShape == genericOp.getShape(operand)) {
+      newInputs.push_back(operand->get());
+      continue;
+    }
+    Type newType = RankedTensorType::get(
+        newShape,
+        cast<RankedTensorType>(operand->get().getType()).getElementType());
+
+    Value newInput = rewriter.create<tensor::ExpandShapeOp>(
+        loc, newType, operand->get(), reassociation);
+    newInputs.push_back(newInput);
+  }
+
+  SmallVector<SmallVector<int64_t>> newOutputShapes;
+  SmallVector<AffineMap> outputMaps;
+  for (int i = 0; i < genericOp.getNumDpsInits(); ++i) {
+    OpOperand *output = genericOp.getDpsInitOperand(i);
+    AffineMap oldOutputMap = genericOp.getMatchingIndexingMap(output);
+    ArrayRef<int64_t> oldShape = genericOp.getShape(output);
+    SmallVector<int64_t> thisOutputShape;
+
+    SmallVector<AffineExpr> outputExpr;
+    for (unsigned idx = 0; idx <= oldShape.size(); ++idx) {
+      if (idx == insertSplitIndex) {
+        thisOutputShape.push_back(output_dimsize);
+        outputExpr.push_back(rewriter.getAffineDimExpr(insertSplitDimension));
+      }
+      if (idx < oldShape.size()) {
+        thisOutputShape.push_back(oldShape[idx]);
+        unsigned dim = oldOutputMap.getDimPosition(idx);
+        outputExpr.push_back(rewriter.getAffineDimExpr(
+            dim < insertSplitDimension ? dim : dim + 1));
+      }
+    }
+
+    AffineMap newOutputMap = AffineMap::get(oldOutputMap.getNumDims() + 1, 0,
+                                            outputExpr, rewriter.getContext());
+    newMaps.push_back(newOutputMap);
+    newOutputShapes.push_back(thisOutputShape);
+  }
+
+  // Handle dynamic dimensions for identity value tensor.
+  SmallVector<Value> dynValDims;
+  SmallVector<int64_t> newOutputShape = newOutputShapes[0];
+  for (size_t i = 0; i < newOutputShape.size(); ++i) {
+    if (ShapedType::isDynamic(newOutputShape[i])) {
+      dynValDims.push_back(rewriter.create<tensor::DimOp>(
+          loc, genericOp.getDpsInputOperand(0)->get(), i));
+    }
+  }
+
+  Type valueElemType = genericOp.getRegionOutputArgs()[0].getType();
+  Value emptyValTensor = rewriter.create<tensor::EmptyOp>(
+      loc, newOutputShape, valueElemType, dynValDims);
+  Value constantOp = rewriter.create<arith::ConstantOp>(loc, *identity);
+  Value identityVal =
+      rewriter.create<linalg::FillOp>(loc, constantOp, emptyValTensor)
+          .getResult(0);
+
+  // Handle dynamic dimensions for identity index tensor.
+  SmallVector<Value> dynIdxDims;
+  newOutputShape = newOutputShapes[1];
+  for (size_t i = 0; i < newOutputShape.size(); ++i) {
+    if (ShapedType::isDynamic(newOutputShape[i])) {
+      dynIdxDims.push_back(rewriter.create<tensor::DimOp>(
+          loc, genericOp.getDpsInputOperand(0)->get(), i));
+    }
+  }
+  Type idxElemType = genericOp.getRegionOutputArgs()[1].getType();
+  Value zeroIdx = rewriter.create<arith::ConstantOp>(
+      loc, rewriter.getZeroAttr(idxElemType));
+  Value idxInitTensor = rewriter.create<tensor::EmptyOp>(
+      loc, newOutputShape, idxElemType, dynIdxDims);
+  Value identityIndex =
+      rewriter.create<linalg::FillOp>(loc, zeroIdx, idxInitTensor).getResult(0);
+
+  SmallVector<utils::IteratorType> newIteratorTypes;
+  for (auto [index, iteratorType] :
+       llvm::enumerate(genericOp.getIteratorTypesArray())) {
+    if (insertSplitDimension == index)
+      newIteratorTypes.push_back(utils::IteratorType::parallel);
+    newIteratorTypes.push_back(iteratorType);
+  }
+  if (insertSplitDimension == genericOp.getIteratorTypesArray().size()) {
+    newIteratorTypes.push_back(utils::IteratorType::parallel);
+  }
+
+  // Create partial linalg.generic op with global index computation.
+  Value tileSize = rewriter.create<arith::ConstantIndexOp>(loc, ratio);
+  auto partialOp = rewriter.create<linalg::GenericOp>(
+      loc, TypeRange{identityVal.getType(), identityIndex.getType()}, newInputs,
+      ValueRange{identityVal, identityIndex}, newMaps, newIteratorTypes);
+
+  rewriter.inlineRegionBefore(genericOp.getRegion(), partialOp.getRegion(),
+                              partialOp.getRegion().begin());
+
+  Block &body = partialOp.getRegion().front();
+  rewriter.setInsertionPointToStart(&body);
+
+  unsigned innerIdxDim = reductionDim + 1;
+  unsigned outerIdxDim = insertSplitDimension;
+
+  // Compute global index (gidx) for reduction when the original reduction
+  // dimension is split into [outerIdx, innerIdx] using `ratio`. This is used to
+  // correctly compute the global index for comparisons and index selection.
+  Value outerIdx = rewriter.create<linalg::IndexOp>(loc, outerIdxDim);
+  Value innerIdx = rewriter.create<linalg::IndexOp>(loc, innerIdxDim);
+  Value offset = rewriter.create<arith::MulIOp>(loc, outerIdx, tileSize);
+  Value gidx = rewriter.create<arith::AddIOp>(loc, offset, innerIdx);
+
+  auto selectOp = dyn_cast<arith::SelectOp>(combinerOps[1]);
+  Value oldIdx = selectOp.getTrueValue();
+  Value newIdx = gidx;
+  if (oldIdx.getType() != gidx.getType()) {
+    newIdx = rewriter.create<arith::IndexCastOp>(loc, oldIdx.getType(), gidx);
+  }
+  selectOp.setOperand(1, newIdx);
+  rewriter.setInsertionPointAfter(partialOp);
+
+  unsigned intermRank = newOutputShape.size();
+  AffineMap valueMap = rewriter.getMultiDimIdentityMap(intermRank);
+  AffineMap indexMap = valueMap;
+  SmallVector<utils::IteratorType> reductionIteratorTypes;
+  SmallVector<AffineExpr> resultExprs;
+  for (unsigned i : llvm::seq<unsigned>(0, intermRank)) {
+    if (insertSplitIndex == i) {
+      reductionIteratorTypes.push_back(utils::IteratorType::reduction);
+    } else {
+      resultExprs.push_back(rewriter.getAffineDimExpr(i));
+      reductionIteratorTypes.push_back(utils::IteratorType::parallel);
+    }
+  }
+
+  AffineMap outputMap =
+      AffineMap::get(intermRank, 0, resultExprs, rewriter.getContext());
+  SmallVector<AffineMap> finalReductionMaps = {valueMap, indexMap, outputMap,
+                                               outputMap};
+
+  assert(combinerOps.size() == 3 &&
+         "Expected exactly 3 combiner ops for argmax (max, select, cmp)");
+  assert(isa<arith::MaximumFOp>(combinerOps[0]) &&
+         "First combiner op must be arith.maximumf");
+  assert(isa<arith::SelectOp>(combinerOps[1]) &&
+         "Second combiner op must be arith.select");
+  assert(isa<arith::CmpFOp>(combinerOps[2]) &&
+         "Third combiner op must be arith.cmpf");
+
+  // Create block for final reduction region.
+  auto finalReduction = rewriter.create<linalg::GenericOp>(
+      loc, genericOp.getResultTypes(),
+      ValueRange{partialOp.getResult(0), partialOp.getResult(1)},
+      genericOp.getDpsInits(), finalReductionMaps, reductionIteratorTypes,
+      [combinerOps](OpBuilder &b, Location loc, ValueRange inputs) {
+        Operation *clonedMax = b.clone(*combinerOps[0]);
+        clonedMax->setOperands({inputs[0], inputs[2]});
+        Operation *clonedCmp = b.clone(*combinerOps[2]);
+        clonedCmp->setOperands({inputs[0], inputs[2]});
+        Operation *clonedSel = b.clone(*combinerOps[1]);
+        clonedSel->setOperands({clonedCmp->getResult(0), inputs[1], inputs[3]});
+        b.create<linalg::YieldOp>(
+            loc, ValueRange{clonedMax->getResult(0), clonedSel->getResult(0)});
+      });
+
+  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,
path: compiler/src/iree/compiler/DispatchCreation/SplitReduction.cpp
user: login: hanhanW
body: style nit:

```suggestion
      /*initOrAlloc=*/nullptr, /*fillOp=*/nullptr,
```
created_at: 2025-05-07T17:41:46Z
author_association: CONTRIBUTOR
diff_hunk: @@ -24,14 +25,355 @@ namespace mlir::iree_compiler::DispatchCreation {
 
 // TODO(thomasraoux): Move to attributes.
 static llvm::cl::opt<int64_t>
-    splitReductionRatio("iree-dispatch-creation-split-matmul-reduction",
-                        llvm::cl::desc("split ratio"), llvm::cl::init(1));
+    splitMatmulReductionRatio("iree-dispatch-creation-split-matmul-reduction",
+                              llvm::cl::desc("split ratio"), llvm::cl::init(1));
+
+static llvm::cl::opt<bool> enableDefaultArgmaxSplitPattern(
+    "iree-dispatch-creation-enable-default-argmax-split-pattern",
+    llvm::cl::desc("Enable default argmax split-k for known reduction patterns "
+                   "with a reduction dimension of 128."),
+    llvm::cl::init(true));
path: compiler/src/iree/compiler/DispatchCreation/SplitReduction.cpp
user: login: hanhanW
body: You don't need to address this comment now. It is a question when I saw the code. Default is ambiguous in the splitReduction context.
created_at: 2025-05-07T17:59:31Z
author_association: CONTRIBUTOR
path: compiler/src/iree/compiler/DispatchCreation/test/split_argmax_reduction.mlir
user: login: hanhanW
body: I'd put the test to `split_reduction.mlir`. I only split the files when it has too many test cases. When it happens, I start thinking if the pass has more responsibilities than it should. There are exceptions, though. This case does not fit my mental model about the split. So I suggest moving the test to split_reduction.mlir, and it is easier for others to find the test. (I think it is also how we structure the tests.)
created_at: 2025-05-09T17:30:54Z
author_association: CONTRIBUTOR
diff_hunk: @@ -63,28 +81,72 @@ namespace {
 struct SplitReductionPass final
     : public impl::SplitReductionPassBase<SplitReductionPass> {
   void runOnOperation() override {
-    if (splitReductionRatio.getValue() <= 1 &&
-        topkSplitReductionRatio.empty()) {
+    if (splitMatmulReductionRatio.getValue() <= 1 &&
+        topkSplitReductionRatio.empty() &&
+        (splitArgmaxReductionRatio.getValue() <= 1 &&
+         !enableStaticArgmaxSplit)) {
       return;
     }
 
+    if (enableStaticArgmaxSplit) {
+      // Use default split-k pattern for argmax ops: split the reduction dim
+      // (e.g., 131072) into tiles of size 128, resulting in 1024 tiles (131072
+      // / 128). So, the split ratio refers to the tile size of the reduction
+      // dimension.
+      splitArgmaxReductionRatio = 128;
+    }
path: compiler/src/iree/compiler/DispatchCreation/SplitReduction.cpp
user: login: hanhanW
body: Why do we need the additional flag, but not just making `iree-dispatch-creation-split-argmax-reduction` default to 128? The other question is that do we want it happening by default? If so, we should document how 128 is derived. If it is an experimental number, we can also just mention it in the comment. People will ask why when they see the magic number, IMO.

My expectation is that the default value is 1 (and we don't need the `enableStaticArgmaxSplit` variable). You use `--iree-dispatch-creation-split-argmax-reduction=128` in your compilation config. We can set the default value to a reasonable value when we are confident for this.

(I don't work on this area, so maybe we are already confident about this. Then it is okay to set it to 128 by default.)
created_at: 2025-05-09T17:42:19Z
author_association: CONTRIBUTOR
diff_hunk: @@ -63,28 +81,72 @@ namespace {
 struct SplitReductionPass final
     : public impl::SplitReductionPassBase<SplitReductionPass> {
   void runOnOperation() override {
-    if (splitReductionRatio.getValue() <= 1 &&
-        topkSplitReductionRatio.empty()) {
+    if (splitMatmulReductionRatio.getValue() <= 1 &&
+        topkSplitReductionRatio.empty() &&
+        (splitArgmaxReductionRatio.getValue() <= 1 &&
+         !enableStaticArgmaxSplit)) {
       return;
     }
 
+    if (enableStaticArgmaxSplit) {
+      // Use default split-k pattern for argmax ops: split the reduction dim
+      // (e.g., 131072) into tiles of size 128, resulting in 1024 tiles (131072
+      // / 128). So, the split ratio refers to the tile size of the reduction
+      // dimension.
+      splitArgmaxReductionRatio = 128;
+    }
+
     MLIRContext *context = &getContext();
     auto funcOp = getOperation();
 
+    SmallVector<linalg::MatmulOp> matmulCandidates;
+    SmallVector<IREE::LinalgExt::TopkOp> topkCandidates;
+    SmallVector<linalg::GenericOp> argmaxCandidates;
+
+    IRRewriter rewriter(context);
+    funcOp->walk([&](Operation *op) {
+      TypeSwitch<Operation *>(op)
+          .Case<linalg::MatmulOp>([&](auto matmulOp) {
+            if (splitMatmulReductionRatio > 1) {
+              matmulCandidates.push_back(matmulOp);
+            }
+          })
+          .Case<IREE::LinalgExt::TopkOp>([&](auto topkOp) {
+            if (!topkSplitReductionRatio.empty()) {
+              topkCandidates.push_back(topkOp);
+            }
+          })
+          .Case<linalg::GenericOp>([&](auto genericOp) {
+            if (splitArgmaxReductionRatio > 1 &&
+                IREE::LinalgExt::isArgmaxOp(genericOp)) {
+              argmaxCandidates.push_back(genericOp);
+            }
+          });
+    });
+
+    // Split matmul ops.
     auto matmulSplitReductionControlFn =
         [&](linalg::LinalgOp op) -> linalg::SplitReductionOptions {
       // For matmul make the new parallel dimension first so that it looks
       // like a batch_matmul and can follow the same codegen.
-      return {int64_t(splitReductionRatio), 0, /*innerParallel=*/false};
+      return {int64_t(splitMatmulReductionRatio), 0, /*innerParallel=*/false};
path: compiler/src/iree/compiler/DispatchCreation/SplitReduction.cpp
user: login: hanhanW
body: It is not your responsibility because the code was bad in the first place. Could you help remove the `int64_t()` cast since you're touching the code? I think we don't need the cast at all. The type is already int64_t.

(side note: ideally, if we really need the cast, we should use c++ style like static_cast. We should not use c style casting.)
created_at: 2025-05-09T17:45:56Z
author_association: CONTRIBUTOR
diff_hunk: @@ -389,4 +390,315 @@ splitReduction(RewriterBase &rewriter, LinalgExt::TopkOp topkOp,
   return success();
 }
 
+struct ArgmaxCombinerOps {
+  Operation *maxOp = nullptr;    // arith.maximumf
+  Operation *selectOp = nullptr; // arith.select
+  Operation *cmpOp = nullptr;    // arith.cmpf
+};
+
+// Matches the combiner pattern in a linalg.generic argmax-style reduction:
+// Example MLIR:
+// %4:2 = linalg.generic {
+//     indexing_maps = [...],
+//     iterator_types = ["parallel", "reduction"]
+//   } ins(%arg0 : tensor<?x128xbf16>) outs(%1, %3 : tensor<?xbf16>,
+//   tensor<?xi64>) {
+// ^bb0(%in: bf16, %out: bf16, %out_0: i64):
+//   %5 = linalg.index 1 : index
+//   %6 = arith.index_cast %5 : index to i64
+//   %7 = arith.maximumf %in, %out : bf16
+//   %8 = arith.cmpf ogt, %in, %out : bf16
+//   %9 = arith.select %8, %6, %out_0 : i64
+//   linalg.yield %7, %9 : bf16, i64
+// } -> (tensor<?xbf16>, tensor<?xi64>)
+//
+// This function extracts the `arith.maximumf`, `arith.cmpf`, and `arith.select`
+// operations from the body to facilitate transformations such as split
+// reduction.
+static FailureOr<ArgmaxCombinerOps>
+collectArgmaxCombinerOps(linalg::GenericOp genericOp) {
+
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  ArgmaxCombinerOps ops;
+
+  auto yieldOp = cast<linalg::YieldOp>(genericOp.getBody()->getTerminator());
+
+  // Extract max value producer: arith.maximumf.
+  Value maxResult = yieldOp.getOperand(0);
+  auto maxOp = dyn_cast<arith::MaximumFOp>(maxResult.getDefiningOp());
+
+  // Extract index result producer: arith.select.
+  Value indexResult = yieldOp.getOperand(1);
+  auto selectOp = dyn_cast<arith::SelectOp>(indexResult.getDefiningOp());
+
+  // Extract the condition of the select, expected to be arith.cmpf with
+  // predicate OGT.
+  auto cmpOp = dyn_cast<arith::CmpFOp>(selectOp.getCondition().getDefiningOp());
+
+  ops.maxOp = maxOp;
+  ops.selectOp = selectOp;
+  ops.cmpOp = cmpOp;
+
+  return ops;
+}
path: compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/SplitReduction.cpp
user: login: hanhanW
body: style nit 1: don't use blank lines when you don't have to. resist starting functions with a blank line. https://google.github.io/styleguide/cppguide.html#Vertical_Whitespace
nit 2: ops is not used until the end. we should move the declaration down.
style nit 3: I'd use `cast` because you are not handling the 'failure'. We expect it is castable and we already have the assertion in the first line.

```suggestion
static FailureOr<ArgmaxCombinerOps>
collectArgmaxCombinerOps(linalg::GenericOp genericOp) {
  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
  auto yieldOp = cast<linalg::YieldOp>(genericOp.getBody()->getTerminator());

  // Extract max value producer: arith.maximumf.
  Value maxResult = yieldOp.getOperand(0);
  auto maxOp = dyn_cast<arith::MaximumFOp>(maxResult.getDefiningOp());

  // Extract index result producer: arith.select.
  Value indexResult = yieldOp.getOperand(1);
  auto selectOp = dyn_cast<arith::SelectOp>(indexResult.getDefiningOp());

  // Extract the condition of the select, expected to be arith.cmpf with
  // predicate OGT.
  auto cmpOp = dyn_cast<arith::CmpFOp>(selectOp.getCondition().getDefiningOp());

  ArgmaxCombinerOps ops;
  ops.maxOp = maxOp;
  ops.selectOp = selectOp;
  ops.cmpOp = cmpOp;
  return ops;
}
```

(I'd suggest aggregate initialization, e.g., `return {.maxOp  = maxOp, ...}`, if IREE is built with c++20. Unfortunately, we are not using c++20 atm..)
created_at: 2025-05-09T17:54:03Z
author_association: CONTRIBUTOR
diff_hunk: @@ -118,4 +118,64 @@ FailureOr<std::pair<Value, Value>> rewriteFft(Operation *op, Value operand,
                                               int64_t fftLength,
                                               PatternRewriter &rewriter);
 
+/// Apply transformation to split a linalg.generic argmax reduction
+/// into a two-stage reduction using an additional parallel dimension.
+/// The transformation first computes a partial argmax over tiles (parallel),
+/// then reduces those results into a final result (reduction).
+///
+/// This pattern is specialized for reductions that yield both the maximum
+/// value and its index, using the combination of `arith.maximumf`,
+/// `arith.cmpf`, and `arith.select` ops. It assumes a known structure of the
+/// region and injects index computations to track global indices.
+///
+/// Returns the resulting partial and final linalg.generic ops, or failure
+/// if the pattern does not match or cannot be split.
+///
+/// Example:
+/// ```
+/// // Original argmax op reducing over dim=512
+/// %4:2 = linalg.generic {
+///   indexing_maps = [...],
+///   iterator_types = ["parallel", "reduction"]
+/// } ins(%arg0 : tensor<?x512xbf16>)
+///   outs(%out_val, %out_idx : tensor<?xbf16>, tensor<?xi64>) {
+/// ^bb0(%in: bf16, %out: bf16, %out_0: i64):
+///   %idx = linalg.index 1 : index
+///   %cast = arith.index_cast %idx : index to i64
+///   %max = arith.maximumf %in, %out : bf16
+///   %cmp = arith.cmpf ogt, %in, %out : bf16
+///   %sel = arith.select %cmp, %cast, %out_0 : i64
+///   linalg.yield %max, %sel : bf16, i64
+/// } -> (tensor<?xbf16>, tensor<?xi64>)
+/// ```
+/// To:
+/// ```
+/// // After splitting K=512 into 4 x 128
+/// %expanded = tensor.expand_shape %arg0 [[0], [1, 2]] : tensor<?x512xbf16>
+///     into tensor<?x4x128xbf16>
+/// %init_val = linalg.fill ins(%cst : bf16) outs(%empty : tensor<?x4xbf16>)
+/// %init_idx = linalg.fill ins(%zero : i64) outs(%empty : tensor<?x4xi64>)
+/// %partial:2 = linalg.generic {
+///   indexing_maps = [...],
+///   iterator_types = ["parallel", "reduction"]
+/// } ins(%expanded : tensor<?x4x128xbf16>)
+///   outs(%init_val, %init_idx : tensor<?x4xbf16>, tensor<?x4xi64>) {
+///   // compute global index: outer_idx * 128 + inner_idx
+///   ...
+/// }
+///
+/// // Final argmax over the tile dimension (dim=1 of ?x4)
+/// %final:2 = linalg.generic {
+///   indexing_maps = [...],
+///   iterator_types = ["reduction"]
+/// } ins(%partial#0, %partial#1)
+///   outs(%out_val, %out_idx : tensor<?xbf16>, tensor<?xi64>) {
+///   // same combiner: maximumf, cmpf, select
+///   ...
+/// }
+/// ```
path: compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/Transforms.h
user: login: hanhanW
body: Optional nit: remove the markdown style block. I've seen that people use it in IREE and MLIR, but I really don't get the point. I'm not a fan of using it because it adds visual noise to me. Marking it optional because there are no convention, IMO. We can just replace ``` with blank, and have indents for examples. I took a stab at the formatting a bit, please make the change if you think it looks better.

```suggestion
/// Example: original argmax op reducing over dim=512
///
///   %4:2 = linalg.generic {
///     indexing_maps = [...],
///     iterator_types = ["parallel", "reduction"]
///   } ins(%arg0 : tensor<?x512xbf16>)
///     outs(%out_val, %out_idx : tensor<?xbf16>, tensor<?xi64>) {
///   ^bb0(%in: bf16, %out: bf16, %out_0: i64):
///     %idx = linalg.index 1 : index
///     %cast = arith.index_cast %idx : index to i64
///     %max = arith.maximumf %in, %out : bf16
///     %cmp = arith.cmpf ogt, %in, %out : bf16
///     %sel = arith.select %cmp, %cast, %out_0 : i64
///     linalg.yield %max, %sel : bf16, i64
///   } -> (tensor<?xbf16>, tensor<?xi64>)
///
/// To: splitting K=512 into 4 x 128 + final argmax over the tile dimension
///     (dim=1 of ?x4)
///
///   %expanded = tensor.expand_shape %arg0 [[0], [1, 2]] : tensor<?x512xbf16>
///       into tensor<?x4x128xbf16>
///   %init_val = linalg.fill ins(%cst : bf16) outs(%empty : tensor<?x4xbf16>)
///   %init_idx = linalg.fill ins(%zero : i64) outs(%empty : tensor<?x4xi64>)
///   %partial:2 = linalg.generic {
///     indexing_maps = [...],
///     iterator_types = ["parallel", "reduction"]
///   } ins(%expanded : tensor<?x4x128xbf16>)
///     outs(%init_val, %init_idx : tensor<?x4xbf16>, tensor<?x4xi64>) {
///     // compute global index: outer_idx * 128 + inner_idx
///     ...
///   }
///   %final:2 = linalg.generic {
///     indexing_maps = [...],
///     iterator_types = ["reduction"]
///   } ins(%partial#0, %partial#1)
///     outs(%out_val, %out_idx : tensor<?xbf16>, tensor<?xi64>) {
///     // same combiner: maximumf, cmpf, select
///     ...
///   }
```


created_at: 2025-05-09T18:07:26Z
author_association: CONTRIBUTOR
diff_hunk: @@ -389,4 +390,315 @@ splitReduction(RewriterBase &rewriter, LinalgExt::TopkOp topkOp,
   return success();
 }
 
+struct ArgmaxCombinerOps {
+  Operation *maxOp = nullptr;    // arith.maximumf
+  Operation *selectOp = nullptr; // arith.select
+  Operation *cmpOp = nullptr;    // arith.cmpf
+};
+
+// Matches the combiner pattern in a linalg.generic argmax-style reduction:
+// Example MLIR:
+// %4:2 = linalg.generic {
+//     indexing_maps = [...],
+//     iterator_types = ["parallel", "reduction"]
+//   } ins(%arg0 : tensor<?x128xbf16>) outs(%1, %3 : tensor<?xbf16>,
+//   tensor<?xi64>) {
+// ^bb0(%in: bf16, %out: bf16, %out_0: i64):
+//   %5 = linalg.index 1 : index
+//   %6 = arith.index_cast %5 : index to i64
+//   %7 = arith.maximumf %in, %out : bf16
+//   %8 = arith.cmpf ogt, %in, %out : bf16
+//   %9 = arith.select %8, %6, %out_0 : i64
+//   linalg.yield %7, %9 : bf16, i64
+// } -> (tensor<?xbf16>, tensor<?xi64>)
+//
+// This function extracts the `arith.maximumf`, `arith.cmpf`, and `arith.select`
+// operations from the body to facilitate transformations such as split
+// reduction.
+static FailureOr<ArgmaxCombinerOps>
path: compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/SplitReduction.cpp
user: login: hanhanW
body: nit: remove `FailureOr`, because the assumption is that it is always a argmax-style op. It is just a local method and we already have the sanity check. We don't need to handle the failure for this, and it makes code easier.
created_at: 2025-05-09T18:20:05Z
author_association: CONTRIBUTOR
diff_hunk: @@ -389,4 +390,315 @@ splitReduction(RewriterBase &rewriter, LinalgExt::TopkOp topkOp,
   return success();
 }
 
+struct ArgmaxCombinerOps {
+  Operation *maxOp = nullptr;    // arith.maximumf
+  Operation *selectOp = nullptr; // arith.select
+  Operation *cmpOp = nullptr;    // arith.cmpf
+};
+
+// Matches the combiner pattern in a linalg.generic argmax-style reduction:
+// Example MLIR:
+// %4:2 = linalg.generic {
+//     indexing_maps = [...],
+//     iterator_types = ["parallel", "reduction"]
+//   } ins(%arg0 : tensor<?x128xbf16>) outs(%1, %3 : tensor<?xbf16>,
+//   tensor<?xi64>) {
+// ^bb0(%in: bf16, %out: bf16, %out_0: i64):
+//   %5 = linalg.index 1 : index
+//   %6 = arith.index_cast %5 : index to i64
+//   %7 = arith.maximumf %in, %out : bf16
+//   %8 = arith.cmpf ogt, %in, %out : bf16
+//   %9 = arith.select %8, %6, %out_0 : i64
+//   linalg.yield %7, %9 : bf16, i64
+// } -> (tensor<?xbf16>, tensor<?xi64>)
+//
+// This function extracts the `arith.maximumf`, `arith.cmpf`, and `arith.select`
+// operations from the body to facilitate transformations such as split
+// reduction.
+static FailureOr<ArgmaxCombinerOps>
+collectArgmaxCombinerOps(linalg::GenericOp genericOp) {
+
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  ArgmaxCombinerOps ops;
+
+  auto yieldOp = cast<linalg::YieldOp>(genericOp.getBody()->getTerminator());
+
+  // Extract max value producer: arith.maximumf.
+  Value maxResult = yieldOp.getOperand(0);
+  auto maxOp = dyn_cast<arith::MaximumFOp>(maxResult.getDefiningOp());
+
+  // Extract index result producer: arith.select.
+  Value indexResult = yieldOp.getOperand(1);
+  auto selectOp = dyn_cast<arith::SelectOp>(indexResult.getDefiningOp());
+
+  // Extract the condition of the select, expected to be arith.cmpf with
+  // predicate OGT.
+  auto cmpOp = dyn_cast<arith::CmpFOp>(selectOp.getCondition().getDefiningOp());
+
+  ops.maxOp = maxOp;
+  ops.selectOp = selectOp;
+  ops.cmpOp = cmpOp;
+
+  return ops;
+}
+
+FailureOr<linalg::SplitReductionResult>
+splitArgmaxReduction(RewriterBase &rewriter, linalg::GenericOp genericOp,
+                     linalg::ControlSplitReductionFn controlSplitReductionFn) {
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  OpBuilder::InsertionGuard guard(rewriter);
+  rewriter.setInsertionPoint(genericOp);
+  Location loc = genericOp->getLoc();
+
+  linalg::SplitReductionOptions control = controlSplitReductionFn(genericOp);
+  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, 4> loopRanges = genericOp.getStaticLoopRanges();
+  int64_t reductionDimSize = loopRanges[reductionDim];
+
+  //  The total number of output elements along this new dimension is
+  //  reductionDimSize / ratio.
+  int64_t outputDimsize = reductionDimSize / ratio;
+
+  if (reductionDimSize == ShapedType::kDynamic ||
+      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");
+  }
+
+  FailureOr<ArgmaxCombinerOps> maybeOps = collectArgmaxCombinerOps(genericOp);
path: compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/SplitReduction.cpp
user: login: hanhanW
body: How about renaming it to `combinerOps`? We dont need the `maybe` prefix after we remove `FailureOr`, and `ops` seem to be too vague.
created_at: 2025-05-09T18:36:10Z
author_association: CONTRIBUTOR
diff_hunk: @@ -389,4 +390,315 @@ splitReduction(RewriterBase &rewriter, LinalgExt::TopkOp topkOp,
   return success();
 }
 
+struct ArgmaxCombinerOps {
+  Operation *maxOp = nullptr;    // arith.maximumf
+  Operation *selectOp = nullptr; // arith.select
+  Operation *cmpOp = nullptr;    // arith.cmpf
+};
+
+// Matches the combiner pattern in a linalg.generic argmax-style reduction:
+// Example MLIR:
+// %4:2 = linalg.generic {
+//     indexing_maps = [...],
+//     iterator_types = ["parallel", "reduction"]
+//   } ins(%arg0 : tensor<?x128xbf16>) outs(%1, %3 : tensor<?xbf16>,
+//   tensor<?xi64>) {
+// ^bb0(%in: bf16, %out: bf16, %out_0: i64):
+//   %5 = linalg.index 1 : index
+//   %6 = arith.index_cast %5 : index to i64
+//   %7 = arith.maximumf %in, %out : bf16
+//   %8 = arith.cmpf ogt, %in, %out : bf16
+//   %9 = arith.select %8, %6, %out_0 : i64
+//   linalg.yield %7, %9 : bf16, i64
+// } -> (tensor<?xbf16>, tensor<?xi64>)
+//
+// This function extracts the `arith.maximumf`, `arith.cmpf`, and `arith.select`
+// operations from the body to facilitate transformations such as split
+// reduction.
+static FailureOr<ArgmaxCombinerOps>
+collectArgmaxCombinerOps(linalg::GenericOp genericOp) {
+
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  ArgmaxCombinerOps ops;
+
+  auto yieldOp = cast<linalg::YieldOp>(genericOp.getBody()->getTerminator());
+
+  // Extract max value producer: arith.maximumf.
+  Value maxResult = yieldOp.getOperand(0);
+  auto maxOp = dyn_cast<arith::MaximumFOp>(maxResult.getDefiningOp());
+
+  // Extract index result producer: arith.select.
+  Value indexResult = yieldOp.getOperand(1);
+  auto selectOp = dyn_cast<arith::SelectOp>(indexResult.getDefiningOp());
+
+  // Extract the condition of the select, expected to be arith.cmpf with
+  // predicate OGT.
+  auto cmpOp = dyn_cast<arith::CmpFOp>(selectOp.getCondition().getDefiningOp());
+
+  ops.maxOp = maxOp;
+  ops.selectOp = selectOp;
+  ops.cmpOp = cmpOp;
+
+  return ops;
+}
+
+FailureOr<linalg::SplitReductionResult>
+splitArgmaxReduction(RewriterBase &rewriter, linalg::GenericOp genericOp,
+                     linalg::ControlSplitReductionFn controlSplitReductionFn) {
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  OpBuilder::InsertionGuard guard(rewriter);
+  rewriter.setInsertionPoint(genericOp);
+  Location loc = genericOp->getLoc();
+
+  linalg::SplitReductionOptions control = controlSplitReductionFn(genericOp);
+  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, 4> loopRanges = genericOp.getStaticLoopRanges();
+  int64_t reductionDimSize = loopRanges[reductionDim];
+
+  //  The total number of output elements along this new dimension is
+  //  reductionDimSize / ratio.
+  int64_t outputDimsize = reductionDimSize / ratio;
+
+  if (reductionDimSize == ShapedType::kDynamic ||
+      reductionDimSize % ratio != 0) {
+    return rewriter.notifyMatchFailure(
+        genericOp, "Reduction dimension not divisible by split ratio");
+  }
path: compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/SplitReduction.cpp
user: login: hanhanW
body: nit 1: I think we should move the `outputDimSize` to the chunk that it is used, i.e., the place that you compute the output shapes/maps/etc. It is weird that we do the computation before the check. It could be dynamic value. In this case, we don't need to compute `outputDimSize` at all, right?

nit 2: switch to new preferred style, `ShapedType::isDynamic(reductionDimSize)`

picky nit for the error message: We typically start the first sentence with a lower-case letter, because it matches error message styles commonly produced by other tools

https://llvm.org/docs/CodingStandards.html#error-and-warning-messages

> Also, to match error message styles commonly produced by other tools, start the first sentence with a lower-case letter, and finish the last sentence without a period, if it would end in one otherwise.
created_at: 2025-05-09T18:45:47Z
author_association: CONTRIBUTOR
diff_hunk: @@ -389,4 +390,315 @@ splitReduction(RewriterBase &rewriter, LinalgExt::TopkOp topkOp,
   return success();
 }
 
+struct ArgmaxCombinerOps {
+  Operation *maxOp = nullptr;    // arith.maximumf
+  Operation *selectOp = nullptr; // arith.select
+  Operation *cmpOp = nullptr;    // arith.cmpf
+};
+
+// Matches the combiner pattern in a linalg.generic argmax-style reduction:
+// Example MLIR:
+// %4:2 = linalg.generic {
+//     indexing_maps = [...],
+//     iterator_types = ["parallel", "reduction"]
+//   } ins(%arg0 : tensor<?x128xbf16>) outs(%1, %3 : tensor<?xbf16>,
+//   tensor<?xi64>) {
+// ^bb0(%in: bf16, %out: bf16, %out_0: i64):
+//   %5 = linalg.index 1 : index
+//   %6 = arith.index_cast %5 : index to i64
+//   %7 = arith.maximumf %in, %out : bf16
+//   %8 = arith.cmpf ogt, %in, %out : bf16
+//   %9 = arith.select %8, %6, %out_0 : i64
+//   linalg.yield %7, %9 : bf16, i64
+// } -> (tensor<?xbf16>, tensor<?xi64>)
+//
+// This function extracts the `arith.maximumf`, `arith.cmpf`, and `arith.select`
+// operations from the body to facilitate transformations such as split
+// reduction.
+static FailureOr<ArgmaxCombinerOps>
+collectArgmaxCombinerOps(linalg::GenericOp genericOp) {
+
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  ArgmaxCombinerOps ops;
+
+  auto yieldOp = cast<linalg::YieldOp>(genericOp.getBody()->getTerminator());
+
+  // Extract max value producer: arith.maximumf.
+  Value maxResult = yieldOp.getOperand(0);
+  auto maxOp = dyn_cast<arith::MaximumFOp>(maxResult.getDefiningOp());
+
+  // Extract index result producer: arith.select.
+  Value indexResult = yieldOp.getOperand(1);
+  auto selectOp = dyn_cast<arith::SelectOp>(indexResult.getDefiningOp());
+
+  // Extract the condition of the select, expected to be arith.cmpf with
+  // predicate OGT.
+  auto cmpOp = dyn_cast<arith::CmpFOp>(selectOp.getCondition().getDefiningOp());
+
+  ops.maxOp = maxOp;
+  ops.selectOp = selectOp;
+  ops.cmpOp = cmpOp;
+
+  return ops;
+}
+
+FailureOr<linalg::SplitReductionResult>
+splitArgmaxReduction(RewriterBase &rewriter, linalg::GenericOp genericOp,
+                     linalg::ControlSplitReductionFn controlSplitReductionFn) {
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  OpBuilder::InsertionGuard guard(rewriter);
+  rewriter.setInsertionPoint(genericOp);
+  Location loc = genericOp->getLoc();
+
+  linalg::SplitReductionOptions control = controlSplitReductionFn(genericOp);
+  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, 4> loopRanges = genericOp.getStaticLoopRanges();
+  int64_t reductionDimSize = loopRanges[reductionDim];
+
+  //  The total number of output elements along this new dimension is
+  //  reductionDimSize / ratio.
+  int64_t outputDimsize = reductionDimSize / ratio;
+
+  if (reductionDimSize == ShapedType::kDynamic ||
+      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");
+  }
+
+  FailureOr<ArgmaxCombinerOps> maybeOps = collectArgmaxCombinerOps(genericOp);
+  if (failed(maybeOps))
+    return rewriter.notifyMatchFailure(genericOp,
+                                       "invalid combiner for argmax");
+
+  ArgmaxCombinerOps combinerOps = *maybeOps;
+  Operation *reductionOp = combinerOps.maxOp;
+
+  std::optional<TypedAttr> identity = arith::getNeutralElement(reductionOp);
+  if (!identity.has_value())
+    return rewriter.notifyMatchFailure(
+        genericOp, "Unknown identity value for the reduction");
+
+  SmallVector<Value> newInputs;
+  SmallVector<AffineMap> newMaps;
+  // Calculate the new shapes and indexing maps of the input operands.
+  for (OpOperand *operand : genericOp.getDpsInputOperands()) {
+    AffineMap map = genericOp.getMatchingIndexingMap(operand);
+    SmallVector<int64_t> newShape;
+    SmallVector<AffineExpr> exprs;
+    SmallVector<ReassociationIndices> reassociation;
+    unsigned index = 0;
+    for (unsigned idx : llvm::seq<unsigned>(0, map.getNumResults())) {
+      unsigned dim = map.getDimPosition(idx);
+      if (reductionDim == dim) {
+        if (control.innerParallel) {
+          newShape.push_back(ratio); // reduce
+          newShape.push_back(genericOp.getShape(operand)[idx] /
+                             ratio); // parallel (insert)
+          exprs.push_back(rewriter.getAffineDimExpr(
+              dim < insertSplitDimension ? dim : dim + 1));
+          exprs.push_back(rewriter.getAffineDimExpr(insertSplitDimension));
+        } else {
+          newShape.push_back(genericOp.getShape(operand)[idx] /
+                             ratio); // parallel (insert)
+          newShape.push_back(ratio); // reduce
+          exprs.push_back(rewriter.getAffineDimExpr(insertSplitDimension));
+          exprs.push_back(rewriter.getAffineDimExpr(
+              dim < insertSplitDimension ? dim : dim + 1));
+        }
+        reassociation.push_back({index++, index++});
+        continue;
+      }
+      newShape.push_back(genericOp.getShape(operand)[idx]);
+      exprs.push_back(rewriter.getAffineDimExpr(
+          dim < insertSplitDimension ? dim : dim + 1));
+      reassociation.push_back({index++});
+    }
+    newMaps.push_back(
+        AffineMap::get(map.getNumDims() + 1, 0, exprs, genericOp.getContext()));
+    // If the shape is unchanged the input doesn't change.
+    if (newShape == genericOp.getShape(operand)) {
+      newInputs.push_back(operand->get());
+      continue;
+    }
+    Type newType = RankedTensorType::get(
+        newShape,
+        cast<RankedTensorType>(operand->get().getType()).getElementType());
+
+    Value newInput = rewriter.create<tensor::ExpandShapeOp>(
+        loc, newType, operand->get(), reassociation);
+    newInputs.push_back(newInput);
+  }
+
+  SmallVector<SmallVector<int64_t>> newOutputShapes;
+  SmallVector<AffineMap> outputMaps;
+  for (int i = 0; i < genericOp.getNumDpsInits(); ++i) {
+    OpOperand *output = genericOp.getDpsInitOperand(i);
+    AffineMap oldOutputMap = genericOp.getMatchingIndexingMap(output);
+    ArrayRef<int64_t> oldShape = genericOp.getShape(output);
+    SmallVector<int64_t> thisOutputShape;
+
+    SmallVector<AffineExpr> outputExpr;
+    for (unsigned idx = 0; idx <= oldShape.size(); ++idx) {
+      if (idx == insertSplitIndex) {
+        thisOutputShape.push_back(outputDimsize);
+        outputExpr.push_back(rewriter.getAffineDimExpr(insertSplitDimension));
+      }
+      if (idx < oldShape.size()) {
+        thisOutputShape.push_back(oldShape[idx]);
+        unsigned dim = oldOutputMap.getDimPosition(idx);
+        outputExpr.push_back(rewriter.getAffineDimExpr(
+            dim < insertSplitDimension ? dim : dim + 1));
+      }
+    }
+
+    AffineMap newOutputMap = AffineMap::get(oldOutputMap.getNumDims() + 1, 0,
+                                            outputExpr, rewriter.getContext());
+    newMaps.push_back(newOutputMap);
+    newOutputShapes.push_back(thisOutputShape);
+  }
+
+  // Handle dynamic dimensions for identity value tensor.
+  SmallVector<Value> dynValDims;
+  SmallVector<int64_t> newOutputShape = newOutputShapes[0];
+  for (size_t i = 0; i < newOutputShape.size(); ++i) {
+    if (ShapedType::isDynamic(newOutputShape[i])) {
+      dynValDims.push_back(rewriter.create<tensor::DimOp>(
+          loc, genericOp.getDpsInputOperand(0)->get(), i));
+    }
+  }
+
+  Type valueElemType = genericOp.getRegionOutputArgs()[0].getType();
+  Value emptyValTensor = rewriter.create<tensor::EmptyOp>(
+      loc, newOutputShape, valueElemType, dynValDims);
+  Value constantOp = rewriter.create<arith::ConstantOp>(loc, *identity);
+  Value identityVal =
+      rewriter.create<linalg::FillOp>(loc, constantOp, emptyValTensor)
+          .getResult(0);
+
+  // Handle dynamic dimensions for identity index tensor.
+  SmallVector<Value> dynIdxDims;
+  newOutputShape = newOutputShapes[1];
+  for (size_t i = 0; i < newOutputShape.size(); ++i) {
+    if (ShapedType::isDynamic(newOutputShape[i])) {
+      dynIdxDims.push_back(rewriter.create<tensor::DimOp>(
+          loc, genericOp.getDpsInputOperand(0)->get(), i));
+    }
+  }
path: compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/SplitReduction.cpp
user: login: hanhanW
body: IIUC, this part is generating the expanded operands and the properties of the partial reduction op. Furthermore, it creates a "folded" linalg.fill op on the expanded domain.

I think the logic can be much simpler because we have few more patterns help a lot. First, you can declare a `expandValue()` method that applies the expand_shape for all the operands, and the partial reduction op takes them as inputs and outputs. Here, I think you'd have a question about how to fold the expand shape away.

1. We have [canonicalization patterns](https://github.com/llvm/llvm-project/blob/31fd77aa51a643245f8eb277483554509b771832/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp#L976-L983) in FillOp that swaps the (fill, reshape) pair.
2. We can explicitly populate the [patterns](https://github.com/llvm/llvm-project/blob/65a6cbde5bb074ce377cb1aa6145241ad1719c17/mlir/lib/Dialect/Tensor/Transforms/EmptyOpPatterns.cpp#L130C20-L138) that fold reshapes into empty ops. I don't know why upstream people do not put them EmptyOp canonicalization patterns, but they are valid to use IMO.
3. For creating the expand_shape op, you don't need to compute the dynamic values yourself. You can use [this builder](https://github.com/llvm/llvm-project/blob/65a6cbde5bb074ce377cb1aa6145241ad1719c17/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td#L1144-L1146) that takes result type and reassociation map. The new type can be constructed easily, so I'm not going to dig into details. After you get the type, you can use [getReassociationIndicesForReshape](https://github.com/llvm/llvm-project/blob/0d5124775cace200f6f99905989ebeb65853b16e/mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h#L67-L71C1) method to get the map, and use them to create the expand shape op. I checked that the upstream builder infers the output shape for you and they build the dynamic values for you.

This way, we can generate the expanded operands cleanly.

The next problem is about the generic op properties. For indexing maps, I think you can use `argmaxOp.getIndexingMapsArray()` to get all the indexing maps. It returns `SmallVector<AffineMap>` vector. Then you can declare a `getExpandedIndexingMap()` and iterate on all the indexing map to get the expanded indexing maps.

Then you get all the properties, except iterators. The iterator change is in below, so I'll leave the comment there.

I played a bit with the IR, I think you will generate something like below. (please ignore the transform script, it is just for kicking in the patterns that fold expand_shape away for empty ops.)

```mlir
#map = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
#map1 = affine_map<(d0, d1, d2) -> (d0, d1)>
#map2 = affine_map<(d0, d1) -> (d0, d1)>
#map3 = affine_map<(d0, d1) -> (d0)>
module {
  util.func public @argmax(%arg0: tensor<?x131072xbf16>, %arg1: index) -> tensor<?xi64> {
    %cst = arith.constant 0xFF80 : bf16
    %c0_i64 = arith.constant 0 : i64
    %0 = tensor.empty(%arg1) : tensor<?xbf16>
    %1 = linalg.fill ins(%cst : bf16) outs(%0 : tensor<?xbf16>) -> tensor<?xbf16>
    %2 = tensor.empty(%arg1) : tensor<?xi64>
    %3 = linalg.fill ins(%c0_i64 : i64) outs(%2 : tensor<?xi64>) -> tensor<?xi64>
    %c128 = arith.constant 128 : index
    %c0 = arith.constant 0 : index
    %dim = tensor.dim %arg0, %c0 : tensor<?x131072xbf16>
    %expanded = tensor.expand_shape %arg0 [[0], [1, 2]] output_shape [%dim, 1024, 128] : tensor<?x131072xbf16> into tensor<?x1024x128xbf16>
    %c1024 = arith.constant 1024 : index
    %4 = arith.divsi %dim, %c1024 : index
    %expanded_0 = tensor.expand_shape %1 [[0, 1]] output_shape [%dim, 1024] : tensor<?xbf16> into tensor<?x1024xbf16>
    %expanded_1 = tensor.expand_shape %3 [[0, 1]] output_shape [%dim, 1024] : tensor<?xi64> into tensor<?x1024xi64>
    %5:2 = linalg.generic {indexing_maps = [#map, #map1, #map1], iterator_types = ["parallel", "parallel", "reduction"]} ins(%expanded : tensor<?x1024x128xbf16>) outs(%expanded_0, %expanded_1 : tensor<?x1024xbf16>, tensor<?x1024xi64>) {
    ^bb0(%in: bf16, %out: bf16, %out_2: i64):
      %7 = linalg.index 1 : index
      %8 = linalg.index 2 : index
      %9 = arith.muli %7, %c128 : index
      %10 = arith.addi %9, %8 : index
      %11 = arith.index_cast %10 : index to i64
      %12 = arith.maximumf %in, %out : bf16
      %13 = arith.cmpf ogt, %in, %out : bf16
      %14 = arith.select %13, %11, %out_2 : i64
      linalg.yield %12, %14 : bf16, i64
    } -> (tensor<?x1024xbf16>, tensor<?x1024xi64>)
    %6:2 = linalg.generic {indexing_maps = [#map2, #map2, #map3, #map3], iterator_types = ["parallel", "reduction"]} ins(%5#0, %5#1 : tensor<?x1024xbf16>, tensor<?x1024xi64>) outs(%1, %3 : tensor<?xbf16>, tensor<?xi64>) {
    ^bb0(%in: bf16, %in_2: i64, %out: bf16, %out_3: i64):
      %7 = arith.maximumf %in, %out : bf16
      %8 = arith.cmpf ogt, %in, %out : bf16
      %9 = arith.select %8, %in_2, %out_3 : i64
      linalg.yield %7, %9 : bf16, i64
    } -> (tensor<?xbf16>, tensor<?xi64>)
    util.return %6#1 : tensor<?xi64>
  }
  module attributes {transform.with_named_sequence} {
    transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
      %0 = transform.structured.match ops{["util.func"]} in %arg0 : (!transform.any_op) -> !transform.op<"util.func">
      transform.apply_patterns to %0 {
        transform.apply_patterns.tensor.fold_tensor_empty
      } : !transform.op<"util.func">
      transform.yield
    }
  }
}
```

And I verified that the expand ops are folded away, if you run `iree-opt --canonicalize -transform-interpreter ~/repro.mlir`.

Does it make sense?
created_at: 2025-05-09T20:37:03Z
author_association: CONTRIBUTOR
diff_hunk: @@ -389,4 +390,315 @@ splitReduction(RewriterBase &rewriter, LinalgExt::TopkOp topkOp,
   return success();
 }
 
+struct ArgmaxCombinerOps {
+  Operation *maxOp = nullptr;    // arith.maximumf
+  Operation *selectOp = nullptr; // arith.select
+  Operation *cmpOp = nullptr;    // arith.cmpf
+};
+
+// Matches the combiner pattern in a linalg.generic argmax-style reduction:
+// Example MLIR:
+// %4:2 = linalg.generic {
+//     indexing_maps = [...],
+//     iterator_types = ["parallel", "reduction"]
+//   } ins(%arg0 : tensor<?x128xbf16>) outs(%1, %3 : tensor<?xbf16>,
+//   tensor<?xi64>) {
+// ^bb0(%in: bf16, %out: bf16, %out_0: i64):
+//   %5 = linalg.index 1 : index
+//   %6 = arith.index_cast %5 : index to i64
+//   %7 = arith.maximumf %in, %out : bf16
+//   %8 = arith.cmpf ogt, %in, %out : bf16
+//   %9 = arith.select %8, %6, %out_0 : i64
+//   linalg.yield %7, %9 : bf16, i64
+// } -> (tensor<?xbf16>, tensor<?xi64>)
+//
+// This function extracts the `arith.maximumf`, `arith.cmpf`, and `arith.select`
+// operations from the body to facilitate transformations such as split
+// reduction.
+static FailureOr<ArgmaxCombinerOps>
+collectArgmaxCombinerOps(linalg::GenericOp genericOp) {
+
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  ArgmaxCombinerOps ops;
+
+  auto yieldOp = cast<linalg::YieldOp>(genericOp.getBody()->getTerminator());
+
+  // Extract max value producer: arith.maximumf.
+  Value maxResult = yieldOp.getOperand(0);
+  auto maxOp = dyn_cast<arith::MaximumFOp>(maxResult.getDefiningOp());
+
+  // Extract index result producer: arith.select.
+  Value indexResult = yieldOp.getOperand(1);
+  auto selectOp = dyn_cast<arith::SelectOp>(indexResult.getDefiningOp());
+
+  // Extract the condition of the select, expected to be arith.cmpf with
+  // predicate OGT.
+  auto cmpOp = dyn_cast<arith::CmpFOp>(selectOp.getCondition().getDefiningOp());
+
+  ops.maxOp = maxOp;
+  ops.selectOp = selectOp;
+  ops.cmpOp = cmpOp;
+
+  return ops;
+}
+
+FailureOr<linalg::SplitReductionResult>
+splitArgmaxReduction(RewriterBase &rewriter, linalg::GenericOp genericOp,
+                     linalg::ControlSplitReductionFn controlSplitReductionFn) {
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  OpBuilder::InsertionGuard guard(rewriter);
+  rewriter.setInsertionPoint(genericOp);
+  Location loc = genericOp->getLoc();
+
+  linalg::SplitReductionOptions control = controlSplitReductionFn(genericOp);
+  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, 4> loopRanges = genericOp.getStaticLoopRanges();
+  int64_t reductionDimSize = loopRanges[reductionDim];
+
+  //  The total number of output elements along this new dimension is
+  //  reductionDimSize / ratio.
+  int64_t outputDimsize = reductionDimSize / ratio;
+
+  if (reductionDimSize == ShapedType::kDynamic ||
+      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");
+  }
+
+  FailureOr<ArgmaxCombinerOps> maybeOps = collectArgmaxCombinerOps(genericOp);
+  if (failed(maybeOps))
+    return rewriter.notifyMatchFailure(genericOp,
+                                       "invalid combiner for argmax");
+
+  ArgmaxCombinerOps combinerOps = *maybeOps;
+  Operation *reductionOp = combinerOps.maxOp;
+
+  std::optional<TypedAttr> identity = arith::getNeutralElement(reductionOp);
+  if (!identity.has_value())
+    return rewriter.notifyMatchFailure(
+        genericOp, "Unknown identity value for the reduction");
+
+  SmallVector<Value> newInputs;
+  SmallVector<AffineMap> newMaps;
+  // Calculate the new shapes and indexing maps of the input operands.
+  for (OpOperand *operand : genericOp.getDpsInputOperands()) {
+    AffineMap map = genericOp.getMatchingIndexingMap(operand);
+    SmallVector<int64_t> newShape;
+    SmallVector<AffineExpr> exprs;
+    SmallVector<ReassociationIndices> reassociation;
+    unsigned index = 0;
+    for (unsigned idx : llvm::seq<unsigned>(0, map.getNumResults())) {
+      unsigned dim = map.getDimPosition(idx);
+      if (reductionDim == dim) {
+        if (control.innerParallel) {
+          newShape.push_back(ratio); // reduce
+          newShape.push_back(genericOp.getShape(operand)[idx] /
+                             ratio); // parallel (insert)
+          exprs.push_back(rewriter.getAffineDimExpr(
+              dim < insertSplitDimension ? dim : dim + 1));
+          exprs.push_back(rewriter.getAffineDimExpr(insertSplitDimension));
+        } else {
+          newShape.push_back(genericOp.getShape(operand)[idx] /
+                             ratio); // parallel (insert)
+          newShape.push_back(ratio); // reduce
+          exprs.push_back(rewriter.getAffineDimExpr(insertSplitDimension));
+          exprs.push_back(rewriter.getAffineDimExpr(
+              dim < insertSplitDimension ? dim : dim + 1));
+        }
+        reassociation.push_back({index++, index++});
+        continue;
+      }
+      newShape.push_back(genericOp.getShape(operand)[idx]);
+      exprs.push_back(rewriter.getAffineDimExpr(
+          dim < insertSplitDimension ? dim : dim + 1));
+      reassociation.push_back({index++});
+    }
+    newMaps.push_back(
+        AffineMap::get(map.getNumDims() + 1, 0, exprs, genericOp.getContext()));
+    // If the shape is unchanged the input doesn't change.
+    if (newShape == genericOp.getShape(operand)) {
+      newInputs.push_back(operand->get());
+      continue;
+    }
+    Type newType = RankedTensorType::get(
+        newShape,
+        cast<RankedTensorType>(operand->get().getType()).getElementType());
+
+    Value newInput = rewriter.create<tensor::ExpandShapeOp>(
+        loc, newType, operand->get(), reassociation);
+    newInputs.push_back(newInput);
+  }
+
+  SmallVector<SmallVector<int64_t>> newOutputShapes;
+  SmallVector<AffineMap> outputMaps;
+  for (int i = 0; i < genericOp.getNumDpsInits(); ++i) {
+    OpOperand *output = genericOp.getDpsInitOperand(i);
+    AffineMap oldOutputMap = genericOp.getMatchingIndexingMap(output);
+    ArrayRef<int64_t> oldShape = genericOp.getShape(output);
+    SmallVector<int64_t> thisOutputShape;
+
+    SmallVector<AffineExpr> outputExpr;
+    for (unsigned idx = 0; idx <= oldShape.size(); ++idx) {
+      if (idx == insertSplitIndex) {
+        thisOutputShape.push_back(outputDimsize);
+        outputExpr.push_back(rewriter.getAffineDimExpr(insertSplitDimension));
+      }
+      if (idx < oldShape.size()) {
+        thisOutputShape.push_back(oldShape[idx]);
+        unsigned dim = oldOutputMap.getDimPosition(idx);
+        outputExpr.push_back(rewriter.getAffineDimExpr(
+            dim < insertSplitDimension ? dim : dim + 1));
+      }
+    }
+
+    AffineMap newOutputMap = AffineMap::get(oldOutputMap.getNumDims() + 1, 0,
+                                            outputExpr, rewriter.getContext());
+    newMaps.push_back(newOutputMap);
+    newOutputShapes.push_back(thisOutputShape);
+  }
+
+  // Handle dynamic dimensions for identity value tensor.
+  SmallVector<Value> dynValDims;
+  SmallVector<int64_t> newOutputShape = newOutputShapes[0];
+  for (size_t i = 0; i < newOutputShape.size(); ++i) {
+    if (ShapedType::isDynamic(newOutputShape[i])) {
+      dynValDims.push_back(rewriter.create<tensor::DimOp>(
+          loc, genericOp.getDpsInputOperand(0)->get(), i));
+    }
+  }
+
+  Type valueElemType = genericOp.getRegionOutputArgs()[0].getType();
+  Value emptyValTensor = rewriter.create<tensor::EmptyOp>(
+      loc, newOutputShape, valueElemType, dynValDims);
+  Value constantOp = rewriter.create<arith::ConstantOp>(loc, *identity);
+  Value identityVal =
+      rewriter.create<linalg::FillOp>(loc, constantOp, emptyValTensor)
+          .getResult(0);
+
+  // Handle dynamic dimensions for identity index tensor.
+  SmallVector<Value> dynIdxDims;
+  newOutputShape = newOutputShapes[1];
+  for (size_t i = 0; i < newOutputShape.size(); ++i) {
+    if (ShapedType::isDynamic(newOutputShape[i])) {
+      dynIdxDims.push_back(rewriter.create<tensor::DimOp>(
+          loc, genericOp.getDpsInputOperand(0)->get(), i));
+    }
+  }
+  Type idxElemType = genericOp.getRegionOutputArgs()[1].getType();
+  Value zeroIdx = rewriter.create<arith::ConstantOp>(
+      loc, rewriter.getZeroAttr(idxElemType));
+  Value idxInitTensor = rewriter.create<tensor::EmptyOp>(
+      loc, newOutputShape, idxElemType, dynIdxDims);
+  Value identityIndex =
+      rewriter.create<linalg::FillOp>(loc, zeroIdx, idxInitTensor).getResult(0);
+
+  SmallVector<utils::IteratorType> newIteratorTypes;
+  for (auto [index, iteratorType] :
+       llvm::enumerate(genericOp.getIteratorTypesArray())) {
+    if (insertSplitDimension == index)
+      newIteratorTypes.push_back(utils::IteratorType::parallel);
+    newIteratorTypes.push_back(iteratorType);
+  }
+  if (insertSplitDimension == genericOp.getIteratorTypesArray().size()) {
+    newIteratorTypes.push_back(utils::IteratorType::parallel);
+  }
path: compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/SplitReduction.cpp
user: login: hanhanW
body: I think we can use the `insert` method from SmallVector. Something like

```cpp
SmallVector<utils::IteratorType> newIteratorTypes = genericOp.getIteratorTypesArray();
newIteratorTypes.insert(newIteratorTypes.begin() + index, utils::IteratorType::parallel);
```
created_at: 2025-05-09T20:42:15Z
author_association: CONTRIBUTOR
diff_hunk: @@ -389,4 +390,315 @@ splitReduction(RewriterBase &rewriter, LinalgExt::TopkOp topkOp,
   return success();
 }
 
+struct ArgmaxCombinerOps {
+  Operation *maxOp = nullptr;    // arith.maximumf
+  Operation *selectOp = nullptr; // arith.select
+  Operation *cmpOp = nullptr;    // arith.cmpf
+};
+
+// Matches the combiner pattern in a linalg.generic argmax-style reduction:
+// Example MLIR:
+// %4:2 = linalg.generic {
+//     indexing_maps = [...],
+//     iterator_types = ["parallel", "reduction"]
+//   } ins(%arg0 : tensor<?x128xbf16>) outs(%1, %3 : tensor<?xbf16>,
+//   tensor<?xi64>) {
+// ^bb0(%in: bf16, %out: bf16, %out_0: i64):
+//   %5 = linalg.index 1 : index
+//   %6 = arith.index_cast %5 : index to i64
+//   %7 = arith.maximumf %in, %out : bf16
+//   %8 = arith.cmpf ogt, %in, %out : bf16
+//   %9 = arith.select %8, %6, %out_0 : i64
+//   linalg.yield %7, %9 : bf16, i64
+// } -> (tensor<?xbf16>, tensor<?xi64>)
+//
+// This function extracts the `arith.maximumf`, `arith.cmpf`, and `arith.select`
+// operations from the body to facilitate transformations such as split
+// reduction.
+static FailureOr<ArgmaxCombinerOps>
+collectArgmaxCombinerOps(linalg::GenericOp genericOp) {
+
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  ArgmaxCombinerOps ops;
+
+  auto yieldOp = cast<linalg::YieldOp>(genericOp.getBody()->getTerminator());
+
+  // Extract max value producer: arith.maximumf.
+  Value maxResult = yieldOp.getOperand(0);
+  auto maxOp = dyn_cast<arith::MaximumFOp>(maxResult.getDefiningOp());
+
+  // Extract index result producer: arith.select.
+  Value indexResult = yieldOp.getOperand(1);
+  auto selectOp = dyn_cast<arith::SelectOp>(indexResult.getDefiningOp());
+
+  // Extract the condition of the select, expected to be arith.cmpf with
+  // predicate OGT.
+  auto cmpOp = dyn_cast<arith::CmpFOp>(selectOp.getCondition().getDefiningOp());
+
+  ops.maxOp = maxOp;
+  ops.selectOp = selectOp;
+  ops.cmpOp = cmpOp;
+
+  return ops;
+}
+
+FailureOr<linalg::SplitReductionResult>
+splitArgmaxReduction(RewriterBase &rewriter, linalg::GenericOp genericOp,
+                     linalg::ControlSplitReductionFn controlSplitReductionFn) {
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  OpBuilder::InsertionGuard guard(rewriter);
+  rewriter.setInsertionPoint(genericOp);
+  Location loc = genericOp->getLoc();
+
+  linalg::SplitReductionOptions control = controlSplitReductionFn(genericOp);
+  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, 4> loopRanges = genericOp.getStaticLoopRanges();
+  int64_t reductionDimSize = loopRanges[reductionDim];
+
+  //  The total number of output elements along this new dimension is
+  //  reductionDimSize / ratio.
+  int64_t outputDimsize = reductionDimSize / ratio;
+
+  if (reductionDimSize == ShapedType::kDynamic ||
+      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");
+  }
+
+  FailureOr<ArgmaxCombinerOps> maybeOps = collectArgmaxCombinerOps(genericOp);
+  if (failed(maybeOps))
+    return rewriter.notifyMatchFailure(genericOp,
+                                       "invalid combiner for argmax");
+
+  ArgmaxCombinerOps combinerOps = *maybeOps;
+  Operation *reductionOp = combinerOps.maxOp;
+
+  std::optional<TypedAttr> identity = arith::getNeutralElement(reductionOp);
+  if (!identity.has_value())
+    return rewriter.notifyMatchFailure(
+        genericOp, "Unknown identity value for the reduction");
+
+  SmallVector<Value> newInputs;
+  SmallVector<AffineMap> newMaps;
+  // Calculate the new shapes and indexing maps of the input operands.
+  for (OpOperand *operand : genericOp.getDpsInputOperands()) {
+    AffineMap map = genericOp.getMatchingIndexingMap(operand);
+    SmallVector<int64_t> newShape;
+    SmallVector<AffineExpr> exprs;
+    SmallVector<ReassociationIndices> reassociation;
+    unsigned index = 0;
+    for (unsigned idx : llvm::seq<unsigned>(0, map.getNumResults())) {
+      unsigned dim = map.getDimPosition(idx);
+      if (reductionDim == dim) {
+        if (control.innerParallel) {
+          newShape.push_back(ratio); // reduce
+          newShape.push_back(genericOp.getShape(operand)[idx] /
+                             ratio); // parallel (insert)
+          exprs.push_back(rewriter.getAffineDimExpr(
+              dim < insertSplitDimension ? dim : dim + 1));
+          exprs.push_back(rewriter.getAffineDimExpr(insertSplitDimension));
+        } else {
+          newShape.push_back(genericOp.getShape(operand)[idx] /
+                             ratio); // parallel (insert)
+          newShape.push_back(ratio); // reduce
+          exprs.push_back(rewriter.getAffineDimExpr(insertSplitDimension));
+          exprs.push_back(rewriter.getAffineDimExpr(
+              dim < insertSplitDimension ? dim : dim + 1));
+        }
+        reassociation.push_back({index++, index++});
+        continue;
+      }
+      newShape.push_back(genericOp.getShape(operand)[idx]);
+      exprs.push_back(rewriter.getAffineDimExpr(
+          dim < insertSplitDimension ? dim : dim + 1));
+      reassociation.push_back({index++});
+    }
+    newMaps.push_back(
+        AffineMap::get(map.getNumDims() + 1, 0, exprs, genericOp.getContext()));
+    // If the shape is unchanged the input doesn't change.
+    if (newShape == genericOp.getShape(operand)) {
+      newInputs.push_back(operand->get());
+      continue;
+    }
+    Type newType = RankedTensorType::get(
+        newShape,
+        cast<RankedTensorType>(operand->get().getType()).getElementType());
+
+    Value newInput = rewriter.create<tensor::ExpandShapeOp>(
+        loc, newType, operand->get(), reassociation);
+    newInputs.push_back(newInput);
+  }
+
+  SmallVector<SmallVector<int64_t>> newOutputShapes;
+  SmallVector<AffineMap> outputMaps;
+  for (int i = 0; i < genericOp.getNumDpsInits(); ++i) {
+    OpOperand *output = genericOp.getDpsInitOperand(i);
+    AffineMap oldOutputMap = genericOp.getMatchingIndexingMap(output);
+    ArrayRef<int64_t> oldShape = genericOp.getShape(output);
+    SmallVector<int64_t> thisOutputShape;
+
+    SmallVector<AffineExpr> outputExpr;
+    for (unsigned idx = 0; idx <= oldShape.size(); ++idx) {
+      if (idx == insertSplitIndex) {
+        thisOutputShape.push_back(outputDimsize);
+        outputExpr.push_back(rewriter.getAffineDimExpr(insertSplitDimension));
+      }
+      if (idx < oldShape.size()) {
+        thisOutputShape.push_back(oldShape[idx]);
+        unsigned dim = oldOutputMap.getDimPosition(idx);
+        outputExpr.push_back(rewriter.getAffineDimExpr(
+            dim < insertSplitDimension ? dim : dim + 1));
+      }
+    }
+
+    AffineMap newOutputMap = AffineMap::get(oldOutputMap.getNumDims() + 1, 0,
+                                            outputExpr, rewriter.getContext());
+    newMaps.push_back(newOutputMap);
+    newOutputShapes.push_back(thisOutputShape);
+  }
+
+  // Handle dynamic dimensions for identity value tensor.
+  SmallVector<Value> dynValDims;
+  SmallVector<int64_t> newOutputShape = newOutputShapes[0];
+  for (size_t i = 0; i < newOutputShape.size(); ++i) {
+    if (ShapedType::isDynamic(newOutputShape[i])) {
+      dynValDims.push_back(rewriter.create<tensor::DimOp>(
+          loc, genericOp.getDpsInputOperand(0)->get(), i));
+    }
+  }
+
+  Type valueElemType = genericOp.getRegionOutputArgs()[0].getType();
+  Value emptyValTensor = rewriter.create<tensor::EmptyOp>(
+      loc, newOutputShape, valueElemType, dynValDims);
+  Value constantOp = rewriter.create<arith::ConstantOp>(loc, *identity);
+  Value identityVal =
+      rewriter.create<linalg::FillOp>(loc, constantOp, emptyValTensor)
+          .getResult(0);
+
+  // Handle dynamic dimensions for identity index tensor.
+  SmallVector<Value> dynIdxDims;
+  newOutputShape = newOutputShapes[1];
+  for (size_t i = 0; i < newOutputShape.size(); ++i) {
+    if (ShapedType::isDynamic(newOutputShape[i])) {
+      dynIdxDims.push_back(rewriter.create<tensor::DimOp>(
+          loc, genericOp.getDpsInputOperand(0)->get(), i));
+    }
+  }
+  Type idxElemType = genericOp.getRegionOutputArgs()[1].getType();
+  Value zeroIdx = rewriter.create<arith::ConstantOp>(
+      loc, rewriter.getZeroAttr(idxElemType));
+  Value idxInitTensor = rewriter.create<tensor::EmptyOp>(
+      loc, newOutputShape, idxElemType, dynIdxDims);
+  Value identityIndex =
+      rewriter.create<linalg::FillOp>(loc, zeroIdx, idxInitTensor).getResult(0);
+
+  SmallVector<utils::IteratorType> newIteratorTypes;
+  for (auto [index, iteratorType] :
+       llvm::enumerate(genericOp.getIteratorTypesArray())) {
+    if (insertSplitDimension == index)
+      newIteratorTypes.push_back(utils::IteratorType::parallel);
+    newIteratorTypes.push_back(iteratorType);
+  }
+  if (insertSplitDimension == genericOp.getIteratorTypesArray().size()) {
+    newIteratorTypes.push_back(utils::IteratorType::parallel);
+  }
+
+  // Create partial linalg.generic op with global index computation.
+  Value tileSize = rewriter.create<arith::ConstantIndexOp>(loc, ratio);
+  auto partialOp = rewriter.create<linalg::GenericOp>(
+      loc, TypeRange{identityVal.getType(), identityIndex.getType()}, newInputs,
+      ValueRange{identityVal, identityIndex}, newMaps, newIteratorTypes);
+
+  rewriter.inlineRegionBefore(genericOp.getRegion(), partialOp.getRegion(),
+                              partialOp.getRegion().begin());
+
+  Block &body = partialOp.getRegion().front();
+  rewriter.setInsertionPointToStart(&body);
+
+  unsigned innerIdxDim = reductionDim + 1;
+  unsigned outerIdxDim = insertSplitDimension;
+
+  // Compute global index (gidx) for reduction when the original reduction
+  // dimension is split into [outerIdx, innerIdx] using `ratio`. This is used to
+  // correctly compute the global index for comparisons and index selection.
+  Value outerIdx = rewriter.create<linalg::IndexOp>(loc, outerIdxDim);
+  Value innerIdx = rewriter.create<linalg::IndexOp>(loc, innerIdxDim);
+  Value offset = rewriter.create<arith::MulIOp>(loc, outerIdx, tileSize);
+  Value gidx = rewriter.create<arith::AddIOp>(loc, offset, innerIdx);
+
+  auto selectOp = dyn_cast<arith::SelectOp>(combinerOps.selectOp);
+  Value oldIdx = selectOp.getTrueValue();
+  Value newIdx = gidx;
+  if (oldIdx.getType() != gidx.getType()) {
+    newIdx = rewriter.create<arith::IndexCastOp>(loc, oldIdx.getType(), gidx);
+  }
+  selectOp.setOperand(1, newIdx);
+  rewriter.setInsertionPointAfter(partialOp);
+
+  unsigned intermRank = newOutputShape.size();
+  AffineMap valueMap = rewriter.getMultiDimIdentityMap(intermRank);
+  AffineMap indexMap = valueMap;
+  SmallVector<utils::IteratorType> reductionIteratorTypes;
path: compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/SplitReduction.cpp
user: login: hanhanW
body: I think the construction of `reductionIteratorTypes` can be easier, and it makes the below loop simpler.

```cpp
SmallVector<utils::IteratorType> reductionIteratorTypes(intermRank, utils::IteratorType::parallel);
reductionIteratorTypes[insertSplitIndex] = utils::IteratorType::reduction;
```

Then you dont need the else statement in the below loop.
created_at: 2025-05-09T20:45:19Z
author_association: CONTRIBUTOR
diff_hunk: @@ -389,4 +390,315 @@ splitReduction(RewriterBase &rewriter, LinalgExt::TopkOp topkOp,
   return success();
 }
 
+struct ArgmaxCombinerOps {
+  Operation *maxOp = nullptr;    // arith.maximumf
+  Operation *selectOp = nullptr; // arith.select
+  Operation *cmpOp = nullptr;    // arith.cmpf
+};
+
+// Matches the combiner pattern in a linalg.generic argmax-style reduction:
+// Example MLIR:
+// %4:2 = linalg.generic {
+//     indexing_maps = [...],
+//     iterator_types = ["parallel", "reduction"]
+//   } ins(%arg0 : tensor<?x128xbf16>) outs(%1, %3 : tensor<?xbf16>,
+//   tensor<?xi64>) {
+// ^bb0(%in: bf16, %out: bf16, %out_0: i64):
+//   %5 = linalg.index 1 : index
+//   %6 = arith.index_cast %5 : index to i64
+//   %7 = arith.maximumf %in, %out : bf16
+//   %8 = arith.cmpf ogt, %in, %out : bf16
+//   %9 = arith.select %8, %6, %out_0 : i64
+//   linalg.yield %7, %9 : bf16, i64
+// } -> (tensor<?xbf16>, tensor<?xi64>)
+//
+// This function extracts the `arith.maximumf`, `arith.cmpf`, and `arith.select`
+// operations from the body to facilitate transformations such as split
+// reduction.
+static FailureOr<ArgmaxCombinerOps>
+collectArgmaxCombinerOps(linalg::GenericOp genericOp) {
+
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  ArgmaxCombinerOps ops;
+
+  auto yieldOp = cast<linalg::YieldOp>(genericOp.getBody()->getTerminator());
+
+  // Extract max value producer: arith.maximumf.
+  Value maxResult = yieldOp.getOperand(0);
+  auto maxOp = dyn_cast<arith::MaximumFOp>(maxResult.getDefiningOp());
+
+  // Extract index result producer: arith.select.
+  Value indexResult = yieldOp.getOperand(1);
+  auto selectOp = dyn_cast<arith::SelectOp>(indexResult.getDefiningOp());
+
+  // Extract the condition of the select, expected to be arith.cmpf with
+  // predicate OGT.
+  auto cmpOp = dyn_cast<arith::CmpFOp>(selectOp.getCondition().getDefiningOp());
+
+  ops.maxOp = maxOp;
+  ops.selectOp = selectOp;
+  ops.cmpOp = cmpOp;
+
+  return ops;
+}
+
+FailureOr<linalg::SplitReductionResult>
+splitArgmaxReduction(RewriterBase &rewriter, linalg::GenericOp genericOp,
+                     linalg::ControlSplitReductionFn controlSplitReductionFn) {
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  OpBuilder::InsertionGuard guard(rewriter);
+  rewriter.setInsertionPoint(genericOp);
+  Location loc = genericOp->getLoc();
+
+  linalg::SplitReductionOptions control = controlSplitReductionFn(genericOp);
+  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;
+  }
path: compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/SplitReduction.cpp
user: login: hanhanW
body: `innerParallel` seems to be always false for argmax case, should we bail out when it is true?

Also, I'd remove the blank line in the between. They belong the same block to me.
created_at: 2025-05-09T20:46:49Z
author_association: CONTRIBUTOR
diff_hunk: @@ -63,28 +81,72 @@ namespace {
 struct SplitReductionPass final
     : public impl::SplitReductionPassBase<SplitReductionPass> {
   void runOnOperation() override {
-    if (splitReductionRatio.getValue() <= 1 &&
-        topkSplitReductionRatio.empty()) {
+    if (splitMatmulReductionRatio.getValue() <= 1 &&
+        topkSplitReductionRatio.empty() &&
+        (splitArgmaxReductionRatio.getValue() <= 1 &&
+         !enableStaticArgmaxSplit)) {
       return;
     }
 
+    if (enableStaticArgmaxSplit) {
+      // Use default split-k pattern for argmax ops: split the reduction dim
+      // (e.g., 131072) into tiles of size 128, resulting in 1024 tiles (131072
+      // / 128). So, the split ratio refers to the tile size of the reduction
+      // dimension.
+      splitArgmaxReductionRatio = 128;
+    }
path: compiler/src/iree/compiler/DispatchCreation/SplitReduction.cpp
user: login: bangtianliu
body: > Why do we need the additional flag, but not just making `iree-dispatch-creation-split-argmax-reduction` default to 128? The other question is that do we want it happening by default? If so, we should document how 128 is derived. If it is an experimental number, we can also just mention it in the comment. People will ask why when they see the magic number, IMO.
> 
> My expectation is that the default value is 1 (and we don't need the `enableStaticArgmaxSplit` variable). You use `--iree-dispatch-creation-split-argmax-reduction=128` in your compilation config. We can set the default value to a reasonable value when we are confident for this.
> 
> (I don't work on this area, so maybe we are already confident about this. Then it is okay to set it to 128 by default.)

Here is the context: https://github.com/iree-org/iree/pull/20717#issuecomment-2856229221

We want to enable it by default so that the performance issue in #20650 can be solved. 
created_at: 2025-05-09T21:09:06Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,71 @@
+// RUN: iree-opt --pass-pipeline='builtin.module(util.func(iree-dispatch-creation-split-reduction-ops,cse, canonicalize))' %s | FileCheck %s
+
+util.func public @argmax(%arg0: tensor<?x131072xbf16>, %arg1: index) -> tensor<?xi64> {
+  %cst = arith.constant 0xFF80 : bf16
+  %c0_i64 = arith.constant 0 : i64
+  %0 = tensor.empty(%arg1) : tensor<?xbf16>
+  %1 = linalg.fill ins(%cst : bf16) outs(%0 : tensor<?xbf16>) -> tensor<?xbf16>
+  %2 = tensor.empty(%arg1) : tensor<?xi64>
+  %3 = linalg.fill ins(%c0_i64 : i64) outs(%2 : tensor<?xi64>) -> tensor<?xi64>
+  %4:2 = linalg.generic {
+      indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
+                       affine_map<(d0, d1) -> (d0)>,
+                       affine_map<(d0, d1) -> (d0)>],
+      iterator_types = ["parallel", "reduction"]}
+      ins(%arg0 : tensor<?x131072xbf16>) outs(%1, %3 : tensor<?xbf16>, tensor<?xi64>) {
+  ^bb0(%in: bf16, %out: bf16, %out_0: i64):
+    %5 = linalg.index 1 : index
+    %6 = arith.index_cast %5 : index to i64
+    %7 = arith.maximumf %in, %out : bf16
+    %8 = arith.cmpf ogt, %in, %out : bf16
+    %9 = arith.select %8, %6, %out_0 : i64
+    linalg.yield %7, %9 : bf16, i64
+  } -> (tensor<?xbf16>, tensor<?xi64>)
+  util.return %4#1 : tensor<?xi64>
+}
+
+// CHECK-LABEL:   util.func public @argmax
+
+// Check identity value preparation
+// CHECK: %[[CST:.+]] = arith.constant 0xFF80 : bf16
+// CHECK: %[[ZERO:.+]] = arith.constant 0 : i64
+// CHECK: %[[FINALVAL_EMPTY:.+]] = tensor.empty(%[[ARG1:.+]]) : tensor<?xbf16>
+// CHECK: %[[FINALVAL:.+]] = linalg.fill ins(%[[CST]] : bf16) outs(%[[FINALVAL_EMPTY]] : tensor<?xbf16>) -> tensor<?xbf16>
+// CHECK: %[[FINALIDX_EMPTY:.+]] = tensor.empty(%[[ARG1]]) : tensor<?xi64>
+// CHECK: %[[FINALIDX:.+]] = linalg.fill ins(%[[ZERO]] : i64) outs(%[[FINALIDX_EMPTY]] : tensor<?xi64>) -> tensor<?xi64>
+
+// Check partial reduction
+// CHECK: %[[EXPAND:.+]] = tensor.expand_shape %arg0 {{\[}}[0], [1, 2]] output_shape [%{{.+}}, 1024, 128] : tensor<?x131072xbf16> into tensor<?x1024x128xbf16>
+// CHECK: %[[INITVAL:.+]] = tensor.empty(%{{.+}}) : tensor<?x1024xbf16>
+// CHECK: %[[FILLVAL:.+]] = linalg.fill ins(%{{.+}} : bf16) outs(%[[INITVAL]] : tensor<?x1024xbf16>) -> tensor<?x1024xbf16>
+// CHECK: %[[INITIDX:.+]] = tensor.empty(%{{.+}}) : tensor<?x1024xi64>
+// CHECK: %[[FILLIDX:.+]] = linalg.fill ins(%{{.+}} : i64) outs(%[[INITIDX]] : tensor<?x1024xi64>) -> tensor<?x1024xi64>
+
+// CHECK: %[[PARTIAL:.+]]:2 = linalg.generic
+// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]
path: compiler/src/iree/compiler/DispatchCreation/test/split_argmax_reduction.mlir
user: login: hanhanW
body: I missed a comment I had in my mind, we may want to check indexing maps as well.
created_at: 2025-05-09T21:11:45Z
author_association: CONTRIBUTOR
diff_hunk: @@ -63,28 +81,72 @@ namespace {
 struct SplitReductionPass final
     : public impl::SplitReductionPassBase<SplitReductionPass> {
   void runOnOperation() override {
-    if (splitReductionRatio.getValue() <= 1 &&
-        topkSplitReductionRatio.empty()) {
+    if (splitMatmulReductionRatio.getValue() <= 1 &&
+        topkSplitReductionRatio.empty() &&
+        (splitArgmaxReductionRatio.getValue() <= 1 &&
+         !enableStaticArgmaxSplit)) {
       return;
     }
 
+    if (enableStaticArgmaxSplit) {
+      // Use default split-k pattern for argmax ops: split the reduction dim
+      // (e.g., 131072) into tiles of size 128, resulting in 1024 tiles (131072
+      // / 128). So, the split ratio refers to the tile size of the reduction
+      // dimension.
+      splitArgmaxReductionRatio = 128;
+    }
path: compiler/src/iree/compiler/DispatchCreation/SplitReduction.cpp
user: login: MaheshRavishankar
body: I think we need to have two knobs, the size of the reduction at which to apply split-k on and one to select the split to apply.
created_at: 2025-05-10T17:32:15Z
author_association: COLLABORATOR
diff_hunk: @@ -24,8 +26,24 @@ namespace mlir::iree_compiler::DispatchCreation {
 
 // TODO(thomasraoux): Move to attributes.
 static llvm::cl::opt<int64_t>
-    splitReductionRatio("iree-dispatch-creation-split-matmul-reduction",
-                        llvm::cl::desc("split ratio"), llvm::cl::init(1));
+    splitMatmulReductionRatio("iree-dispatch-creation-split-matmul-reduction",
+                              llvm::cl::desc("split ratio"), llvm::cl::init(1));
+
+static llvm::cl::opt<bool> enableStaticArgmaxSplit(
path: compiler/src/iree/compiler/DispatchCreation/SplitReduction.cpp
user: login: MaheshRavishankar
body: I don't think we want these options if we are trying to turn it on by default. You need to make it two values

1. The minimum size of the reduction dimension that triggers the split
2. The size of (one of the) the split dimension, I think the size of the reduction dimension in the first operation.

You can set the value of 1 to default to 128*1024 and the second to be 128.
created_at: 2025-05-10T17:47:17Z
author_association: COLLABORATOR
diff_hunk: @@ -24,8 +26,24 @@ namespace mlir::iree_compiler::DispatchCreation {
 
 // TODO(thomasraoux): Move to attributes.
 static llvm::cl::opt<int64_t>
-    splitReductionRatio("iree-dispatch-creation-split-matmul-reduction",
-                        llvm::cl::desc("split ratio"), llvm::cl::init(1));
+    splitMatmulReductionRatio("iree-dispatch-creation-split-matmul-reduction",
+                              llvm::cl::desc("split ratio"), llvm::cl::init(1));
+
+static llvm::cl::opt<bool> enableStaticArgmaxSplit(
path: compiler/src/iree/compiler/DispatchCreation/SplitReduction.cpp
user: login: hanhanW
body: I see the reason of having two flags now, it makes sense to me. Thanks!
created_at: 2025-05-13T00:22:48Z
author_association: CONTRIBUTOR
diff_hunk: @@ -389,4 +390,315 @@ splitReduction(RewriterBase &rewriter, LinalgExt::TopkOp topkOp,
   return success();
 }
 
+struct ArgmaxCombinerOps {
+  Operation *maxOp = nullptr;    // arith.maximumf
+  Operation *selectOp = nullptr; // arith.select
+  Operation *cmpOp = nullptr;    // arith.cmpf
+};
+
+// Matches the combiner pattern in a linalg.generic argmax-style reduction:
+// Example MLIR:
+// %4:2 = linalg.generic {
+//     indexing_maps = [...],
+//     iterator_types = ["parallel", "reduction"]
+//   } ins(%arg0 : tensor<?x128xbf16>) outs(%1, %3 : tensor<?xbf16>,
+//   tensor<?xi64>) {
+// ^bb0(%in: bf16, %out: bf16, %out_0: i64):
+//   %5 = linalg.index 1 : index
+//   %6 = arith.index_cast %5 : index to i64
+//   %7 = arith.maximumf %in, %out : bf16
+//   %8 = arith.cmpf ogt, %in, %out : bf16
+//   %9 = arith.select %8, %6, %out_0 : i64
+//   linalg.yield %7, %9 : bf16, i64
+// } -> (tensor<?xbf16>, tensor<?xi64>)
+//
+// This function extracts the `arith.maximumf`, `arith.cmpf`, and `arith.select`
+// operations from the body to facilitate transformations such as split
+// reduction.
+static FailureOr<ArgmaxCombinerOps>
+collectArgmaxCombinerOps(linalg::GenericOp genericOp) {
+
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  ArgmaxCombinerOps ops;
+
+  auto yieldOp = cast<linalg::YieldOp>(genericOp.getBody()->getTerminator());
+
+  // Extract max value producer: arith.maximumf.
+  Value maxResult = yieldOp.getOperand(0);
+  auto maxOp = dyn_cast<arith::MaximumFOp>(maxResult.getDefiningOp());
+
+  // Extract index result producer: arith.select.
+  Value indexResult = yieldOp.getOperand(1);
+  auto selectOp = dyn_cast<arith::SelectOp>(indexResult.getDefiningOp());
+
+  // Extract the condition of the select, expected to be arith.cmpf with
+  // predicate OGT.
+  auto cmpOp = dyn_cast<arith::CmpFOp>(selectOp.getCondition().getDefiningOp());
+
+  ops.maxOp = maxOp;
+  ops.selectOp = selectOp;
+  ops.cmpOp = cmpOp;
+
+  return ops;
+}
+
+FailureOr<linalg::SplitReductionResult>
+splitArgmaxReduction(RewriterBase &rewriter, linalg::GenericOp genericOp,
+                     linalg::ControlSplitReductionFn controlSplitReductionFn) {
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  OpBuilder::InsertionGuard guard(rewriter);
+  rewriter.setInsertionPoint(genericOp);
+  Location loc = genericOp->getLoc();
+
+  linalg::SplitReductionOptions control = controlSplitReductionFn(genericOp);
+  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, 4> loopRanges = genericOp.getStaticLoopRanges();
+  int64_t reductionDimSize = loopRanges[reductionDim];
+
+  //  The total number of output elements along this new dimension is
+  //  reductionDimSize / ratio.
+  int64_t outputDimsize = reductionDimSize / ratio;
+
+  if (reductionDimSize == ShapedType::kDynamic ||
+      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");
+  }
+
+  FailureOr<ArgmaxCombinerOps> maybeOps = collectArgmaxCombinerOps(genericOp);
+  if (failed(maybeOps))
+    return rewriter.notifyMatchFailure(genericOp,
+                                       "invalid combiner for argmax");
+
+  ArgmaxCombinerOps combinerOps = *maybeOps;
+  Operation *reductionOp = combinerOps.maxOp;
+
+  std::optional<TypedAttr> identity = arith::getNeutralElement(reductionOp);
+  if (!identity.has_value())
+    return rewriter.notifyMatchFailure(
+        genericOp, "Unknown identity value for the reduction");
+
+  SmallVector<Value> newInputs;
+  SmallVector<AffineMap> newMaps;
+  // Calculate the new shapes and indexing maps of the input operands.
+  for (OpOperand *operand : genericOp.getDpsInputOperands()) {
+    AffineMap map = genericOp.getMatchingIndexingMap(operand);
+    SmallVector<int64_t> newShape;
+    SmallVector<AffineExpr> exprs;
+    SmallVector<ReassociationIndices> reassociation;
+    unsigned index = 0;
+    for (unsigned idx : llvm::seq<unsigned>(0, map.getNumResults())) {
+      unsigned dim = map.getDimPosition(idx);
+      if (reductionDim == dim) {
+        if (control.innerParallel) {
+          newShape.push_back(ratio); // reduce
+          newShape.push_back(genericOp.getShape(operand)[idx] /
+                             ratio); // parallel (insert)
+          exprs.push_back(rewriter.getAffineDimExpr(
+              dim < insertSplitDimension ? dim : dim + 1));
+          exprs.push_back(rewriter.getAffineDimExpr(insertSplitDimension));
+        } else {
+          newShape.push_back(genericOp.getShape(operand)[idx] /
+                             ratio); // parallel (insert)
+          newShape.push_back(ratio); // reduce
+          exprs.push_back(rewriter.getAffineDimExpr(insertSplitDimension));
+          exprs.push_back(rewriter.getAffineDimExpr(
+              dim < insertSplitDimension ? dim : dim + 1));
+        }
+        reassociation.push_back({index++, index++});
+        continue;
+      }
+      newShape.push_back(genericOp.getShape(operand)[idx]);
+      exprs.push_back(rewriter.getAffineDimExpr(
+          dim < insertSplitDimension ? dim : dim + 1));
+      reassociation.push_back({index++});
+    }
+    newMaps.push_back(
+        AffineMap::get(map.getNumDims() + 1, 0, exprs, genericOp.getContext()));
+    // If the shape is unchanged the input doesn't change.
+    if (newShape == genericOp.getShape(operand)) {
+      newInputs.push_back(operand->get());
+      continue;
+    }
+    Type newType = RankedTensorType::get(
+        newShape,
+        cast<RankedTensorType>(operand->get().getType()).getElementType());
+
+    Value newInput = rewriter.create<tensor::ExpandShapeOp>(
+        loc, newType, operand->get(), reassociation);
+    newInputs.push_back(newInput);
+  }
+
+  SmallVector<SmallVector<int64_t>> newOutputShapes;
+  SmallVector<AffineMap> outputMaps;
+  for (int i = 0; i < genericOp.getNumDpsInits(); ++i) {
+    OpOperand *output = genericOp.getDpsInitOperand(i);
+    AffineMap oldOutputMap = genericOp.getMatchingIndexingMap(output);
+    ArrayRef<int64_t> oldShape = genericOp.getShape(output);
+    SmallVector<int64_t> thisOutputShape;
+
+    SmallVector<AffineExpr> outputExpr;
+    for (unsigned idx = 0; idx <= oldShape.size(); ++idx) {
+      if (idx == insertSplitIndex) {
+        thisOutputShape.push_back(outputDimsize);
+        outputExpr.push_back(rewriter.getAffineDimExpr(insertSplitDimension));
+      }
+      if (idx < oldShape.size()) {
+        thisOutputShape.push_back(oldShape[idx]);
+        unsigned dim = oldOutputMap.getDimPosition(idx);
+        outputExpr.push_back(rewriter.getAffineDimExpr(
+            dim < insertSplitDimension ? dim : dim + 1));
+      }
+    }
+
+    AffineMap newOutputMap = AffineMap::get(oldOutputMap.getNumDims() + 1, 0,
+                                            outputExpr, rewriter.getContext());
+    newMaps.push_back(newOutputMap);
+    newOutputShapes.push_back(thisOutputShape);
+  }
+
+  // Handle dynamic dimensions for identity value tensor.
+  SmallVector<Value> dynValDims;
+  SmallVector<int64_t> newOutputShape = newOutputShapes[0];
+  for (size_t i = 0; i < newOutputShape.size(); ++i) {
+    if (ShapedType::isDynamic(newOutputShape[i])) {
+      dynValDims.push_back(rewriter.create<tensor::DimOp>(
+          loc, genericOp.getDpsInputOperand(0)->get(), i));
+    }
+  }
+
+  Type valueElemType = genericOp.getRegionOutputArgs()[0].getType();
+  Value emptyValTensor = rewriter.create<tensor::EmptyOp>(
+      loc, newOutputShape, valueElemType, dynValDims);
+  Value constantOp = rewriter.create<arith::ConstantOp>(loc, *identity);
+  Value identityVal =
+      rewriter.create<linalg::FillOp>(loc, constantOp, emptyValTensor)
+          .getResult(0);
+
+  // Handle dynamic dimensions for identity index tensor.
+  SmallVector<Value> dynIdxDims;
+  newOutputShape = newOutputShapes[1];
+  for (size_t i = 0; i < newOutputShape.size(); ++i) {
+    if (ShapedType::isDynamic(newOutputShape[i])) {
+      dynIdxDims.push_back(rewriter.create<tensor::DimOp>(
+          loc, genericOp.getDpsInputOperand(0)->get(), i));
+    }
+  }
path: compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/SplitReduction.cpp
user: login: bangtianliu
body: I am confused about how this 
```
    %expanded_0 = tensor.expand_shape %1 [[0, 1]] output_shape [%dim, 1024] : tensor<?xbf16> into tensor<?x1024xbf16>
```
is generated in your provided mlir since from my end I noticed that association can not be computed from <?xbf16> to <?x1024xbf16>. 

created_at: 2025-05-13T02:40:13Z
author_association: CONTRIBUTOR
diff_hunk: @@ -389,4 +390,315 @@ splitReduction(RewriterBase &rewriter, LinalgExt::TopkOp topkOp,
   return success();
 }
 
+struct ArgmaxCombinerOps {
+  Operation *maxOp = nullptr;    // arith.maximumf
+  Operation *selectOp = nullptr; // arith.select
+  Operation *cmpOp = nullptr;    // arith.cmpf
+};
+
+// Matches the combiner pattern in a linalg.generic argmax-style reduction:
+// Example MLIR:
+// %4:2 = linalg.generic {
+//     indexing_maps = [...],
+//     iterator_types = ["parallel", "reduction"]
+//   } ins(%arg0 : tensor<?x128xbf16>) outs(%1, %3 : tensor<?xbf16>,
+//   tensor<?xi64>) {
+// ^bb0(%in: bf16, %out: bf16, %out_0: i64):
+//   %5 = linalg.index 1 : index
+//   %6 = arith.index_cast %5 : index to i64
+//   %7 = arith.maximumf %in, %out : bf16
+//   %8 = arith.cmpf ogt, %in, %out : bf16
+//   %9 = arith.select %8, %6, %out_0 : i64
+//   linalg.yield %7, %9 : bf16, i64
+// } -> (tensor<?xbf16>, tensor<?xi64>)
+//
+// This function extracts the `arith.maximumf`, `arith.cmpf`, and `arith.select`
+// operations from the body to facilitate transformations such as split
+// reduction.
+static FailureOr<ArgmaxCombinerOps>
+collectArgmaxCombinerOps(linalg::GenericOp genericOp) {
+
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  ArgmaxCombinerOps ops;
+
+  auto yieldOp = cast<linalg::YieldOp>(genericOp.getBody()->getTerminator());
+
+  // Extract max value producer: arith.maximumf.
+  Value maxResult = yieldOp.getOperand(0);
+  auto maxOp = dyn_cast<arith::MaximumFOp>(maxResult.getDefiningOp());
+
+  // Extract index result producer: arith.select.
+  Value indexResult = yieldOp.getOperand(1);
+  auto selectOp = dyn_cast<arith::SelectOp>(indexResult.getDefiningOp());
+
+  // Extract the condition of the select, expected to be arith.cmpf with
+  // predicate OGT.
+  auto cmpOp = dyn_cast<arith::CmpFOp>(selectOp.getCondition().getDefiningOp());
+
+  ops.maxOp = maxOp;
+  ops.selectOp = selectOp;
+  ops.cmpOp = cmpOp;
+
+  return ops;
+}
+
+FailureOr<linalg::SplitReductionResult>
+splitArgmaxReduction(RewriterBase &rewriter, linalg::GenericOp genericOp,
+                     linalg::ControlSplitReductionFn controlSplitReductionFn) {
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  OpBuilder::InsertionGuard guard(rewriter);
+  rewriter.setInsertionPoint(genericOp);
+  Location loc = genericOp->getLoc();
+
+  linalg::SplitReductionOptions control = controlSplitReductionFn(genericOp);
+  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, 4> loopRanges = genericOp.getStaticLoopRanges();
+  int64_t reductionDimSize = loopRanges[reductionDim];
+
+  //  The total number of output elements along this new dimension is
+  //  reductionDimSize / ratio.
+  int64_t outputDimsize = reductionDimSize / ratio;
+
+  if (reductionDimSize == ShapedType::kDynamic ||
+      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");
+  }
+
+  FailureOr<ArgmaxCombinerOps> maybeOps = collectArgmaxCombinerOps(genericOp);
+  if (failed(maybeOps))
+    return rewriter.notifyMatchFailure(genericOp,
+                                       "invalid combiner for argmax");
+
+  ArgmaxCombinerOps combinerOps = *maybeOps;
+  Operation *reductionOp = combinerOps.maxOp;
+
+  std::optional<TypedAttr> identity = arith::getNeutralElement(reductionOp);
+  if (!identity.has_value())
+    return rewriter.notifyMatchFailure(
+        genericOp, "Unknown identity value for the reduction");
+
+  SmallVector<Value> newInputs;
+  SmallVector<AffineMap> newMaps;
+  // Calculate the new shapes and indexing maps of the input operands.
+  for (OpOperand *operand : genericOp.getDpsInputOperands()) {
+    AffineMap map = genericOp.getMatchingIndexingMap(operand);
+    SmallVector<int64_t> newShape;
+    SmallVector<AffineExpr> exprs;
+    SmallVector<ReassociationIndices> reassociation;
+    unsigned index = 0;
+    for (unsigned idx : llvm::seq<unsigned>(0, map.getNumResults())) {
+      unsigned dim = map.getDimPosition(idx);
+      if (reductionDim == dim) {
+        if (control.innerParallel) {
+          newShape.push_back(ratio); // reduce
+          newShape.push_back(genericOp.getShape(operand)[idx] /
+                             ratio); // parallel (insert)
+          exprs.push_back(rewriter.getAffineDimExpr(
+              dim < insertSplitDimension ? dim : dim + 1));
+          exprs.push_back(rewriter.getAffineDimExpr(insertSplitDimension));
+        } else {
+          newShape.push_back(genericOp.getShape(operand)[idx] /
+                             ratio); // parallel (insert)
+          newShape.push_back(ratio); // reduce
+          exprs.push_back(rewriter.getAffineDimExpr(insertSplitDimension));
+          exprs.push_back(rewriter.getAffineDimExpr(
+              dim < insertSplitDimension ? dim : dim + 1));
+        }
+        reassociation.push_back({index++, index++});
+        continue;
+      }
+      newShape.push_back(genericOp.getShape(operand)[idx]);
+      exprs.push_back(rewriter.getAffineDimExpr(
+          dim < insertSplitDimension ? dim : dim + 1));
+      reassociation.push_back({index++});
+    }
+    newMaps.push_back(
+        AffineMap::get(map.getNumDims() + 1, 0, exprs, genericOp.getContext()));
+    // If the shape is unchanged the input doesn't change.
+    if (newShape == genericOp.getShape(operand)) {
+      newInputs.push_back(operand->get());
+      continue;
+    }
+    Type newType = RankedTensorType::get(
+        newShape,
+        cast<RankedTensorType>(operand->get().getType()).getElementType());
+
+    Value newInput = rewriter.create<tensor::ExpandShapeOp>(
+        loc, newType, operand->get(), reassociation);
+    newInputs.push_back(newInput);
+  }
+
+  SmallVector<SmallVector<int64_t>> newOutputShapes;
+  SmallVector<AffineMap> outputMaps;
+  for (int i = 0; i < genericOp.getNumDpsInits(); ++i) {
+    OpOperand *output = genericOp.getDpsInitOperand(i);
+    AffineMap oldOutputMap = genericOp.getMatchingIndexingMap(output);
+    ArrayRef<int64_t> oldShape = genericOp.getShape(output);
+    SmallVector<int64_t> thisOutputShape;
+
+    SmallVector<AffineExpr> outputExpr;
+    for (unsigned idx = 0; idx <= oldShape.size(); ++idx) {
+      if (idx == insertSplitIndex) {
+        thisOutputShape.push_back(outputDimsize);
+        outputExpr.push_back(rewriter.getAffineDimExpr(insertSplitDimension));
+      }
+      if (idx < oldShape.size()) {
+        thisOutputShape.push_back(oldShape[idx]);
+        unsigned dim = oldOutputMap.getDimPosition(idx);
+        outputExpr.push_back(rewriter.getAffineDimExpr(
+            dim < insertSplitDimension ? dim : dim + 1));
+      }
+    }
+
+    AffineMap newOutputMap = AffineMap::get(oldOutputMap.getNumDims() + 1, 0,
+                                            outputExpr, rewriter.getContext());
+    newMaps.push_back(newOutputMap);
+    newOutputShapes.push_back(thisOutputShape);
+  }
+
+  // Handle dynamic dimensions for identity value tensor.
+  SmallVector<Value> dynValDims;
+  SmallVector<int64_t> newOutputShape = newOutputShapes[0];
+  for (size_t i = 0; i < newOutputShape.size(); ++i) {
+    if (ShapedType::isDynamic(newOutputShape[i])) {
+      dynValDims.push_back(rewriter.create<tensor::DimOp>(
+          loc, genericOp.getDpsInputOperand(0)->get(), i));
+    }
+  }
+
+  Type valueElemType = genericOp.getRegionOutputArgs()[0].getType();
+  Value emptyValTensor = rewriter.create<tensor::EmptyOp>(
+      loc, newOutputShape, valueElemType, dynValDims);
+  Value constantOp = rewriter.create<arith::ConstantOp>(loc, *identity);
+  Value identityVal =
+      rewriter.create<linalg::FillOp>(loc, constantOp, emptyValTensor)
+          .getResult(0);
+
+  // Handle dynamic dimensions for identity index tensor.
+  SmallVector<Value> dynIdxDims;
+  newOutputShape = newOutputShapes[1];
+  for (size_t i = 0; i < newOutputShape.size(); ++i) {
+    if (ShapedType::isDynamic(newOutputShape[i])) {
+      dynIdxDims.push_back(rewriter.create<tensor::DimOp>(
+          loc, genericOp.getDpsInputOperand(0)->get(), i));
+    }
+  }
path: compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/SplitReduction.cpp
user: login: bangtianliu
body: Especially its is converted from %1
```
    %1 = linalg.fill ins(%cst : bf16) outs(%0 : tensor<?xbf16>) -> tensor<?xbf16>
```

For example, assume the case in which %1 is of <1xbf16>. cc @hanhanW 
created_at: 2025-05-13T03:19:46Z
author_association: CONTRIBUTOR
diff_hunk: @@ -389,4 +390,315 @@ splitReduction(RewriterBase &rewriter, LinalgExt::TopkOp topkOp,
   return success();
 }
 
+struct ArgmaxCombinerOps {
+  Operation *maxOp = nullptr;    // arith.maximumf
+  Operation *selectOp = nullptr; // arith.select
+  Operation *cmpOp = nullptr;    // arith.cmpf
+};
+
+// Matches the combiner pattern in a linalg.generic argmax-style reduction:
+// Example MLIR:
+// %4:2 = linalg.generic {
+//     indexing_maps = [...],
+//     iterator_types = ["parallel", "reduction"]
+//   } ins(%arg0 : tensor<?x128xbf16>) outs(%1, %3 : tensor<?xbf16>,
+//   tensor<?xi64>) {
+// ^bb0(%in: bf16, %out: bf16, %out_0: i64):
+//   %5 = linalg.index 1 : index
+//   %6 = arith.index_cast %5 : index to i64
+//   %7 = arith.maximumf %in, %out : bf16
+//   %8 = arith.cmpf ogt, %in, %out : bf16
+//   %9 = arith.select %8, %6, %out_0 : i64
+//   linalg.yield %7, %9 : bf16, i64
+// } -> (tensor<?xbf16>, tensor<?xi64>)
+//
+// This function extracts the `arith.maximumf`, `arith.cmpf`, and `arith.select`
+// operations from the body to facilitate transformations such as split
+// reduction.
+static FailureOr<ArgmaxCombinerOps>
+collectArgmaxCombinerOps(linalg::GenericOp genericOp) {
+
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  ArgmaxCombinerOps ops;
+
+  auto yieldOp = cast<linalg::YieldOp>(genericOp.getBody()->getTerminator());
+
+  // Extract max value producer: arith.maximumf.
+  Value maxResult = yieldOp.getOperand(0);
+  auto maxOp = dyn_cast<arith::MaximumFOp>(maxResult.getDefiningOp());
+
+  // Extract index result producer: arith.select.
+  Value indexResult = yieldOp.getOperand(1);
+  auto selectOp = dyn_cast<arith::SelectOp>(indexResult.getDefiningOp());
+
+  // Extract the condition of the select, expected to be arith.cmpf with
+  // predicate OGT.
+  auto cmpOp = dyn_cast<arith::CmpFOp>(selectOp.getCondition().getDefiningOp());
+
+  ops.maxOp = maxOp;
+  ops.selectOp = selectOp;
+  ops.cmpOp = cmpOp;
+
+  return ops;
+}
+
+FailureOr<linalg::SplitReductionResult>
+splitArgmaxReduction(RewriterBase &rewriter, linalg::GenericOp genericOp,
+                     linalg::ControlSplitReductionFn controlSplitReductionFn) {
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  OpBuilder::InsertionGuard guard(rewriter);
+  rewriter.setInsertionPoint(genericOp);
+  Location loc = genericOp->getLoc();
+
+  linalg::SplitReductionOptions control = controlSplitReductionFn(genericOp);
+  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, 4> loopRanges = genericOp.getStaticLoopRanges();
+  int64_t reductionDimSize = loopRanges[reductionDim];
+
+  //  The total number of output elements along this new dimension is
+  //  reductionDimSize / ratio.
+  int64_t outputDimsize = reductionDimSize / ratio;
+
+  if (reductionDimSize == ShapedType::kDynamic ||
+      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");
+  }
+
+  FailureOr<ArgmaxCombinerOps> maybeOps = collectArgmaxCombinerOps(genericOp);
+  if (failed(maybeOps))
+    return rewriter.notifyMatchFailure(genericOp,
+                                       "invalid combiner for argmax");
+
+  ArgmaxCombinerOps combinerOps = *maybeOps;
+  Operation *reductionOp = combinerOps.maxOp;
+
+  std::optional<TypedAttr> identity = arith::getNeutralElement(reductionOp);
+  if (!identity.has_value())
+    return rewriter.notifyMatchFailure(
+        genericOp, "Unknown identity value for the reduction");
+
+  SmallVector<Value> newInputs;
+  SmallVector<AffineMap> newMaps;
+  // Calculate the new shapes and indexing maps of the input operands.
+  for (OpOperand *operand : genericOp.getDpsInputOperands()) {
+    AffineMap map = genericOp.getMatchingIndexingMap(operand);
+    SmallVector<int64_t> newShape;
+    SmallVector<AffineExpr> exprs;
+    SmallVector<ReassociationIndices> reassociation;
+    unsigned index = 0;
+    for (unsigned idx : llvm::seq<unsigned>(0, map.getNumResults())) {
+      unsigned dim = map.getDimPosition(idx);
+      if (reductionDim == dim) {
+        if (control.innerParallel) {
+          newShape.push_back(ratio); // reduce
+          newShape.push_back(genericOp.getShape(operand)[idx] /
+                             ratio); // parallel (insert)
+          exprs.push_back(rewriter.getAffineDimExpr(
+              dim < insertSplitDimension ? dim : dim + 1));
+          exprs.push_back(rewriter.getAffineDimExpr(insertSplitDimension));
+        } else {
+          newShape.push_back(genericOp.getShape(operand)[idx] /
+                             ratio); // parallel (insert)
+          newShape.push_back(ratio); // reduce
+          exprs.push_back(rewriter.getAffineDimExpr(insertSplitDimension));
+          exprs.push_back(rewriter.getAffineDimExpr(
+              dim < insertSplitDimension ? dim : dim + 1));
+        }
+        reassociation.push_back({index++, index++});
+        continue;
+      }
+      newShape.push_back(genericOp.getShape(operand)[idx]);
+      exprs.push_back(rewriter.getAffineDimExpr(
+          dim < insertSplitDimension ? dim : dim + 1));
+      reassociation.push_back({index++});
+    }
+    newMaps.push_back(
+        AffineMap::get(map.getNumDims() + 1, 0, exprs, genericOp.getContext()));
+    // If the shape is unchanged the input doesn't change.
+    if (newShape == genericOp.getShape(operand)) {
+      newInputs.push_back(operand->get());
+      continue;
+    }
+    Type newType = RankedTensorType::get(
+        newShape,
+        cast<RankedTensorType>(operand->get().getType()).getElementType());
+
+    Value newInput = rewriter.create<tensor::ExpandShapeOp>(
+        loc, newType, operand->get(), reassociation);
+    newInputs.push_back(newInput);
+  }
+
+  SmallVector<SmallVector<int64_t>> newOutputShapes;
+  SmallVector<AffineMap> outputMaps;
+  for (int i = 0; i < genericOp.getNumDpsInits(); ++i) {
+    OpOperand *output = genericOp.getDpsInitOperand(i);
+    AffineMap oldOutputMap = genericOp.getMatchingIndexingMap(output);
+    ArrayRef<int64_t> oldShape = genericOp.getShape(output);
+    SmallVector<int64_t> thisOutputShape;
+
+    SmallVector<AffineExpr> outputExpr;
+    for (unsigned idx = 0; idx <= oldShape.size(); ++idx) {
+      if (idx == insertSplitIndex) {
+        thisOutputShape.push_back(outputDimsize);
+        outputExpr.push_back(rewriter.getAffineDimExpr(insertSplitDimension));
+      }
+      if (idx < oldShape.size()) {
+        thisOutputShape.push_back(oldShape[idx]);
+        unsigned dim = oldOutputMap.getDimPosition(idx);
+        outputExpr.push_back(rewriter.getAffineDimExpr(
+            dim < insertSplitDimension ? dim : dim + 1));
+      }
+    }
+
+    AffineMap newOutputMap = AffineMap::get(oldOutputMap.getNumDims() + 1, 0,
+                                            outputExpr, rewriter.getContext());
+    newMaps.push_back(newOutputMap);
+    newOutputShapes.push_back(thisOutputShape);
+  }
+
+  // Handle dynamic dimensions for identity value tensor.
+  SmallVector<Value> dynValDims;
+  SmallVector<int64_t> newOutputShape = newOutputShapes[0];
+  for (size_t i = 0; i < newOutputShape.size(); ++i) {
+    if (ShapedType::isDynamic(newOutputShape[i])) {
+      dynValDims.push_back(rewriter.create<tensor::DimOp>(
+          loc, genericOp.getDpsInputOperand(0)->get(), i));
+    }
+  }
+
+  Type valueElemType = genericOp.getRegionOutputArgs()[0].getType();
+  Value emptyValTensor = rewriter.create<tensor::EmptyOp>(
+      loc, newOutputShape, valueElemType, dynValDims);
+  Value constantOp = rewriter.create<arith::ConstantOp>(loc, *identity);
+  Value identityVal =
+      rewriter.create<linalg::FillOp>(loc, constantOp, emptyValTensor)
+          .getResult(0);
+
+  // Handle dynamic dimensions for identity index tensor.
+  SmallVector<Value> dynIdxDims;
+  newOutputShape = newOutputShapes[1];
+  for (size_t i = 0; i < newOutputShape.size(); ++i) {
+    if (ShapedType::isDynamic(newOutputShape[i])) {
+      dynIdxDims.push_back(rewriter.create<tensor::DimOp>(
+          loc, genericOp.getDpsInputOperand(0)->get(), i));
+    }
+  }
path: compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/SplitReduction.cpp
user: login: hanhanW
body: Sorry @bangtianliu and good catch! I misunderstood the outs operands part in the algorithm, and here is the suggestion and the reasons behind the suggestion.

In the beginning, I thought that we could reuse the destination tensors (i.e., inits) because we should reuse the inits as much as possible. They could be `tensor.empty()` or other buffers. However, I was wrong on this. Fundamentally, it is not possible because we always need a storage to hold the data until we do the final reduction. (Vectorization is a different story, which could claim that generating `tensor.empty` is always valid. But I'm not going to expand the context here.)

Thus, we always need to create a `tensor.empty()` op and fill in the corresponding values. We can do it with few lines to achieve what you're doing here. Your implementation is correct, and my suggestion is making the code simpler.

My suggestion is to have `getSplitReductionInit(Value origInit, TypeAttr identity, int64_t insertDimSize, unsigned insertSplitIndex)` method and reuse it for all the outs operands. I'm assuming that all the outs need the same behavior. This is true for argmax, but it may be false for some linalg ops that I'm not aware of. Can always generalize it later when we see a need.

In the implementation, we create `arith.constant`, `tensor.empty`, and `linalg.fill ins(%cst) outs(%empty)` ops. This assumption is valid to me because it is a reduction-based generic op. The `arith.constant` op creation looks good to me. For `tensor.empty` op, what matters is the shape and the element type. You can use [tensor::getMixedSizes()](https://github.com/llvm/llvm-project/blob/27983696a6d6caf8f90d77745598aeaec88b7009/mlir/include/mlir/Dialect/Tensor/IR/Tensor.h#L133-L135) to get the shape with `SmallVector<OpFoldResult>` type, use the insert trick to insert the dimension, and create the `tensor.empty()` op with the [builder](https://github.com/llvm/llvm-project/blob/50316c1eb8ed189dc3a979d552b8b04b0730b287/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td#L311-L313). This builder takes `SmallVector<OpFoldResult>` shape and the element type, so you don't need to create the shapes like the current implementation. The internal implementation is very similar to what you have. The final implementation could be something like:

```cpp
// The getElementTypeOrSelf util looks a little silly as it is always a RankedTensorType, but it is helpful.
Type elemType = getElementTypeOrSelf(origInit.getType());
SmallVector<OpFoldResult> shape = tensor::getMixedSizes(builder, loc, origInit);
shape.insert(shape.begin() + insertSplitIndex, builder.getIndexAttr(insertDimSize));
Value emptyValue =
    rewriter.create<tensor::EmptyOp>(loc, shape, elemType);
Value identityValue =
    rewriter.create<linalg::FillOp>(loc, valueConst, emptyValue).getResult(0);
```

I will review the other parts tomorrow morning, and thanks for addressing my comments!
created_at: 2025-05-13T23:04:31Z
author_association: CONTRIBUTOR
diff_hunk: @@ -389,4 +390,352 @@ splitReduction(RewriterBase &rewriter, LinalgExt::TopkOp topkOp,
   return success();
 }
 
+struct ArgmaxCombinerOps {
+  Operation *maxOp = nullptr;    // arith.maximumf
+  Operation *selectOp = nullptr; // arith.select
+  Operation *cmpOp = nullptr;    // arith.cmpf
+};
+
+// Matches the combiner pattern in a linalg.generic argmax-style reduction:
+// Example MLIR:
+// %4:2 = linalg.generic {
+//     indexing_maps = [...],
+//     iterator_types = ["parallel", "reduction"]
+//   } ins(%arg0 : tensor<?x128xbf16>) outs(%1, %3 : tensor<?xbf16>,
+//   tensor<?xi64>) {
+// ^bb0(%in: bf16, %out: bf16, %out_0: i64):
+//   %5 = linalg.index 1 : index
+//   %6 = arith.index_cast %5 : index to i64
+//   %7 = arith.maximumf %in, %out : bf16
+//   %8 = arith.cmpf ogt, %in, %out : bf16
+//   %9 = arith.select %8, %6, %out_0 : i64
+//   linalg.yield %7, %9 : bf16, i64
+// } -> (tensor<?xbf16>, tensor<?xi64>)
+//
+// This function extracts the `arith.maximumf`, `arith.cmpf`, and `arith.select`
+// operations from the body to facilitate transformations such as split
+// reduction.
+static ArgmaxCombinerOps collectArgmaxCombinerOps(linalg::GenericOp genericOp) {
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  auto yieldOp = cast<linalg::YieldOp>(genericOp.getBody()->getTerminator());
+
+  // Extract max value producer: arith.maximumf.
+  Value maxResult = yieldOp.getOperand(0);
+  auto maxOp = cast<arith::MaximumFOp>(maxResult.getDefiningOp());
+
+  // Extract index result producer: arith.select.
+  Value indexResult = yieldOp.getOperand(1);
+  auto selectOp = cast<arith::SelectOp>(indexResult.getDefiningOp());
+
+  // Extract the condition of the select, expected to be arith.cmpf with
+  // predicate OGT.
+  auto cmpOp = cast<arith::CmpFOp>(selectOp.getCondition().getDefiningOp());
+
+  ArgmaxCombinerOps ops;
+  ops.maxOp = maxOp;
+  ops.selectOp = selectOp;
+  ops.cmpOp = cmpOp;
+  return ops;
+}
+
+static Value expandValue(OpBuilder &builder, Location loc, Value value,
+                         RankedTensorType expandedType) {
+  RankedTensorType originalType = cast<RankedTensorType>(value.getType());
+  if (originalType == expandedType)
+    return value;
+
+  auto reassociation =
+      mlir::getReassociationIndicesForReshape(originalType, expandedType);
+  assert(reassociation && "invalid reassociation for expand_shape");
path: compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/SplitReduction.cpp
user: login: hanhanW
body: nit: I'd replace the error message with something like `failed to infer reassociation indices from types`.
created_at: 2025-05-14T03:44:48Z
author_association: CONTRIBUTOR
diff_hunk: @@ -389,4 +390,352 @@ splitReduction(RewriterBase &rewriter, LinalgExt::TopkOp topkOp,
   return success();
 }
 
+struct ArgmaxCombinerOps {
+  Operation *maxOp = nullptr;    // arith.maximumf
+  Operation *selectOp = nullptr; // arith.select
+  Operation *cmpOp = nullptr;    // arith.cmpf
+};
+
+// Matches the combiner pattern in a linalg.generic argmax-style reduction:
+// Example MLIR:
+// %4:2 = linalg.generic {
+//     indexing_maps = [...],
+//     iterator_types = ["parallel", "reduction"]
+//   } ins(%arg0 : tensor<?x128xbf16>) outs(%1, %3 : tensor<?xbf16>,
+//   tensor<?xi64>) {
+// ^bb0(%in: bf16, %out: bf16, %out_0: i64):
+//   %5 = linalg.index 1 : index
+//   %6 = arith.index_cast %5 : index to i64
+//   %7 = arith.maximumf %in, %out : bf16
+//   %8 = arith.cmpf ogt, %in, %out : bf16
+//   %9 = arith.select %8, %6, %out_0 : i64
+//   linalg.yield %7, %9 : bf16, i64
+// } -> (tensor<?xbf16>, tensor<?xi64>)
+//
+// This function extracts the `arith.maximumf`, `arith.cmpf`, and `arith.select`
+// operations from the body to facilitate transformations such as split
+// reduction.
+static ArgmaxCombinerOps collectArgmaxCombinerOps(linalg::GenericOp genericOp) {
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  auto yieldOp = cast<linalg::YieldOp>(genericOp.getBody()->getTerminator());
+
+  // Extract max value producer: arith.maximumf.
+  Value maxResult = yieldOp.getOperand(0);
+  auto maxOp = cast<arith::MaximumFOp>(maxResult.getDefiningOp());
+
+  // Extract index result producer: arith.select.
+  Value indexResult = yieldOp.getOperand(1);
+  auto selectOp = cast<arith::SelectOp>(indexResult.getDefiningOp());
+
+  // Extract the condition of the select, expected to be arith.cmpf with
+  // predicate OGT.
+  auto cmpOp = cast<arith::CmpFOp>(selectOp.getCondition().getDefiningOp());
+
+  ArgmaxCombinerOps ops;
+  ops.maxOp = maxOp;
+  ops.selectOp = selectOp;
+  ops.cmpOp = cmpOp;
+  return ops;
+}
+
+static Value expandValue(OpBuilder &builder, Location loc, Value value,
+                         RankedTensorType expandedType) {
+  RankedTensorType originalType = cast<RankedTensorType>(value.getType());
+  if (originalType == expandedType)
+    return value;
+
+  auto reassociation =
+      mlir::getReassociationIndicesForReshape(originalType, expandedType);
+  assert(reassociation && "invalid reassociation for expand_shape");
+  return builder.create<tensor::ExpandShapeOp>(loc, expandedType, value,
+                                               *reassociation);
+}
+
+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) {
+      // 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));
+      }
+    } else {
+      unsigned shifted = (dim < insertSplitDim) ? dim : dim + 1;
+      exprs.push_back(getAffineDimExpr(shifted, ctx));
+    }
path: compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/SplitReduction.cpp
user: login: hanhanW
body: style nit: prefer early-exit to simplify the code and reduce levels of indents. https://llvm.org/docs/CodingStandards.html#use-early-exits-and-continue-to-simplify-code

E.g.,

```cpp
if (dim != reductionDim) {
  unsigned shifted = (dim < insertSplitDim) ? dim : dim + 1;
  exprs.push_back(getAffineDimExpr(shifted, ctx));
  continue;
}
// other code
...
```
created_at: 2025-05-14T03:46:34Z
author_association: CONTRIBUTOR
diff_hunk: @@ -389,4 +390,352 @@ splitReduction(RewriterBase &rewriter, LinalgExt::TopkOp topkOp,
   return success();
 }
 
+struct ArgmaxCombinerOps {
+  Operation *maxOp = nullptr;    // arith.maximumf
+  Operation *selectOp = nullptr; // arith.select
+  Operation *cmpOp = nullptr;    // arith.cmpf
+};
+
+// Matches the combiner pattern in a linalg.generic argmax-style reduction:
+// Example MLIR:
+// %4:2 = linalg.generic {
+//     indexing_maps = [...],
+//     iterator_types = ["parallel", "reduction"]
+//   } ins(%arg0 : tensor<?x128xbf16>) outs(%1, %3 : tensor<?xbf16>,
+//   tensor<?xi64>) {
+// ^bb0(%in: bf16, %out: bf16, %out_0: i64):
+//   %5 = linalg.index 1 : index
+//   %6 = arith.index_cast %5 : index to i64
+//   %7 = arith.maximumf %in, %out : bf16
+//   %8 = arith.cmpf ogt, %in, %out : bf16
+//   %9 = arith.select %8, %6, %out_0 : i64
+//   linalg.yield %7, %9 : bf16, i64
+// } -> (tensor<?xbf16>, tensor<?xi64>)
+//
+// This function extracts the `arith.maximumf`, `arith.cmpf`, and `arith.select`
+// operations from the body to facilitate transformations such as split
+// reduction.
+static ArgmaxCombinerOps collectArgmaxCombinerOps(linalg::GenericOp genericOp) {
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  auto yieldOp = cast<linalg::YieldOp>(genericOp.getBody()->getTerminator());
+
+  // Extract max value producer: arith.maximumf.
+  Value maxResult = yieldOp.getOperand(0);
+  auto maxOp = cast<arith::MaximumFOp>(maxResult.getDefiningOp());
+
+  // Extract index result producer: arith.select.
+  Value indexResult = yieldOp.getOperand(1);
+  auto selectOp = cast<arith::SelectOp>(indexResult.getDefiningOp());
+
+  // Extract the condition of the select, expected to be arith.cmpf with
+  // predicate OGT.
+  auto cmpOp = cast<arith::CmpFOp>(selectOp.getCondition().getDefiningOp());
+
+  ArgmaxCombinerOps ops;
+  ops.maxOp = maxOp;
+  ops.selectOp = selectOp;
+  ops.cmpOp = cmpOp;
+  return ops;
+}
+
+static Value expandValue(OpBuilder &builder, Location loc, Value value,
+                         RankedTensorType expandedType) {
+  RankedTensorType originalType = cast<RankedTensorType>(value.getType());
+  if (originalType == expandedType)
+    return value;
+
+  auto reassociation =
+      mlir::getReassociationIndicesForReshape(originalType, expandedType);
+  assert(reassociation && "invalid reassociation for expand_shape");
+  return builder.create<tensor::ExpandShapeOp>(loc, expandedType, value,
+                                               *reassociation);
+}
+
+static AffineMap getExpandedInputIndexingMap(AffineMap oldMap,
path: compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/SplitReduction.cpp
user: login: hanhanW
body: style nit: please add a doc, it is not trivial to understand what it returns based on names.

> Almost every function declaration should have comments immediately preceding it that describe what the function does and how to use it. These comments may be omitted only if the function is simple and obvious (e.g., simple accessors for obvious properties of the class). 


https://google.github.io/styleguide/cppguide.html#Function_Comments
created_at: 2025-05-14T03:48:15Z
author_association: CONTRIBUTOR
diff_hunk: @@ -389,4 +390,352 @@ splitReduction(RewriterBase &rewriter, LinalgExt::TopkOp topkOp,
   return success();
 }
 
+struct ArgmaxCombinerOps {
+  Operation *maxOp = nullptr;    // arith.maximumf
+  Operation *selectOp = nullptr; // arith.select
+  Operation *cmpOp = nullptr;    // arith.cmpf
+};
+
+// Matches the combiner pattern in a linalg.generic argmax-style reduction:
+// Example MLIR:
+// %4:2 = linalg.generic {
+//     indexing_maps = [...],
+//     iterator_types = ["parallel", "reduction"]
+//   } ins(%arg0 : tensor<?x128xbf16>) outs(%1, %3 : tensor<?xbf16>,
+//   tensor<?xi64>) {
+// ^bb0(%in: bf16, %out: bf16, %out_0: i64):
+//   %5 = linalg.index 1 : index
+//   %6 = arith.index_cast %5 : index to i64
+//   %7 = arith.maximumf %in, %out : bf16
+//   %8 = arith.cmpf ogt, %in, %out : bf16
+//   %9 = arith.select %8, %6, %out_0 : i64
+//   linalg.yield %7, %9 : bf16, i64
+// } -> (tensor<?xbf16>, tensor<?xi64>)
+//
+// This function extracts the `arith.maximumf`, `arith.cmpf`, and `arith.select`
+// operations from the body to facilitate transformations such as split
+// reduction.
+static ArgmaxCombinerOps collectArgmaxCombinerOps(linalg::GenericOp genericOp) {
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  auto yieldOp = cast<linalg::YieldOp>(genericOp.getBody()->getTerminator());
+
+  // Extract max value producer: arith.maximumf.
+  Value maxResult = yieldOp.getOperand(0);
+  auto maxOp = cast<arith::MaximumFOp>(maxResult.getDefiningOp());
+
+  // Extract index result producer: arith.select.
+  Value indexResult = yieldOp.getOperand(1);
+  auto selectOp = cast<arith::SelectOp>(indexResult.getDefiningOp());
+
+  // Extract the condition of the select, expected to be arith.cmpf with
+  // predicate OGT.
+  auto cmpOp = cast<arith::CmpFOp>(selectOp.getCondition().getDefiningOp());
+
+  ArgmaxCombinerOps ops;
+  ops.maxOp = maxOp;
+  ops.selectOp = selectOp;
+  ops.cmpOp = cmpOp;
+  return ops;
+}
+
+static Value expandValue(OpBuilder &builder, Location loc, Value value,
+                         RankedTensorType expandedType) {
+  RankedTensorType originalType = cast<RankedTensorType>(value.getType());
+  if (originalType == expandedType)
+    return value;
+
+  auto reassociation =
+      mlir::getReassociationIndicesForReshape(originalType, expandedType);
+  assert(reassociation && "invalid reassociation for expand_shape");
+  return builder.create<tensor::ExpandShapeOp>(loc, expandedType, value,
+                                               *reassociation);
+}
+
+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) {
+      // 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));
+      }
+    } else {
+      unsigned shifted = (dim < insertSplitDim) ? dim : dim + 1;
+      exprs.push_back(getAffineDimExpr(shifted, ctx));
+    }
+  }
+
+  return AffineMap::get(oldMap.getNumDims() + 1, oldMap.getNumSymbols(), exprs,
+                        ctx);
+}
+
+static AffineMap getExpandedOutputIndexingMap(AffineMap oldMap,
path: compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/SplitReduction.cpp
user: login: hanhanW
body: style nit: please add a doc. This is not obvious to me.
created_at: 2025-05-14T03:48:27Z
author_association: CONTRIBUTOR
diff_hunk: @@ -389,4 +390,352 @@ splitReduction(RewriterBase &rewriter, LinalgExt::TopkOp topkOp,
   return success();
 }
 
+struct ArgmaxCombinerOps {
+  Operation *maxOp = nullptr;    // arith.maximumf
+  Operation *selectOp = nullptr; // arith.select
+  Operation *cmpOp = nullptr;    // arith.cmpf
+};
+
+// Matches the combiner pattern in a linalg.generic argmax-style reduction:
+// Example MLIR:
+// %4:2 = linalg.generic {
+//     indexing_maps = [...],
+//     iterator_types = ["parallel", "reduction"]
+//   } ins(%arg0 : tensor<?x128xbf16>) outs(%1, %3 : tensor<?xbf16>,
+//   tensor<?xi64>) {
+// ^bb0(%in: bf16, %out: bf16, %out_0: i64):
+//   %5 = linalg.index 1 : index
+//   %6 = arith.index_cast %5 : index to i64
+//   %7 = arith.maximumf %in, %out : bf16
+//   %8 = arith.cmpf ogt, %in, %out : bf16
+//   %9 = arith.select %8, %6, %out_0 : i64
+//   linalg.yield %7, %9 : bf16, i64
+// } -> (tensor<?xbf16>, tensor<?xi64>)
+//
+// This function extracts the `arith.maximumf`, `arith.cmpf`, and `arith.select`
+// operations from the body to facilitate transformations such as split
+// reduction.
+static ArgmaxCombinerOps collectArgmaxCombinerOps(linalg::GenericOp genericOp) {
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  auto yieldOp = cast<linalg::YieldOp>(genericOp.getBody()->getTerminator());
+
+  // Extract max value producer: arith.maximumf.
+  Value maxResult = yieldOp.getOperand(0);
+  auto maxOp = cast<arith::MaximumFOp>(maxResult.getDefiningOp());
+
+  // Extract index result producer: arith.select.
+  Value indexResult = yieldOp.getOperand(1);
+  auto selectOp = cast<arith::SelectOp>(indexResult.getDefiningOp());
+
+  // Extract the condition of the select, expected to be arith.cmpf with
+  // predicate OGT.
+  auto cmpOp = cast<arith::CmpFOp>(selectOp.getCondition().getDefiningOp());
+
+  ArgmaxCombinerOps ops;
+  ops.maxOp = maxOp;
+  ops.selectOp = selectOp;
+  ops.cmpOp = cmpOp;
+  return ops;
+}
+
+static Value expandValue(OpBuilder &builder, Location loc, Value value,
+                         RankedTensorType expandedType) {
+  RankedTensorType originalType = cast<RankedTensorType>(value.getType());
+  if (originalType == expandedType)
+    return value;
+
+  auto reassociation =
+      mlir::getReassociationIndicesForReshape(originalType, expandedType);
+  assert(reassociation && "invalid reassociation for expand_shape");
+  return builder.create<tensor::ExpandShapeOp>(loc, expandedType, value,
+                                               *reassociation);
+}
+
+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) {
+      // 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));
+      }
+    } else {
+      unsigned shifted = (dim < insertSplitDim) ? dim : dim + 1;
+      exprs.push_back(getAffineDimExpr(shifted, ctx));
+    }
+  }
+
+  return AffineMap::get(oldMap.getNumDims() + 1, oldMap.getNumSymbols(), exprs,
+                        ctx);
+}
+
+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);
+}
+
+FailureOr<linalg::SplitReductionResult>
+splitArgmaxReduction(RewriterBase &rewriter, linalg::GenericOp genericOp,
+                     linalg::ControlSplitReductionFn controlSplitReductionFn) {
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  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, 4> 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");
+  }
path: compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/SplitReduction.cpp
user: login: hanhanW
body: optional style nit: I'd remove the blank line in the between. They are reasonable to belong to the same chunk to me. (Mark it optional because maybe you have different thought.)
created_at: 2025-05-14T03:50:18Z
author_association: CONTRIBUTOR
diff_hunk: @@ -389,4 +390,352 @@ splitReduction(RewriterBase &rewriter, LinalgExt::TopkOp topkOp,
   return success();
 }
 
+struct ArgmaxCombinerOps {
+  Operation *maxOp = nullptr;    // arith.maximumf
+  Operation *selectOp = nullptr; // arith.select
+  Operation *cmpOp = nullptr;    // arith.cmpf
+};
+
+// Matches the combiner pattern in a linalg.generic argmax-style reduction:
+// Example MLIR:
+// %4:2 = linalg.generic {
+//     indexing_maps = [...],
+//     iterator_types = ["parallel", "reduction"]
+//   } ins(%arg0 : tensor<?x128xbf16>) outs(%1, %3 : tensor<?xbf16>,
+//   tensor<?xi64>) {
+// ^bb0(%in: bf16, %out: bf16, %out_0: i64):
+//   %5 = linalg.index 1 : index
+//   %6 = arith.index_cast %5 : index to i64
+//   %7 = arith.maximumf %in, %out : bf16
+//   %8 = arith.cmpf ogt, %in, %out : bf16
+//   %9 = arith.select %8, %6, %out_0 : i64
+//   linalg.yield %7, %9 : bf16, i64
+// } -> (tensor<?xbf16>, tensor<?xi64>)
+//
+// This function extracts the `arith.maximumf`, `arith.cmpf`, and `arith.select`
+// operations from the body to facilitate transformations such as split
+// reduction.
+static ArgmaxCombinerOps collectArgmaxCombinerOps(linalg::GenericOp genericOp) {
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  auto yieldOp = cast<linalg::YieldOp>(genericOp.getBody()->getTerminator());
+
+  // Extract max value producer: arith.maximumf.
+  Value maxResult = yieldOp.getOperand(0);
+  auto maxOp = cast<arith::MaximumFOp>(maxResult.getDefiningOp());
+
+  // Extract index result producer: arith.select.
+  Value indexResult = yieldOp.getOperand(1);
+  auto selectOp = cast<arith::SelectOp>(indexResult.getDefiningOp());
+
+  // Extract the condition of the select, expected to be arith.cmpf with
+  // predicate OGT.
+  auto cmpOp = cast<arith::CmpFOp>(selectOp.getCondition().getDefiningOp());
+
+  ArgmaxCombinerOps ops;
+  ops.maxOp = maxOp;
+  ops.selectOp = selectOp;
+  ops.cmpOp = cmpOp;
+  return ops;
+}
+
+static Value expandValue(OpBuilder &builder, Location loc, Value value,
+                         RankedTensorType expandedType) {
+  RankedTensorType originalType = cast<RankedTensorType>(value.getType());
+  if (originalType == expandedType)
+    return value;
+
+  auto reassociation =
+      mlir::getReassociationIndicesForReshape(originalType, expandedType);
+  assert(reassociation && "invalid reassociation for expand_shape");
+  return builder.create<tensor::ExpandShapeOp>(loc, expandedType, value,
+                                               *reassociation);
+}
+
+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) {
+      // 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));
+      }
+    } else {
+      unsigned shifted = (dim < insertSplitDim) ? dim : dim + 1;
+      exprs.push_back(getAffineDimExpr(shifted, ctx));
+    }
+  }
+
+  return AffineMap::get(oldMap.getNumDims() + 1, oldMap.getNumSymbols(), exprs,
+                        ctx);
+}
+
+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);
+}
+
+FailureOr<linalg::SplitReductionResult>
+splitArgmaxReduction(RewriterBase &rewriter, linalg::GenericOp genericOp,
+                     linalg::ControlSplitReductionFn controlSplitReductionFn) {
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  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, 4> 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");
+  }
+
+  ArgmaxCombinerOps combinerOps = collectArgmaxCombinerOps(genericOp);
+  Operation *reductionOp = combinerOps.maxOp;
+
+  std::optional<TypedAttr> identity = arith::getNeutralElement(reductionOp);
+  if (!identity.has_value())
+    return rewriter.notifyMatchFailure(
+        genericOp, "Unknown identity value for the reduction");
path: compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/SplitReduction.cpp
user: login: hanhanW
body: optional nit: Same here, I'd remove the blank line.

nit 2: Different from LLVM, IREE always wraps single statement with braces. IREE follows mixed LLVM/Google style,  and this is one of few difference from LLVM style.
created_at: 2025-05-14T03:52:00Z
author_association: CONTRIBUTOR
diff_hunk: @@ -389,4 +390,352 @@ splitReduction(RewriterBase &rewriter, LinalgExt::TopkOp topkOp,
   return success();
 }
 
+struct ArgmaxCombinerOps {
+  Operation *maxOp = nullptr;    // arith.maximumf
+  Operation *selectOp = nullptr; // arith.select
+  Operation *cmpOp = nullptr;    // arith.cmpf
+};
+
+// Matches the combiner pattern in a linalg.generic argmax-style reduction:
+// Example MLIR:
+// %4:2 = linalg.generic {
+//     indexing_maps = [...],
+//     iterator_types = ["parallel", "reduction"]
+//   } ins(%arg0 : tensor<?x128xbf16>) outs(%1, %3 : tensor<?xbf16>,
+//   tensor<?xi64>) {
+// ^bb0(%in: bf16, %out: bf16, %out_0: i64):
+//   %5 = linalg.index 1 : index
+//   %6 = arith.index_cast %5 : index to i64
+//   %7 = arith.maximumf %in, %out : bf16
+//   %8 = arith.cmpf ogt, %in, %out : bf16
+//   %9 = arith.select %8, %6, %out_0 : i64
+//   linalg.yield %7, %9 : bf16, i64
+// } -> (tensor<?xbf16>, tensor<?xi64>)
+//
+// This function extracts the `arith.maximumf`, `arith.cmpf`, and `arith.select`
+// operations from the body to facilitate transformations such as split
+// reduction.
+static ArgmaxCombinerOps collectArgmaxCombinerOps(linalg::GenericOp genericOp) {
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  auto yieldOp = cast<linalg::YieldOp>(genericOp.getBody()->getTerminator());
+
+  // Extract max value producer: arith.maximumf.
+  Value maxResult = yieldOp.getOperand(0);
+  auto maxOp = cast<arith::MaximumFOp>(maxResult.getDefiningOp());
+
+  // Extract index result producer: arith.select.
+  Value indexResult = yieldOp.getOperand(1);
+  auto selectOp = cast<arith::SelectOp>(indexResult.getDefiningOp());
+
+  // Extract the condition of the select, expected to be arith.cmpf with
+  // predicate OGT.
+  auto cmpOp = cast<arith::CmpFOp>(selectOp.getCondition().getDefiningOp());
+
+  ArgmaxCombinerOps ops;
+  ops.maxOp = maxOp;
+  ops.selectOp = selectOp;
+  ops.cmpOp = cmpOp;
+  return ops;
+}
+
+static Value expandValue(OpBuilder &builder, Location loc, Value value,
+                         RankedTensorType expandedType) {
+  RankedTensorType originalType = cast<RankedTensorType>(value.getType());
+  if (originalType == expandedType)
+    return value;
+
+  auto reassociation =
+      mlir::getReassociationIndicesForReshape(originalType, expandedType);
+  assert(reassociation && "invalid reassociation for expand_shape");
+  return builder.create<tensor::ExpandShapeOp>(loc, expandedType, value,
+                                               *reassociation);
+}
+
+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) {
+      // 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));
+      }
+    } else {
+      unsigned shifted = (dim < insertSplitDim) ? dim : dim + 1;
+      exprs.push_back(getAffineDimExpr(shifted, ctx));
+    }
+  }
+
+  return AffineMap::get(oldMap.getNumDims() + 1, oldMap.getNumSymbols(), exprs,
+                        ctx);
+}
+
+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);
+}
+
+FailureOr<linalg::SplitReductionResult>
+splitArgmaxReduction(RewriterBase &rewriter, linalg::GenericOp genericOp,
+                     linalg::ControlSplitReductionFn controlSplitReductionFn) {
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  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, 4> 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");
+  }
+
+  ArgmaxCombinerOps combinerOps = collectArgmaxCombinerOps(genericOp);
+  Operation *reductionOp = combinerOps.maxOp;
+
+  std::optional<TypedAttr> identity = arith::getNeutralElement(reductionOp);
+  if (!identity.has_value())
+    return rewriter.notifyMatchFailure(
+        genericOp, "Unknown identity value for the reduction");
+
+  SmallVector<Value> newInputs;
+  for (OpOperand *operand : genericOp.getDpsInputOperands()) {
+    AffineMap map = genericOp.getMatchingIndexingMap(operand);
+    ArrayRef<int64_t> oldShape = genericOp.getShape(operand);
+    Type elementType =
+        cast<RankedTensorType>(operand->get().getType()).getElementType();
path: compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/SplitReduction.cpp
user: login: hanhanW
body: optional nit: we can use `getElementTypeOrSelf`. Mark it optional because there are no convention. I personally like it better in this place because it could be one-line code.
created_at: 2025-05-14T03:53:31Z
author_association: CONTRIBUTOR
diff_hunk: @@ -389,4 +390,352 @@ splitReduction(RewriterBase &rewriter, LinalgExt::TopkOp topkOp,
   return success();
 }
 
+struct ArgmaxCombinerOps {
+  Operation *maxOp = nullptr;    // arith.maximumf
+  Operation *selectOp = nullptr; // arith.select
+  Operation *cmpOp = nullptr;    // arith.cmpf
+};
+
+// Matches the combiner pattern in a linalg.generic argmax-style reduction:
+// Example MLIR:
+// %4:2 = linalg.generic {
+//     indexing_maps = [...],
+//     iterator_types = ["parallel", "reduction"]
+//   } ins(%arg0 : tensor<?x128xbf16>) outs(%1, %3 : tensor<?xbf16>,
+//   tensor<?xi64>) {
+// ^bb0(%in: bf16, %out: bf16, %out_0: i64):
+//   %5 = linalg.index 1 : index
+//   %6 = arith.index_cast %5 : index to i64
+//   %7 = arith.maximumf %in, %out : bf16
+//   %8 = arith.cmpf ogt, %in, %out : bf16
+//   %9 = arith.select %8, %6, %out_0 : i64
+//   linalg.yield %7, %9 : bf16, i64
+// } -> (tensor<?xbf16>, tensor<?xi64>)
+//
+// This function extracts the `arith.maximumf`, `arith.cmpf`, and `arith.select`
+// operations from the body to facilitate transformations such as split
+// reduction.
+static ArgmaxCombinerOps collectArgmaxCombinerOps(linalg::GenericOp genericOp) {
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  auto yieldOp = cast<linalg::YieldOp>(genericOp.getBody()->getTerminator());
+
+  // Extract max value producer: arith.maximumf.
+  Value maxResult = yieldOp.getOperand(0);
+  auto maxOp = cast<arith::MaximumFOp>(maxResult.getDefiningOp());
+
+  // Extract index result producer: arith.select.
+  Value indexResult = yieldOp.getOperand(1);
+  auto selectOp = cast<arith::SelectOp>(indexResult.getDefiningOp());
+
+  // Extract the condition of the select, expected to be arith.cmpf with
+  // predicate OGT.
+  auto cmpOp = cast<arith::CmpFOp>(selectOp.getCondition().getDefiningOp());
+
+  ArgmaxCombinerOps ops;
+  ops.maxOp = maxOp;
+  ops.selectOp = selectOp;
+  ops.cmpOp = cmpOp;
+  return ops;
+}
+
+static Value expandValue(OpBuilder &builder, Location loc, Value value,
+                         RankedTensorType expandedType) {
+  RankedTensorType originalType = cast<RankedTensorType>(value.getType());
+  if (originalType == expandedType)
+    return value;
+
+  auto reassociation =
+      mlir::getReassociationIndicesForReshape(originalType, expandedType);
+  assert(reassociation && "invalid reassociation for expand_shape");
+  return builder.create<tensor::ExpandShapeOp>(loc, expandedType, value,
+                                               *reassociation);
+}
+
+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) {
+      // 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));
+      }
+    } else {
+      unsigned shifted = (dim < insertSplitDim) ? dim : dim + 1;
+      exprs.push_back(getAffineDimExpr(shifted, ctx));
+    }
+  }
+
+  return AffineMap::get(oldMap.getNumDims() + 1, oldMap.getNumSymbols(), exprs,
+                        ctx);
+}
+
+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);
+}
+
+FailureOr<linalg::SplitReductionResult>
+splitArgmaxReduction(RewriterBase &rewriter, linalg::GenericOp genericOp,
+                     linalg::ControlSplitReductionFn controlSplitReductionFn) {
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  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, 4> 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");
+  }
+
+  ArgmaxCombinerOps combinerOps = collectArgmaxCombinerOps(genericOp);
+  Operation *reductionOp = combinerOps.maxOp;
+
+  std::optional<TypedAttr> identity = arith::getNeutralElement(reductionOp);
+  if (!identity.has_value())
+    return rewriter.notifyMatchFailure(
+        genericOp, "Unknown identity value for the reduction");
+
+  SmallVector<Value> newInputs;
+  for (OpOperand *operand : genericOp.getDpsInputOperands()) {
+    AffineMap map = genericOp.getMatchingIndexingMap(operand);
+    ArrayRef<int64_t> oldShape = genericOp.getShape(operand);
+    Type elementType =
+        cast<RankedTensorType>(operand->get().getType()).getElementType();
+
+    SmallVector<int64_t> expandedShape;
+    for (unsigned idx : llvm::seq<unsigned>(0u, map.getNumResults())) {
+      unsigned dim = map.getDimPosition(idx);
+      if (dim == reductionDim) {
+        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.
+        }
+      } else {
+        expandedShape.push_back(oldShape[idx]);
+      }
path: compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/SplitReduction.cpp
user: login: hanhanW
body: ditto: prefer early-exit style and try to save levels of indents. (you'll use `if (dim != reductionDim)` if you apply the change.)
created_at: 2025-05-14T03:54:28Z
author_association: CONTRIBUTOR
diff_hunk: @@ -389,4 +390,352 @@ splitReduction(RewriterBase &rewriter, LinalgExt::TopkOp topkOp,
   return success();
 }
 
+struct ArgmaxCombinerOps {
+  Operation *maxOp = nullptr;    // arith.maximumf
+  Operation *selectOp = nullptr; // arith.select
+  Operation *cmpOp = nullptr;    // arith.cmpf
+};
+
+// Matches the combiner pattern in a linalg.generic argmax-style reduction:
+// Example MLIR:
+// %4:2 = linalg.generic {
+//     indexing_maps = [...],
+//     iterator_types = ["parallel", "reduction"]
+//   } ins(%arg0 : tensor<?x128xbf16>) outs(%1, %3 : tensor<?xbf16>,
+//   tensor<?xi64>) {
+// ^bb0(%in: bf16, %out: bf16, %out_0: i64):
+//   %5 = linalg.index 1 : index
+//   %6 = arith.index_cast %5 : index to i64
+//   %7 = arith.maximumf %in, %out : bf16
+//   %8 = arith.cmpf ogt, %in, %out : bf16
+//   %9 = arith.select %8, %6, %out_0 : i64
+//   linalg.yield %7, %9 : bf16, i64
+// } -> (tensor<?xbf16>, tensor<?xi64>)
+//
+// This function extracts the `arith.maximumf`, `arith.cmpf`, and `arith.select`
+// operations from the body to facilitate transformations such as split
+// reduction.
+static ArgmaxCombinerOps collectArgmaxCombinerOps(linalg::GenericOp genericOp) {
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  auto yieldOp = cast<linalg::YieldOp>(genericOp.getBody()->getTerminator());
+
+  // Extract max value producer: arith.maximumf.
+  Value maxResult = yieldOp.getOperand(0);
+  auto maxOp = cast<arith::MaximumFOp>(maxResult.getDefiningOp());
+
+  // Extract index result producer: arith.select.
+  Value indexResult = yieldOp.getOperand(1);
+  auto selectOp = cast<arith::SelectOp>(indexResult.getDefiningOp());
+
+  // Extract the condition of the select, expected to be arith.cmpf with
+  // predicate OGT.
+  auto cmpOp = cast<arith::CmpFOp>(selectOp.getCondition().getDefiningOp());
+
+  ArgmaxCombinerOps ops;
+  ops.maxOp = maxOp;
+  ops.selectOp = selectOp;
+  ops.cmpOp = cmpOp;
+  return ops;
+}
+
+static Value expandValue(OpBuilder &builder, Location loc, Value value,
+                         RankedTensorType expandedType) {
+  RankedTensorType originalType = cast<RankedTensorType>(value.getType());
+  if (originalType == expandedType)
+    return value;
+
+  auto reassociation =
+      mlir::getReassociationIndicesForReshape(originalType, expandedType);
+  assert(reassociation && "invalid reassociation for expand_shape");
+  return builder.create<tensor::ExpandShapeOp>(loc, expandedType, value,
+                                               *reassociation);
+}
+
+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) {
+      // 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));
+      }
+    } else {
+      unsigned shifted = (dim < insertSplitDim) ? dim : dim + 1;
+      exprs.push_back(getAffineDimExpr(shifted, ctx));
+    }
+  }
+
+  return AffineMap::get(oldMap.getNumDims() + 1, oldMap.getNumSymbols(), exprs,
+                        ctx);
+}
+
+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);
+}
+
+FailureOr<linalg::SplitReductionResult>
+splitArgmaxReduction(RewriterBase &rewriter, linalg::GenericOp genericOp,
+                     linalg::ControlSplitReductionFn controlSplitReductionFn) {
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  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, 4> 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");
+  }
+
+  ArgmaxCombinerOps combinerOps = collectArgmaxCombinerOps(genericOp);
+  Operation *reductionOp = combinerOps.maxOp;
+
+  std::optional<TypedAttr> identity = arith::getNeutralElement(reductionOp);
+  if (!identity.has_value())
+    return rewriter.notifyMatchFailure(
+        genericOp, "Unknown identity value for the reduction");
+
+  SmallVector<Value> newInputs;
+  for (OpOperand *operand : genericOp.getDpsInputOperands()) {
+    AffineMap map = genericOp.getMatchingIndexingMap(operand);
+    ArrayRef<int64_t> oldShape = genericOp.getShape(operand);
+    Type elementType =
+        cast<RankedTensorType>(operand->get().getType()).getElementType();
+
+    SmallVector<int64_t> expandedShape;
+    for (unsigned idx : llvm::seq<unsigned>(0u, map.getNumResults())) {
+      unsigned dim = map.getDimPosition(idx);
+      if (dim == reductionDim) {
+        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.
+        }
+      } else {
+        expandedShape.push_back(oldShape[idx]);
+      }
+    }
+
+    RankedTensorType expandedType =
+        RankedTensorType::get(expandedShape, elementType);
path: compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/SplitReduction.cpp
user: login: hanhanW
body: style nit: use auto because `XXX::get` already spells the type.

https://llvm.org/docs/CodingStandards.html#use-auto-type-deduction-to-make-code-more-readable
created_at: 2025-05-14T03:55:01Z
author_association: CONTRIBUTOR
diff_hunk: @@ -389,4 +390,352 @@ splitReduction(RewriterBase &rewriter, LinalgExt::TopkOp topkOp,
   return success();
 }
 
+struct ArgmaxCombinerOps {
+  Operation *maxOp = nullptr;    // arith.maximumf
+  Operation *selectOp = nullptr; // arith.select
+  Operation *cmpOp = nullptr;    // arith.cmpf
+};
+
+// Matches the combiner pattern in a linalg.generic argmax-style reduction:
+// Example MLIR:
+// %4:2 = linalg.generic {
+//     indexing_maps = [...],
+//     iterator_types = ["parallel", "reduction"]
+//   } ins(%arg0 : tensor<?x128xbf16>) outs(%1, %3 : tensor<?xbf16>,
+//   tensor<?xi64>) {
+// ^bb0(%in: bf16, %out: bf16, %out_0: i64):
+//   %5 = linalg.index 1 : index
+//   %6 = arith.index_cast %5 : index to i64
+//   %7 = arith.maximumf %in, %out : bf16
+//   %8 = arith.cmpf ogt, %in, %out : bf16
+//   %9 = arith.select %8, %6, %out_0 : i64
+//   linalg.yield %7, %9 : bf16, i64
+// } -> (tensor<?xbf16>, tensor<?xi64>)
+//
+// This function extracts the `arith.maximumf`, `arith.cmpf`, and `arith.select`
+// operations from the body to facilitate transformations such as split
+// reduction.
+static ArgmaxCombinerOps collectArgmaxCombinerOps(linalg::GenericOp genericOp) {
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  auto yieldOp = cast<linalg::YieldOp>(genericOp.getBody()->getTerminator());
+
+  // Extract max value producer: arith.maximumf.
+  Value maxResult = yieldOp.getOperand(0);
+  auto maxOp = cast<arith::MaximumFOp>(maxResult.getDefiningOp());
+
+  // Extract index result producer: arith.select.
+  Value indexResult = yieldOp.getOperand(1);
+  auto selectOp = cast<arith::SelectOp>(indexResult.getDefiningOp());
+
+  // Extract the condition of the select, expected to be arith.cmpf with
+  // predicate OGT.
+  auto cmpOp = cast<arith::CmpFOp>(selectOp.getCondition().getDefiningOp());
+
+  ArgmaxCombinerOps ops;
+  ops.maxOp = maxOp;
+  ops.selectOp = selectOp;
+  ops.cmpOp = cmpOp;
+  return ops;
+}
+
+static Value expandValue(OpBuilder &builder, Location loc, Value value,
+                         RankedTensorType expandedType) {
+  RankedTensorType originalType = cast<RankedTensorType>(value.getType());
+  if (originalType == expandedType)
+    return value;
+
+  auto reassociation =
+      mlir::getReassociationIndicesForReshape(originalType, expandedType);
+  assert(reassociation && "invalid reassociation for expand_shape");
+  return builder.create<tensor::ExpandShapeOp>(loc, expandedType, value,
+                                               *reassociation);
+}
+
+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) {
+      // 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));
+      }
+    } else {
+      unsigned shifted = (dim < insertSplitDim) ? dim : dim + 1;
+      exprs.push_back(getAffineDimExpr(shifted, ctx));
+    }
+  }
+
+  return AffineMap::get(oldMap.getNumDims() + 1, oldMap.getNumSymbols(), exprs,
+                        ctx);
+}
+
+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);
+}
+
+FailureOr<linalg::SplitReductionResult>
+splitArgmaxReduction(RewriterBase &rewriter, linalg::GenericOp genericOp,
+                     linalg::ControlSplitReductionFn controlSplitReductionFn) {
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  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, 4> 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");
+  }
+
+  ArgmaxCombinerOps combinerOps = collectArgmaxCombinerOps(genericOp);
+  Operation *reductionOp = combinerOps.maxOp;
+
+  std::optional<TypedAttr> identity = arith::getNeutralElement(reductionOp);
+  if (!identity.has_value())
+    return rewriter.notifyMatchFailure(
+        genericOp, "Unknown identity value for the reduction");
+
+  SmallVector<Value> newInputs;
+  for (OpOperand *operand : genericOp.getDpsInputOperands()) {
+    AffineMap map = genericOp.getMatchingIndexingMap(operand);
+    ArrayRef<int64_t> oldShape = genericOp.getShape(operand);
+    Type elementType =
+        cast<RankedTensorType>(operand->get().getType()).getElementType();
+
+    SmallVector<int64_t> expandedShape;
+    for (unsigned idx : llvm::seq<unsigned>(0u, map.getNumResults())) {
+      unsigned dim = map.getDimPosition(idx);
+      if (dim == reductionDim) {
+        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.
+        }
+      } else {
+        expandedShape.push_back(oldShape[idx]);
+      }
+    }
+
+    RankedTensorType 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;
path: compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/SplitReduction.cpp
user: login: hanhanW
body: style nit: I think you accidentally add one more leading whitespace in the comment. :)
nit 2: can we rename it to `outputDimSize` for consistency?

```suggestion
  // The total number of output elements along this new dimension is
  // reductionDimSize / ratio.
  int64_t outputDimSize = reductionDimSize / ratio;
```
created_at: 2025-05-14T03:56:49Z
author_association: CONTRIBUTOR
diff_hunk: @@ -389,4 +390,352 @@ splitReduction(RewriterBase &rewriter, LinalgExt::TopkOp topkOp,
   return success();
 }
 
+struct ArgmaxCombinerOps {
+  Operation *maxOp = nullptr;    // arith.maximumf
+  Operation *selectOp = nullptr; // arith.select
+  Operation *cmpOp = nullptr;    // arith.cmpf
+};
+
+// Matches the combiner pattern in a linalg.generic argmax-style reduction:
+// Example MLIR:
+// %4:2 = linalg.generic {
+//     indexing_maps = [...],
+//     iterator_types = ["parallel", "reduction"]
+//   } ins(%arg0 : tensor<?x128xbf16>) outs(%1, %3 : tensor<?xbf16>,
+//   tensor<?xi64>) {
+// ^bb0(%in: bf16, %out: bf16, %out_0: i64):
+//   %5 = linalg.index 1 : index
+//   %6 = arith.index_cast %5 : index to i64
+//   %7 = arith.maximumf %in, %out : bf16
+//   %8 = arith.cmpf ogt, %in, %out : bf16
+//   %9 = arith.select %8, %6, %out_0 : i64
+//   linalg.yield %7, %9 : bf16, i64
+// } -> (tensor<?xbf16>, tensor<?xi64>)
+//
+// This function extracts the `arith.maximumf`, `arith.cmpf`, and `arith.select`
+// operations from the body to facilitate transformations such as split
+// reduction.
+static ArgmaxCombinerOps collectArgmaxCombinerOps(linalg::GenericOp genericOp) {
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  auto yieldOp = cast<linalg::YieldOp>(genericOp.getBody()->getTerminator());
+
+  // Extract max value producer: arith.maximumf.
+  Value maxResult = yieldOp.getOperand(0);
+  auto maxOp = cast<arith::MaximumFOp>(maxResult.getDefiningOp());
+
+  // Extract index result producer: arith.select.
+  Value indexResult = yieldOp.getOperand(1);
+  auto selectOp = cast<arith::SelectOp>(indexResult.getDefiningOp());
+
+  // Extract the condition of the select, expected to be arith.cmpf with
+  // predicate OGT.
+  auto cmpOp = cast<arith::CmpFOp>(selectOp.getCondition().getDefiningOp());
+
+  ArgmaxCombinerOps ops;
+  ops.maxOp = maxOp;
+  ops.selectOp = selectOp;
+  ops.cmpOp = cmpOp;
+  return ops;
+}
+
+static Value expandValue(OpBuilder &builder, Location loc, Value value,
+                         RankedTensorType expandedType) {
+  RankedTensorType originalType = cast<RankedTensorType>(value.getType());
+  if (originalType == expandedType)
+    return value;
+
+  auto reassociation =
+      mlir::getReassociationIndicesForReshape(originalType, expandedType);
+  assert(reassociation && "invalid reassociation for expand_shape");
+  return builder.create<tensor::ExpandShapeOp>(loc, expandedType, value,
+                                               *reassociation);
+}
+
+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) {
+      // 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));
+      }
+    } else {
+      unsigned shifted = (dim < insertSplitDim) ? dim : dim + 1;
+      exprs.push_back(getAffineDimExpr(shifted, ctx));
+    }
+  }
+
+  return AffineMap::get(oldMap.getNumDims() + 1, oldMap.getNumSymbols(), exprs,
+                        ctx);
+}
+
+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);
+}
+
+FailureOr<linalg::SplitReductionResult>
+splitArgmaxReduction(RewriterBase &rewriter, linalg::GenericOp genericOp,
+                     linalg::ControlSplitReductionFn controlSplitReductionFn) {
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  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, 4> 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");
+  }
+
+  ArgmaxCombinerOps combinerOps = collectArgmaxCombinerOps(genericOp);
+  Operation *reductionOp = combinerOps.maxOp;
+
+  std::optional<TypedAttr> identity = arith::getNeutralElement(reductionOp);
+  if (!identity.has_value())
+    return rewriter.notifyMatchFailure(
+        genericOp, "Unknown identity value for the reduction");
+
+  SmallVector<Value> newInputs;
+  for (OpOperand *operand : genericOp.getDpsInputOperands()) {
+    AffineMap map = genericOp.getMatchingIndexingMap(operand);
+    ArrayRef<int64_t> oldShape = genericOp.getShape(operand);
+    Type elementType =
+        cast<RankedTensorType>(operand->get().getType()).getElementType();
+
+    SmallVector<int64_t> expandedShape;
+    for (unsigned idx : llvm::seq<unsigned>(0u, map.getNumResults())) {
+      unsigned dim = map.getDimPosition(idx);
+      if (dim == reductionDim) {
+        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.
+        }
+      } else {
+        expandedShape.push_back(oldShape[idx]);
+      }
+    }
+
+    RankedTensorType 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.
+  Type valueElemType = genericOp.getRegionOutputArgs()[0].getType();
+  ArrayRef<int64_t> origValueShape = genericOp.getShape(valueInit);
+  SmallVector<int64_t> expandedValueShape;
+  for (unsigned i = 0; i <= origValueShape.size(); ++i) {
+    if (i == insertSplitIndex)
+      expandedValueShape.push_back(outputDimsize);
+    if (i < origValueShape.size())
+      expandedValueShape.push_back(origValueShape[i]);
+  }
+  SmallVector<Value> dynValueDims;
+  for (size_t i = 0; i < expandedValueShape.size(); ++i) {
+    if (ShapedType::isDynamic(expandedValueShape[i])) {
+      unsigned sourceIdx = (i < insertSplitIndex) ? i : i - 1;
+      dynValueDims.push_back(
+          rewriter.create<tensor::DimOp>(loc, valueInit->get(), sourceIdx));
+    }
+  }
+  Value valueConst = rewriter.create<arith::ConstantOp>(loc, *identity);
+  Value emptyValue = rewriter.create<tensor::EmptyOp>(
+      loc, expandedValueShape, valueElemType, dynValueDims);
+  Value identityValue =
+      rewriter.create<linalg::FillOp>(loc, valueConst, emptyValue).getResult(0);
+
+  // Index identity.
+  Type indexElemType = genericOp.getRegionOutputArgs()[1].getType();
+  ArrayRef<int64_t> origIndexShape = genericOp.getShape(indexInit);
+  SmallVector<int64_t> expandedIndexShape;
+  for (unsigned i = 0; i <= origIndexShape.size(); ++i) {
+    if (i == insertSplitIndex)
+      expandedIndexShape.push_back(outputDimsize);
+    if (i < origIndexShape.size())
+      expandedIndexShape.push_back(origIndexShape[i]);
+  }
+  SmallVector<Value> dynIdxDims;
+  for (unsigned i = 0; i < expandedIndexShape.size(); ++i) {
+    if (ShapedType::isDynamic(expandedIndexShape[i])) {
+      unsigned sourceIdx = (i < insertSplitIndex) ? i : i - 1;
+      dynIdxDims.push_back(
+          rewriter.create<tensor::DimOp>(loc, indexInit->get(), sourceIdx));
+    }
+  }
+  Value zeroIdx = rewriter.create<arith::ConstantOp>(
+      loc, rewriter.getZeroAttr(indexElemType));
+  Value emptyIdx = rewriter.create<tensor::EmptyOp>(loc, expandedIndexShape,
+                                                    indexElemType, dynIdxDims);
+  Value identityIndex =
+      rewriter.create<linalg::FillOp>(loc, zeroIdx, emptyIdx).getResult(0);
+
+  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()));
+    }
+  }
+
+  // Create partial linalg.generic op with global index computation.
+  Value tileSize = rewriter.create<arith::ConstantIndexOp>(loc, ratio);
+  auto partialOp = rewriter.create<linalg::GenericOp>(
+      loc, TypeRange{identityValue.getType(), identityIndex.getType()},
+      newInputs, ValueRange{identityValue, identityIndex}, newMaps,
+      newIteratorTypes);
+
+  rewriter.inlineRegionBefore(genericOp.getRegion(), partialOp.getRegion(),
+                              partialOp.getRegion().begin());
+
+  Block &body = partialOp.getRegion().front();
+  rewriter.setInsertionPointToStart(&body);
+
+  unsigned innerIdxDim = reductionDim + 1;
+  unsigned outerIdxDim = insertSplitDimension;
+
+  // Compute global index (gidx) for reduction when the original reduction
+  // dimension is split into [outerIdx, innerIdx] using `ratio`. This is used to
+  // correctly compute the global index for comparisons and index selection.
+  Value outerIdx = rewriter.create<linalg::IndexOp>(loc, outerIdxDim);
+  Value innerIdx = rewriter.create<linalg::IndexOp>(loc, innerIdxDim);
+  Value offset = rewriter.create<arith::MulIOp>(loc, outerIdx, tileSize);
+  Value gidx = rewriter.create<arith::AddIOp>(loc, offset, innerIdx);
+
+  auto selectOp = dyn_cast<arith::SelectOp>(combinerOps.selectOp);
+  Value oldIdx = selectOp.getTrueValue();
+  Value newIdx = gidx;
+  if (oldIdx.getType() != gidx.getType()) {
+    newIdx = rewriter.create<arith::IndexCastOp>(loc, oldIdx.getType(), gidx);
+  }
+  selectOp.setOperand(1, newIdx);
+  rewriter.setInsertionPointAfter(partialOp);
+
+  unsigned intermRank = expandedValueShape.size();
+  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 (unsigned i : llvm::seq<unsigned>(0, intermRank)) {
+    if (i != insertSplitIndex)
+      resultExprs.push_back(rewriter.getAffineDimExpr(i));
+  }
path: compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/SplitReduction.cpp
user: login: hanhanW
body: style nit: use auto because `llvm::seq<>` already spells the type.
style nit 2: wrap simple single statement with braces.

```suggestion
  for (auto i : llvm::seq<unsigned>(0, intermRank)) {
    if (i != insertSplitIndex) {
      resultExprs.push_back(rewriter.getAffineDimExpr(i));
    }
  }
```
created_at: 2025-05-14T03:57:59Z
author_association: CONTRIBUTOR
diff_hunk: @@ -389,4 +390,332 @@ splitReduction(RewriterBase &rewriter, LinalgExt::TopkOp topkOp,
   return success();
 }
 
+struct ArgmaxCombinerOps {
+  Operation *maxOp = nullptr;    // arith.maximumf
+  Operation *selectOp = nullptr; // arith.select
+  Operation *cmpOp = nullptr;    // arith.cmpf
+};
+
+// Matches the combiner pattern in a linalg.generic argmax-style reduction:
+// Example MLIR:
+// %4:2 = linalg.generic {
+//     indexing_maps = [...],
+//     iterator_types = ["parallel", "reduction"]
+//   } ins(%arg0 : tensor<?x128xbf16>) outs(%1, %3 : tensor<?xbf16>,
+//   tensor<?xi64>) {
+// ^bb0(%in: bf16, %out: bf16, %out_0: i64):
+//   %5 = linalg.index 1 : index
+//   %6 = arith.index_cast %5 : index to i64
+//   %7 = arith.maximumf %in, %out : bf16
+//   %8 = arith.cmpf ogt, %in, %out : bf16
+//   %9 = arith.select %8, %6, %out_0 : i64
+//   linalg.yield %7, %9 : bf16, i64
+// } -> (tensor<?xbf16>, tensor<?xi64>)
+//
+// This function extracts the `arith.maximumf`, `arith.cmpf`, and `arith.select`
+// operations from the body to facilitate transformations such as split
+// reduction.
+static ArgmaxCombinerOps collectArgmaxCombinerOps(linalg::GenericOp genericOp) {
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  auto yieldOp = cast<linalg::YieldOp>(genericOp.getBody()->getTerminator());
+
+  // Extract max value producer: arith.maximumf.
+  Value maxResult = yieldOp.getOperand(0);
+  auto maxOp = cast<arith::MaximumFOp>(maxResult.getDefiningOp());
+
+  // Extract index result producer: arith.select.
+  Value indexResult = yieldOp.getOperand(1);
+  auto selectOp = cast<arith::SelectOp>(indexResult.getDefiningOp());
+
+  // Extract the condition of the select, expected to be arith.cmpf with
+  // predicate OGT.
+  auto cmpOp = cast<arith::CmpFOp>(selectOp.getCondition().getDefiningOp());
+
+  ArgmaxCombinerOps ops;
+  ops.maxOp = maxOp;
+  ops.selectOp = selectOp;
+  ops.cmpOp = cmpOp;
+  return ops;
+}
+
+static Value expandValue(OpBuilder &builder, Location loc, Value value,
+                         RankedTensorType expandedType) {
+  RankedTensorType originalType = cast<RankedTensorType>(value.getType());
+  if (originalType == expandedType)
+    return value;
+
+  auto reassociation =
+      mlir::getReassociationIndicesForReshape(originalType, expandedType);
+  assert(reassociation && "failed to infer reassociation indices from types");
+  return builder.create<tensor::ExpandShapeOp>(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 = builder.create<arith::ConstantOp>(
+      loc, elemType, llvm::cast<TypedAttr>(identityAttr));
+  Value empty = builder.create<tensor::EmptyOp>(loc, shape, elemType);
+  return builder.create<linalg::FillOp>(loc, identityVal, empty).getResult(0);
+}
+
+FailureOr<linalg::SplitReductionResult>
+splitArgmaxReduction(RewriterBase &rewriter, linalg::GenericOp genericOp,
+                     linalg::ControlSplitReductionFn controlSplitReductionFn) {
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  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, 4> 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");
+  }
+
+  ArgmaxCombinerOps combinerOps = collectArgmaxCombinerOps(genericOp);
+  Operation *reductionOp = combinerOps.maxOp;
+  std::optional<TypedAttr> identity = arith::getNeutralElement(reductionOp);
+  if (!identity.has_value())
+    return rewriter.notifyMatchFailure(
+        genericOp, "Unknown identity value for the reduction");
path: compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/SplitReduction.cpp
user: login: hanhanW
body: Wrap with braces.
created_at: 2025-05-14T23:41:05Z
author_association: CONTRIBUTOR
diff_hunk: @@ -389,4 +390,332 @@ splitReduction(RewriterBase &rewriter, LinalgExt::TopkOp topkOp,
   return success();
 }
 
+struct ArgmaxCombinerOps {
+  Operation *maxOp = nullptr;    // arith.maximumf
+  Operation *selectOp = nullptr; // arith.select
+  Operation *cmpOp = nullptr;    // arith.cmpf
+};
+
+// Matches the combiner pattern in a linalg.generic argmax-style reduction:
+// Example MLIR:
+// %4:2 = linalg.generic {
+//     indexing_maps = [...],
+//     iterator_types = ["parallel", "reduction"]
+//   } ins(%arg0 : tensor<?x128xbf16>) outs(%1, %3 : tensor<?xbf16>,
+//   tensor<?xi64>) {
+// ^bb0(%in: bf16, %out: bf16, %out_0: i64):
+//   %5 = linalg.index 1 : index
+//   %6 = arith.index_cast %5 : index to i64
+//   %7 = arith.maximumf %in, %out : bf16
+//   %8 = arith.cmpf ogt, %in, %out : bf16
+//   %9 = arith.select %8, %6, %out_0 : i64
+//   linalg.yield %7, %9 : bf16, i64
+// } -> (tensor<?xbf16>, tensor<?xi64>)
+//
+// This function extracts the `arith.maximumf`, `arith.cmpf`, and `arith.select`
+// operations from the body to facilitate transformations such as split
+// reduction.
+static ArgmaxCombinerOps collectArgmaxCombinerOps(linalg::GenericOp genericOp) {
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  auto yieldOp = cast<linalg::YieldOp>(genericOp.getBody()->getTerminator());
+
+  // Extract max value producer: arith.maximumf.
+  Value maxResult = yieldOp.getOperand(0);
+  auto maxOp = cast<arith::MaximumFOp>(maxResult.getDefiningOp());
+
+  // Extract index result producer: arith.select.
+  Value indexResult = yieldOp.getOperand(1);
+  auto selectOp = cast<arith::SelectOp>(indexResult.getDefiningOp());
+
+  // Extract the condition of the select, expected to be arith.cmpf with
+  // predicate OGT.
+  auto cmpOp = cast<arith::CmpFOp>(selectOp.getCondition().getDefiningOp());
+
+  ArgmaxCombinerOps ops;
+  ops.maxOp = maxOp;
+  ops.selectOp = selectOp;
+  ops.cmpOp = cmpOp;
+  return ops;
+}
+
+static Value expandValue(OpBuilder &builder, Location loc, Value value,
+                         RankedTensorType expandedType) {
+  RankedTensorType originalType = cast<RankedTensorType>(value.getType());
+  if (originalType == expandedType)
+    return value;
path: compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/SplitReduction.cpp
user: login: hanhanW
body: style: wrap with braces.
created_at: 2025-05-14T23:41:17Z
author_association: CONTRIBUTOR
diff_hunk: @@ -389,4 +390,332 @@ splitReduction(RewriterBase &rewriter, LinalgExt::TopkOp topkOp,
   return success();
 }
 
+struct ArgmaxCombinerOps {
+  Operation *maxOp = nullptr;    // arith.maximumf
+  Operation *selectOp = nullptr; // arith.select
+  Operation *cmpOp = nullptr;    // arith.cmpf
+};
+
+// Matches the combiner pattern in a linalg.generic argmax-style reduction:
+// Example MLIR:
+// %4:2 = linalg.generic {
+//     indexing_maps = [...],
+//     iterator_types = ["parallel", "reduction"]
+//   } ins(%arg0 : tensor<?x128xbf16>) outs(%1, %3 : tensor<?xbf16>,
+//   tensor<?xi64>) {
+// ^bb0(%in: bf16, %out: bf16, %out_0: i64):
+//   %5 = linalg.index 1 : index
+//   %6 = arith.index_cast %5 : index to i64
+//   %7 = arith.maximumf %in, %out : bf16
+//   %8 = arith.cmpf ogt, %in, %out : bf16
+//   %9 = arith.select %8, %6, %out_0 : i64
+//   linalg.yield %7, %9 : bf16, i64
+// } -> (tensor<?xbf16>, tensor<?xi64>)
+//
+// This function extracts the `arith.maximumf`, `arith.cmpf`, and `arith.select`
+// operations from the body to facilitate transformations such as split
+// reduction.
+static ArgmaxCombinerOps collectArgmaxCombinerOps(linalg::GenericOp genericOp) {
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  auto yieldOp = cast<linalg::YieldOp>(genericOp.getBody()->getTerminator());
+
+  // Extract max value producer: arith.maximumf.
+  Value maxResult = yieldOp.getOperand(0);
+  auto maxOp = cast<arith::MaximumFOp>(maxResult.getDefiningOp());
+
+  // Extract index result producer: arith.select.
+  Value indexResult = yieldOp.getOperand(1);
+  auto selectOp = cast<arith::SelectOp>(indexResult.getDefiningOp());
+
+  // Extract the condition of the select, expected to be arith.cmpf with
+  // predicate OGT.
+  auto cmpOp = cast<arith::CmpFOp>(selectOp.getCondition().getDefiningOp());
+
+  ArgmaxCombinerOps ops;
+  ops.maxOp = maxOp;
+  ops.selectOp = selectOp;
+  ops.cmpOp = cmpOp;
+  return ops;
+}
+
+static Value expandValue(OpBuilder &builder, Location loc, Value value,
+                         RankedTensorType expandedType) {
+  RankedTensorType originalType = cast<RankedTensorType>(value.getType());
+  if (originalType == expandedType)
+    return value;
+
+  auto reassociation =
+      mlir::getReassociationIndicesForReshape(originalType, expandedType);
path: compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/SplitReduction.cpp
user: login: hanhanW
body: Sorry, I missed this in the previous reviews. I think we usually drop `mlir::` prefix in transformation implementation. This is why we have `iree_compiler` under `mlir` namespace.
created_at: 2025-05-14T23:42:17Z
author_association: CONTRIBUTOR
diff_hunk: @@ -389,4 +390,332 @@ splitReduction(RewriterBase &rewriter, LinalgExt::TopkOp topkOp,
   return success();
 }
 
+struct ArgmaxCombinerOps {
+  Operation *maxOp = nullptr;    // arith.maximumf
+  Operation *selectOp = nullptr; // arith.select
+  Operation *cmpOp = nullptr;    // arith.cmpf
+};
+
+// Matches the combiner pattern in a linalg.generic argmax-style reduction:
+// Example MLIR:
+// %4:2 = linalg.generic {
+//     indexing_maps = [...],
+//     iterator_types = ["parallel", "reduction"]
+//   } ins(%arg0 : tensor<?x128xbf16>) outs(%1, %3 : tensor<?xbf16>,
+//   tensor<?xi64>) {
+// ^bb0(%in: bf16, %out: bf16, %out_0: i64):
+//   %5 = linalg.index 1 : index
+//   %6 = arith.index_cast %5 : index to i64
+//   %7 = arith.maximumf %in, %out : bf16
+//   %8 = arith.cmpf ogt, %in, %out : bf16
+//   %9 = arith.select %8, %6, %out_0 : i64
+//   linalg.yield %7, %9 : bf16, i64
+// } -> (tensor<?xbf16>, tensor<?xi64>)
+//
+// This function extracts the `arith.maximumf`, `arith.cmpf`, and `arith.select`
+// operations from the body to facilitate transformations such as split
+// reduction.
+static ArgmaxCombinerOps collectArgmaxCombinerOps(linalg::GenericOp genericOp) {
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  auto yieldOp = cast<linalg::YieldOp>(genericOp.getBody()->getTerminator());
+
+  // Extract max value producer: arith.maximumf.
+  Value maxResult = yieldOp.getOperand(0);
+  auto maxOp = cast<arith::MaximumFOp>(maxResult.getDefiningOp());
+
+  // Extract index result producer: arith.select.
+  Value indexResult = yieldOp.getOperand(1);
+  auto selectOp = cast<arith::SelectOp>(indexResult.getDefiningOp());
+
+  // Extract the condition of the select, expected to be arith.cmpf with
+  // predicate OGT.
+  auto cmpOp = cast<arith::CmpFOp>(selectOp.getCondition().getDefiningOp());
+
+  ArgmaxCombinerOps ops;
+  ops.maxOp = maxOp;
+  ops.selectOp = selectOp;
+  ops.cmpOp = cmpOp;
+  return ops;
+}
+
+static Value expandValue(OpBuilder &builder, Location loc, Value value,
+                         RankedTensorType expandedType) {
+  RankedTensorType originalType = cast<RankedTensorType>(value.getType());
+  if (originalType == expandedType)
+    return value;
+
+  auto reassociation =
+      mlir::getReassociationIndicesForReshape(originalType, expandedType);
+  assert(reassociation && "failed to infer reassociation indices from types");
+  return builder.create<tensor::ExpandShapeOp>(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 = builder.create<arith::ConstantOp>(
+      loc, elemType, llvm::cast<TypedAttr>(identityAttr));
path: compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/SplitReduction.cpp
user: login: hanhanW
body: nit: drop `llvm::` prefix for the casting.
created_at: 2025-05-14T23:42:40Z
author_association: CONTRIBUTOR
diff_hunk: @@ -389,4 +390,332 @@ splitReduction(RewriterBase &rewriter, LinalgExt::TopkOp topkOp,
   return success();
 }
 
+struct ArgmaxCombinerOps {
+  Operation *maxOp = nullptr;    // arith.maximumf
+  Operation *selectOp = nullptr; // arith.select
+  Operation *cmpOp = nullptr;    // arith.cmpf
+};
+
+// Matches the combiner pattern in a linalg.generic argmax-style reduction:
+// Example MLIR:
+// %4:2 = linalg.generic {
+//     indexing_maps = [...],
+//     iterator_types = ["parallel", "reduction"]
+//   } ins(%arg0 : tensor<?x128xbf16>) outs(%1, %3 : tensor<?xbf16>,
+//   tensor<?xi64>) {
+// ^bb0(%in: bf16, %out: bf16, %out_0: i64):
+//   %5 = linalg.index 1 : index
+//   %6 = arith.index_cast %5 : index to i64
+//   %7 = arith.maximumf %in, %out : bf16
+//   %8 = arith.cmpf ogt, %in, %out : bf16
+//   %9 = arith.select %8, %6, %out_0 : i64
+//   linalg.yield %7, %9 : bf16, i64
+// } -> (tensor<?xbf16>, tensor<?xi64>)
+//
+// This function extracts the `arith.maximumf`, `arith.cmpf`, and `arith.select`
+// operations from the body to facilitate transformations such as split
+// reduction.
+static ArgmaxCombinerOps collectArgmaxCombinerOps(linalg::GenericOp genericOp) {
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  auto yieldOp = cast<linalg::YieldOp>(genericOp.getBody()->getTerminator());
+
+  // Extract max value producer: arith.maximumf.
+  Value maxResult = yieldOp.getOperand(0);
+  auto maxOp = cast<arith::MaximumFOp>(maxResult.getDefiningOp());
+
+  // Extract index result producer: arith.select.
+  Value indexResult = yieldOp.getOperand(1);
+  auto selectOp = cast<arith::SelectOp>(indexResult.getDefiningOp());
+
+  // Extract the condition of the select, expected to be arith.cmpf with
+  // predicate OGT.
+  auto cmpOp = cast<arith::CmpFOp>(selectOp.getCondition().getDefiningOp());
+
+  ArgmaxCombinerOps ops;
+  ops.maxOp = maxOp;
+  ops.selectOp = selectOp;
+  ops.cmpOp = cmpOp;
+  return ops;
+}
+
+static Value expandValue(OpBuilder &builder, Location loc, Value value,
+                         RankedTensorType expandedType) {
+  RankedTensorType originalType = cast<RankedTensorType>(value.getType());
+  if (originalType == expandedType)
+    return value;
+
+  auto reassociation =
+      mlir::getReassociationIndicesForReshape(originalType, expandedType);
+  assert(reassociation && "failed to infer reassociation indices from types");
+  return builder.create<tensor::ExpandShapeOp>(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 = builder.create<arith::ConstantOp>(
+      loc, elemType, llvm::cast<TypedAttr>(identityAttr));
+  Value empty = builder.create<tensor::EmptyOp>(loc, shape, elemType);
+  return builder.create<linalg::FillOp>(loc, identityVal, empty).getResult(0);
+}
+
+FailureOr<linalg::SplitReductionResult>
+splitArgmaxReduction(RewriterBase &rewriter, linalg::GenericOp genericOp,
+                     linalg::ControlSplitReductionFn controlSplitReductionFn) {
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  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, 4> 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");
+  }
+
+  ArgmaxCombinerOps combinerOps = collectArgmaxCombinerOps(genericOp);
+  Operation *reductionOp = combinerOps.maxOp;
+  std::optional<TypedAttr> identity = arith::getNeutralElement(reductionOp);
+  if (!identity.has_value())
+    return rewriter.notifyMatchFailure(
+        genericOp, "Unknown identity value for the reduction");
+
+  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 (unsigned idx : llvm::seq<unsigned>(0u, map.getNumResults())) {
+      unsigned dim = map.getDimPosition(idx);
+      if (dim != reductionDim) {
+        expandedShape.push_back(oldShape[idx]);
+        continue;
+      }
+      if (dim == reductionDim) {
path: compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/SplitReduction.cpp
user: login: hanhanW
body: I think this `if` is not necessary, right?
created_at: 2025-05-14T23:43:35Z
author_association: CONTRIBUTOR
diff_hunk: @@ -389,4 +390,332 @@ splitReduction(RewriterBase &rewriter, LinalgExt::TopkOp topkOp,
   return success();
 }
 
+struct ArgmaxCombinerOps {
+  Operation *maxOp = nullptr;    // arith.maximumf
+  Operation *selectOp = nullptr; // arith.select
+  Operation *cmpOp = nullptr;    // arith.cmpf
+};
+
+// Matches the combiner pattern in a linalg.generic argmax-style reduction:
+// Example MLIR:
+// %4:2 = linalg.generic {
+//     indexing_maps = [...],
+//     iterator_types = ["parallel", "reduction"]
+//   } ins(%arg0 : tensor<?x128xbf16>) outs(%1, %3 : tensor<?xbf16>,
+//   tensor<?xi64>) {
+// ^bb0(%in: bf16, %out: bf16, %out_0: i64):
+//   %5 = linalg.index 1 : index
+//   %6 = arith.index_cast %5 : index to i64
+//   %7 = arith.maximumf %in, %out : bf16
+//   %8 = arith.cmpf ogt, %in, %out : bf16
+//   %9 = arith.select %8, %6, %out_0 : i64
+//   linalg.yield %7, %9 : bf16, i64
+// } -> (tensor<?xbf16>, tensor<?xi64>)
+//
+// This function extracts the `arith.maximumf`, `arith.cmpf`, and `arith.select`
+// operations from the body to facilitate transformations such as split
+// reduction.
+static ArgmaxCombinerOps collectArgmaxCombinerOps(linalg::GenericOp genericOp) {
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  auto yieldOp = cast<linalg::YieldOp>(genericOp.getBody()->getTerminator());
+
+  // Extract max value producer: arith.maximumf.
+  Value maxResult = yieldOp.getOperand(0);
+  auto maxOp = cast<arith::MaximumFOp>(maxResult.getDefiningOp());
+
+  // Extract index result producer: arith.select.
+  Value indexResult = yieldOp.getOperand(1);
+  auto selectOp = cast<arith::SelectOp>(indexResult.getDefiningOp());
+
+  // Extract the condition of the select, expected to be arith.cmpf with
+  // predicate OGT.
+  auto cmpOp = cast<arith::CmpFOp>(selectOp.getCondition().getDefiningOp());
+
+  ArgmaxCombinerOps ops;
+  ops.maxOp = maxOp;
+  ops.selectOp = selectOp;
+  ops.cmpOp = cmpOp;
+  return ops;
+}
+
+static Value expandValue(OpBuilder &builder, Location loc, Value value,
+                         RankedTensorType expandedType) {
+  RankedTensorType originalType = cast<RankedTensorType>(value.getType());
+  if (originalType == expandedType)
+    return value;
+
+  auto reassociation =
+      mlir::getReassociationIndicesForReshape(originalType, expandedType);
+  assert(reassociation && "failed to infer reassociation indices from types");
+  return builder.create<tensor::ExpandShapeOp>(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 = builder.create<arith::ConstantOp>(
+      loc, elemType, llvm::cast<TypedAttr>(identityAttr));
+  Value empty = builder.create<tensor::EmptyOp>(loc, shape, elemType);
+  return builder.create<linalg::FillOp>(loc, identityVal, empty).getResult(0);
+}
+
+FailureOr<linalg::SplitReductionResult>
+splitArgmaxReduction(RewriterBase &rewriter, linalg::GenericOp genericOp,
+                     linalg::ControlSplitReductionFn controlSplitReductionFn) {
+  assert(isArgmaxOp(genericOp) && "expected operation to be an argmax op");
+
+  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, 4> 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");
+  }
+
+  ArgmaxCombinerOps combinerOps = collectArgmaxCombinerOps(genericOp);
+  Operation *reductionOp = combinerOps.maxOp;
+  std::optional<TypedAttr> identity = arith::getNeutralElement(reductionOp);
+  if (!identity.has_value())
+    return rewriter.notifyMatchFailure(
+        genericOp, "Unknown identity value for the reduction");
+
+  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 (unsigned idx : llvm::seq<unsigned>(0u, map.getNumResults())) {
path: compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/SplitReduction.cpp
user: login: hanhanW
body: nit: I missed it in the previous reviews, use `auto`.
created_at: 2025-05-14T23:44:12Z
author_association: CONTRIBUTOR
diff_hunk: @@ -27,3 +27,110 @@ util.func public @matmul_with_encoding(%arg0: tensor<100x200xf32, #encoding>, %a
 }
 // CHECK-LABEL: util.func public @matmul_with_encoding(
 // CHECK:         linalg.matmul
+
+// -----
+
+util.func public @argmax(%arg0: tensor<?x131072xbf16>, %arg1: index) -> tensor<?xi64> {
+  %cst = arith.constant 0xFF80 : bf16
+  %c0_i64 = arith.constant 0 : i64
+  %0 = tensor.empty(%arg1) : tensor<?xbf16>
+  %1 = linalg.fill ins(%cst : bf16) outs(%0 : tensor<?xbf16>) -> tensor<?xbf16>
+  %2 = tensor.empty(%arg1) : tensor<?xi64>
+  %3 = linalg.fill ins(%c0_i64 : i64) outs(%2 : tensor<?xi64>) -> tensor<?xi64>
+  %4:2 = linalg.generic {
+      indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
+                       affine_map<(d0, d1) -> (d0)>,
+                       affine_map<(d0, d1) -> (d0)>],
+      iterator_types = ["parallel", "reduction"]}
+      ins(%arg0 : tensor<?x131072xbf16>) outs(%1, %3 : tensor<?xbf16>, tensor<?xi64>) {
+  ^bb0(%in: bf16, %out: bf16, %out_0: i64):
+    %5 = linalg.index 1 : index
+    %6 = arith.index_cast %5 : index to i64
+    %7 = arith.maximumf %in, %out : bf16
+    %8 = arith.cmpf ogt, %in, %out : bf16
+    %9 = arith.select %8, %6, %out_0 : i64
+    linalg.yield %7, %9 : bf16, i64
+  } -> (tensor<?xbf16>, tensor<?xi64>)
+  util.return %4#1 : tensor<?xi64>
+}
+
+// CHECK: #[[$MAP0:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+// CHECK: #[[$MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+// CHECK: #[[$MAP2:.+]] = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK: #[[$MAP3:.+]] = affine_map<(d0, d1) -> (d0)>
+
+// CHECK-LABEL:   util.func public @argmax
+// Check identity value preparation
+// CHECK: %[[CST:.+]] = arith.constant 0xFF80 : bf16
+// CHECK: %[[ZERO:.+]] = arith.constant 0 : i64
+// CHECK: %[[FINALVAL_EMPTY:.+]] = tensor.empty(%[[ARG1:.+]]) : tensor<?xbf16>
+// CHECK: %[[FINALVAL:.+]] = linalg.fill ins(%[[CST]] : bf16) outs(%[[FINALVAL_EMPTY]] : tensor<?xbf16>) -> tensor<?xbf16>
+// CHECK: %[[FINALIDX_EMPTY:.+]] = tensor.empty(%[[ARG1]]) : tensor<?xi64>
+// CHECK: %[[FINALIDX:.+]] = linalg.fill ins(%[[ZERO]] : i64) outs(%[[FINALIDX_EMPTY]] : tensor<?xi64>) -> tensor<?xi64>
+
+// Check partial reduction
path: compiler/src/iree/compiler/DispatchCreation/test/split_reduction.mlir
user: login: hanhanW
body: ```suggestion
// Check partial reduction.
```
created_at: 2025-05-14T23:47:37Z
author_association: CONTRIBUTOR
diff_hunk: @@ -27,3 +27,110 @@ util.func public @matmul_with_encoding(%arg0: tensor<100x200xf32, #encoding>, %a
 }
 // CHECK-LABEL: util.func public @matmul_with_encoding(
 // CHECK:         linalg.matmul
+
+// -----
+
+util.func public @argmax(%arg0: tensor<?x131072xbf16>, %arg1: index) -> tensor<?xi64> {
+  %cst = arith.constant 0xFF80 : bf16
+  %c0_i64 = arith.constant 0 : i64
+  %0 = tensor.empty(%arg1) : tensor<?xbf16>
+  %1 = linalg.fill ins(%cst : bf16) outs(%0 : tensor<?xbf16>) -> tensor<?xbf16>
+  %2 = tensor.empty(%arg1) : tensor<?xi64>
+  %3 = linalg.fill ins(%c0_i64 : i64) outs(%2 : tensor<?xi64>) -> tensor<?xi64>
+  %4:2 = linalg.generic {
+      indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
+                       affine_map<(d0, d1) -> (d0)>,
+                       affine_map<(d0, d1) -> (d0)>],
+      iterator_types = ["parallel", "reduction"]}
+      ins(%arg0 : tensor<?x131072xbf16>) outs(%1, %3 : tensor<?xbf16>, tensor<?xi64>) {
+  ^bb0(%in: bf16, %out: bf16, %out_0: i64):
+    %5 = linalg.index 1 : index
+    %6 = arith.index_cast %5 : index to i64
+    %7 = arith.maximumf %in, %out : bf16
+    %8 = arith.cmpf ogt, %in, %out : bf16
+    %9 = arith.select %8, %6, %out_0 : i64
+    linalg.yield %7, %9 : bf16, i64
+  } -> (tensor<?xbf16>, tensor<?xi64>)
+  util.return %4#1 : tensor<?xi64>
+}
+
+// CHECK: #[[$MAP0:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+// CHECK: #[[$MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+// CHECK: #[[$MAP2:.+]] = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK: #[[$MAP3:.+]] = affine_map<(d0, d1) -> (d0)>
+
+// CHECK-LABEL:   util.func public @argmax
+// Check identity value preparation
+// CHECK: %[[CST:.+]] = arith.constant 0xFF80 : bf16
+// CHECK: %[[ZERO:.+]] = arith.constant 0 : i64
+// CHECK: %[[FINALVAL_EMPTY:.+]] = tensor.empty(%[[ARG1:.+]]) : tensor<?xbf16>
+// CHECK: %[[FINALVAL:.+]] = linalg.fill ins(%[[CST]] : bf16) outs(%[[FINALVAL_EMPTY]] : tensor<?xbf16>) -> tensor<?xbf16>
+// CHECK: %[[FINALIDX_EMPTY:.+]] = tensor.empty(%[[ARG1]]) : tensor<?xi64>
+// CHECK: %[[FINALIDX:.+]] = linalg.fill ins(%[[ZERO]] : i64) outs(%[[FINALIDX_EMPTY]] : tensor<?xi64>) -> tensor<?xi64>
+
+// Check partial reduction
+// CHECK: %[[EXPAND:.+]] = tensor.expand_shape %arg0 {{\[}}[0], [1, 2]] output_shape [%{{.+}}, 1024, 128] : tensor<?x131072xbf16> into tensor<?x1024x128xbf16>
+// CHECK: %[[INITVAL:.+]] = tensor.empty(%{{.+}}) : tensor<?x1024xbf16>
+// CHECK: %[[FILLVAL:.+]] = linalg.fill ins(%{{.+}} : bf16) outs(%[[INITVAL]] : tensor<?x1024xbf16>) -> tensor<?x1024xbf16>
+// CHECK: %[[INITIDX:.+]] = tensor.empty(%{{.+}}) : tensor<?x1024xi64>
+// CHECK: %[[FILLIDX:.+]] = linalg.fill ins(%{{.+}} : i64) outs(%[[INITIDX]] : tensor<?x1024xi64>) -> tensor<?x1024xi64>
+
+// CHECK: %[[PARTIAL:.+]]:2 = linalg.generic
+// CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP1]]]
+// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]
+// CHECK-SAME: ins(%[[EXPAND]] : tensor<?x1024x128xbf16>)
+// CHECK-SAME: outs(%[[FILLVAL]], %[[FILLIDX]] : tensor<?x1024xbf16>, tensor<?x1024xi64>)
+// CHECK: ^bb0(%[[VAL:.+]]: bf16, %[[ACC:.+]]: bf16, %[[IDX:.+]]: i64)
+// CHECK: %[[OUTER:.+]] = linalg.index 1 : index
+// CHECK: %[[INNER:.+]] = linalg.index 2 : index
+// CHECK: %[[OFFSET:.+]] = arith.muli %[[OUTER]], %{{.+}} : index
+// CHECK: %[[GIDX:.+]] = arith.addi %[[OFFSET]], %[[INNER]] : index
+// CHECK: %[[CAST:.+]] = arith.index_cast %[[GIDX]] : index to i64
+// CHECK: %[[MAX:.+]] = arith.maximumf %[[VAL]], %[[ACC]] : bf16
+// CHECK: %[[CMP:.+]] = arith.cmpf ogt, %[[VAL]], %[[ACC]] : bf16
+// CHECK: %[[SEL:.+]] = arith.select %[[CMP]], %[[CAST]], %[[IDX]] : i64
+// CHECK: linalg.yield %[[MAX]], %[[SEL]] : bf16, i64
+
+// Final reduction
path: compiler/src/iree/compiler/DispatchCreation/test/split_reduction.mlir
user: login: hanhanW
body: ```suggestion
// Check final reduction.
```
created_at: 2025-05-14T23:47:56Z
author_association: CONTRIBUTOR

PR comments:
user: login: bangtianliu
created_at: 2025-05-06T18:52:22Z
body: > I want to review the template part if you don't mind wait me a bit. I can do the review later today or tomorrow. No needs to wait for me if others approve the change. I can send post review comments or raise a PR myself after you land the PR.

Thanks! No worries — it doesn’t seem like this PR is urgent to land.
author_association: CONTRIBUTOR
user: login: bangtianliu
created_at: 2025-05-06T22:09:05Z
body: Per Mahesh's suggestion, this adds a new flag `--iree-dispatch-creation-enable-default-argmax-split-pattern` to enable argmax split-k when the reduction dimension is 128.

Additionally, the meaning of `splitArgmaxReductionRatio` is updated to represent the number of elements per split along the reduction dimension (i.e., the tile size), rather than the number of splits.

In my local experimental test, the runtime for a <1x131072xbf16> input was 9.50 ms with the default flag disabled. When the default split pattern (with a reduction dimension of 128) was enabled, the runtime decreased significantly to 0.187 ms.

author_association: CONTRIBUTOR
user: login: bangtianliu
created_at: 2025-05-06T22:32:03Z
body: It seems that enabling this default pattern will cause some CI errors. I will look into them.
author_association: CONTRIBUTOR
user: login: bangtianliu
created_at: 2025-05-07T22:44:09Z
body: > I did not do a detailed review, but here are a few high-level comments so far:
> 
> 1. We already have separation in the main function, i.e., `SplitReductionPass::runOnOperation`. Can we avoid the template? Template is good for some cases, but the downside is hard to debug and maintain. It is also hard to trace code. I'd prefer to avoid it as much as possible. This falls to the case that we can avoid template. I don't see a need to use template here. Naming the method as splitArgmaxReduction looks much easier to me. https://google.github.io/styleguide/cppguide.html#Template_metaprogramming
> 2. Can you elaborate why the upstream method can't be reused? I can see some reasons when I browse the code, but it is easier if you can list down why and what the difference is in PR description. It is also helpful for reviewers and future contributors, as they would track back to the commit.
> 3. There are too many dup code between your implementation and upstream splitReduction. Can you think a way to refactor the code? I can help if you need some guidance.
> 4. I think it is better to move the `splitArgmaxReduction()` to `LinalgExt/Transforms/Transforms.[h|cpp]`. Because (a) I feel that we will upstream it in the future, (b) this makes the pass file cleaner, (c) it belongs to transformations about LinalgExt (linalg extention), IMO, (d) it can be reused by other passes when needed. RE(d): we do have a good example of reusing splitReduction in CPU codegen to perform innermost reduction. I think it is better to expose the implementation outside this pass in the first place.
> 5. Similar to (3), please add more details to the PR description. The PR is not trivial, we should have informative PR description. Also, it'd be a plus if you mention that there is a refactoring about `isArgmaxOp` in the PR description. https://google.github.io/eng-practices/review/developer/cl-descriptions.html
> 
> The topK variant locates at wrong files, I'll help fix it in a separate PR.

Thanks for your time and feedback during the code reviews. I've removed the template, updated the PR description, and reformatted the code according to your suggestions—except for point 3. Could you please provide further guidance on that part?
author_association: CONTRIBUTOR
user: login: bangtianliu
created_at: 2025-05-07T23:20:01Z
body: oh, I missed this one:
>  I think it is better to move the splitArgmaxReduction() to LinalgExt/Transforms/Transforms.[h|cpp]

I will send another commit to address it. 
author_association: CONTRIBUTOR
user: login: bangtianliu
created_at: 2025-05-13T05:32:59Z
body: Most of the comments have been addressed, except for the following two:
- The flag options: https://github.com/iree-org/iree/pull/20717#discussion_r2083256177 (DONE)
- It's still unclear whether `tensor.expand_shape` can be used for creating the identity value and index, per https://github.com/iree-org/iree/pull/20717#discussion_r2085851140

I'll continue working on these tomorrow.
author_association: CONTRIBUTOR