[Mlir-commits] [mlir] [mlir][OpenACC] Keep ThreadY active for inner-combine-fed worker reductions (PR #211696)

[llvmlistbot at llvm.org]

https://github.com/khaki3 updated https://github.com/llvm/llvm-project/pull/211696

>From 9c99893bd0cf15ccb8e7e7afb9c21a8b4266822e Mon Sep 17 00:00:00 2001
From: Kazuaki Matsumura <kmatsumura at nvidia.com>
Date: Thu, 23 Jul 2026 16:41:01 -0700
Subject: [PATCH 1/2] [mlir][OpenACC] Keep ThreadY active for inner-combine-fed
 worker reductions

A block_y + thread_y reduction accumulator that is filled by an inner
block-scoped acc.reduction_combine holds one distinct partial per worker
row (each row's shared slot is written by the inner combine). Reducing it
into the destination must therefore run on every worker row. The previous
classification kept ThreadY active only for "proven worker-private"
combines (isThreadYPrivate with hasPrivateDest), which excludes the
block_y+thread_y -> grid/global shape produced by e.g. a CUF
`kernel do(2)` reduction, so it fell back to the row-zero legacy path and
dropped the other workers' partials (result -1 instead of -4 for a 2D SUM).

Add a discriminator: a block_y+thread_y accumulator keeps ThreadY active
when it is fed by an inner block-scoped combine (distinct per-worker
partials). A plain worker accumulate, which lowers to a worker-wide
all_reduce that already broadcasts the total, correctly stays row-zero to
avoid multiplying by the worker count.
---
 .../Dialect/OpenACC/Transforms/ACCCGToGPU.cpp | 41 +++++++++++++++++++
 1 file changed, 41 insertions(+)

diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
index 51db77332c102..ede68126f4cae 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
@@ -1240,6 +1240,32 @@ static bool hasUnsafeEffectsWhenBroadening(Operation *op) {
   return !op->hasTrait<OpTrait::HasRecursiveMemoryEffects>();
 }
 
+/// True when \p accumulator is the destination of a separate block-scoped
+/// acc.reduction_combine (an inner-level combine that fills its per-worker
+/// shared slot), so it holds one distinct partial per worker row rather than a
+/// worker-wide broadcast.
+static bool isFedByInnerBlockCombine(acc::PrivateLocalOp accumulator,
+                                     Operation *selfCombine) {
+  if (!accumulator)
+    return false;
+  for (Operation *user : accumulator.getResult().getUsers()) {
+    if (user == selfCombine)
+      continue;
+    auto combineOp = dyn_cast<acc::ReductionCombineOp>(user);
+    if (!combineOp ||
+        unwrapMemRefConversion(combineOp.getDestMemref()).getDefiningOp() !=
+            accumulator.getOperation())
+      continue;
+    SmallVector<mlir::acc::GPUParallelDimAttr> parDims =
+        getReductionCombineParDims(combineOp);
+    if (llvm::any_of(parDims, [](mlir::acc::GPUParallelDimAttr d) {
+          return d.isAnyBlock();
+        }))
+      return true;
+  }
+  return false;
+}
+
 /// Records whether \p combineOp requires ThreadY to remain active.
 static void classifyThreadYCombine(ThreadYBroadeningInfo &info,
                                    Operation *combineOp, Value src, Value dest,
@@ -1263,6 +1289,21 @@ static void classifyThreadYCombine(ThreadYBroadeningInfo &info,
     return;
   }
 
+  // A block_y + thread_y accumulator that is itself fed by an inner
+  // block-scoped combine holds one distinct partial per worker row (each row's
+  // shared slot was filled by the inner combine). Reducing it into the
+  // destination must therefore run on every worker row; predicating it to
+  // ThreadY row zero would drop the other workers' partials. This differs from
+  // a plain worker accumulate, which lowers to a worker-wide all_reduce that
+  // already broadcasts the total, and so must stay row-zero to avoid
+  // multiplying by the worker count.
+  if (hasThreadY && hasBlock &&
+      isThreadYPrivate(srcPrivate, /*allowBlock=*/true, computeRegion) &&
+      isFedByInnerBlockCombine(srcPrivate, combineOp)) {
+    info.hasActiveWorkerCombine = true;
+    return;
+  }
+
   info.hasExplicitInactiveCombine = true;
   if (!info.diagnosticOp)
     info.diagnosticOp = combineOp;

>From 090ae0ed6eea7f170da9ece392e8d28bb718e49c Mon Sep 17 00:00:00 2001
From: Kazuaki Matsumura <kmatsumura at nvidia.com>
Date: Thu, 23 Jul 2026 22:08:49 -0700
Subject: [PATCH 2/2] [mlir][OpenACC] NFC: shorten worker-reduction comment

---
 .../lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp | 15 ++++-----------
 1 file changed, 4 insertions(+), 11 deletions(-)

diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
index ede68126f4cae..969b782a7460c 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
@@ -1241,9 +1241,7 @@ static bool hasUnsafeEffectsWhenBroadening(Operation *op) {
 }
 
 /// True when \p accumulator is the destination of a separate block-scoped
-/// acc.reduction_combine (an inner-level combine that fills its per-worker
-/// shared slot), so it holds one distinct partial per worker row rather than a
-/// worker-wide broadcast.
+/// combine, so it holds a distinct per-worker partial rather than a broadcast.
 static bool isFedByInnerBlockCombine(acc::PrivateLocalOp accumulator,
                                      Operation *selfCombine) {
   if (!accumulator)
@@ -1289,14 +1287,9 @@ static void classifyThreadYCombine(ThreadYBroadeningInfo &info,
     return;
   }
 
-  // A block_y + thread_y accumulator that is itself fed by an inner
-  // block-scoped combine holds one distinct partial per worker row (each row's
-  // shared slot was filled by the inner combine). Reducing it into the
-  // destination must therefore run on every worker row; predicating it to
-  // ThreadY row zero would drop the other workers' partials. This differs from
-  // a plain worker accumulate, which lowers to a worker-wide all_reduce that
-  // already broadcasts the total, and so must stay row-zero to avoid
-  // multiplying by the worker count.
+  // A block_y+thread_y accumulator fed by an inner block-scoped combine holds
+  // a distinct partial per worker row, so its combine runs on every row; a
+  // plain worker accumulate broadcasts via all_reduce and stays row-zero.
   if (hasThreadY && hasBlock &&
       isThreadYPrivate(srcPrivate, /*allowBlock=*/true, computeRegion) &&
       isFedByInnerBlockCombine(srcPrivate, combineOp)) {