Repository: iree-org/iree
Pull request #20751: [GPU] Enable vector distribute on reduction operations by default
Author: login: pashu123
Created: 2025-05-07T23:55:26Z | Merged: 2025-05-13T21:41:57Z

-- Set only reduction tile sizes on the parallel operation (the
workgroup tile sizes are dominated by the reduction operation).
-- The bitwidth now selects max(operands_bitwidth) for an operation,
this is just to match warp distribution numerics.


Diff:
diff --git a/compiler/src/iree/compiler/Codegen/Common/DecomposeSoftmax.cpp b/compiler/src/iree/compiler/Codegen/Common/DecomposeSoftmax.cpp
index 25bc161eea7d..4aca280e6bd1 100644
--- a/compiler/src/iree/compiler/Codegen/Common/DecomposeSoftmax.cpp
+++ b/compiler/src/iree/compiler/Codegen/Common/DecomposeSoftmax.cpp
@@ -11,6 +11,7 @@
 #include "mlir/Dialect/Math/IR/Math.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 namespace mlir::iree_compiler {
 
@@ -18,6 +19,46 @@ namespace mlir::iree_compiler {
 #include "iree/compiler/Codegen/Common/Passes.h.inc"
 
 namespace {
+
+static bool isScalarOrTensorOfSizeOne(Type t) {
+  if (auto tensorType = dyn_cast<RankedTensorType>(t)) {
+    return tensorType.hasStaticShape() && tensorType.getNumElements() == 1;
+  }
+  return t.isIntOrIndexOrFloat();
+}
+
+struct FuseElementWiseGenericOps : public OpRewritePattern<linalg::GenericOp> {
+  using OpRewritePattern<linalg::GenericOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(linalg::GenericOp genericOp,
+                                PatternRewriter &rewriter) const override {
+    // Avoid doing this for scalar operations.
+    auto isScalarValue = [](Value v) {
+      return isScalarOrTensorOfSizeOne(v.getType());
+    };
+    if (llvm::all_of(genericOp.getOperands(), isScalarValue) &&
+        llvm::all_of(genericOp.getResults(), isScalarValue)) {
+      return failure();
+    }
+
+    // Find the first operand that is defined by another generic op on tensors.
+    for (OpOperand &opOperand : genericOp->getOpOperands()) {
+      if (!linalg::areElementwiseOpsFusable(&opOperand))
+        continue;
+
+      FailureOr<linalg::ElementwiseOpFusionResult> fusionResult =
+          linalg::fuseElementwiseOps(rewriter, &opOperand);
+      if (succeeded(fusionResult)) {
+        auto replacements = fusionResult->fusedOp->getResults().take_back(
+            genericOp.getNumResults());
+        rewriter.replaceOp(genericOp, replacements);
+        return success();
+      }
+    }
+    return failure();
+  }
+};
+
 /// Given an N-dimensional tensor x, this op converts
 /// softmax(x) to the following sequence of operations:
 ///
@@ -37,6 +78,7 @@ namespace {
 ///    softmax = z / l
 ///
 static LogicalResult convertSoftmaxToGenerics(mlir::FunctionOpInterface funcOp,
+                                              RewritePatternSet &patterns,
                                               bool useFusion) {
   IRRewriter rewriter(funcOp.getContext());
   SmallVector<Operation *> toDelete;
@@ -44,7 +86,6 @@ static LogicalResult convertSoftmaxToGenerics(mlir::FunctionOpInterface funcOp,
   funcOp.walk([&](linalg::SoftmaxOp softmaxOp) {
     softmaxOpsToDecompose.push_back(softmaxOp);
   });
-
   OpBuilder::InsertionGuard guard(rewriter);
   for (Operation *softmaxOp : softmaxOpsToDecompose) {
     // Cast linalg::softmax to AggregatedOpInterface since this where
@@ -59,42 +100,13 @@ static LogicalResult convertSoftmaxToGenerics(mlir::FunctionOpInterface funcOp,
           softmaxOp, "linalg::SoftmaxOp could not be decomposed"));
       return failure();
     }
-
     // Replace the result of linalg::softmax with the `result` generated via
     // the decomposition above.
     rewriter.replaceOp(decomposableSoftmaxOp, *result);
-
-    if (useFusion) {
-      // Fusion later depends on couple of Ops/Values - we try to obtain the
-      // same by backtracking through the generated value's def-chain.
-      Operation *resultOp = (*result)[0].getDefiningOp();
-      Value numerator = resultOp->getOperand(0);
-      Operation *numeratorOp = numerator.getDefiningOp();
-
-      // Rematerialize operands that are marked for this.
-      SmallVector<OpOperand *> uses = llvm::to_vector(llvm::map_range(
-          numerator.getUses(), [](OpOperand &use) { return &use; }));
-      for (OpOperand *use : uses) {
-        Operation *consumer = use->getOwner();
-        OpBuilder::InsertionGuard g(rewriter);
-        rewriter.setInsertionPoint(consumer);
-        FailureOr<linalg::ElementwiseOpFusionResult> fusionResult =
-            linalg::fuseElementwiseOps(rewriter, use);
-        if (succeeded(fusionResult)) {
-          SmallVector<Value> replacements = llvm::to_vector(
-              llvm::map_range(consumer->getResults(), [&](Value oldValue) {
-                return fusionResult->replacements.lookup(oldValue);
-              }));
-          rewriter.replaceOp(consumer, replacements);
-        }
-      }
-      toDelete.push_back(numeratorOp);
-    }
   }
-  for (Operation *op : toDelete) {
-    rewriter.eraseOp(op);
+  if (!softmaxOpsToDecompose.empty() && useFusion) {
+    patterns.insert<FuseElementWiseGenericOps>(funcOp.getContext());
   }
-
   return success();
 }
 
@@ -110,8 +122,14 @@ struct DecomposeSoftmaxPass
   void runOnOperation() override {
     MLIRContext *context = &getContext();
     IRRewriter rewriter(context);
-    if (failed(convertSoftmaxToGenerics(getOperation(), useFusion)))
+    RewritePatternSet patterns(context);
+    if (failed(convertSoftmaxToGenerics(getOperation(), patterns, useFusion))) {
       return signalPassFailure();
+    }
+    linalg::populateEraseUnusedOperandsAndResultsPatterns(patterns);
+    if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) {
+      return signalPassFailure();
+    }
   }
 };
 
diff --git a/compiler/src/iree/compiler/Codegen/Common/test/decompose_softmax.mlir b/compiler/src/iree/compiler/Codegen/Common/test/decompose_softmax.mlir
index 43da04ba6e72..4075a3b14a17 100644
--- a/compiler/src/iree/compiler/Codegen/Common/test/decompose_softmax.mlir
+++ b/compiler/src/iree/compiler/Codegen/Common/test/decompose_softmax.mlir
@@ -11,9 +11,10 @@ func.func @softmax(%arg0: tensor<2x16x32xf32>) -> tensor<2x16x32xf32> {
 // CHECK-DAG:  #[[MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
 // CHECK-DAG:  #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
 // CHECK:      func.func @softmax(%[[ARG0:[a-zA-Z0-9_]+]]: tensor<2x16x32xf32>) -> tensor<2x16x32xf32> {
-// CHECK:        %[[D0:.+]] = tensor.empty() : tensor<2x16x32xf32>
-// CHECK:        %[[D1:.+]] = tensor.empty() : tensor<2x16xf32>
-// CHECK:        %[[CST:.+]] = arith.constant 0xFFC00000 : f32
+// CHECK-DAG:    %[[D0:.+]] = tensor.empty() : tensor<2x16x32xf32>
+// CHECK-DAG:    %[[D1:.+]] = tensor.empty() : tensor<2x16xf32>
+// CHECK-DAG:    %[[CST:.+]] = arith.constant 0xFFC00000 : f32
+// CHECK-DAG:    %[[CST0:.+]] = arith.constant 0.0
 // CHECK:        %[[D2:.+]] = linalg.fill ins(%[[CST]] : f32) outs(%[[D1]] : tensor<2x16xf32>) -> tensor<2x16xf32>
 // CHECK:        %[[D3:.+]] = linalg.generic {indexing_maps = [#[[MAP]], #[[MAP1]]], iterator_types = ["parallel",
 // CHECK-SAME:     "parallel", "reduction"]} ins(%[[ARG0]] : tensor<2x16x32xf32>) outs(%[[D2]] : tensor<2x16xf32>) {
@@ -21,7 +22,6 @@ func.func @softmax(%arg0: tensor<2x16x32xf32>) -> tensor<2x16x32xf32> {
 // CHECK:          %[[D8:.+]] = arith.maxnumf %[[IN]], %[[OUT]] : f32
 // CHECK:          linalg.yield %[[D8]] : f32
 // CHECK:        } -> tensor<2x16xf32>
-// CHECK:        %[[CST0:.+]] = arith.constant 0.0
 // CHECK:        %[[D4:.+]] = linalg.fill ins(%[[CST0]] : f32) outs(%[[D1]] : tensor<2x16xf32>) -> tensor<2x16xf32>
 // CHECK:        %[[D5:.+]] = linalg.generic {indexing_maps = [#[[MAP]], #[[MAP1]], #[[MAP1]]], iterator_types =
 // CHECK-SAME:     ["parallel", "parallel", "reduction"]} ins(%[[ARG0]], %[[D3]] : tensor<2x16x32xf32>, tensor<2x16xf32>)
@@ -32,24 +32,25 @@ func.func @softmax(%arg0: tensor<2x16x32xf32>) -> tensor<2x16x32xf32> {
 // CHECK:          %[[D10:.+]] = arith.addf %[[D9]], %[[OUT]]
 // CHECK:          linalg.yield %[[D10]] : f32
 // CHECK:        } -> tensor<2x16xf32>
-// CHECK:        %[[D7:.+]]:2 = linalg.generic {indexing_maps = [#[[MAP]], #[[MAP1]], #[[MAP1]], #[[MAP]], #[[MAP]]], iterator_types =
+// CHECK:        %[[D7:.+]] = linalg.generic {indexing_maps = [#[[MAP]], #[[MAP1]], #[[MAP1]], #[[MAP]]], iterator_types =
 // CHECK-SAME:     ["parallel", "parallel", "parallel"]} ins(%[[ARG0]], %[[D3]], %[[D5]] : tensor<2x16x32xf32>, tensor<2x16xf32>, tensor<2x16xf32>)
-// CHECK-SAME:     outs(%[[D0]], %[[D0]] : tensor<2x16x32xf32>, tensor<2x16x32xf32>) {
-// CHECK:        ^bb0(%[[IN:.+]]: f32, %[[IN_1:.+]]: f32, %[[IN_2:.+]]: f32, %[[OUT0:.+]]: f32, %[[OUT1:.+]]: f32):
+// CHECK-SAME:     outs(%[[D0]] : tensor<2x16x32xf32>) {
+// CHECK:        ^bb0(%[[IN:.+]]: f32, %[[IN_1:.+]]: f32, %[[IN_2:.+]]: f32, %[[OUT0:.+]]: f32):
 // CHECK:          %[[D8:.+]] = arith.subf %[[IN]], %[[IN_1]] : f32
 // CHECK:          %[[D9:.+]] = math.exp %[[D8]] : f32
 // CHECK:          %[[D10:.+]] = arith.divf %[[D9]], %[[IN_2]] : f32
-// CHECK:          linalg.yield %[[D9]], %[[D10]] : f32
-// CHECK:        } -> (tensor<2x16x32xf32>, tensor<2x16x32xf32>)
-// CHECK:        return %[[D7]]#1 : tensor<2x16x32xf32>
+// CHECK:          linalg.yield %[[D10]] : f32
+// CHECK:        } -> tensor<2x16x32xf32>
+// CHECK:        return %[[D7]] : tensor<2x16x32xf32>
 // CHECK:      }
 
 // CHECK-NO-FUSE-DAG:  #[[MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
 // CHECK-NO-FUSE-DAG:  #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
 // CHECK-NO-FUSE:      func.func @softmax(%[[ARG0:[a-zA-Z0-9_]+]]: tensor<2x16x32xf32>) -> tensor<2x16x32xf32> {
-// CHECK-NO-FUSE:        %[[D0:.+]] = tensor.empty() : tensor<2x16x32xf32>
-// CHECK-NO-FUSE:        %[[D1:.+]] = tensor.empty() : tensor<2x16xf32>
-// CHECK-NO-FUSE:        %[[CST:.+]] = arith.constant 0xFFC00000 : f32
+// CHECK-NO-FUSE-DAG:     %[[D0:.+]] = tensor.empty() : tensor<2x16x32xf32>
+// CHECK-NO-FUSE-DAG:     %[[D1:.+]] = tensor.empty() : tensor<2x16xf32>
+// CHECK-NO-FUSE-DAG:     %[[CST:.+]] = arith.constant 0xFFC00000 : f32
+// CHECK-NO-FUSE-DAG:     %[[CST0:.+]] = arith.constant 0.000000e+00 : f32
 // CHECK-NO-FUSE:        %[[D2:.+]] = linalg.fill ins(%[[CST]] : f32) outs(%[[D1]] : tensor<2x16xf32>) -> tensor<2x16xf32>
 // CHECK-NO-FUSE:        %[[D3:.+]] = linalg.generic {indexing_maps = [#[[MAP]], #[[MAP1]]], iterator_types = ["parallel",
 // CHECK-NO-FUSE-SAME:     "parallel", "reduction"]} ins(%[[ARG0]] : tensor<2x16x32xf32>) outs(%[[D2]] : tensor<2x16xf32>) {
@@ -64,7 +65,6 @@ func.func @softmax(%arg0: tensor<2x16x32xf32>) -> tensor<2x16x32xf32> {
 // CHECK-NO-FUSE:          %[[D9:.+]] = math.exp %[[D8]] : f32
 // CHECK-NO-FUSE:          linalg.yield %[[D9]] : f32
 // CHECK-NO-FUSE:        } -> tensor<2x16x32xf32>
-// CHECK-NO-FUSE:        %[[CST0:.+]] = arith.constant 0.000000e+00 : f32
 // CHECK-NO-FUSE:        %[[D5:.+]] = linalg.fill ins(%[[CST0]] : f32) outs(%[[D1]] : tensor<2x16xf32>) -> tensor<2x16xf32>
 // CHECK-NO-FUSE:        %[[D6:.+]] = linalg.generic {indexing_maps = [#[[MAP]], #[[MAP1]]], iterator_types =
 // CHECK-NO-FUSE-SAME:     ["parallel", "parallel", "reduction"]} ins(%[[D4]] : tensor<2x16x32xf32>)
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp b/compiler/src/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp
index 4557e7629f46..38c376a29c8d 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp
@@ -68,11 +68,6 @@ llvm::cl::opt<bool> clLLVMGPUVectorizePipeline(
     llvm::cl::desc("forces use of the legacy LLVMGPU vectorize pipeline"),
     llvm::cl::init(false));
 
-llvm::cl::opt<bool> clGPUTestVectorDistributeOnReduction(
-    "iree-codegen-llvmgpu-test-vector-distribution-on-reduction",
-    llvm::cl::desc("test vector distribute on reduction."),
-    llvm::cl::init(false));
-
 llvm::cl::opt<bool> clGPUEnableVectorDistribution(
     "iree-codegen-llvmgpu-use-vector-distribution",
     llvm::cl::desc("enable the usage of the vector distribution pipeline"),
@@ -290,7 +285,6 @@ static bool hasReductionIterator(linalg::LinalgOp &op) {
 /// bitwidth of the operation. The bitwidth is minimum of the init and src
 /// operands.
 static FailureOr<int64_t> getBitWidth(linalg::LinalgOp op) {
-
   Value init = op.getDpsInitOperand(0)->get();
   Type initElemType = getElementTypeOrSelf(init);
   Value src = op.getDpsInputOperand(0)->get();
@@ -312,9 +306,8 @@ static FailureOr<int64_t> getBitWidth(linalg::LinalgOp op) {
     return failure();
   }
 
-  int64_t bitWidth = std::min(initElemType.getIntOrFloatBitWidth(),
+  int64_t bitWidth = std::max(initElemType.getIntOrFloatBitWidth(),
                               srcElemType.getIntOrFloatBitWidth());
-
   return bitWidth;
 }
 
@@ -357,58 +350,60 @@ getVectorDistributeReductionConfig(linalg::LinalgOp op,
 
   SmallVector<int64_t, 4> bounds = op.getStaticLoopRanges();
 
-  SmallVector<int64_t> workgroupTileSizes(op.getNumLoops(), 0);
   SmallVector<int64_t> threadTileSizes(op.getNumLoops(), 0);
-  SmallVector<int64_t> reductionTileSizes(op.getNumLoops(), 0);
   SmallVector<int64_t> threadCounts(op.getNumLoops(), 1);
   SmallVector<int64_t> subGroupCounts(op.getNumLoops(), 1);
   SmallVector<int64_t> mapping(op.getNumLoops());
   std::iota(mapping.begin(), mapping.end(), 0);
 
-  // Set the shared wgp dims to 1.
-  for (int64_t i : sharedWgpDims) {
-    workgroupTileSizes[i] = 1;
-  }
-
   // Set the configuration for the operation with no reduction dims.
+  // The workgroup tile sizes are set by the reduction operation.
   if (reductionDims.empty()) {
+    SmallVector<int64_t> reductionTileSizes(op.getNumLoops(), 1);
+
+    // For the shared wgp dimension, set the reduction tile sizes to be zero.
+    for (auto i : sharedWgpDims) {
+      reductionTileSizes[i] = 0;
+    }
+
     int64_t parallelSize = bounds[parallelDims.back()];
-    if (parallelSize % threadLoads != 0) {
+    if (ShapedType::isDynamic(parallelSize) ||
+        parallelSize % threadLoads != 0) {
       return failure();
     }
-
-    int64_t lastDimWgpTileSize = workgroupSize * threadLoads;
+    int64_t lastDimReductionTileSize = workgroupSize * threadLoads;
 
     // Setting subgroupBasis to minimum i.e., 1 and threadBasis
     // to maximum i.e., subgroupSize.
     int64_t subgroupBasis = 1;
     int64_t threadBasis = subgroupSize;
 
-    lastDimWgpTileSize = llvm::APIntOps::GreatestCommonDivisor(
-                             {64, static_cast<uint64_t>(parallelSize)},
-                             {64, static_cast<uint64_t>(lastDimWgpTileSize)})
-                             .getZExtValue();
+    lastDimReductionTileSize =
+        llvm::APIntOps::GreatestCommonDivisor(
+            {64, static_cast<uint64_t>(parallelSize)},
+            {64, static_cast<uint64_t>(lastDimReductionTileSize)})
+            .getZExtValue();
 
     if (!(sharedWgpDims.size() == op.getNumLoops())) {
-      while (lastDimWgpTileSize % (threadBasis * threadLoads) != 0) {
+      int subgroupStride = threadBasis * threadLoads;
+      while (lastDimReductionTileSize % subgroupStride != 0) {
         threadBasis >>= 1;
+        subgroupStride = threadBasis * threadLoads;
       }
-      subgroupBasis = lastDimWgpTileSize / (threadLoads * threadBasis) == 0
-                          ? 1
-                          : lastDimWgpTileSize / (threadLoads * threadBasis);
+      int subgroup = lastDimReductionTileSize / subgroupStride;
+      subgroupBasis = (subgroup == 0) ? 1 : subgroup;
     }
 
     // Since all the dimensions are contained within the shared parallel
     // dimension, set the tile sizes to 1.
     if (sharedWgpDims.size() == op.getNumLoops()) {
-      lastDimWgpTileSize = 1;
+      lastDimReductionTileSize = 1;
       threadLoads = 1;
       threadBasis = 1;
       subgroupBasis = 1;
     }
 
-    workgroupTileSizes[parallelDims.back()] = lastDimWgpTileSize;
-    reductionTileSizes[parallelDims.back()] = lastDimWgpTileSize;
+    reductionTileSizes[parallelDims.back()] = lastDimReductionTileSize;
     threadTileSizes[parallelDims.back()] = threadLoads;
     threadCounts[parallelDims.back()] = threadBasis;
     subGroupCounts[parallelDims.back()] = subgroupBasis;
@@ -420,7 +415,7 @@ getVectorDistributeReductionConfig(linalg::LinalgOp op,
         {b.getI64ArrayAttr(threadCounts), b.getI64ArrayAttr(mapping)});
 
     SmallVector<NamedAttribute, 4> configAttrs = {
-        NamedAttribute("workgroup", b.getI64ArrayAttr(workgroupTileSizes)),
+        NamedAttribute("reduction", b.getI64ArrayAttr(reductionTileSizes)),
         NamedAttribute("thread", b.getI64ArrayAttr(threadTileSizes)),
         NamedAttribute("thread_basis", threadBasisAttr),
         NamedAttribute("subgroup_basis", subgroupBasisAttr)};
@@ -431,6 +426,13 @@ getVectorDistributeReductionConfig(linalg::LinalgOp op,
     return loweringConfig;
   }
 
+  SmallVector<int64_t> workgroupTileSizes(op.getNumLoops(), 0);
+
+  // Set the shared wgp dims to 1.
+  for (int64_t i : sharedWgpDims) {
+    workgroupTileSizes[i] = 1;
+  }
+
   // Setting the config for operation with atleast one reduction dimension.
   SmallVector<int64_t> partialReductionTileSizes(op.getNumLoops(), 0);
   int64_t lastReductionDim = reductionDims.back();
@@ -473,13 +475,13 @@ getVectorDistributeReductionConfig(linalg::LinalgOp op,
                              .getZExtValue();
 
   int64_t threadBasis = subgroupSize;
-  while (partialReductionSize % (threadBasis * threadLoads) != 0)
+  int subgroupStride = threadBasis * threadLoads;
+  while (partialReductionSize % subgroupStride != 0) {
     threadBasis >>= 1;
-
-  int64_t subgroupBasis =
-      partialReductionSize / (threadLoads * threadBasis) == 0
-          ? 1
-          : partialReductionSize / (threadLoads * threadBasis);
+    subgroupStride = threadBasis * threadLoads;
+  }
+  int subgroup = partialReductionSize / subgroupStride;
+  int64_t subgroupBasis = (subgroup == 0) ? 1 : subgroup;
 
   partialReductionTileSizes[lastReductionDim] = partialReductionSize;
   threadTileSizes[lastReductionDim] = threadLoads;
@@ -543,6 +545,9 @@ populateConfigInfo(const llvm::SetVector<linalg::LinalgOp> &computeOps,
           return failure();
         }
         setLoweringConfig(linalgOp, *loweringConfig);
+        // If the config is set on the op, no need to iterate
+        // over other operands.
+        break;
       }
     }
   }
@@ -2988,13 +2993,10 @@ static LogicalResult setRootConfig(IREE::GPU::TargetAttr target,
       LDBG("Contract Config");
       return success();
     }
-    // Skinny matmuls go through the vectordistribute pipeline by default.
-    if (clGPUTestVectorDistributeOnReduction || isMatmulLike(linalgOp)) {
-      if (succeeded(setReductionVectorDistributionConfig(target, entryPointFn,
-                                                         linalgOp))) {
-        LDBG("Vector Distribution Subgroup Reduction Config");
-        return success();
-      }
+    if (succeeded(setReductionVectorDistributionConfig(target, entryPointFn,
+                                                       linalgOp))) {
+      LDBG("Vector Distribution Subgroup Reduction Config");
+      return success();
     }
     if (succeeded(setWarpReductionConfig(target, entryPointFn, linalgOp))) {
       LDBG("Warp Reduction Config");
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/config_vector_distribute_reduction_gfx942.mlir b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/config_vector_distribute_reduction_gfx942.mlir
index 2654018f1cb5..7d434ce8b597 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/config_vector_distribute_reduction_gfx942.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/config_vector_distribute_reduction_gfx942.mlir
@@ -2,10 +2,6 @@
 // RUN:   --iree-codegen-llvmgpu-use-igemm=false \
 // RUN:   --pass-pipeline="builtin.module(iree-llvmgpu-select-lowering-strategy)" %s | FileCheck %s
 
-// RUN: iree-opt --split-input-file --iree-gpu-test-target=gfx942 \
-// RUN:   --iree-codegen-llvmgpu-test-vector-distribution-on-reduction \
-// RUN:   --pass-pipeline="builtin.module(iree-llvmgpu-select-lowering-strategy)" %s | FileCheck %s --check-prefix=CHECK-REDUCTION
-
 // CHECK:      #iree_codegen.translation_info<pipeline = LLVMGPUVectorDistribute
 
 #pipeline_layout = #hal.pipeline.layout<bindings = [
@@ -87,14 +83,14 @@ func.func @reduction_with_no_consumer() {
     iree_tensor_ext.dispatch.tensor.store %7, %1, offsets = [0, 0], sizes = [2, 32], strides = [1, 1] : tensor<2x32xf32> -> !iree_tensor_ext.dispatch.tensor<writeonly:tensor<2x32xf32>>
     return
 }
-// CHECK-REDUCTION:      #iree_codegen.translation_info<pipeline = LLVMGPUVectorDistribute
+// CHECK:      #iree_codegen.translation_info<pipeline = LLVMGPUVectorDistribute
 
-// CHECK-REDUCTION-LABEL: func.func @reduction_with_no_consumer
-// CHECK-REDUCTION:           lowering_config = #iree_gpu.lowering_config
-// CHECK-REDUCTION-SAME:      partial_reduction = [0, 0, 1, 4096]
-// CHECK-REDUCTION-SAME:      subgroup_basis = {{\[}}[1, 1, 1, 8], [0, 1, 2, 3]
-// CHECK-REDUCTION-SAME:      thread = [0, 0, 1, 8], thread_basis = {{\[}}[1, 1, 1, 64], [0, 1, 2, 3]
-// CHECK-REDUCTION-SAME:      workgroup = [1, 1, 0, 0]
+// CHECK-LABEL: func.func @reduction_with_no_consumer
+// CHECK:           lowering_config = #iree_gpu.lowering_config
+// CHECK-SAME:      partial_reduction = [0, 0, 1, 4096]
+// CHECK-SAME:      subgroup_basis = {{\[}}[1, 1, 1, 16], [0, 1, 2, 3]
+// CHECK-SAME:      thread = [0, 0, 1, 4], thread_basis = {{\[}}[1, 1, 1, 64], [0, 1, 2, 3]
+// CHECK-SAME:      workgroup = [1, 1, 0, 0]
 
 // -----
 
@@ -150,33 +146,32 @@ func.func @test_multiple_reduction() {
 // Attaches lowering config to the operation with reduction iterator or
 // parallel iterator with new dimensions.
 
-// CHECK-REDUCTION-LABEL: func.func @test_multiple_reduction
-// CHECK-REDUCTION:       %{{.*}} = linalg.generic {indexing_maps = [#map, #map1],
-// CHECK-REDUCTION-SAME:    iterator_types = ["parallel", "parallel", "reduction", "reduction"]}
-// CHECK-REDUCTION-SAME:    ins(%{{.*}} : tensor<2x32x10x16384xf32>)
-// CHECK-REDUCTION-SAME:    outs({{.*}}: tensor<2x32xf32>)
-// CHECK-REDUCTION-SAME:    attrs =  {lowering_config = #iree_gpu.lowering_config<{
-// CHECK-REDUCTION-SAME:               partial_reduction = [0, 0, 1, 8192],
-// CHECK-REDUCTION-SAME:               subgroup_basis = {{\[}}[1, 1, 1, 16], [0, 1, 2, 3]],
-// CHECK-REDUCTION-SAME:               thread = [0, 0, 1, 8], thread_basis = {{\[}}[1, 1, 1, 64], [0, 1, 2, 3]],
-// CHECK-REDUCTION-SAME:               workgroup = [1, 1, 0, 0]
-// CHECK-REDUCTION:       %{{.*}} = linalg.generic {indexing_maps = [#map, #map1, #map1],
-// CHECK-REDUCTION-SAME:    iterator_types = ["parallel", "parallel", "reduction", "reduction"]}
-// CHECK-REDUCTION-SAME:    ins{{.*}}, {{.*}} : tensor<2x32x10x16384xf32>, tensor<2x32xf32>)
-// CHECK-REDUCTION-SAME:    outs(%{{.*}} : tensor<2x32xf32>)
-// CHECK-REDUCTION-SAME:    attrs =  {lowering_config = #iree_gpu.lowering_config<{
-// CHECK-REDUCTION-SAME:              partial_reduction = [0, 0, 1, 8192],
-// CHECK-REDUCTION-SAME:              subgroup_basis = {{\[}}[1, 1, 1, 16], [0, 1, 2, 3]],
-// CHECK-REDUCTION-SAME:              thread = [0, 0, 1, 8], thread_basis = {{\[}}[1, 1, 1, 64], [0, 1, 2, 3]],
-// CHECK-REDUCTION-SAME:              workgroup = [1, 1, 0, 0]
-// CHECK-REDUCTION:       %{{.*}} = linalg.generic {indexing_maps = [#map, #map1, #map1, #map],
-// CHECK-REDUCTION-SAME:    iterator_types = ["parallel", "parallel", "parallel", "parallel"]}
-// CHECK-REDUCTION-SAME:    ins({{.*}}, %{{.*}}, {{.*}} : tensor<2x32x10x16384xf16>, tensor<2x32xf32>, tensor<2x32xf32>)
-// CHECK-REDUCTION-SAME:    outs(%{{.*}} : tensor<2x32x10x16384xf32>)
-// CHECK-REDUCTION-SAME:    attrs =  {lowering_config = #iree_gpu.lowering_config<{
-// CHECK-REDUCTION-SAME:              subgroup_basis = {{\[}}[1, 1, 1, 16], [0, 1, 2, 3]],
-// CHECK-REDUCTION-SAME:              thread = [0, 0, 0, 8], thread_basis = {{\[}}[1, 1, 1, 64], [0, 1, 2, 3]],
-// CHECK-REDUCTION-SAME:              workgroup = [1, 1, 0, 8192]
+// CHECK-LABEL: func.func @test_multiple_reduction
+// CHECK:       %{{.*}} = linalg.generic {indexing_maps = [#map, #map1],
+// CHECK-SAME:    iterator_types = ["parallel", "parallel", "reduction", "reduction"]}
+// CHECK-SAME:    ins(%{{.*}} : tensor<2x32x10x16384xf32>)
+// CHECK-SAME:    outs({{.*}}: tensor<2x32xf32>)
+// CHECK-SAME:    attrs =  {lowering_config = #iree_gpu.lowering_config<{
+// CHECK-SAME:               partial_reduction = [0, 0, 1, 4096],
+// CHECK-SAME:               subgroup_basis = {{\[}}[1, 1, 1, 16], [0, 1, 2, 3]],
+// CHECK-SAME:               thread = [0, 0, 1, 4], thread_basis = {{\[}}[1, 1, 1, 64], [0, 1, 2, 3]],
+// CHECK-SAME:               workgroup = [1, 1, 0, 0]
+// CHECK:       %{{.*}} = linalg.generic {indexing_maps = [#map, #map1, #map1],
+// CHECK-SAME:    iterator_types = ["parallel", "parallel", "reduction", "reduction"]}
+// CHECK-SAME:    ins{{.*}}, {{.*}} : tensor<2x32x10x16384xf32>, tensor<2x32xf32>)
+// CHECK-SAME:    outs(%{{.*}} : tensor<2x32xf32>)
+// CHECK-SAME:    attrs =  {lowering_config = #iree_gpu.lowering_config<{
+// CHECK-SAME:              partial_reduction = [0, 0, 1, 4096],
+// CHECK-SAME:              subgroup_basis = {{\[}}[1, 1, 1, 16], [0, 1, 2, 3]],
+// CHECK-SAME:              thread = [0, 0, 1, 4], thread_basis = {{\[}}[1, 1, 1, 64], [0, 1, 2, 3]],
+// CHECK:       %{{.*}} = linalg.generic {indexing_maps = [#map, #map1, #map1, #map],
+// CHECK-SAME:    iterator_types = ["parallel", "parallel", "parallel", "parallel"]}
+// CHECK-SAME:    ins({{.*}}, %{{.*}}, {{.*}} : tensor<2x32x10x16384xf16>, tensor<2x32xf32>, tensor<2x32xf32>)
+// CHECK-SAME:    outs(%{{.*}} : tensor<2x32x10x16384xf32>)
+// CHECK-SAME:    attrs =  {lowering_config = #iree_gpu.lowering_config<{
+// CHECK-SAME:              reduction = [0, 0, 1, 4096],
+// CHECK-SAME:              subgroup_basis = {{\[}}[1, 1, 1, 16], [0, 1, 2, 3]],
+// CHECK-SAME:              thread = [0, 0, 0, 4], thread_basis = {{\[}}[1, 1, 1, 64], [0, 1, 2, 3]]
 
 // -----
 
@@ -202,12 +197,12 @@ func.func @test_dyn_reduction(%arg0: !iree_tensor_ext.dispatch.tensor<readwrite:
   iree_tensor_ext.dispatch.tensor.store %5, %arg0, offsets = [0, 0], sizes = [128, 128], strides = [1, 1] : tensor<128x128xf32> -> !iree_tensor_ext.dispatch.tensor<readwrite:tensor<128x128xf32>>
   return
 }
-//       CHECK-REDUCTION: #[[$TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = LLVMGPUVectorDistribute workgroup_size = [2, 1, 1] subgroup_size = 64
-//       CHECK-REDUCTION: func.func @test_dyn_reduction
-//  CHECK-REDUCTION-SAME:     translation_info = #[[$TRANSLATION]]
-//       CHECK-REDUCTION:   linalg.generic
-//  CHECK-REDUCTION-SAME:      attrs =  {lowering_config = #iree_gpu.lowering_config<{
-//  CHECK-REDUCTION-SAME:               partial_reduction = [0, 0, 1, 32],
-//  CHECK-REDUCTION-SAME:               subgroup_basis = {{\[}}[1, 1, 1, 1], [0, 1, 2, 3]],
-//  CHECK-REDUCTION-SAME:               thread = [0, 0, 1, 16], thread_basis = {{\[}}[1, 1, 1, 2], [0, 1, 2, 3]],
-//  CHECK-REDUCTION-SAME:               workgroup = [1, 1, 0, 0]
+//       CHECK: #[[$TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = LLVMGPUVectorDistribute workgroup_size = [8, 1, 1] subgroup_size = 64
+//       CHECK: func.func @test_dyn_reduction
+//  CHECK-SAME:     translation_info = #[[$TRANSLATION]]
+//       CHECK:   linalg.generic
+//  CHECK-SAME:      attrs =  {lowering_config = #iree_gpu.lowering_config<{
+//  CHECK-SAME:               partial_reduction = [0, 0, 1, 32],
+//  CHECK-SAME:               subgroup_basis = {{\[}}[1, 1, 1, 1], [0, 1, 2, 3]],
+//  CHECK-SAME:               thread = [0, 0, 1, 4], thread_basis = {{\[}}[1, 1, 1, 8], [0, 1, 2, 3]],
+//  CHECK-SAME:               workgroup = [1, 1, 0, 0]
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/nvvm_pipeline_test.mlir b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/nvvm_pipeline_test.mlir
index c2b68882bb55..0e01e4c920a2 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/nvvm_pipeline_test.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/nvvm_pipeline_test.mlir
@@ -719,14 +719,14 @@ hal.executable @mma_fused_fp16 {
   #hal.pipeline.binding<storage_buffer>,
   #hal.pipeline.binding<storage_buffer>
 ]>
-hal.executable @warp_reduction_dispatch {
+hal.executable @vector_distribute_dispatch {
 hal.executable.variant @cuda target(<"cuda", "cuda-nvptx-fb">) {
-  hal.executable.export public @warp_reduction_dispatch layout(#pipeline_layout) count(%arg0: !hal.device, %arg1: index, %arg2 : index) -> (index, index, index) {
+  hal.executable.export public @vector_distribute_dispatch layout(#pipeline_layout) count(%arg0: !hal.device, %arg1: index, %arg2 : index) -> (index, index, index) {
       %x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root %arg1, %arg2
       hal.return %x, %y, %z : index, index, index
     }
   builtin.module {
-    func.func @warp_reduction_dispatch() {
+    func.func @vector_distribute_dispatch() {
       %c0 = arith.constant 0 : index
       %c1024 = arith.constant 1024 : index
       %cst = arith.constant 1.000000e+00 : f32
@@ -751,13 +751,13 @@ hal.executable.variant @cuda target(<"cuda", "cuda-nvptx-fb">) {
 }
 }
 
-// Check that we generate a warp reduce code sequence.
-//   CHECK-LABEL: hal.executable public @warp_reduction_dispatch
+// Check that we generate a vector distribute code sequence.
+//   CHECK-LABEL: hal.executable public @vector_distribute_dispatch
 //         CHECK:   hal.executable.variant public @cuda
 // CHECK-COUNT-5:     nvvm.shfl.sync  bfly
-//         CHECK:     llvm.store %{{.*}}, %{{.*}} : f32, !llvm.ptr<3>
+//         CHECK:     llvm.store %{{.*}}, %{{.*}} : vector<1xf32>, !llvm.ptr<3>
 //         CHECK:     nvvm.barrier0
-//         CHECK:     llvm.load {{.*}} : !llvm.ptr<3> -> f32
+//         CHECK:     llvm.load {{.*}} : !llvm.ptr<3> -> vector<1xf32>
 // CHECK-COUNT-2:     nvvm.shfl.sync  bfly
 
 // -----
@@ -769,14 +769,14 @@ hal.executable.variant @cuda target(<"cuda", "cuda-nvptx-fb">) {
   #hal.pipeline.binding<storage_buffer>,
   #hal.pipeline.binding<storage_buffer>
 ]>
-hal.executable @warp_reduction_broadcast_dispatch {
+hal.executable @vector_distribution_broadcast_dispatch {
 hal.executable.variant @cuda target(<"cuda", "cuda-nvptx-fb">) {
-  hal.executable.export public @warp_reduction_broadcast_dispatch layout(#pipeline_layout) count(%arg0: !hal.device, %arg1: index, %arg2 : index) -> (index, index, index) {
+  hal.executable.export public @vector_distribution_broadcast_dispatch layout(#pipeline_layout) count(%arg0: !hal.device, %arg1: index, %arg2 : index) -> (index, index, index) {
       %x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root %arg1, %arg2
       hal.return %x, %y, %z : index, index, index
     }
   builtin.module {
-    func.func @warp_reduction_broadcast_dispatch() {
+    func.func @vector_distribution_broadcast_dispatch() {
       %c0 = arith.constant 0 : index
       %c1024 = arith.constant 1024 : index
       %cst_0 = arith.constant 3.840000e+02 : f32
@@ -811,15 +811,15 @@ hal.executable.variant @cuda target(<"cuda", "cuda-nvptx-fb">) {
 }
 
 // Check that we generate a group reduce fused with broadcast + elementwise.
-//   CHECK-LABEL: hal.executable public @warp_reduction_broadcast_dispatch
+//   CHECK-LABEL: hal.executable public @vector_distribution_broadcast_dispatch
 //         CHECK:   hal.executable.variant public @cuda
 // CHECK-COUNT-5:     nvvm.shfl.sync  bfly
-//         CHECK:     llvm.store %{{.*}}, %{{.*}} : f32, !llvm.ptr<3>
+//         CHECK:     llvm.store %{{.*}}, %{{.*}} : vector<1xf32>, !llvm.ptr<3>
 //         CHECK:     nvvm.barrier0
-//         CHECK:     llvm.load {{.*}} : !llvm.ptr<3> -> f32
+//         CHECK:     llvm.load {{.*}} : !llvm.ptr<3> -> vector<1xf32>
 // CHECK-COUNT-2:     nvvm.shfl.sync  bfly
 //         CHECK:     llvm.fdiv %{{.*}}, %{{.*}}
-//         CHECK:     llvm.store %{{.*}}, %{{.*}} {alignment = 4 : i64} : vector<4xf32>, !llvm.ptr<1>
+//         CHECK:     llvm.store %{{.*}}, %{{.*}} {alignment = 4 : i64} : vector<4xf32>, !llvm.ptr<3>
 
 // -----
 
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/reduction_pipeline_cuda.mlir b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/reduction_pipeline_cuda.mlir
index 26e3258abeb4..25a7c6d6de47 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/reduction_pipeline_cuda.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/reduction_pipeline_cuda.mlir
@@ -36,28 +36,25 @@ hal.executable.variant @cuda target(<"cuda", "cuda-nvptx-fb">) {
 }
 }
 
-//         CHECK: #[[TRANSLATION_INFO:.+]] = #iree_codegen.translation_info<pipeline = LLVMGPUWarpReduction workgroup_size = [256, 1, 1] subgroup_size = 32>
+//         CHECK: #[[TRANSLATION_INFO:.+]] = #iree_codegen.translation_info<pipeline = LLVMGPUVectorDistribute workgroup_size = [256, 1, 1] subgroup_size = 32
 //         CHECK:  func.func @warp_reduction_dispatch()
 //    CHECK-SAME:      translation_info = #[[TRANSLATION_INFO]]
-//     CHECK-DAG:    %[[CST:.+]] = arith.constant dense<0.000000e+00> : vector<4xf32>
+//     CHECK-DAG:    %[[CST:.+]] = arith.constant dense<0.000000e+00> : vector<1x1x4xf32>
 //     CHECK-DAG:    %[[TID:.+]] = gpu.thread_id  x
-//         CHECK:    %[[R0:.+]] = scf.for %{{.*}} = %c0 to %c10240 step %c1024 iter_args(%[[A0:.+]] = %[[CST]]) -> (vector<4xf32>) {
+//         CHECK:    %[[R0:.+]] = scf.for %{{.*}} = %c0 to %c10240 step %c1024 iter_args(%[[A0:.+]] = %[[CST]]) -> (vector<1x1x4xf32>) {
 //         CHECK:      %[[V:.+]] = vector.transfer_read {{.*}} : memref<512x10240xf32, #hal.descriptor_type<storage_buffer>>, vector<4xf32>
-//         CHECK:      %[[ADD:.+]] = arith.addf %[[V]], %[[A0]] : vector<4xf32>
-//         CHECK:      scf.yield %[[ADD]] : vector<4xf32>
-//         CHECK:    }
-// CHECK-COUNT-5:    gpu.shuffle  xor {{.*}} : f32
-//         CHECK:    %[[ALLOC:.+]] = memref.alloc() : memref<8xf32, #gpu.address_space<workgroup>>
-//         CHECK:    scf.if %{{.*}} {
-//         CHECK:      memref.store %{{.*}}, %[[ALLOC]]{{.*}} : memref<8xf32, #gpu.address_space<workgroup>>
+//         CHECK:      %[[STRIDED:.+]] = vector.insert_strided_slice %[[V]], {{.*}} : vector<4xf32> into vector<1x1x4xf32>
+//         CHECK:      %[[ADD:.+]] = arith.addf %[[STRIDED]], %[[A0]] : vector<1x1x4xf32>
+//         CHECK:      scf.yield %[[ADD]] : vector<1x1x4xf32>
 //         CHECK:    }
+//         CHECK:    gpu.subgroup_reduce  add {{.*}} cluster(size = 32) : (f32) -> f32
+//         CHECK:    %[[ALLOC:.+]] = memref.alloc() : memref<10xf32, #gpu.address_space<workgroup>>
+//         CHECK:    %[[SVIEW:.+]] = memref.subview {{.*}} : memref<10xf32, #gpu.address_space<workgroup>> to memref<8xf32, strided<[1]>, #gpu.address_space<workgroup>>
+//         CHECK:    vector.transfer_write %{{.*}}, %[[SVIEW]]{{.*}} : vector<1xf32>, memref<8xf32, strided<[1]>, #gpu.address_space<workgroup>>
 //         CHECK:    gpu.barrier
-//         CHECK:    memref.load %[[ALLOC]]{{.*}} : memref<8xf32, #gpu.address_space<workgroup>>
-// CHECK-COUNT-3:    gpu.shuffle  xor {{.*}} : f32
-//         CHECK:    gpu.shuffle  idx {{.*}} : f32
-//         CHECK:    scf.if %{{.*}} {
-//         CHECK:      vector.transfer_write {{.*}} : vector<1xf32>, memref<512xf32, #hal.descriptor_type<storage_buffer>>
-//         CHECK:    }
+//         CHECK:    vector.transfer_read %[[SVIEW]]{{.*}} : memref<8xf32,
+//         CHECK:    gpu.subgroup_reduce  add {{.*}} cluster(size = 8) : (f32) -> f32
+//         CHECK:    vector.transfer_write {{.*}} : vector<f32>, memref<512xf32, #hal.descriptor_type<storage_buffer>>
 
 // -----
 
@@ -108,23 +105,18 @@ hal.executable.variant @cuda target(<"cuda", "cuda-nvptx-fb">) {
 }
 }
 
-//         CHECK: #[[TRANSLATION_INFO:.+]] = #iree_codegen.translation_info<pipeline = LLVMGPUWarpReduction workgroup_size = [256, 1, 1] subgroup_size = 32>
+//         CHECK: #[[TRANSLATION_INFO:.+]] = #iree_codegen.translation_info<pipeline = LLVMGPUVectorDistribute workgroup_size = [256, 1, 1] subgroup_size = 32
 //         CHECK:  func.func @warp_reduction_broadcast_dispatch()
 //    CHECK-SAME:      translation_info = #[[TRANSLATION_INFO]]
-//         CHECK:    scf.for {{.*}} -> (vector<4xf32>) {
-//         CHECK:      vector.transfer_read {{.*}} : memref<512x10240xf32, #hal.descriptor_type<storage_buffer>>, vector<4xf32>
-//         CHECK:      arith.addf {{.*}} : vector<4xf32>
+//         CHECK:    scf.for {{.*}} -> (vector<1x1x4xf32>) {
+//         CHECK:      vector.transfer_read {{.*}} : memref<512x10240xf32,
+//         CHECK:      arith.addf {{.*}} : vector<1x1x4xf32>
 //         CHECK:      scf.yield
-// CHECK-COUNT-5:    gpu.shuffle  xor
-//         CHECK:    scf.if
-//         CHECK:      memref.store {{.*}} : memref<8xf32, #gpu.address_space<workgroup>>
-//         CHECK:    }
-// CHECK-COUNT-3:    gpu.shuffle  xor
-//         CHECK:    gpu.shuffle  idx
-//         CHECK:    arith.divf {{.*}} : vector<4xf32>
-//         CHECK:    scf.for
-//         CHECK:      vector.transfer_write {{.*}} : vector<4xf32>, memref<512x10240xf32, #hal.descriptor_type<storage_buffer>>
-//         CHECK:    }
+//         CHECK:    gpu.subgroup_reduce
+//         CHECK:    vector.transfer_write {{.*}} : vector<1xf32
+//         CHECK:    gpu.subgroup_reduce
+//         CHECK:    arith.divf {{.*}} : vector<1x1x4xf32>
+//         CHECK:    vector.transfer_write {{.*}} : vector<4xf32>, {{.*}}
 //         CHECK:    return
 
 // -----
@@ -157,75 +149,31 @@ hal.executable.variant @cuda target(<"cuda", "cuda-nvptx-fb">) {
 }
 }
 
-//         CHECK: #[[TRANSLATION_INFO:.+]] = #iree_codegen.translation_info<pipeline = LLVMGPUWarpReduction workgroup_size = [1024, 1, 1] subgroup_size = 32>
+//         CHECK: #[[TRANSLATION_INFO:.+]] = #iree_codegen.translation_info<pipeline = LLVMGPUVectorDistribute workgroup_size = [1024, 1, 1] subgroup_size = 32
 //         CHECK:  func.func @softmax()
 //    CHECK-SAME:      translation_info = #[[TRANSLATION_INFO]]
-//         CHECK:    scf.for {{.*}} -> (vector<4xf32>) {
-//         CHECK:      vector.transfer_read {{.*}} : memref<12x128x40960xf32, #hal.descriptor_type<storage_buffer>>, vector<4xf32>
-//         CHECK:      arith.maxnumf {{.*}} : vector<4xf32>
+//         CHECK:    scf.for {{.*}} -> (vector<1x1x4xf32>) {
+//         CHECK:      vector.transfer_read {{.*}} : memref<12x128x40960xf32,
+//         CHECK:      arith.maxnumf {{.*}} : vector<1x1x4xf32>
 //         CHECK:      scf.yield
-//         CHECK:    vector.reduction <maxnumf>, %{{.*}} : vector<4xf32> into f32
-//         CHECK:    gpu.shuffle  xor
-//         CHECK:    arith.maxnumf
-//         CHECK:    gpu.shuffle  xor
-//         CHECK:    arith.maxnumf
-//         CHECK:    gpu.shuffle  xor
-//         CHECK:    arith.maxnumf
-//         CHECK:    gpu.shuffle  xor
-//         CHECK:    arith.maxnumf
-//         CHECK:    gpu.shuffle  xor
-//         CHECK:    arith.maxnumf
-//         CHECK:    arith.remui
-//         CHECK:    scf.if
-//         CHECK:      memref.store {{.*}} : memref<32xf32, #gpu.address_space<workgroup>>
-//         CHECK:    }
+//         CHECK:    vector.multi_reduction <maxnumf>
+//         CHECK:    gpu.subgroup_reduce  maxnumf
+//         CHECK:    vector.transfer_write
 //         CHECK:    gpu.barrier
-//         CHECK:    arith.minui
-//         CHECK:    memref.load
-//         CHECK:    gpu.shuffle  xor
-//         CHECK:    arith.maxnumf
-//         CHECK:    gpu.shuffle  xor
-//         CHECK:    arith.maxnumf
-//         CHECK:    gpu.shuffle  xor
-//         CHECK:    arith.maxnumf
-//         CHECK:    gpu.shuffle  xor
-//         CHECK:    arith.maxnumf
-//         CHECK:    gpu.shuffle  xor
-//         CHECK:    arith.maxnumf
-//         CHECK:    vector.broadcast %{{.*}} : f32 to vector<4xf32>
-//         CHECK:    scf.for {{.*}} -> (vector<4xf32>) {
+//         CHECK:    gpu.subgroup_reduce  maxnumf
+//         CHECK:    vector.broadcast %{{.*}} : f32 to vector<1x1x4xf32>
+//         CHECK:    scf.for {{.*}} -> (vector<1x1x4xf32>) {
 //         CHECK:      vector.transfer_read
 //         CHECK:      arith.subf
 //         CHECK:      math.exp
 //         CHECK:      arith.addf
 //         CHECK:      scf.yield
-//         CHECK:    vector.reduction <add>, %{{.*}} : vector<4xf32> into f32
-//         CHECK:    gpu.shuffle  xor
-//         CHECK:    arith.addf
-//         CHECK:    gpu.shuffle  xor
-//         CHECK:    arith.addf
-//         CHECK:    gpu.shuffle  xor
-//         CHECK:    arith.addf
-//         CHECK:    gpu.shuffle  xor
-//         CHECK:    arith.addf
-//         CHECK:    gpu.shuffle  xor
-//         CHECK:    arith.addf
-//         CHECK:    scf.if
-//         CHECK:      memref.store {{.*}} : memref<32xf32, #gpu.address_space<workgroup>>
-//         CHECK:    }
+//         CHECK:    vector.multi_reduction <add>
+//         CHECK:    gpu.subgroup_reduce  add
+//         CHECK:    vector.transfer_write
 //         CHECK:    gpu.barrier
-//         CHECK:    memref.load
-//         CHECK:    gpu.shuffle  xor
-//         CHECK:    arith.addf
-//         CHECK:    gpu.shuffle  xor
-//         CHECK:    arith.addf
-//         CHECK:    gpu.shuffle  xor
-//         CHECK:    arith.addf
-//         CHECK:    gpu.shuffle  xor
-//         CHECK:    arith.addf
-//         CHECK:    gpu.shuffle  xor
-//         CHECK:    arith.addf
-//         CHECK:    arith.addf
+//         CHECK:    vector.transfer_read
+//         CHECK:    gpu.subgroup_reduce  add
 //         CHECK:    vector.broadcast
 //         CHECK:    scf.for
 //         CHECK:      vector.transfer_read
@@ -308,12 +256,7 @@ hal.executable.variant public @cuda_nvptx_fb target(<"cuda", "cuda-nvptx-fb">) {
 }
 
 //   CHECK-LABEL: func.func @group_reduction
-//         CHECK:   %[[RD:.+]] = vector.transfer_read {{.*}} memref<8x64xf32, #hal.descriptor_type<storage_buffer>>, vector<2xf32>
-//         CHECK:   %[[ADD:.+]] = arith.addf %[[RD]]
-//         CHECK:   vector.reduction <add>, %[[ADD]]
-// CHECK-COUNT-5:   gpu.shuffle  xor{{.*}}{{[[:space:]].*}}{{.*}} arith.addf
-//         CHECK:   scf.if
-//         CHECK:     vector.transfer_write {{.*}} memref<8xf32, #hal.descriptor_type<storage_buffer>>
+//       CHECK:    gpu.subgroup_reduce  add {{.*}} cluster(size = 32) : (f32) -> f32
 
 // -----
 
@@ -356,14 +299,7 @@ hal.executable.variant public @cuda_nvptx_fb target(<"cuda", "cuda-nvptx-fb">) {
 }
 
 //   CHECK-LABEL: func.func @group_elementwise_reduction_elementwise
-//         CHECK:   vector.transfer_read {{.*}} vector<2xf32>
-//         CHECK:   arith.addf{{.*}} : vector<2xf32>
-//         CHECK:   arith.addf{{.*}} : vector<2xf32>
-//         CHECK:   arith.addf{{.*}} : vector<2xf32>
-// CHECK-COUNT-5:   gpu.shuffle  xor{{.*}}{{[[:space:]].*}}{{.*}} arith.addf
-//         CHECK:   %[[SQRT_VEC:.+]] = math.sqrt
-//         CHECK:   scf.if
-//         CHECK:     vector.transfer_write %[[SQRT_VEC]], {{.*}} : vector<1xf32>, memref<8xf32, #hal.descriptor_type<storage_buffer>>
+//       CHECK:    gpu.subgroup_reduce  add {{.*}} cluster(size = 32) : (f32) -> f32
 
 // -----
 
@@ -399,18 +335,8 @@ hal.executable.variant public @cuda_nvptx_fb target(<"cuda", "cuda-nvptx-fb">) {
 }
 
 //   CHECK-LABEL: func.func @group_reduction_larger
-// CHECK-COUNT-5: gpu.shuffle  xor{{.*}}{{[[:space:]].*}}{{.*}} arith.addf
-//         CHECK: %[[ALLOC:.*]] = memref.alloc() : memref<8xf32, #gpu.address_space<workgroup>>
-//         CHECK: scf.if
-//         CHECK:   memref.store %{{.*}}, %[[ALLOC]][%{{.*}}] : memref<8xf32, #gpu.address_space<workgroup>>
-//         CHECK: }
-//         CHECK: arith.minui
-//         CHECK: memref.load
-// CHECK-COUNT-3: gpu.shuffle  xor{{.*}}{{[[:space:]].*}}{{.*}} arith.addf
-//         CHECK: %[[RES:.*]], %{{.*}} = gpu.shuffle  idx
-//         CHECK: %[[RES_VEC:.*]] = vector.broadcast %[[RES]] : f32 to vector<1xf32>
-//         CHECK: scf.if
-//         CHECK:   vector.transfer_write %[[RES_VEC]]
+//   CHECK:    gpu.subgroup_reduce  add {{.*}} cluster(size = 32) : (f32) -> f32
+//   CHECK:    gpu.subgroup_reduce  add {{.*}}
 
 // -----
 
@@ -446,7 +372,7 @@ hal.executable.variant public @cuda_nvptx_fb target(<"cuda", "cuda-nvptx-fb">) {
 }
 
 //   CHECK-LABEL: func.func @group_reduction_1d
-// CHECK-COUNT-5: gpu.shuffle  xor{{.*}}{{[[:space:]].*}}{{.*}} arith.addf
+//   CHECK:    gpu.subgroup_reduce  add {{.*}} cluster(size = 32) : (f32) -> f32
 
 // -----
 
@@ -491,7 +417,7 @@ hal.executable.variant public @cuda_nvptx_fb target(<"cuda", "cuda-nvptx-fb">) {
 }
 
 //   CHECK-LABEL: func.func @group_elementwise_reduction_elementwise_4d
-// CHECK-COUNT-5: gpu.shuffle  xor{{.*}}{{[[:space:]].*}}{{.*}} arith.addf
+//   CHECK:    gpu.subgroup_reduce  add {{.*}} cluster(size = 32) : (f32) -> f32
 
 // -----
 
@@ -550,12 +476,8 @@ hal.executable private @i4_dequant_matvec {
 //     CHECK-DAG:   %[[C32:.+]] = arith.constant 32 : index
 //     CHECK-DAG:   %[[C1:.+]] = arith.constant 1 : index
 //     CHECK-DAG:   %[[CST:.+]] = arith.constant dense<0.000000e+00> : vector<1x1x4xf16>
-//         CHECK:   %[[FOR:.+]] = scf.for %{{.+}} = %[[C0]] to %[[C32]] step %[[C1]] iter_args(%{{.*}} = %[[CST]]) -> (vector<1x1x4xf16>)
-//         CHECK:   %{{.*}} = arith.extui %{{.*}} : vector<1x1x4xi4> to vector<1x1x4xi32>
-//         CHECK:   %{{.*}} = arith.uitofp %{{.*}} : vector<1x1x4xi32> to vector<1x1x4xf16>
-//         CHECK:   %{{.*}} = arith.subf %{{.*}}, %{{.*}} : vector<1x1x4xf16>
-//         CHECK:   %{{.*}} = arith.mulf %{{.*}}, %{{.*}} : vector<1x1x4xf16>
-//         CHECK:   %{{.*}} = arith.mulf %{{.*}}, %{{.*}} : vector<1x1x4xf16>
-//         CHECK:   %{{.*}} = arith.addf %{{.*}}, %{{.*}} : vector<1x1x4xf16>
+//         CHECK:   scf.for %{{.+}} = %[[C0]] to %[[C32]] step %[[C1]] iter_args(%{{.*}} = %[[CST]]) -> (vector<1x1x4xf16>)
+//         CHECK:     arith.mulf %{{.*}}, %{{.*}} : vector<1x1x4xf16>
+//         CHECK:     arith.addf %{{.*}}, %{{.*}} : vector<1x1x4xf16>
 
-//         CHECK:   %{{.*}} = vector.multi_reduction <add>, %{{.*}}, %{{.*}} [0, 1, 2] : vector<1x1x4xf16> to f16
+//         CHECK:   vector.multi_reduction <add>, %{{.*}}, %{{.*}} [0, 1, 2] : vector<1x1x4xf16> to f16
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/reduction_pipeline_rocm.mlir b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/reduction_pipeline_rocm.mlir
index b78058a8e2c3..07e21aa3b196 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/reduction_pipeline_rocm.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/reduction_pipeline_rocm.mlir
@@ -36,10 +36,10 @@ hal.executable.variant public @rocm_hsaco_fb target(<"rocm", "rocm-hsaco-fb">) {
 }
 }
 
-//         CDNA3: #[[$TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = LLVMGPUWarpReduction workgroup_size = [32, 1, 1] subgroup_size = 32>
+//         CDNA3: #[[$TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = LLVMGPUVectorDistribute workgroup_size = [32, 1, 1] subgroup_size = 32
 //         CDNA3: func.func @group_reduction_1d()
 //    CDNA3-SAME:    translation_info = #[[$TRANSLATION]]
-// CDNA3-COUNT-5:     gpu.shuffle  xor{{.*}}{{[[:space:]].*}}{{.*}} arith.addf
+//         CDNA3:        gpu.subgroup_reduce  add {{.*}} cluster(size = 32) : (f32) -> f32
 
 // -----
 
@@ -76,9 +76,9 @@ hal.executable.variant public @rocm_hsaco_fb target(<"rocm", "rocm-hsaco-fb">) {
 
 // On CDNA, we prefer wave64 with subgroup size of 64.
 
-//        CHECK: #[[$TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = LLVMGPUWarpReduction workgroup_size = [64, 1, 1] subgroup_size = 64>
+//        CHECK: #[[$TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = LLVMGPUVectorDistribute workgroup_size = [64, 1, 1] subgroup_size = 64
 //        CHECK: func.func @group_reduction_1d
-// CHECK-COUNT-5:     gpu.shuffle  xor{{.*}}{{[[:space:]].*}}{{.*}} arith.addf
+//        CHECK:     gpu.subgroup_reduce  add {{.*}} cluster(size = 64) : (f32) -> f32
 
 // -----
 
@@ -300,7 +300,7 @@ hal.executable private @matvec_fp16 {
 // Multi-row matvec with wave32.
 // TODO(kuhar): We should reduce the number of `gpu.shuffles` performed.
 
-//          CDNA3: #[[$TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = LLVMGPUVectorDistribute workgroup_size = [32, 1, 1] subgroup_size = 32
+//          CDNA3: #[[$TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = LLVMGPUVectorDistribute workgroup_size = [64, 1, 1] subgroup_size = 32
 //          CDNA3: func.func @matvec_fp16()
 //     CDNA3-SAME:     translation_info = #[[$TRANSLATION]]
 //      CDNA3-DAG:   %[[C0:.+]] = arith.constant 0 : index
@@ -379,15 +379,15 @@ hal.executable public @multi_reduction {
 
 // Check that all loops are singly nested.
 //
-//          CHECK: #[[$TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = LLVMGPUWarpReduction workgroup_size = [256, 1, 1] subgroup_size = 64>
+//          CHECK: #[[$TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = LLVMGPUVectorDistribute workgroup_size = [960, 1, 1] subgroup_size = 64
 //          CHECK: func.func @multi_reduction()
 //     CHECK-SAME:     translation_info = #[[$TRANSLATION]]
 //      CHECK-DAG:   %[[C0:.+]] = arith.constant 0 : index
 //      CHECK-DAG:   %[[C1:.+]] = arith.constant 1 : index
-//      CHECK-DAG:   %[[C225:.+]] = arith.constant 225 : index
-//          CHECK:   %[[RES0:.+]] = scf.for %[[ARG0:[a-zA-Z0-9]+]] = %[[C0]] to %[[C225]] step %[[C1]]
-//     CHECK-NEXT:     %[[DELIN:.+]]:2 = affine.delinearize_index %[[ARG0]] into (15, 15) : index, index
-//          CHECK:   %[[RES1:.+]] = scf.for %[[ARG0:[a-zA-Z0-9]+]] = %[[C0]] to %[[C225]] step %[[C1]]
-//     CHECK-NEXT:     %[[DELIN:.+]]:2 = affine.delinearize_index %[[ARG0]] into (15, 15) : index, index
-//          CHECK:   scf.for %[[ARG0:[a-zA-Z0-9]+]] = %[[C0]] to %[[C225]] step %[[C1]]
-//     CHECK-NEXT:     %[[DELIN:.+]]:2 = affine.delinearize_index %[[ARG0]] into (15, 15) : index, index
+//      CHECK-DAG:   %[[C60:.+]] = arith.constant 60 : index
+//          CHECK:   %[[RES0:.+]] = scf.for %[[ARG0:[a-zA-Z0-9]+]] = %[[C0]] to %[[C60]] step %[[C1]]
+//      CHECK-NOT:      scf.for
+//          CHECK:    scf.yield
+//          CHECK:   %[[RES1:.+]] = scf.for %[[ARG0:[a-zA-Z0-9]+]] = %[[C0]] to %[[C60]] step %[[C1]]
+//      CHECK-NOT:      scf.for
+//          CHECK:    scf.yield
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/reduction_pipeline_softmax_rocm.mlir b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/reduction_pipeline_softmax_rocm.mlir
index 548bce71f20d..2d3d78ae2a38 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/reduction_pipeline_softmax_rocm.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/reduction_pipeline_softmax_rocm.mlir
@@ -19,37 +19,21 @@ func.func @softmax() {
   return
 }
 
-//          CHECK: #[[$TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = LLVMGPUWarpReduction workgroup_size = [1024, 1, 1] subgroup_size = 32>
+//          CHECK: #[[$TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = LLVMGPUVectorDistribute workgroup_size = [1024, 1, 1] subgroup_size = 32
 //    CHECK-LABEL: func.func @softmax
 //     CHECK-SAME:     translation_info = #[[$TRANSLATION]]
-// CHECK-COUNT-20:   gpu.shuffle  xor{{.*}}{{[[:space:]].*}}{{.*}}
-
-// -----
-
-#pipeline_layout = #hal.pipeline.layout<bindings = [
-  #hal.pipeline.binding<storage_buffer>,
-  #hal.pipeline.binding<storage_buffer>
-]>
-func.func @softmax() {
-  %c0 = arith.constant 0 : index
-  %cst = arith.constant -3.40282347E+38 : f32
-  %cst_0 = arith.constant 0.000000e+00 : f32
-  %cst_1 = arith.constant 1.000000e+00 : f32
-  %0 = hal.interface.binding.subspan layout(#pipeline_layout) binding(0) alignment(64) offset(%c0) : !iree_tensor_ext.dispatch.tensor<readonly:tensor<12x128x40960xf32>>
-  %1 = hal.interface.binding.subspan layout(#pipeline_layout) binding(1) alignment(64) offset(%c0) : !iree_tensor_ext.dispatch.tensor<writeonly:tensor<12x128x40960xf32>>
-  %2 = iree_tensor_ext.dispatch.tensor.load %0, offsets = [0, 0, 0], sizes = [12, 128, 40960], strides = [1, 1, 1] : !iree_tensor_ext.dispatch.tensor<readonly:tensor<12x128x40960xf32>> -> tensor<12x128x40960xf32>
-  %3 = tensor.empty() : tensor<12x128x40960xf32>
-  %4 = linalg.softmax dimension(2) ins(%2 : tensor<12x128x40960xf32>) outs(%3 : tensor<12x128x40960xf32>) -> tensor<12x128x40960xf32>
-  iree_tensor_ext.dispatch.tensor.store %4, %1, offsets = [0, 0, 0], sizes = [12, 128, 40960], strides = [1, 1, 1] : tensor<12x128x40960xf32> -> !iree_tensor_ext.dispatch.tensor<writeonly:tensor<12x128x40960xf32>>
-  return
-}
+//     CHECK:    gpu.subgroup_reduce  maxnumf {{.*}} cluster(size = 32) : (f32) -> f32
+//     CHECK:    gpu.subgroup_reduce  maxnumf {{.*}} cluster(size = 32) : (f32) -> f32
+//     CHECK:    gpu.subgroup_reduce  add {{.*}} cluster(size = 32) : (f32) -> f32
 
 // On CDNA, we prefer wave64 with subgroup size 64.
 
-//          CDNA3: #[[$TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = LLVMGPUWarpReduction workgroup_size = [1024, 1, 1] subgroup_size = 64>
+//          CDNA3: #[[$TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = LLVMGPUVectorDistribute workgroup_size = [1024, 1, 1] subgroup_size = 64
 //          CDNA3: func.func @softmax
 //     CDNA3-SAME:      translation_info = #[[$TRANSLATION]]
-// CDNA3-COUNT-20:   gpu.shuffle  xor{{.*}}{{[[:space:]].*}}{{.*}}
+//     CDNA3:    gpu.subgroup_reduce  maxnumf {{.*}} cluster(size = 64) : (f32) -> f32
+//     CDNA3:    gpu.subgroup_reduce  maxnumf {{.*}} cluster(size = 16) : (f32) -> f32
+//     CDNA3:    gpu.subgroup_reduce  add {{.*}} cluster(size = 16) : (f32) -> f32
 
 // -----
 


Reviews:
user: login: Groverkss
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-13T09:02:28Z
user: login: pashu123
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-13T13:02:03Z
user: login: Groverkss
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-13T13:05:25Z
user: login: Groverkss
body: LGTM
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-13T13:14:15Z

Review comments:
diff_hunk: @@ -543,6 +545,7 @@ populateConfigInfo(const llvm::SetVector<linalg::LinalgOp> &computeOps,
           return failure();
         }
         setLoweringConfig(linalgOp, *loweringConfig);
+        break;
path: compiler/src/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp
user: login: Groverkss
body: what's happening here?
created_at: 2025-05-13T09:02:05Z
author_association: CONTRIBUTOR
diff_hunk: @@ -543,6 +545,7 @@ populateConfigInfo(const llvm::SetVector<linalg::LinalgOp> &computeOps,
           return failure();
         }
         setLoweringConfig(linalgOp, *loweringConfig);
+        break;
path: compiler/src/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp
user: login: pashu123
body: Making sure that once we attach the lowering config, we can just bail out. Since we are iterating over the operand in the inner most loop.
created_at: 2025-05-13T13:02:03Z
author_association: CONTRIBUTOR
diff_hunk: @@ -543,6 +545,7 @@ populateConfigInfo(const llvm::SetVector<linalg::LinalgOp> &computeOps,
           return failure();
         }
         setLoweringConfig(linalgOp, *loweringConfig);
+        break;
path: compiler/src/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp
user: login: Groverkss
body: Can you add a comment here?
created_at: 2025-05-13T13:05:25Z
author_association: CONTRIBUTOR
diff_hunk: @@ -357,58 +350,60 @@ getVectorDistributeReductionConfig(linalg::LinalgOp op,
 
   SmallVector<int64_t, 4> bounds = op.getStaticLoopRanges();
 
-  SmallVector<int64_t> workgroupTileSizes(op.getNumLoops(), 0);
   SmallVector<int64_t> threadTileSizes(op.getNumLoops(), 0);
-  SmallVector<int64_t> reductionTileSizes(op.getNumLoops(), 0);
   SmallVector<int64_t> threadCounts(op.getNumLoops(), 1);
   SmallVector<int64_t> subGroupCounts(op.getNumLoops(), 1);
   SmallVector<int64_t> mapping(op.getNumLoops());
   std::iota(mapping.begin(), mapping.end(), 0);
 
-  // Set the shared wgp dims to 1.
-  for (int64_t i : sharedWgpDims) {
-    workgroupTileSizes[i] = 1;
-  }
-
   // Set the configuration for the operation with no reduction dims.
+  // The workgroup tile sizes are set by the reduction operation.
   if (reductionDims.empty()) {
+    SmallVector<int64_t> reductionTileSizes(op.getNumLoops(), 1);
+
+    // For the shared wgp dimension, set the reduction tile sizes to be zero.
+    for (auto i : sharedWgpDims) {
+      reductionTileSizes[i] = 0;
+    }
+
     int64_t parallelSize = bounds[parallelDims.back()];
-    if (parallelSize % threadLoads != 0) {
+    if (ShapedType::isDynamic(parallelSize) ||
+        parallelSize % threadLoads != 0) {
       return failure();
     }
-
-    int64_t lastDimWgpTileSize = workgroupSize * threadLoads;
+    int64_t lastDimReductionTileSize = workgroupSize * threadLoads;
 
     // Setting subgroupBasis to minimum i.e., 1 and threadBasis
     // to maximum i.e., subgroupSize.
     int64_t subgroupBasis = 1;
     int64_t threadBasis = subgroupSize;
 
-    lastDimWgpTileSize = llvm::APIntOps::GreatestCommonDivisor(
-                             {64, static_cast<uint64_t>(parallelSize)},
-                             {64, static_cast<uint64_t>(lastDimWgpTileSize)})
-                             .getZExtValue();
+    lastDimReductionTileSize =
+        llvm::APIntOps::GreatestCommonDivisor(
+            {64, static_cast<uint64_t>(parallelSize)},
+            {64, static_cast<uint64_t>(lastDimReductionTileSize)})
+            .getZExtValue();
 
     if (!(sharedWgpDims.size() == op.getNumLoops())) {
-      while (lastDimWgpTileSize % (threadBasis * threadLoads) != 0) {
+      while (lastDimReductionTileSize % (threadBasis * threadLoads) != 0) {
         threadBasis >>= 1;
       }
-      subgroupBasis = lastDimWgpTileSize / (threadLoads * threadBasis) == 0
-                          ? 1
-                          : lastDimWgpTileSize / (threadLoads * threadBasis);
+      subgroupBasis =
+          lastDimReductionTileSize / (threadLoads * threadBasis) == 0
+              ? 1
+              : lastDimReductionTileSize / (threadLoads * threadBasis);
path: compiler/src/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp
user: login: Groverkss
body: can you split this on multiple lines?
created_at: 2025-05-13T13:07:05Z
author_association: CONTRIBUTOR

PR comments:
user: login: benvanik
created_at: 2025-05-08T00:48:20Z
body: Every kernel, huh?
(Please provide useful descriptions on PRs ;)
author_association: COLLABORATOR
user: login: pashu123
created_at: 2025-05-08T01:35:29Z
body: > Every kernel, huh? (Please provide useful descriptions on PRs ;)

Apologies, I was just testing on the CI.
author_association: CONTRIBUTOR
user: login: benvanik
created_at: 2025-05-08T02:04:55Z
body: (that's cool - in the future, please note that everyone watching the repository sees these and they are enshrined in the repo forever - write _something_ :)
author_association: COLLABORATOR
user: login: benvanik
created_at: 2025-05-13T22:28:40Z
body: going to try a rerun, but fyi this _may have_ caused a regression:
https://github.com/iree-org/iree/actions/runs/15007320793/job/42169379980
```
ERROR iree-test-suites/sharktank_models/benchmarks/model_benchmark_run.py::sdxl :: vae_rocm - check 69.94073186302559 <= 68.2: sdxl vae benchmark time should not regress more than a factor of 1.1
```
author_association: COLLABORATOR
user: login: pashu123
created_at: 2025-05-13T23:07:24Z
body: > going to try a rerun, but fyi this _may have_ caused a regression: https://github.com/iree-org/iree/actions/runs/15007320793/job/42169379980
> 
> ```
> ERROR iree-test-suites/sharktank_models/benchmarks/model_benchmark_run.py::sdxl :: vae_rocm - check 69.94073186302559 <= 68.2: sdxl vae benchmark time should not regress more than a factor of 1.1
> ```
Sure, if it persists, I'll revert.
author_association: CONTRIBUTOR
user: login: benvanik
created_at: 2025-05-13T23:13:06Z
body: Looks like it's right on the edge - subsequent merges on main are failing: https://github.com/iree-org/iree/actions/runs/15007458464/job/42169770312

you may need to up the tolerance (or figure out why it regressed)
author_association: COLLABORATOR
user: login: pashu123
created_at: 2025-05-13T23:20:43Z
body: > Looks like it's right on the edge - subsequent merges on main are failing: https://github.com/iree-org/iree/actions/runs/15007458464/job/42169770312
> 
> you may need to up the tolerance (or figure out why it regressed)

The latest one is passing though: https://github.com/iree-org/iree/actions/runs/15007816465/job/42171017802
author_association: CONTRIBUTOR
user: login: pashu123
created_at: 2025-05-13T23:24:55Z
body: > > Looks like it's right on the edge - subsequent merges on main are failing: https://github.com/iree-org/iree/actions/runs/15007458464/job/42169770312
> > you may need to up the tolerance (or figure out why it regressed)
> 
> The latest one is passing though: https://github.com/iree-org/iree/actions/runs/15007816465/job/42171017802

This PR also goes through: https://github.com/iree-org/iree/actions/runs/15007320793/job/42172666060 . The CI would be flaky.
author_association: CONTRIBUTOR
user: login: pashu123
created_at: 2025-05-13T23:26:00Z
body: Oh! I missed it. Yes, it's right on edge. 
author_association: CONTRIBUTOR
user: login: MaheshRavishankar
created_at: 2025-05-14T01:27:19Z
body: Actually looks like only VAE regressed. Would be good to triage that, but almost everything else got better. Maybe we should change the golden time for others as well
author_association: COLLABORATOR
user: login: benvanik
created_at: 2025-05-14T01:30:54Z
body: yeah, if timings went down the lower so we don't backslide!
author_association: COLLABORATOR