[Mlir-commits] [mlir] [mlir][gpu][nvvm] Add subgroup_reduce shuffle fallback for clustered and non-i32 cases (PR #209098)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Johnny Lin (johnny19436)

<details>
<summary>Changes</summary>

`gpu.subgroup_reduce` lowering in GPUToNVVM currently relies on the `nvvm.redux.sync`
fast path and rejects clustered reductions and non-i32 payloads.

This patch keeps the existing redux fast path unchanged for supported uniform i32
cases, and adds a fallback path for unsupported uniform subgroup reductions by
reusing existing GPU subgroup-reduce rewrite patterns:
- `populateGpuBreakDownSubgroupReducePatterns`
- `populateGpuLowerSubgroupReduceToShufflePatterns`
- `populateGpuLowerClusteredSubgroupReduceToShufflePatterns`

The fallback is applied before LLVM conversion and lowers these cases to
`gpu.shuffle`-based reductions, which are then legalized by existing NVVM
shuffle lowering.

Behavior after this patch:
- uniform + i32 + non-clustered + redux-supported op: still lowers to `nvvm.redux.sync`
- uniform + non-i32: now lowers via shuffle fallback
- uniform + clustered: now lowers via shuffle fallback
- non-uniform subgroup_reduce behavior is unchanged in this patch

Tests:
- Updated `mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir` with focused cases for:
  - existing redux fast path
  - non-i32 fallback
  - clustered fallback

Commands run:
- `mlir-opt mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir -convert-gpu-to-nvvm='has-redux=1' -split-input-file | FileCheck mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir`
- `mlir-opt mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir -convert-gpu-to-nvvm='has-redux=1 allow-pattern-rollback=0' -split-input-file | FileCheck mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir`

---
Full diff: https://github.com/llvm/llvm-project/pull/209098.diff


2 Files Affected:

- (modified) mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp (+35) 
- (modified) mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir (+22) 


``````````diff
diff --git a/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp b/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp
index 80420c26537c3..59b150deb546d 100644
--- a/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp
+++ b/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp
@@ -96,6 +96,14 @@ convertToNVVMReductionKind(gpu::AllReduceOperation mode) {
   return std::nullopt;
 }
 
+static bool canLowerSubgroupReduceToNVVMRedux(gpu::SubgroupReduceOp op) {
+  if (op.getClusterSize() || !op.getUniform())
+    return false;
+  if (!op.getValue().getType().isInteger(32))
+    return false;
+  return convertToNVVMReductionKind(op.getOp()).has_value();
+}
+
 static constexpr llvm::StringLiteral kNVVMNamedBarrierIdPrefix =
     "__named_barrier_id";
 static constexpr int32_t kNVVMFirstNamedBarrierId = 1;
@@ -529,6 +537,33 @@ struct LowerGpuOpsToNVVMOpsPass final
     // ops which need to be lowered further, which is not supported by a
     // single conversion pass.
     {
+      // Lower subgroup reductions that cannot use nvvm.redux to shuffles
+      // before conversion. Keep redux-compatible cases untouched so the
+      // dedicated conversion pattern still applies.
+      SmallVector<Operation *> subgroupReduceOpsToLower;
+      m.walk([&](gpu::SubgroupReduceOp op) {
+        if (!op.getUniform())
+          return;
+        if (!this->hasRedux || !canLowerSubgroupReduceToNVVMRedux(op))
+          subgroupReduceOpsToLower.push_back(op.getOperation());
+      });
+      if (!subgroupReduceOpsToLower.empty()) {
+        RewritePatternSet subgroupReducePatterns(m.getContext());
+        populateGpuBreakDownSubgroupReducePatterns(
+            subgroupReducePatterns, /*maxShuffleBitwidth=*/kNVVMWarpSize);
+        populateGpuLowerSubgroupReduceToShufflePatterns(
+            subgroupReducePatterns,
+            /*subgroupSize=*/kNVVMWarpSize,
+            /*shuffleBitwidth=*/kNVVMWarpSize);
+        populateGpuLowerClusteredSubgroupReduceToShufflePatterns(
+            subgroupReducePatterns,
+            /*subgroupSize=*/kNVVMWarpSize,
+            /*shuffleBitwidth=*/kNVVMWarpSize);
+        if (failed(applyOpPatternsGreedily(subgroupReduceOpsToLower,
+                                           std::move(subgroupReducePatterns))))
+          return signalPassFailure();
+      }
+
       RewritePatternSet patterns(m.getContext());
       populateGpuRewritePatterns(patterns);
       // Transform N-D vector.from_elements to 1-D vector.from_elements before
diff --git a/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir b/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir
index b96069ac41a44..8c76285b3e47e 100644
--- a/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir
+++ b/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir
@@ -739,6 +739,28 @@ gpu.module @test_module_30 {
   }
 }
 
+gpu.module @test_module_30_fallback {
+  // CHECK-LABEL: func @subgroup_reduce_add_f32_fallback
+  gpu.func @subgroup_reduce_add_f32_fallback(%arg0 : f32, %buf : memref<f32>) {
+    // CHECK-NOT: nvvm.redux.sync
+    // CHECK: nvvm.shfl.sync bfly
+    %result = gpu.subgroup_reduce add %arg0 uniform {} : (f32) -> (f32)
+    memref.store %result, %buf[] : memref<f32>
+    gpu.return
+  }
+
+  // CHECK-LABEL: func @subgroup_reduce_clustered_i32_fallback
+  gpu.func @subgroup_reduce_clustered_i32_fallback(%arg0 : i32,
+                                                   %buf : memref<i32>) {
+    // CHECK-NOT: nvvm.redux.sync
+    // CHECK: nvvm.shfl.sync bfly
+    %result = gpu.subgroup_reduce add %arg0 uniform cluster(size = 8) :
+        (i32) -> (i32)
+    memref.store %result, %buf[] : memref<i32>
+    gpu.return
+  }
+}
+
 gpu.module @test_module_31 {
   // CHECK: llvm.func @__nv_fmodf(f32, f32) -> f32
   // CHECK: llvm.func @__nv_fmod(f64, f64) -> f64

``````````

</details>


https://github.com/llvm/llvm-project/pull/209098