[Mlir-commits] [mlir] [mlir][acc] Give gang routine privates a workgroup slot of their own (PR #236264)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Oct 27 15:46:53 PDT 2026
https://github.com/khaki3 created https://github.com/llvm/llvm-project/pull/236264
A private hoisted to the top of a gang routine is shared by all threads of the block. ACCCGToGPU lowered it like any block-level private: thread 0 allocated it with memref.alloc, published the pointer through a workgroup attribution, and every thread loaded it back after a barrier; the matching memref.dealloc ran at the end of the routine. That is one heap allocation and free per routine call per block, even for a scalar.
A gang routine runs once per block, so a statically shaped private can live in a workgroup attribution of the routine itself. The private is now cast from that slot to the generic address space, and a workgroup barrier at the end of the routine keeps a following call from overwriting the slot while threads still read it.
The heap path stays for runtime-sized privates, privates whose element size is unknown, worker and vector routines, routines that call themselves (their activations would share one slot), and once the slots of a routine would exceed 256 bytes.
Limitations: recursion through another routine is not detected, and the calling kernels' shared memory budget does not account for the slots, as it already did not for the pointer slot.
>From 034ae9d03361d219128784d9dad594e7f7ce0767 Mon Sep 17 00:00:00 2001
From: Kazuaki Matsumura <kmatsumura at nvidia.com>
Date: Tue, 27 Oct 2026 15:38:55 -0700
Subject: [PATCH] [mlir][acc] Give gang routine privates a workgroup slot of
their own
A private hoisted to the top of a gang routine is shared by all threads of
the block. ACCCGToGPU lowered it like any block-level private: thread 0
allocated it with memref.alloc, published the pointer through a workgroup
attribution, and every thread loaded it back after a barrier; the matching
memref.dealloc ran at the end of the routine. That is one heap allocation and
free per routine call per block, even for a scalar.
A gang routine runs once per block, so a statically shaped private can live
in a workgroup attribution of the routine itself. The private is now cast from
that slot to the generic address space, and a workgroup barrier at the end of
the routine keeps a following call from overwriting the slot while threads
still read it.
The heap path stays for runtime-sized privates, privates whose element size is
unknown, worker and vector routines, routines that call themselves (their
activations would share one slot), and once the slots of a routine would
exceed 256 bytes.
Limitations: recursion through another routine is not detected, and the
calling kernels' shared memory budget does not account for the slots, as it
already did not for the pointer slot.
---
.../Dialect/OpenACC/Transforms/ACCCGToGPU.cpp | 78 +++++++
...-to-gpu-privatize-routine-gang-shared.mlir | 194 ++++++++++++++++++
2 files changed, 272 insertions(+)
create mode 100644 mlir/test/Dialect/OpenACC/acc-cg-to-gpu-privatize-routine-gang-shared.mlir
diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
index 06d1cb3a1e0625..ae8b9f7fd5abf5 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
@@ -496,6 +496,10 @@ class ACCCGToGPULowering {
Value allocatePrivatizeHeapStorage(acc::PrivatizeOp privatize,
MemRefType baseTy,
ValueRange mappedDynamicSizes);
+ /// Place static privatize storage of a block routine in a workgroup
+ /// attribution of its own; returns null if over the routine slot budget.
+ Value allocateRoutineSharedSlot(acc::PrivatizeOp privatize, MemRefType baseTy,
+ gpu::GPUFuncOp gpuFuncOp);
/// Clone and lower an `scf.execute_region`.
void processExecuteRegion(scf::ExecuteRegionOp op);
/// Lower `acc.reduction_accumulate`.
@@ -721,6 +725,11 @@ class ACCCGToGPULowering {
// Reusable privatize broadcast slots per type; disabled for kernels.
llvm::DenseMap<Type, Value> privatizeBroadcastCache;
+ // Routine workgroup slots add to the shared memory of every calling kernel,
+ // which is sized elsewhere, so keep them small.
+ static constexpr int64_t kMaxRoutineSharedSlotBytes = 256;
+ int64_t routineSharedSlotBytes = 0;
+
int64_t staticBlockDimX = 1024;
int64_t staticBlockDimY = 1024;
int64_t staticBlockDimZ = 1024;
@@ -2755,6 +2764,66 @@ Value ACCCGToGPULowering::allocatePrivatizeHeapStorage(
return mem;
}
+Value ACCCGToGPULowering::allocateRoutineSharedSlot(acc::PrivatizeOp privatize,
+ MemRefType baseTy,
+ gpu::GPUFuncOp gpuFuncOp) {
+ Location loc = privatize.getLoc();
+ Type elemTy = baseTy.getElementType();
+ if (!baseTy.hasStaticShape() || !baseTy.getLayout().isIdentity() ||
+ baseTy.getMemorySpace() || isa<BaseMemRefType>(elemTy))
+ return {};
+ // Query without a diagnostic: an unsized type keeps the heap path.
+ std::optional<acc::TypeSizeAndAlignment> sizeAndAlignment =
+ accSupport.getTypeSizeAndAlignment(
+ elemTy, computeRegion->getParentOfType<ModuleOp>());
+ if (!sizeAndAlignment)
+ return {};
+ int64_t bytes =
+ sizeAndAlignment->first.getFixedValue() * baseTy.getNumElements();
+ if (bytes <= 0 || routineSharedSlotBytes + bytes > kMaxRoutineSharedSlotBytes)
+ return {};
+ // Activations of a recursive routine would share the slot.
+ StringRef funcName = gpuFuncOp.getName();
+ bool callsItself =
+ gpuFuncOp
+ .walk([&](CallOpInterface call) {
+ auto callee =
+ dyn_cast_if_present<SymbolRefAttr>(call.getCallableForCallee());
+ return callee && callee.getLeafReference() == funcName
+ ? WalkResult::interrupt()
+ : WalkResult::advance();
+ })
+ .wasInterrupted();
+ if (callsItself)
+ return {};
+ routineSharedSlotBytes += bytes;
+
+ auto workgroupAS = gpu::AddressSpaceAttr::get(
+ rewriter.getContext(), gpu::GPUDialect::getWorkgroupAddressSpace());
+ MemRefType slotTy =
+ MemRefType::get(baseTy.getShape(), baseTy.getElementType(),
+ MemRefLayoutAttrInterface{}, workgroupAS);
+ Value slot = gpuFuncOp.addWorkgroupAttribution(slotTy, loc);
+ gpuFuncOp.setWorkgroupAttributionAttr(
+ gpuFuncOp.getNumWorkgroupAttributions() - 1,
+ LLVM::LLVMDialect::getAlignAttrName(), rewriter.getI32IntegerAttr(16));
+ Value mem = castPointerLikeTypeIfNeeded(rewriter, loc, slot, baseTy);
+ Block *slotBlock = rewriter.getInsertionBlock();
+ if (slotBlock->mightHaveTerminator()) {
+ // The next call of the routine rewrites the slot; wait for every reader.
+ OpBuilder::InsertionGuard guard(rewriter);
+ rewriter.setInsertionPoint(slotBlock->getTerminator());
+ createBarrier(loc, mlir::acc::GPUParallelDimsAttr::get(
+ rewriter.getContext(),
+ {mlir::acc::GPUParallelDimAttr::threadYDim(
+ rewriter.getContext())}));
+ }
+ mapping.map(privatize.getResult(), mem);
+ 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");
@@ -2831,6 +2900,7 @@ Value ACCCGToGPULowering::processPrivatize(acc::PrivatizeOp privatize) {
// variables at routine level are shared across ThreadY threads (single
// slot)
bool needsWorkgroupBarrier = false;
+ bool isBlockRoutine = false;
if (isInsideACCSpecializedRoutine(computeRegion)) {
FunctionOpInterface funcOp =
computeRegion->getParentOfType<FunctionOpInterface>();
@@ -2846,6 +2916,7 @@ Value ACCCGToGPULowering::processPrivatize(acc::PrivatizeOp privatize) {
needsWorkgroupBarrier = true;
} else if (routineParDim.isAnyBlock()) {
needsWorkgroupBarrier = true;
+ isBlockRoutine = true;
}
}
@@ -2876,6 +2947,13 @@ Value ACCCGToGPULowering::processPrivatize(acc::PrivatizeOp privatize) {
return allocatePrivatizeHeapStorage(privatize, baseTy, mappedDynamicSizes);
}
+ // A block routine runs once per block, so a small static private shared by
+ // its threads can live in a workgroup slot instead of a broadcast heap copy.
+ if (isBlockRoutine && !threadYIsActive && mappedDynamicSizes.empty()) {
+ if (Value slot = allocateRoutineSharedSlot(privatize, baseTy, gpuFuncOp))
+ return slot;
+ }
+
// 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-routine-gang-shared.mlir b/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-privatize-routine-gang-shared.mlir
new file mode 100644
index 00000000000000..0ce0f1dc82ad98
--- /dev/null
+++ b/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-privatize-routine-gang-shared.mlir
@@ -0,0 +1,194 @@
+// RUN: mlir-opt %s --pass-pipeline="builtin.module(gpu.module(gpu.func(acc-cg-to-gpu)))" --split-input-file | FileCheck %s
+
+// A gang routine runs once per block, so a small static private shared by the
+// block's threads is placed in a workgroup attribution of its own instead of a
+// thread-0 heap allocation whose pointer is broadcast through one.
+
+module attributes {gpu.container_module} {
+ acc.routine @routine_gang func(@gang_scalar) gang
+
+ gpu.module @cuda_device_mod {
+ func.func private @use(memref<f64>)
+
+ // CHECK-LABEL: gpu.func @gang_scalar
+ // CHECK-SAME: workgroup(%[[SLOT:[^ ]+]] : memref<f64, #gpu.address_space<workgroup>> {llvm.align = 16 : i32})
+ // CHECK-NOT: memref.alloc
+ // CHECK: %[[MEM:.*]] = memref.memory_space_cast %[[SLOT]] : memref<f64, #gpu.address_space<workgroup>> to memref<f64>
+ // CHECK: memref.store %{{.*}}, %[[MEM]][] : memref<f64>
+ // CHECK: func.call @use(%[[MEM]])
+ // CHECK: gpu.barrier
+ // CHECK-NOT: memref.dealloc
+ // CHECK: gpu.return
+ gpu.func @gang_scalar()
+ attributes {acc.specialized_routine = #acc.specialized_routine<@routine_gang, <gang_dim1>, "gang_scalar">} {
+ %c1 = arith.constant 1 : index
+ %bx = acc.par_width %c1 par_dim(#acc.par_dim<block_x>)
+ %tx = acc.par_width %c1 par_dim(#acc.par_dim<thread_x>)
+ %priv = acc.privatize {acc.active_par_dims = #acc<active_par_dims[block_x]>}
+ : () -> !acc.private_type<memref<f64>>
+ acc.compute_region launch(%arg0 = %bx, %arg1 = %tx) ins(%arg2 = %priv)
+ : (!acc.private_type<memref<f64>>) {
+ %cst = arith.constant 2.0 : f64
+ %local = acc.private_local %arg2
+ {acc.active_par_dims = #acc<active_par_dims[block_x]>}
+ : (!acc.private_type<memref<f64>>) -> memref<f64>
+ acc.predicate_region {
+ memref.store %cst, %local[] : memref<f64>
+ }
+ acc.predicate_region {
+ func.call @use(%local) : (memref<f64>) -> ()
+ }
+ acc.yield
+ } <{origin = "acc.routine"}>
+ gpu.return
+ }
+ }
+}
+
+// -----
+
+// The slots of one routine share a small byte budget; a private that does not
+// fit keeps the heap allocation broadcast through a pointer slot.
+
+module attributes {gpu.container_module} {
+ acc.routine @routine_gang func(@gang_budget) gang
+
+ gpu.module @cuda_device_mod {
+ func.func private @use2(memref<24xf64>, memref<16xf64>)
+
+ // CHECK-LABEL: gpu.func @gang_budget
+ // CHECK-SAME: workgroup(%{{[^ ]+}} : memref<24xf64, #gpu.address_space<workgroup>> {llvm.align = 16 : i32}, %{{[^ ]+}} : memref<memref<16xf64>, #gpu.address_space<workgroup>> {llvm.align = 16 : i32})
+ // CHECK-NOT: memref.alloc() : memref<24xf64>
+ // CHECK: memref.alloc() : memref<16xf64>
+ // CHECK: memref.dealloc %{{.*}} : memref<16xf64>
+ gpu.func @gang_budget()
+ attributes {acc.specialized_routine = #acc.specialized_routine<@routine_gang, <gang_dim1>, "gang_budget">} {
+ %c1 = arith.constant 1 : index
+ %bx = acc.par_width %c1 par_dim(#acc.par_dim<block_x>)
+ %tx = acc.par_width %c1 par_dim(#acc.par_dim<thread_x>)
+ %p0 = acc.privatize {acc.active_par_dims = #acc<active_par_dims[block_x]>}
+ : () -> !acc.private_type<memref<24xf64>>
+ %p1 = acc.privatize {acc.active_par_dims = #acc<active_par_dims[block_x]>}
+ : () -> !acc.private_type<memref<16xf64>>
+ acc.compute_region launch(%arg0 = %bx, %arg1 = %tx) ins(%arg2 = %p0, %arg3 = %p1)
+ : (!acc.private_type<memref<24xf64>>, !acc.private_type<memref<16xf64>>) {
+ %l0 = acc.private_local %arg2
+ {acc.active_par_dims = #acc<active_par_dims[block_x]>}
+ : (!acc.private_type<memref<24xf64>>) -> memref<24xf64>
+ %l1 = acc.private_local %arg3
+ {acc.active_par_dims = #acc<active_par_dims[block_x]>}
+ : (!acc.private_type<memref<16xf64>>) -> memref<16xf64>
+ acc.predicate_region {
+ func.call @use2(%l0, %l1) : (memref<24xf64>, memref<16xf64>) -> ()
+ }
+ acc.yield
+ } <{origin = "acc.routine"}>
+ gpu.return
+ }
+ }
+}
+
+// -----
+
+// A runtime-sized private keeps the heap allocation.
+
+module attributes {gpu.container_module} {
+ acc.routine @routine_gang func(@gang_dynamic) gang
+
+ gpu.module @cuda_device_mod {
+ func.func private @use_dyn(memref<?xf64>)
+
+ // CHECK-LABEL: gpu.func @gang_dynamic
+ // CHECK-SAME: workgroup(%{{[^ ]+}} : memref<memref<?xf64>, #gpu.address_space<workgroup>> {llvm.align = 16 : i32})
+ // CHECK: memref.alloc(%{{.*}}) : memref<?xf64>
+ // CHECK: memref.dealloc %{{.*}} : memref<?xf64>
+ gpu.func @gang_dynamic(%n: index)
+ attributes {acc.specialized_routine = #acc.specialized_routine<@routine_gang, <gang_dim1>, "gang_dynamic">} {
+ %c1 = arith.constant 1 : index
+ %bx = acc.par_width %c1 par_dim(#acc.par_dim<block_x>)
+ %tx = acc.par_width %c1 par_dim(#acc.par_dim<thread_x>)
+ %priv = acc.privatize(%n) {acc.active_par_dims = #acc<active_par_dims[block_x]>}
+ : (index) -> !acc.private_type<memref<?xf64>>
+ acc.compute_region launch(%arg0 = %bx, %arg1 = %tx) ins(%arg2 = %priv)
+ : (!acc.private_type<memref<?xf64>>) {
+ %local = acc.private_local %arg2
+ {acc.active_par_dims = #acc<active_par_dims[block_x]>}
+ : (!acc.private_type<memref<?xf64>>) -> memref<?xf64>
+ acc.predicate_region {
+ func.call @use_dyn(%local) : (memref<?xf64>) -> ()
+ }
+ acc.yield
+ } <{origin = "acc.routine"}>
+ gpu.return
+ }
+ }
+}
+
+// -----
+
+// A worker routine keeps the heap allocation broadcast through a pointer slot.
+
+module attributes {gpu.container_module} {
+ acc.routine @routine_worker func(@worker_scalar) worker
+
+ gpu.module @cuda_device_mod {
+ func.func private @use(memref<f64>)
+
+ // CHECK-LABEL: gpu.func @worker_scalar
+ // CHECK-SAME: workgroup(%{{[^ ]+}} : memref<memref<f64>, #gpu.address_space<workgroup>> {llvm.align = 16 : i32})
+ // CHECK: memref.alloc() : memref<f64>
+ // CHECK: memref.dealloc %{{.*}} : memref<f64>
+ gpu.func @worker_scalar()
+ attributes {acc.specialized_routine = #acc.specialized_routine<@routine_worker, <worker>, "worker_scalar">} {
+ %c1 = arith.constant 1 : index
+ %ty = acc.par_width %c1 par_dim(#acc.par_dim<thread_y>)
+ %tx = acc.par_width %c1 par_dim(#acc.par_dim<thread_x>)
+ %priv = acc.privatize : () -> !acc.private_type<memref<f64>>
+ acc.compute_region launch(%arg0 = %ty, %arg1 = %tx) ins(%arg2 = %priv)
+ : (!acc.private_type<memref<f64>>) {
+ %local = acc.private_local %arg2 : (!acc.private_type<memref<f64>>) -> memref<f64>
+ acc.predicate_region {
+ func.call @use(%local) : (memref<f64>) -> ()
+ }
+ acc.yield
+ } <{origin = "acc.routine"}>
+ gpu.return
+ }
+ }
+}
+
+// -----
+
+// A private larger than the slot budget keeps the heap allocation.
+
+module attributes {gpu.container_module} {
+ acc.routine @routine_gang func(@gang_large) gang
+
+ gpu.module @cuda_device_mod {
+ func.func private @use_large(memref<64xf64>)
+
+ // CHECK-LABEL: gpu.func @gang_large
+ // CHECK-SAME: workgroup(%{{[^ ]+}} : memref<memref<64xf64>, #gpu.address_space<workgroup>> {llvm.align = 16 : i32})
+ // CHECK: memref.alloc() : memref<64xf64>
+ // CHECK: memref.dealloc %{{.*}} : memref<64xf64>
+ gpu.func @gang_large()
+ attributes {acc.specialized_routine = #acc.specialized_routine<@routine_gang, <gang_dim1>, "gang_large">} {
+ %c1 = arith.constant 1 : index
+ %bx = acc.par_width %c1 par_dim(#acc.par_dim<block_x>)
+ %tx = acc.par_width %c1 par_dim(#acc.par_dim<thread_x>)
+ %priv = acc.privatize {acc.active_par_dims = #acc<active_par_dims[block_x]>}
+ : () -> !acc.private_type<memref<64xf64>>
+ acc.compute_region launch(%arg0 = %bx, %arg1 = %tx) ins(%arg2 = %priv)
+ : (!acc.private_type<memref<64xf64>>) {
+ %local = acc.private_local %arg2
+ {acc.active_par_dims = #acc<active_par_dims[block_x]>}
+ : (!acc.private_type<memref<64xf64>>) -> memref<64xf64>
+ acc.predicate_region {
+ func.call @use_large(%local) : (memref<64xf64>) -> ()
+ }
+ acc.yield
+ } <{origin = "acc.routine"}>
+ gpu.return
+ }
+ }
+}
More information about the Mlir-commits
mailing list