[mlir][OpenACC] Preserve worker rows when combining reductions

[khaki3]

Example:

!$acc parallel loop gang worker reduction(+:sum)
do j = 1, n
  !$acc loop vector reduction(+:sum)
  do i = 1, n
    sum = sum + a(i, j)
  end do
end do

In this code, each worker has a private partial result, but combine predication previously allowed only ThreadY row zero to contribute.

Fix: keep ThreadY active only for proven atomic worker combines, while preserving legacy predication or reporting NYI for unsafe combinations.

[llvmorg-github-actions]

@llvm/pr-subscribers-openacc
@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-openacc

Author: Matsu (khaki3)

Changes

Example:

!$cuf kernel do(3) <<< *,(32,8,1) >>> reduce(+:res)
do k = 1, n3
  do j = 1, n2
    do i = 1, n1
      res = res + a(i,j,k)
    end do
  end do
end do

In this code, each ThreadY worker row produces a partial result, but the final combine was restricted to threadIdx.y == 0, dropping the other rows.

Fix: keep ThreadY active when combining worker-private reductions and reject mixed scopes requiring incompatible predication.

---

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

3 Files Affected:

• (modified) mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp (+43-10)
• (added) mlir/test/Dialect/OpenACC/acc-cg-to-gpu-worker-reduction-combine-mixed-scope.mlir (+43)
• (added) mlir/test/Dialect/OpenACC/acc-cg-to-gpu-worker-reduction-combine.mlir (+95)

diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
index 24aeed6e23bda..6be79b0dd40e4 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
@@ -1246,7 +1246,30 @@ ACCCGToGPULowering::computeActiveAndInactiveParDims(Operation *op,
   mlir::acc::GPUParallelDimAttr lowestParDim =
       mlir::acc::GPUParallelDimAttr::threadXDim(ctx);
   if (block) {
-    block->walk([&](Operation *op) {
+    std::optional<bool> combineThreadYActive;
+    auto applyCombineParDims =
+        [&](Operation *combineOp, Value src,
+            ArrayRef<mlir::acc::GPUParallelDimAttr> combineParDims) {
+          bool isWorkerPrivate =
+              getPrivateScopeForMemref(src) == PrivateMemScope::Worker;
+          for (mlir::acc::GPUParallelDimAttr parDim : combineParDims) {
+            if (parDim.isThreadY()) {
+              if (combineThreadYActive &&
+                  *combineThreadYActive != isWorkerPrivate) {
+                combineOp->emitError()
+                    << "mixed worker-private and non-worker-private reduction "
+                       "combines require incompatible ThreadY predication";
+                hasFailed = true;
+                return failure();
+              }
+              combineThreadYActive = isWorkerPrivate;
+            } else {
+              mlir::acc::removeParDim(ancestorParDims, parDim);
+            }
+          }
+          return success();
+        };
+    block->walk([&](Operation *op) -> WalkResult {
       // Check stores to acc.private_local - add the privatize's par_dims
       // as active dims so predication is correct for per-worker/gang memory.
       if (memref::StoreOp storeOp = dyn_cast<memref::StoreOp>(op)) {
@@ -1277,17 +1300,17 @@ ACCCGToGPULowering::computeActiveAndInactiveParDims(Operation *op,
       // kernel and loop in combined constructs.
       if (acc::ReductionCombineOp reductionCombineOp =
               dyn_cast<acc::ReductionCombineOp>(op)) {
-        for (mlir::acc::GPUParallelDimAttr parDim :
-             getReductionCombineParDims(reductionCombineOp)) {
-          mlir::acc::removeParDim(ancestorParDims, parDim);
-        }
+        if (failed(applyCombineParDims(
+                reductionCombineOp, reductionCombineOp.getSrcMemref(),
+                getReductionCombineParDims(reductionCombineOp))))
+          return WalkResult::interrupt();
       }
       if (acc::ReductionCombineRegionOp combineRegionOp =
               dyn_cast<acc::ReductionCombineRegionOp>(op)) {
-        for (mlir::acc::GPUParallelDimAttr parDim :
-             getReductionCombineParDims(combineRegionOp)) {
-          mlir::acc::removeParDim(ancestorParDims, parDim);
-        }
+        if (failed(applyCombineParDims(
+                combineRegionOp, combineRegionOp.getSrcVar(),
+                getReductionCombineParDims(combineRegionOp))))
+          return WalkResult::interrupt();
       }
       // An array accumulate reduces across its par_dims via gpu.all_reduce, so
       // all those threads must execute it - treat them as active (unlike the
@@ -1301,6 +1324,14 @@ ACCCGToGPULowering::computeActiveAndInactiveParDims(Operation *op,
       }
       return WalkResult::advance();
     });
+    if (combineThreadYActive) {
+      mlir::acc::GPUParallelDimAttr threadY =
+          mlir::acc::GPUParallelDimAttr::threadYDim(ctx);
+      if (*combineThreadYActive)
+        mlir::acc::insertParDim(ancestorParDims, threadY);
+      else
+        mlir::acc::removeParDim(ancestorParDims, threadY);
+    }
   }
 
   // Obtain launch dimensions
@@ -1747,7 +1778,7 @@ ACCCGToGPULowering::getPrivateMemScope(acc::PrivatizeOp privatizeOp) {
   }
   if (hasThreadX)
     return PrivateMemScope::Thread;
-  if (hasBlock && hasThreadY)
+  if (hasThreadY)
     return PrivateMemScope::Worker;
   if (hasBlock)
     return PrivateMemScope::Gang;
@@ -1871,6 +1902,8 @@ void ACCCGToGPULowering::processPredicateRegion(
             SmallVector<mlir::acc::GPUParallelDimAttr>>
       parDimsPair = computeActiveAndInactiveParDims(
           interOp, &interOp.getRegion().front());
+  if (hasFailed)
+    return;
 
   // If ThreadY reduction exists, subgroup alignment is applied
   // (blockDim.x = subgroupSize), so ThreadX lanes exist even without explicit
diff --git a/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-worker-reduction-combine-mixed-scope.mlir b/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-worker-reduction-combine-mixed-scope.mlir
new file mode 100644
index 0000000000000..a0d381740c034
--- /dev/null
+++ b/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-worker-reduction-combine-mixed-scope.mlir
@@ -0,0 +1,43 @@
+// RUN: mlir-opt %s --pass-pipeline="builtin.module(func.func(acc-cg-to-gpu))" \
+// RUN:   -verify-diagnostics
+
+func.func @mixed_scope_worker_reduction_combine(
+    %other: memref<i32>, %result: memref<i32>) {
+  %c1 = arith.constant 1 : index
+  %c4 = arith.constant 4 : index
+  %c32 = arith.constant 32 : index
+  %block_y = acc.par_width %c1 {par_dim = #acc.par_dim<block_y>}
+  %thread_y = acc.par_width %c4 {par_dim = #acc.par_dim<thread_y>}
+  %thread_x = acc.par_width %c32 {par_dim = #acc.par_dim<thread_x>}
+  acc.kernel_environment {
+    %private = acc.privatize [#acc<par_dims[thread_y]>]
+        : () -> !acc.private_type<memref<i32>>
+    // expected-error@+1 {{failed to legalize operation 'acc.compute_region' that was explicitly marked illegal}}
+    acc.compute_region launch(%by = %block_y, %ty = %thread_y, %tx = %thread_x)
+        ins(%private_arg = %private, %other_arg = %other,
+            %result_arg = %result)
+        : (!acc.private_type<memref<i32>>, memref<i32>, memref<i32>) {
+      %c0 = arith.constant 0 : index
+      %c1_inner = arith.constant 1 : index
+      %c0_i32 = arith.constant 0 : i32
+      scf.parallel (%block_iv) = (%c0) to (%by) step (%c1_inner) {
+        %local = acc.private_local %private_arg
+            : (!acc.private_type<memref<i32>>) -> memref<i32>
+        scf.parallel (%worker_iv) = (%c0) to (%ty) step (%c1_inner) {
+          memref.store %c0_i32, %local[] : memref<i32>
+          scf.reduce
+        } {acc.par_dims = #acc<par_dims[thread_y]>}
+        acc.predicate_region {
+          acc.reduction_combine %local into %result_arg <add> : memref<i32>
+              {acc.par_dims = #acc<par_dims[block_y, thread_y]>}
+          // expected-error@+1 {{mixed worker-private and non-worker-private reduction combines require incompatible ThreadY predication}}
+          acc.reduction_combine %other_arg into %result_arg <add> : memref<i32>
+              {acc.par_dims = #acc<par_dims[block_y, thread_y]>}
+        }
+        scf.reduce
+      } {acc.par_dims = #acc<par_dims[block_y]>}
+      acc.yield
+    } {origin = "acc.parallel"}
+  }
+  return
+}
diff --git a/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-worker-reduction-combine.mlir b/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-worker-reduction-combine.mlir
new file mode 100644
index 0000000000000..b2fd246d33d1b
--- /dev/null
+++ b/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-worker-reduction-combine.mlir
@@ -0,0 +1,95 @@
+// RUN: mlir-opt %s --pass-pipeline="builtin.module(func.func(acc-cg-to-gpu))" | FileCheck %s
+
+// A worker-private reduction has one shared slot per ThreadY row. The combine
+// must keep ThreadY active and predicate only ThreadX.
+
+// CHECK-LABEL: func.func @worker_reduction_combine
+// CHECK: gpu.launch {{.*}} threads([[TID_X:%[^,]+]], [[TID_Y:%[^,]+]],
+// CHECK-NOT: arith.cmpi eq, [[TID_Y]]
+// CHECK: %[[IS_X_ZERO:.*]] = arith.cmpi eq, [[TID_X]],
+// CHECK-NOT: arith.andi
+// CHECK: scf.if %[[IS_X_ZERO]]
+// CHECK: acc.atomic.update
+
+func.func @worker_reduction_combine(%result: memref<i32>) {
+  %c1 = arith.constant 1 : index
+  %c4 = arith.constant 4 : index
+  %c32 = arith.constant 32 : index
+  %block_y = acc.par_width %c1 {par_dim = #acc.par_dim<block_y>}
+  %thread_y = acc.par_width %c4 {par_dim = #acc.par_dim<thread_y>}
+  %thread_x = acc.par_width %c32 {par_dim = #acc.par_dim<thread_x>}
+  acc.kernel_environment {
+    %private = acc.privatize [#acc<par_dims[thread_y]>]
+        : () -> !acc.private_type<memref<i32>>
+    acc.compute_region launch(%by = %block_y, %ty = %thread_y, %tx = %thread_x)
+        ins(%private_arg = %private, %result_arg = %result)
+        : (!acc.private_type<memref<i32>>, memref<i32>) {
+      %c0 = arith.constant 0 : index
+      %c1_inner = arith.constant 1 : index
+      %c0_i32 = arith.constant 0 : i32
+      scf.parallel (%block_iv) = (%c0) to (%by) step (%c1_inner) {
+        %local = acc.private_local %private_arg
+            : (!acc.private_type<memref<i32>>) -> memref<i32>
+        scf.parallel (%worker_iv) = (%c0) to (%ty) step (%c1_inner) {
+          memref.store %c0_i32, %local[] : memref<i32>
+          scf.reduce
+        } {acc.par_dims = #acc<par_dims[thread_y]>}
+        acc.predicate_region {
+          acc.reduction_combine %local into %result_arg <add> : memref<i32>
+              {acc.par_dims = #acc<par_dims[block_y, thread_y]>}
+        }
+        scf.reduce
+      } {acc.par_dims = #acc<par_dims[block_y]>}
+      acc.yield
+    } {origin = "acc.parallel"}
+  }
+  return
+}
+
+// CHECK-LABEL: func.func @worker_reduction_combine_region
+// CHECK: gpu.launch {{.*}} threads([[REGION_TID_X:%[^,]+]], [[REGION_TID_Y:%[^,]+]],
+// CHECK-NOT: arith.cmpi eq, [[REGION_TID_Y]]
+// CHECK: %[[REGION_IS_X_ZERO:.*]] = arith.cmpi eq, [[REGION_TID_X]],
+// CHECK-NOT: arith.andi
+// CHECK: scf.if %[[REGION_IS_X_ZERO]]
+// CHECK: arith.addi
+
+func.func @worker_reduction_combine_region(%result: memref<i32>) {
+  %c1 = arith.constant 1 : index
+  %c4 = arith.constant 4 : index
+  %c32 = arith.constant 32 : index
+  %block_y = acc.par_width %c1 {par_dim = #acc.par_dim<block_y>}
+  %thread_y = acc.par_width %c4 {par_dim = #acc.par_dim<thread_y>}
+  %thread_x = acc.par_width %c32 {par_dim = #acc.par_dim<thread_x>}
+  acc.kernel_environment {
+    %private = acc.privatize [#acc<par_dims[thread_y]>]
+        : () -> !acc.private_type<memref<i32>>
+    acc.compute_region launch(%by = %block_y, %ty = %thread_y, %tx = %thread_x)
+        ins(%private_arg = %private, %result_arg = %result)
+        : (!acc.private_type<memref<i32>>, memref<i32>) {
+      %c0 = arith.constant 0 : index
+      %c1_inner = arith.constant 1 : index
+      %c0_i32 = arith.constant 0 : i32
+      scf.parallel (%block_iv) = (%c0) to (%by) step (%c1_inner) {
+        %local = acc.private_local %private_arg
+            : (!acc.private_type<memref<i32>>) -> memref<i32>
+        scf.parallel (%worker_iv) = (%c0) to (%ty) step (%c1_inner) {
+          memref.store %c0_i32, %local[] : memref<i32>
+          scf.reduce
+        } {acc.par_dims = #acc<par_dims[thread_y]>}
+        acc.predicate_region {
+          acc.reduction_combine_region %local into %result_arg : memref<i32> {
+            %lhs = memref.load %result_arg[] : memref<i32>
+            %rhs = memref.load %local[] : memref<i32>
+            %sum = arith.addi %lhs, %rhs : i32
+            memref.store %sum, %result_arg[] : memref<i32>
+            acc.yield
+          } {acc.par_dims = #acc<par_dims[block_y, thread_y]>}
+        }
+        scf.reduce
+      } {acc.par_dims = #acc<par_dims[block_y]>}
+      acc.yield
+    } {origin = "acc.parallel"}
+  }
+  return
+}

[razvanlupusoru]

Thank you.