[Mlir-commits] [mlir] [OpenACC] Support dynamic worker-private shared memory (PR #210770)

[llvmlistbot at llvm.org]

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

Use the maximum ThreadY width when the runtime worker count prevents exact static sizing.

>From b72758c6ac4c5eff9bff3548ccf56ff648582cb8 Mon Sep 17 00:00:00 2001
From: Kazuaki Matsumura <kmatsumura at nvidia.com>
Date: Mon, 20 Jul 2026 10:45:35 -0700
Subject: [PATCH] [OpenACC] Support dynamic worker-private shared memory

Use the maximum ThreadY width when the runtime worker count prevents exact static sizing.
---
 .../Dialect/OpenACC/Utils/OpenACCUtilsCG.cpp  | 20 ++++++++-----------
 ...c-cg-to-gpu-worker-private-dynamic-nw.mlir | 11 +++++-----
 .../Dialect/OpenACC/OpenACCUtilsCGTest.cpp    | 16 ++++++++++-----
 3 files changed, 25 insertions(+), 22 deletions(-)

diff --git a/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsCG.cpp b/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsCG.cpp
index 5f4af6a5bcc5b..745b6fa348533 100644
--- a/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsCG.cpp
+++ b/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsCG.cpp
@@ -317,9 +317,10 @@ collectPrivateLocalParDims(PrivateLocalOp privateLocal,
   return parDims;
 }
 
-static FailureOr<std::optional<int64_t>> getWorkerPrivateSharedMemoryNumCopies(
-    PrivateLocalOp privateLocal, ComputeRegionOp computeRegion,
-    bool isWorkerPrivate, OpenACCSupport *support) {
+static FailureOr<std::optional<int64_t>>
+getWorkerPrivateSharedMemoryNumCopies(PrivateLocalOp privateLocal,
+                                      ComputeRegionOp computeRegion,
+                                      bool isWorkerPrivate, OpenACCSupport *) {
   if (!isWorkerPrivate)
     return std::optional<int64_t>(1);
 
@@ -330,15 +331,10 @@ static FailureOr<std::optional<int64_t>> getWorkerPrivateSharedMemoryNumCopies(
     return std::optional<int64_t>();
 
   auto workerArgConst = workerArg->getDefiningOp<arith::ConstantIndexOp>();
-  if (!workerArgConst) {
-    if (support) {
-      (void)support->emitNYI(privateLocal.getLoc(),
-                             "worker-private variables in shared memory "
-                             "require compile-time constant num_workers");
-      return failure();
-    }
-    return std::optional<int64_t>();
-  }
+  // ThreadY is bounded by 32 after subgroup alignment. Use that upper bound
+  // when the exact worker count is known only at runtime.
+  if (!workerArgConst)
+    return std::optional<int64_t>(32);
   return std::optional<int64_t>(workerArgConst.value());
 }
 
diff --git a/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-worker-private-dynamic-nw.mlir b/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-worker-private-dynamic-nw.mlir
index 5df3d986fa5b0..3120f8dc675e3 100644
--- a/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-worker-private-dynamic-nw.mlir
+++ b/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-worker-private-dynamic-nw.mlir
@@ -1,7 +1,10 @@
-// RUN: mlir-opt %s --pass-pipeline="builtin.module(func.func(acc-cg-to-gpu))" --verify-diagnostics
+// RUN: mlir-opt %s --pass-pipeline="builtin.module(func.func(acc-cg-to-gpu))" | FileCheck %s
 
-// Verify that worker-private shared memory with dynamic num_workers emits
-// a diagnostic instead of silently miscompiling.
+// Dynamic num_workers uses the maximum ThreadY width for shared memory.
+
+// CHECK-LABEL: func.func @test_worker_private_dynamic_nw
+// CHECK: acc.gpu_shared_memory
+// CHECK-SAME: num_copies = 32
 
 func.func @test_worker_private_dynamic_nw(%nw: index) {
   %c32 = arith.constant 32 : index
@@ -11,13 +14,11 @@ func.func @test_worker_private_dynamic_nw(%nw: index) {
   %thread_x = acc.par_width %c32 {par_dim = #acc.par_dim<thread_x>}
   acc.kernel_environment {
     %priv = acc.privatize : () -> !acc.private_type<memref<2xi32>>
-    // expected-error @below {{failed to legalize operation 'acc.compute_region'}}
     acc.compute_region launch(%arg0 = %block_x, %arg1 = %thread_y, %arg2 = %thread_x) ins(%arg10 = %priv) : (!acc.private_type<memref<2xi32>>) {
       %c0 = arith.constant 0 : index
       %c1 = arith.constant 1 : index
       %c0_i32 = arith.constant 0 : i32
       scf.parallel (%iv) = (%c0) to (%arg1) step (%c1) {
-        // expected-error @below {{worker-private variables in shared memory require compile-time constant num_workers}}
         %local = acc.private_local %arg10 : (!acc.private_type<memref<2xi32>>) -> memref<2xi32>
         memref.store %c0_i32, %local[%c0] : memref<2xi32>
         scf.reduce
diff --git a/mlir/unittests/Dialect/OpenACC/OpenACCUtilsCGTest.cpp b/mlir/unittests/Dialect/OpenACC/OpenACCUtilsCGTest.cpp
index 69f529abf6fa6..9bf78556bf265 100644
--- a/mlir/unittests/Dialect/OpenACC/OpenACCUtilsCGTest.cpp
+++ b/mlir/unittests/Dialect/OpenACC/OpenACCUtilsCGTest.cpp
@@ -793,7 +793,7 @@ TEST_F(OpenACCUtilsCGTest, getSharedMemoryBytesGangWorkerReductionAccumulator) {
 }
 
 TEST_F(OpenACCUtilsCGTest,
-       isPrivateLocalSharedMemoryCandidateWorkerPrivateDynamicFails) {
+       isPrivateLocalSharedMemoryCandidateWorkerPrivateDynamicUsesUpperBound) {
   OwningOpRef<ModuleOp> module = ModuleOp::create(b, loc);
   b.setInsertionPointToStart(module->getBody());
   GPUParallelDimsAttr workerDims = GPUParallelDimsAttr::get(
@@ -801,8 +801,7 @@ TEST_F(OpenACCUtilsCGTest,
   auto c1 = arith::ConstantIndexOp::create(b, loc, 1);
   auto bx =
       ParWidthOp::create(b, loc, c1, GPUParallelDimAttr::blockXDim(&context));
-  // A non-constant num_workers: the launch operand exists but is not an
-  // arith.constant, which is what triggers the diagnostic / failure path.
+  // A non-constant num_workers uses the maximum ThreadY width.
   auto dynNumWorkers = arith::AddIOp::create(b, loc, c1, c1);
   auto ty = ParWidthOp::create(b, loc, dynNumWorkers,
                                GPUParallelDimAttr::threadYDim(&context));
@@ -818,11 +817,18 @@ TEST_F(OpenACCUtilsCGTest,
   FailureOr<bool> silent =
       isPrivateLocalSharedMemoryCandidate(privateLocal, cr, *module, policy);
   ASSERT_TRUE(succeeded(silent));
-  EXPECT_FALSE(*silent);
+  EXPECT_TRUE(*silent);
 
   FailureOr<bool> diagnosed = isPrivateLocalSharedMemoryCandidate(
       privateLocal, cr, *module, policy, &support);
-  EXPECT_TRUE(failed(diagnosed));
+  ASSERT_TRUE(succeeded(diagnosed));
+  EXPECT_TRUE(*diagnosed);
+
+  std::optional<int64_t> upperBound =
+      getPrivateLocalSharedMemoryUpperBoundBytes(privateLocal, cr, *module,
+                                                 policy);
+  ASSERT_TRUE(upperBound.has_value());
+  EXPECT_EQ(*upperBound, 512);
 }
 
 TEST_F(OpenACCUtilsCGTest, getPrivateLocalSharedMemoryUpperBoundBytes) {