[Mlir-commits] [mlir] [mlir][openacc] Keep vector-private storage per-lane in ACCCGToGPU (PR #222134)
Delaram Talaashrafi
llvmlistbot at llvm.org
Wed Sep 9 12:03:11 PDT 2026
https://github.com/delaram-talaashrafi updated https://github.com/llvm/llvm-project/pull/222134
>From 89846e38a4088b65ebc628edac9766880cab582a Mon Sep 17 00:00:00 2001
From: Delaram Talaashrafi <dtalaashrafi at rome5.pgi.net>
Date: Tue, 8 Sep 2026 13:18:38 -0700
Subject: [PATCH 1/2] [mlir][openacc] Keep vector-private storage per-lane in
ACCCGToGPU
Give each lane its own allocation when ThreadX is active and the shape is
too large or too dynamic to be considered for stack alloca. processPrivatize
published a single allocation through a workgroup slot indexed by ThreadY,
so all lanes of a worker ended up sharing one buffer.
---
.../Dialect/OpenACC/Transforms/ACCCGToGPU.cpp | 22 ++++++
.../acc-cg-to-gpu-privatize-threadx-heap.mlir | 70 +++++++++++++++++++
2 files changed, 92 insertions(+)
create mode 100644 mlir/test/Dialect/OpenACC/acc-cg-to-gpu-privatize-threadx-heap.mlir
diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
index 90442294d668c..a858e5e35bf3a 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
@@ -2661,6 +2661,28 @@ Value ACCCGToGPULowering::processPrivatize(acc::PrivatizeOp privatize) {
return alloc.getResult();
}
+ if (llvm::any_of(parDimsPair.first, [](mlir::acc::GPUParallelDimAttr parDim) {
+ return parDim.isThreadX();
+ })) {
+ // Reached only after canUseStackAlloca failed: either the static size is
+ // over the per-thread stack budget, or the shape is dynamic so the size
+ // is unknown. Heap-allocate in both cases so a large runtime extent
+ // cannot overflow the stack. Each lane still owns its own copy.
+ Value mem =
+ memref::AllocOp::create(rewriter, loc, baseTy, mappedDynamicSizes);
+ OpBuilder::InsertionGuard guard(rewriter);
+ Block *allocBlock = mem.getDefiningOp()->getBlock();
+ if (allocBlock->mightHaveTerminator()) {
+ rewriter.setInsertionPoint(allocBlock->getTerminator());
+ memref::DeallocOp::create(rewriter, loc, mem);
+ }
+ mapping.map(privatize.getResult(), mem);
+ rewriter.setInsertionPointAfter(mem.getDefiningOp());
+ if (!privatize->getParentOfType<acc::ComputeRegionOp>())
+ rewriter.replaceOp(privatize, mem);
+ return mem;
+ }
+
// Predication - when threadYIsActive, don't predicate on ThreadY dimension
// since each ThreadY needs to execute the allocation for its own slot
SmallVector<mlir::acc::GPUParallelDimAttr> predicateDims;
diff --git a/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-privatize-threadx-heap.mlir b/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-privatize-threadx-heap.mlir
new file mode 100644
index 0000000000000..fc7067cdbbd9f
--- /dev/null
+++ b/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-privatize-threadx-heap.mlir
@@ -0,0 +1,70 @@
+// RUN: mlir-opt %s --pass-pipeline="builtin.module(gpu.module(gpu.func(acc-cg-to-gpu{max-thread-private-stack=16384})))" --split-input-file | FileCheck %s
+
+// ThreadX-private storage in a vector device routine that cannot use
+// memref.alloca (over the stack budget, or a dynamic shape) must not go
+// through the workgroup pointer-broadcast slot. That slot is indexed only by
+// ThreadY, so vector lanes would race. Each lane gets its own memref.alloc.
+
+module attributes {gpu.container_module} {
+ acc.routine @routine_vector func(@large_static) vector
+
+ gpu.module @cuda_device_mod {
+ // CHECK-LABEL: gpu.func @large_static
+ // CHECK: %[[MEM:.*]] = memref.alloc() : memref<4096xi32>
+ // CHECK-NOT: memref.store %[[MEM]], {{.*}}#gpu.address_space<workgroup>
+ // CHECK-NOT: gpu.barrier
+ // CHECK: memref.store %{{.*}}, %[[MEM]][%{{.*}}] : memref<4096xi32>
+ // CHECK: memref.dealloc %[[MEM]] : memref<4096xi32>
+ gpu.func @large_static()
+ attributes {acc.specialized_routine = #acc.specialized_routine<@routine_vector, <vector>, "large_static">} {
+ %c32 = arith.constant 32 : index
+ %tx = acc.par_width %c32 par_dim(#acc.par_dim<thread_x>)
+ %priv = acc.privatize {acc.active_par_dims = #acc<active_par_dims[thread_x]>}
+ : () -> !acc.private_type<memref<4096xi32>>
+ acc.compute_region launch(%arg0 = %tx) ins(%arg10 = %priv)
+ : (!acc.private_type<memref<4096xi32>>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : i32
+ %local = acc.private_local %arg10
+ {acc.active_par_dims = #acc<active_par_dims[thread_x]>}
+ : (!acc.private_type<memref<4096xi32>>) -> memref<4096xi32>
+ memref.store %c1, %local[%c0] : memref<4096xi32>
+ acc.yield
+ } <{origin = "acc.routine"}>
+ gpu.return
+ }
+ }
+}
+
+// -----
+
+module attributes {gpu.container_module} {
+ acc.routine @routine_vector_dyn func(@dynamic) vector
+
+ gpu.module @cuda_device_mod {
+ // CHECK-LABEL: gpu.func @dynamic
+ // CHECK: %[[MEM:.*]] = memref.alloc(%{{.*}}) : memref<?xi32>
+ // CHECK-NOT: memref.store %[[MEM]], {{.*}}#gpu.address_space<workgroup>
+ // CHECK-NOT: gpu.barrier
+ // CHECK: memref.store %{{.*}}, %[[MEM]][%{{.*}}] : memref<?xi32>
+ // CHECK: memref.dealloc %[[MEM]] : memref<?xi32>
+ gpu.func @dynamic(%n: index)
+ attributes {acc.specialized_routine = #acc.specialized_routine<@routine_vector_dyn, <vector>, "dynamic">} {
+ %c32 = arith.constant 32 : index
+ %tx = acc.par_width %c32 par_dim(#acc.par_dim<thread_x>)
+ %priv = acc.privatize(%n) {acc.active_par_dims = #acc<active_par_dims[thread_x]>}
+ : (index) -> !acc.private_type<memref<?xi32>>
+ acc.compute_region launch(%arg0 = %tx) ins(%arg10 = %priv)
+ : (!acc.private_type<memref<?xi32>>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : i32
+ %local = acc.private_local %arg10
+ {acc.active_par_dims = #acc<active_par_dims[thread_x]>}
+ : (!acc.private_type<memref<?xi32>>) -> memref<?xi32>
+ memref.store %c1, %local[%c0] : memref<?xi32>
+ acc.yield
+ } <{origin = "acc.routine"}>
+ gpu.return
+ }
+ }
+}
>From 2321ccb7aebdeb56af70dd48bd1358a199243abf Mon Sep 17 00:00:00 2001
From: Delaram Talaashrafi <dtalaashrafi at rome5.pgi.net>
Date: Wed, 9 Sep 2026 10:29:34 -0700
Subject: [PATCH 2/2] Fix
---
.../Dialect/OpenACC/Transforms/ACCCGToGPU.cpp | 73 +++++++++----------
1 file changed, 36 insertions(+), 37 deletions(-)
diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
index a858e5e35bf3a..6a46810a56a1c 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
@@ -477,6 +477,10 @@ class ACCCGToGPULowering {
std::optional<int64_t> sharedMemCopies = std::nullopt);
/// Lower an `acc.privatize` to device storage.
Value processPrivatize(acc::PrivatizeOp privatize);
+ /// Heap-allocate privatize storage and insert a matching dealloc.
+ Value allocatePrivatizeHeapStorage(acc::PrivatizeOp privatize,
+ MemRefType baseTy,
+ ValueRange mappedDynamicSizes);
/// Clone and lower an `scf.execute_region`.
void processExecuteRegion(scf::ExecuteRegionOp op);
/// Lower `acc.reduction_accumulate`.
@@ -2541,6 +2545,26 @@ void ACCCGToGPULowering::processPredicateRegion(
//
// clang-format on
+Value ACCCGToGPULowering::allocatePrivatizeHeapStorage(
+ acc::PrivatizeOp privatize, MemRefType baseTy,
+ ValueRange mappedDynamicSizes) {
+ Location loc = privatize.getLoc();
+ Value mem =
+ memref::AllocOp::create(rewriter, loc, baseTy, mappedDynamicSizes);
+ OpBuilder::InsertionGuard guard(rewriter);
+ Block *allocBlock = mem.getDefiningOp()->getBlock();
+ if (allocBlock->mightHaveTerminator()) {
+ rewriter.setInsertionPoint(allocBlock->getTerminator());
+ memref::DeallocOp::create(rewriter, loc, mem);
+ }
+ mapping.map(privatize.getResult(), mem);
+ // Ops inside the kernel are rewritten from scratch. If privatize is
+ // outside the kernel, it must be replaced with the allocated storage.
+ if (!privatize->getParentOfType<acc::ComputeRegionOp>())
+ rewriter.replaceOp(privatize, mem);
+ return mem;
+}
+
Value ACCCGToGPULowering::processPrivatize(acc::PrivatizeOp privatize) {
LLVM_DEBUG(llvm::dbgs() << "processing privatize: ";
privatize->print(llvm::dbgs()); llvm::dbgs() << "\n");
@@ -2578,13 +2602,15 @@ Value ACCCGToGPULowering::processPrivatize(acc::PrivatizeOp privatize) {
return privatize.getResult();
}
- for (mlir::acc::GPUParallelDimAttr parDim : parDimsPair.first) {
- if (parDim.isThreadX() &&
- canUseStackAlloca(baseTy, loc, options.maxThreadPrivateStack)) {
- auto alloca = memref::AllocaOp::create(rewriter, loc, baseTy);
- mapping.map(privatize.getResult(), alloca.getResult());
- return alloca.getResult();
- }
+ bool threadXActive =
+ llvm::any_of(parDimsPair.first, [](mlir::acc::GPUParallelDimAttr parDim) {
+ return parDim.isThreadX();
+ });
+ if (threadXActive &&
+ canUseStackAlloca(baseTy, loc, options.maxThreadPrivateStack)) {
+ auto alloca = memref::AllocaOp::create(rewriter, loc, baseTy);
+ mapping.map(privatize.getResult(), alloca.getResult());
+ return alloca.getResult();
}
if (!gpuFuncOp)
@@ -2645,42 +2671,15 @@ Value ACCCGToGPULowering::processPrivatize(acc::PrivatizeOp privatize) {
return alloca.getResult();
}
// Dynamic sizes: use alloc (heap allocation) with dealloc
- auto alloc = memref::AllocOp::create(rewriter, privatize->getLoc(), baseTy,
- mappedDynamicSizes);
-
- // Insert dealloc (free) before the function return
- OpBuilder::InsertPoint currentInsertPoint = rewriter.saveInsertionPoint();
- Block &parentBlock = *alloc->getBlock();
- if (parentBlock.mightHaveTerminator()) {
- rewriter.setInsertionPoint(parentBlock.getTerminator());
- memref::DeallocOp::create(rewriter, privatize->getLoc(), alloc);
- }
- rewriter.restoreInsertionPoint(currentInsertPoint);
-
- mapping.map(privatize.getResult(), alloc.getResult());
- return alloc.getResult();
+ return allocatePrivatizeHeapStorage(privatize, baseTy, mappedDynamicSizes);
}
- if (llvm::any_of(parDimsPair.first, [](mlir::acc::GPUParallelDimAttr parDim) {
- return parDim.isThreadX();
- })) {
+ if (threadXActive) {
// Reached only after canUseStackAlloca failed: either the static size is
// over the per-thread stack budget, or the shape is dynamic so the size
// is unknown. Heap-allocate in both cases so a large runtime extent
// cannot overflow the stack. Each lane still owns its own copy.
- Value mem =
- memref::AllocOp::create(rewriter, loc, baseTy, mappedDynamicSizes);
- OpBuilder::InsertionGuard guard(rewriter);
- Block *allocBlock = mem.getDefiningOp()->getBlock();
- if (allocBlock->mightHaveTerminator()) {
- rewriter.setInsertionPoint(allocBlock->getTerminator());
- memref::DeallocOp::create(rewriter, loc, mem);
- }
- mapping.map(privatize.getResult(), mem);
- rewriter.setInsertionPointAfter(mem.getDefiningOp());
- if (!privatize->getParentOfType<acc::ComputeRegionOp>())
- rewriter.replaceOp(privatize, mem);
- return mem;
+ return allocatePrivatizeHeapStorage(privatize, baseTy, mappedDynamicSizes);
}
// Predication - when threadYIsActive, don't predicate on ThreadY dimension
More information about the Mlir-commits
mailing list