[Mlir-commits] [flang] [mlir] [mlir][acc] Replicate privatized storage across active parallel dims (PR #235974)

[llvmlistbot at llvm.org]

https://github.com/yebinchon updated https://github.com/llvm/llvm-project/pull/235974

>From b0ecd6aa11a8bb5ed6b1e877f25e61fe7aaf3b52 Mon Sep 17 00:00:00 2001
From: Yebin Chon <ychon at nvidia.com>
Date: Mon, 26 Oct 2026 17:50:50 -0700
Subject: [PATCH 1/2] Replicate privatized storage across active parallel dims

---
 .../acc-fir-map-info-prep-privatize.mlir      | 33 +++++++++++++++++++
 .../Dialect/OpenACC/Utils/OpenACCUtilsCG.cpp  | 22 +++++++++----
 2 files changed, 48 insertions(+), 7 deletions(-)

diff --git a/flang/test/Fir/OpenACC/acc-fir-map-info-prep-privatize.mlir b/flang/test/Fir/OpenACC/acc-fir-map-info-prep-privatize.mlir
index 9c0395f01d6ef0..3e24175984626d 100644
--- a/flang/test/Fir/OpenACC/acc-fir-map-info-prep-privatize.mlir
+++ b/flang/test/Fir/OpenACC/acc-fir-map-info-prep-privatize.mlir
@@ -35,6 +35,39 @@ func.func @private_parallel_levels() {
 
 // -----
 
+// Ownership par_dims can be narrower than the active dimensions the kernel
+// indexes the storage by, e.g. a vector reduction temporary allocated inside a
+// gang loop. The storage is replicated across both.
+
+// CHECK-LABEL: func.func @private_active_wider_than_par_dims
+// CHECK: %[[PRIV:.*]] = acc.privatize
+// CHECK: acc.map_info varPtr(%[[PRIV]] : !acc.private_type<memref<8xi32>>)
+// CHECK-SAME: mapFlags(private,gang_private,vector_private)
+func.func @private_active_wider_than_par_dims() {
+  %priv = acc.privatize par_dims(#acc<par_dims[thread_x]>)
+      {acc.active_par_dims = #acc<active_par_dims[block_x, thread_x]>}
+      : () -> !acc.private_type<memref<8xi32>>
+  return
+}
+
+// -----
+
+// Without par_dims the storage is not replicated per level, even when the op
+// has active dimensions.
+
+// CHECK-LABEL: func.func @private_active_without_par_dims
+// CHECK: %[[PRIV:.*]] = acc.privatize
+// CHECK: acc.map_info varPtr(%[[PRIV]] : !acc.private_type<memref<8xi32>>)
+// CHECK-SAME: mapFlags(private)
+func.func @private_active_without_par_dims() {
+  %priv = acc.privatize
+      {acc.active_par_dims = #acc<active_par_dims[block_x, thread_x]>}
+      : () -> !acc.private_type<memref<8xi32>>
+  return
+}
+
+// -----
+
 // Privatized storage need not name any parallel dimension - storage promoted to
 // shared memory is private without being replicated per level. Such a
 // privatization carries no par_dims at all, not an empty one.
diff --git a/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsCG.cpp b/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsCG.cpp
index 7c90f09e4b37b0..59e50240e03015 100644
--- a/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsCG.cpp
+++ b/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsCG.cpp
@@ -631,13 +631,21 @@ MapFlags computePrivatizeMapFlags(PrivatizeOp privatizeOp,
   if (!parDims)
     return flags;
 
-  for (GPUParallelDimAttr parDim : parDims.getArray()) {
-    if (policy.isGang(parDim))
-      flags = flags | MapFlags::gang_private;
-    else if (policy.isWorker(parDim))
-      flags = flags | MapFlags::worker_private;
-    else if (policy.isVector(parDim))
-      flags = flags | MapFlags::vector_private;
+  auto addLevelFlags = [&](ArrayRef<GPUParallelDimAttr> dims) {
+    for (GPUParallelDimAttr parDim : dims) {
+      if (policy.isGang(parDim))
+        flags = flags | MapFlags::gang_private;
+      else if (policy.isWorker(parDim))
+        flags = flags | MapFlags::worker_private;
+      else if (policy.isVector(parDim))
+        flags = flags | MapFlags::vector_private;
+    }
+  };
+  addLevelFlags(parDims.getArray());
+  // Ownership par_dims may be narrower than the dimensions the kernel indexes
+  // the storage by, so replicate across the active dimensions as well.
+  if(ActiveParDimsAttr activeParDims = getActiveParDimsAttr(privatizeOp)) {
+    addLevelFlags(activeParDims.getArray());
   }
   return flags;
 }

>From 6968bf91f16d0adb5b3ace869eaca0920b48588a Mon Sep 17 00:00:00 2001
From: Yebin Chon <ychon at nvidia.com>
Date: Mon, 26 Oct 2026 18:05:14 -0700
Subject: [PATCH 2/2] format

---
 mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsCG.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsCG.cpp b/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsCG.cpp
index 59e50240e03015..0b6dc7a166dbb6 100644
--- a/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsCG.cpp
+++ b/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsCG.cpp
@@ -644,7 +644,7 @@ MapFlags computePrivatizeMapFlags(PrivatizeOp privatizeOp,
   addLevelFlags(parDims.getArray());
   // Ownership par_dims may be narrower than the dimensions the kernel indexes
   // the storage by, so replicate across the active dimensions as well.
-  if(ActiveParDimsAttr activeParDims = getActiveParDimsAttr(privatizeOp)) {
+  if (ActiveParDimsAttr activeParDims = getActiveParDimsAttr(privatizeOp)) {
     addLevelFlags(activeParDims.getArray());
   }
   return flags;