[Mlir-commits] [mlir] [mlir][OpenACC] Lower single-block thread-only array reductions (PR #212369)

[llvmlistbot at llvm.org]

https://github.com/khaki3 created https://github.com/llvm/llvm-project/pull/212369

Example:
```fortran
!$acc parallel loop vector reduction(+:b)
do i = 1, n
  b(i) = i
end do
```

In this code the array reduction is thread(vector)-only with no gang/block launch dim, so a within-block `all_reduce` is a complete reduction. ACCCGToGPU still rejected it as NYI because the accumulate had no block context.

Fix: keep the NYI only when the region actually launches block dimensions; allow the existing per-element `all_reduce` path for single-block thread-only array accumulates.


>From 453cf9156e9d2a9a70ec18825255d2cfe3a55a22 Mon Sep 17 00:00:00 2001
From: Kazuaki Matsumura <kmatsumura at nvidia.com>
Date: Mon, 27 Jul 2026 14:26:37 -0700
Subject: [PATCH] [mlir][OpenACC] Lower single-block thread-only array
 reductions

A within-block all_reduce is a complete reduction when the region
launches no block dimension, so allow thread-only array accumulate in
that case instead of NYI.
---
 .../Dialect/OpenACC/Transforms/ACCCGToGPU.cpp | 10 +++++--
 .../acc-cg-to-gpu-reduction-array.mlir        | 30 +++++++++++++++++++
 2 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
index d50d346ed9204..0d30f54f12447 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
@@ -3512,9 +3512,13 @@ void ACCCGToGPULowering::processAccumulateArrayOp(
     return;
   }
 
-  // A thread-level reduction with no block owner for its elements cannot merge
-  // the cross-thread partials, so report NYI.
-  if (!reductionHasBlockContext(op)) {
+  // A thread-only accumulate merges with a within-block all_reduce, which is
+  // complete only in one block: block context, or a launch with no block dim.
+  // Multi-block thread-only still grid-strides across blocks, so stays NYI.
+  bool regionLaunchesBlocks = llvm::any_of(
+      computeRegion.getLaunchParDims(),
+      [](mlir::acc::GPUParallelDimAttr d) { return d.isAnyBlock(); });
+  if (!reductionHasBlockContext(op) && regionLaunchesBlocks) {
     (void)accSupport.emitNYI(
         loc, "reduction: thread-only array reduction accumulate");
     return;
diff --git a/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-array.mlir b/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-array.mlir
index 93354321865a7..64eaec673dd1c 100644
--- a/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-array.mlir
+++ b/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-array.mlir
@@ -347,3 +347,33 @@ func.func @thread_x_reduction_with_thread_z_width() {
   } {origin = "acc.parallel"}
   return
 }
+
+// Thread-only array accumulate is well-defined when the region launches no
+// block dim (single block): a within-block all_reduce is a complete reduction.
+// CHECK-LABEL: func.func @thread_only_array_reduction_single_block
+// CHECK: gpu.launch
+// CHECK-NOT: acc.reduction_accumulate_array
+// CHECK: scf.for %[[IV:.*]] = %{{.*}} to %{{.*}} step %{{.*}} {
+// CHECK:   %[[ELT:.*]] = memref.load %[[ALLOCA:.*]][%[[IV]]] : memref<8xi32>
+// CHECK:   %[[RED:.*]] = gpu.all_reduce add %[[ELT]]
+// CHECK:   memref.store %[[RED]], %[[ALLOCA]][%[[IV]]] : memref<8xi32>
+// CHECK: }
+func.func @thread_only_array_reduction_single_block() {
+  %c128 = arith.constant 128 : index
+  %tx = acc.par_width %c128 {par_dim = #acc.par_dim<thread_x>}
+  acc.compute_region launch(%ktx = %tx) {
+    %c0 = arith.constant 0 : index
+    %c1 = arith.constant 1 : index
+    %c8 = arith.constant 8 : index
+    %c0_i32 = arith.constant 0 : i32
+    %local = memref.alloca() : memref<8xi32>
+    scf.for %i = %c0 to %c8 step %c1 {
+      memref.store %c0_i32, %local[%i] : memref<8xi32>
+    }
+    %bounds = acc.bounds extent(%c8 : index)
+    acc.reduction_accumulate_array %local bounds(%bounds) <add>
+        : memref<8xi32> {par_dims = #acc<par_dims[thread_x]>}
+    acc.yield
+  } {origin = "acc.parallel"}
+  return
+}