[Mlir-commits] [mlir] [mlir][gpu] Derive has-redux from the chip in the NVVM pipeline (PR #221393)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Alessandro Potenza (alepot55)

<details>
<summary>Changes</summary>

gpu-lower-to-nvvm-pipeline cannot lower a gpu.subgroup_reduce at all. The only pattern for it in convert-gpu-to-nvvm is the one that emits nvvm.redux.sync, and that pattern is added only when the has-redux option is set, which the pipeline does not expose and therefore never sets. The pipeline runs no shuffle expansion either, so the operation reaches the conversion target illegal and the pipeline fails.

Nothing about that option is a decision of its own: redux.sync exists from sm_80, and the pipeline already knows the chip it serializes for. So derive it from cubin-chip. Below sm_80 the pipeline still cannot lower the operation, for the second reason above; I left that alone here.

The integration test fails without the change (the pipeline errors out on the illegal op) and checks for redux.sync.add.s32 with it, on sm_80.

Assisted-by: Claude Code


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


2 Files Affected:

- (modified) mlir/lib/Dialect/GPU/Pipelines/GPUToNVVMPipeline.cpp (+9) 
- (added) mlir/test/Integration/GPU/CUDA/subgroup-reduce-redux.mlir (+20) 


``````````diff
diff --git a/mlir/lib/Dialect/GPU/Pipelines/GPUToNVVMPipeline.cpp b/mlir/lib/Dialect/GPU/Pipelines/GPUToNVVMPipeline.cpp
index 5462cddd44718..6629d8b065b71 100644
--- a/mlir/lib/Dialect/GPU/Pipelines/GPUToNVVMPipeline.cpp
+++ b/mlir/lib/Dialect/GPU/Pipelines/GPUToNVVMPipeline.cpp
@@ -30,6 +30,7 @@
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Pass/PassOptions.h"
 #include "mlir/Transforms/Passes.h"
+#include "llvm/ADT/StringExtras.h"
 
 using namespace mlir;
 
@@ -73,6 +74,14 @@ void buildGpuPassPipeline(OpPassManager &pm,
   opt.useBarePtrCallConv = options.kernelUseBarePtrCallConv;
   opt.indexBitwidth = options.indexBitWidth;
   opt.allowPatternRollback = options.allowPatternRollback;
+  // redux.sync exists from sm_80, and the chip this pipeline serializes for is
+  // the only thing that decides whether it may be used, so there is no reason
+  // to ask for it as an option of its own.
+  StringRef chip(options.cubinChip);
+  unsigned smVersion = 0;
+  if (chip.consume_front("sm_"))
+    chip.take_while(llvm::isDigit).getAsInteger(10, smVersion);
+  opt.hasRedux = smVersion >= 80;
   pm.addNestedPass<gpu::GPUModuleOp>(createConvertGpuOpsToNVVMOps(opt));
   pm.addNestedPass<gpu::GPUModuleOp>(createCanonicalizerPass());
   pm.addNestedPass<gpu::GPUModuleOp>(createCSEPass());
diff --git a/mlir/test/Integration/GPU/CUDA/subgroup-reduce-redux.mlir b/mlir/test/Integration/GPU/CUDA/subgroup-reduce-redux.mlir
new file mode 100644
index 0000000000000..c3d0770c162c1
--- /dev/null
+++ b/mlir/test/Integration/GPU/CUDA/subgroup-reduce-redux.mlir
@@ -0,0 +1,20 @@
+// RUN: mlir-opt %s \
+// RUN:  | mlir-opt -gpu-lower-to-nvvm-pipeline="cubin-chip=sm_80 cubin-format=isa" \
+// RUN:  | FileCheck %s
+
+// A subgroup reduction the source states is executed by the whole subgroup
+// becomes one redux.sync instruction, once the chip the pipeline serializes
+// for is one that has it. Below sm_80 the pipeline still cannot lower the
+// operation at all, since it runs no shuffle expansion either.
+
+// CHECK: redux.sync.add.s32
+
+gpu.module @kernels {
+  gpu.func @reduce(%out: memref<32xi32>) kernel {
+    %tid = gpu.thread_id x
+    %v = arith.index_cast %tid : index to i32
+    %r = gpu.subgroup_reduce add %v uniform : (i32) -> (i32)
+    memref.store %r, %out[%tid] : memref<32xi32>
+    gpu.return
+  }
+}

``````````

</details>


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