[Mlir-commits] [mlir] [mlir][OpenACC] Support dynamic multi-rank array reductions (PR #211336)

[llvmlistbot at llvm.org]

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

Example:
```fortran
integer :: r(:,:)
!$acc parallel loop reduction(+:r)
```

In this code, the reduction accumulator has dynamic dimensions and rank greater than one.

Fix: use runtime dimension sizes to delinearize flattened array-reduction indices.

>From 7e8a8a0c197a2a15a4c7d7997bedbca889b84769 Mon Sep 17 00:00:00 2001
From: Kazuaki Matsumura <kmatsumura at nvidia.com>
Date: Wed, 22 Jul 2026 10:09:13 -0700
Subject: [PATCH] [mlir][OpenACC] Support dynamic multi-rank array reductions

---
 .../Dialect/OpenACC/Transforms/ACCCGToGPU.cpp | 16 +++++------
 .../acc-cg-to-gpu-reduction-array.mlir        | 27 ++++++++++++-------
 2 files changed, 24 insertions(+), 19 deletions(-)

diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
index f0fd7ae96ad4f..92f9cc90dc73f 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
@@ -3259,12 +3259,6 @@ void ACCCGToGPULowering::processAccumulateArrayOp(
     (void)accSupport.emitNYI(loc, "reduction: non-MemRefTy accumulate array");
     return;
   }
-  if (memrefTy.getRank() > 1 && !memrefTy.hasStaticShape()) {
-    (void)accSupport.emitNYI(loc,
-                             "reduction: dynamic multi-rank accumulate array");
-    return;
-  }
-
   FailureOr<arith::AtomicRMWKind> kindOr = getReductionKind(
       op.getReductionOperator(), memrefTy.getElementType(), loc);
   if (failed(kindOr))
@@ -3373,13 +3367,15 @@ void ACCCGToGPULowering::processAccumulateArrayOp(
     Value iv = forOp.getInductionVar();
     SmallVector<Value> indices{iv};
     if (memrefTy.getRank() > 1) {
-      assert(memrefTy.hasStaticShape() &&
-             "multi-rank array reduction accumulator must be static");
       indices.resize(memrefTy.getRank());
       Value linearIndex = iv;
       for (int64_t dim = memrefTy.getRank() - 1; dim >= 0; --dim) {
-        Value dimSize = arith::ConstantIndexOp::create(
-            rewriter, loc, memrefTy.getDimSize(dim));
+        Value dimSize =
+            memrefTy.isDynamicDim(dim)
+                ? memref::DimOp::create(rewriter, loc, memref, dim).getResult()
+                : arith::ConstantIndexOp::create(rewriter, loc,
+                                                 memrefTy.getDimSize(dim))
+                      .getResult();
         indices[dim] =
             arith::RemUIOp::create(rewriter, loc, linearIndex, dimSize);
         if (dim != 0)
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 ace560d869792..422e172e1ec5b 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
@@ -196,20 +196,29 @@ func.func @rank_three_array_reduction() {
   return
 }
 
-// Unsupported dynamic multi-rank accumulators must not generate invalid
-// indexing operations.
 // CHECK-LABEL: func.func @dynamic_rank_two_array_reduction
-// CHECK: gpu.launch
-// CHECK-NOT: memref.load
-// CHECK: gpu.terminator
-func.func @dynamic_rank_two_array_reduction(%local: memref<?x?xi32>, %extent: index) {
+// CHECK: scf.for %[[LINEAR:.*]] =
+// CHECK: %[[DIM1:.*]] = memref.dim %{{.*}}, %{{.*}} : memref<2x?xi32>
+// CHECK: %[[IDX1:.*]] = arith.remui %[[LINEAR]], %[[DIM1]] : index
+// CHECK: %[[ROW:.*]] = arith.divui %[[LINEAR]], %[[DIM1]] : index
+// CHECK: %[[C2:.*]] = arith.constant 2 : index
+// CHECK: %[[IDX0:.*]] = arith.remui %[[ROW]], %[[C2]] : index
+// CHECK: %[[ELEM:.*]] = memref.load %{{.*}}[%[[IDX0]], %[[IDX1]]] : memref<2x?xi32>
+// CHECK: %[[RESULT:.*]] = gpu.all_reduce add %[[ELEM]]
+// CHECK: memref.store %[[RESULT]], %{{.*}}[%[[IDX0]], %[[IDX1]]] : memref<2x?xi32>
+func.func @dynamic_rank_two_array_reduction(
+    %local: memref<2x?xi32>, %n: index) {
   %c1 = arith.constant 1 : index
   %c128 = arith.constant 128 : index
   %bx = acc.par_width %c1 {par_dim = #acc.par_dim<block_x>}
   %tx = acc.par_width %c128 {par_dim = #acc.par_dim<thread_x>}
-  acc.compute_region launch(%kbx = %bx, %ktx = %tx) ins(%arg0 = %local, %ext = %extent) : (memref<?x?xi32>, index) {
-    %bounds = acc.bounds extent(%ext : index)
-    acc.reduction_accumulate_array %arg0 bounds(%bounds) <add> : memref<?x?xi32> {par_dims = #acc<par_dims[block_x, thread_x]>}
+  acc.compute_region launch(%kbx = %bx, %ktx = %tx)
+      ins(%arg0 = %local, %arg1 = %n)
+      : (memref<2x?xi32>, index) {
+    %c2 = arith.constant 2 : index
+    %extent = arith.muli %c2, %arg1 : index
+    %bounds = acc.bounds extent(%extent : index)
+    acc.reduction_accumulate_array %arg0 bounds(%bounds) <add> : memref<2x?xi32> {par_dims = #acc<par_dims[block_x, thread_x]>}
     acc.yield
   } {origin = "acc.parallel"}
   return