[Mlir-commits] [mlir] [mlir][gpu] Store the all-reduce result from the first lane only (PR #221532)

[Alessandro Potenza]

https://github.com/alepot55 created https://github.com/llvm/llvm-project/pull/221532

The final step of the gpu.all_reduce lowering reduces the per-subgroup partial results with createSubgroupReduce, whose contract says that only the first lane holds the result and every other lane holds an undefined value. All the reducing lanes then stored to the same workgroup buffer slot, so with a number of subgroups that is not a power of two (three warps, say) lanes that missed a butterfly step wrote a wrong partial to the same address as the correct one, and the answer depended on which store the hardware let win. This guards the store with isFirstLane, the way the first phase of the lowering already does.

Found with an executable lock-step semantics of the LLVM dialect that reports two lanes of one store writing different values to the same address: on 501 generated kernels the only reports were the 7 whose block had three warps, all at this store. On current NVIDIA hardware the surviving write happens to be the right one, so the bug is latent there, but it is undefined by the lowering's own contract. The two autogenerated tests are updated by hand, as the last upstream change to them did, since generate-test-checks.py now emits a different naming style.

Assisted-by: Claude Code


>From 03b5247470ddb765e6ea7ebce12961789f74bebb Mon Sep 17 00:00:00 2001
From: Alessandro Potenza <ap.alessandro.potenza at gmail.com>
Date: Sat, 5 Sep 2026 16:49:34 +0200
Subject: [PATCH] [mlir][gpu] Store the all-reduce result from the first lane
 only

The final step of the gpu.all_reduce lowering reduces the per-subgroup partial results with createSubgroupReduce, whose contract says that only the first lane holds the result and every other lane holds an undefined value. All the reducing lanes then stored to the same workgroup buffer slot, so with a number of subgroups that is not a power of two (three warps, say) lanes that missed a butterfly step wrote a wrong partial to the same address as the correct one and the answer depended on which store the hardware let win. Guard the store with isFirstLane like the first phase already does.
---
 .../lib/Dialect/GPU/Transforms/AllReduceLowering.cpp |  8 ++++++--
 mlir/test/Dialect/GPU/all-reduce-add.mlir            | 12 ++++++++----
 mlir/test/Dialect/GPU/all-reduce-maxf.mlir           | 12 ++++++++----
 3 files changed, 22 insertions(+), 10 deletions(-)

diff --git a/mlir/lib/Dialect/GPU/Transforms/AllReduceLowering.cpp b/mlir/lib/Dialect/GPU/Transforms/AllReduceLowering.cpp
index d76c3194f025e..e1956ea982357 100644
--- a/mlir/lib/Dialect/GPU/Transforms/AllReduceLowering.cpp
+++ b/mlir/lib/Dialect/GPU/Transforms/AllReduceLowering.cpp
@@ -124,14 +124,18 @@ struct GpuAllReduceRewriter {
 
     // Use the first numSubgroups invocations to reduce the intermediate results
     // from workgroup memory. The final result is written to workgroup memory
-    // again.
+    // again, by the first lane only: the subgroup reduction leaves the result
+    // in the first lane and an undefined value in every other lane, so letting
+    // all of them store would make the outcome depend on which lane's store
+    // wins.
     Value zero = create<arith::ConstantIndexOp>(0);
     createPredicatedBlock(isValidSubgroup, [&] {
       Value index = create<arith::IndexCastOp>(indexType, invocationIdx);
       Value value = create<memref::LoadOp>(valueType, buffer, index);
       Value result =
           createSubgroupReduce(numSubgroups, laneId, value, accumFactory);
-      create<memref::StoreOp>(result, buffer, zero);
+      createPredicatedBlock(
+          isFirstLane, [&] { create<memref::StoreOp>(result, buffer, zero); });
     });
 
     // Synchronize workgroup and load result from workgroup memory.
diff --git a/mlir/test/Dialect/GPU/all-reduce-add.mlir b/mlir/test/Dialect/GPU/all-reduce-add.mlir
index bca9ea3f50bd5..ec25b52228954 100644
--- a/mlir/test/Dialect/GPU/all-reduce-add.mlir
+++ b/mlir/test/Dialect/GPU/all-reduce-add.mlir
@@ -108,7 +108,7 @@ gpu.module @kernels {
     // CHECK:   [[VAL_72:%.*]] = arith.addi [[VAL_28]], [[VAL_2]] : i32
     // CHECK:   [[VAL_73:%.*]] = arith.divsi [[VAL_72]], [[VAL_5]] : i32
     // CHECK:   [[VAL_74:%.*]] = arith.cmpi slt, [[VAL_27]], [[VAL_73]] : i32
-    // CHECK:   cf.cond_br [[VAL_74]], ^bb22, ^bb41
+    // CHECK:   cf.cond_br [[VAL_74]], ^bb22, ^bb43
     // CHECK: ^bb22:
     // CHECK:   [[VAL_75:%.*]] = arith.index_cast [[VAL_27]] : i32 to index
     // CHECK:   [[VAL_76:%.*]] = memref.load [[VAL_1]]{{\[}}[[VAL_75]]] : memref<32xf32, #gpu.address_space<workgroup>>
@@ -169,11 +169,15 @@ gpu.module @kernels {
     // CHECK:   [[VAL_112:%.*]] = arith.addf [[VAL_109]], [[VAL_110]] : f32
     // CHECK:   cf.br ^bb40([[VAL_112]] : f32)
     // CHECK: ^bb40([[VAL_113:%.*]]: f32):
-    // CHECK:   store [[VAL_113]], [[VAL_1]]{{\[}}[[VAL_4]]] : memref<32xf32, #gpu.address_space<workgroup>>
-    // CHECK:   cf.br ^bb42
+    // CHECK:   cf.cond_br [[VAL_30]], ^bb41, ^bb42
     // CHECK: ^bb41:
-    // CHECK:   cf.br ^bb42
+    // CHECK:   store [[VAL_113]], [[VAL_1]]{{\[}}[[VAL_4]]] : memref<32xf32, #gpu.address_space<workgroup>>
+    // CHECK:   cf.br ^bb43
     // CHECK: ^bb42:
+    // CHECK:   cf.br ^bb43
+    // CHECK: ^bb43:
+    // CHECK:   cf.br ^bb44
+    // CHECK: ^bb44:
     // CHECK:   gpu.barrier memfence [#gpu.address_space<workgroup>]
     %sum = gpu.all_reduce add %arg0 uniform {} : (f32) -> (f32)
     gpu.return
diff --git a/mlir/test/Dialect/GPU/all-reduce-maxf.mlir b/mlir/test/Dialect/GPU/all-reduce-maxf.mlir
index 9c32249b2230f..f368b02d4f339 100644
--- a/mlir/test/Dialect/GPU/all-reduce-maxf.mlir
+++ b/mlir/test/Dialect/GPU/all-reduce-maxf.mlir
@@ -108,7 +108,7 @@ gpu.module @kernels {
     // CHECK:   [[VAL_82:%.*]] = arith.addi [[VAL_28]], [[VAL_2]] : i32
     // CHECK:   [[VAL_83:%.*]] = arith.divsi [[VAL_82]], [[VAL_5]] : i32
     // CHECK:   [[VAL_84:%.*]] = arith.cmpi slt, [[VAL_27]], [[VAL_83]] : i32
-    // CHECK:   cf.cond_br [[VAL_84]], ^bb22, ^bb41
+    // CHECK:   cf.cond_br [[VAL_84]], ^bb22, ^bb43
     // CHECK: ^bb22:
     // CHECK:   [[VAL_85:%.*]] = arith.index_cast [[VAL_27]] : i32 to index
     // CHECK:   [[VAL_86:%.*]] = memref.load [[VAL_1]]{{\[}}[[VAL_85]]] : memref<32xf32, #gpu.address_space<workgroup>>
@@ -169,11 +169,15 @@ gpu.module @kernels {
     // CHECK:   [[VAL_132:%.*]] = arith.maxnumf [[VAL_128]], [[VAL_129]] : f32
     // CHECK:   cf.br ^bb40([[VAL_132]] : f32)
     // CHECK: ^bb40([[VAL_133:%.*]]: f32):
-    // CHECK:   store [[VAL_133]], [[VAL_1]]{{\[}}[[VAL_4]]] : memref<32xf32, #gpu.address_space<workgroup>>
-    // CHECK:   cf.br ^bb42
+    // CHECK:   cf.cond_br [[VAL_30]], ^bb41, ^bb42
     // CHECK: ^bb41:
-    // CHECK:   cf.br ^bb42
+    // CHECK:   store [[VAL_133]], [[VAL_1]]{{\[}}[[VAL_4]]] : memref<32xf32, #gpu.address_space<workgroup>>
+    // CHECK:   cf.br ^bb43
     // CHECK: ^bb42:
+    // CHECK:   cf.br ^bb43
+    // CHECK: ^bb43:
+    // CHECK:   cf.br ^bb44
+    // CHECK: ^bb44:
     // CHECK:   gpu.barrier memfence [#gpu.address_space<workgroup>]
     %sum = gpu.all_reduce maxnumf %arg0 uniform {} : (f32) -> (f32)
     gpu.return