[Mlir-commits] [mlir] [mlir][acc] Add GPU workgroup memory operation for planning (PR #205972)
Razvan Lupusoru
llvmlistbot at llvm.org
Fri Jun 26 08:22:28 PDT 2026
https://github.com/razvanlupusoru updated https://github.com/llvm/llvm-project/pull/205972
>From 693a45b2600f758b838e510952a5d932f3ce9406 Mon Sep 17 00:00:00 2001
From: Razvan Lupusoru <rlupusoru at nvidia.com>
Date: Thu, 25 Jun 2026 17:42:25 -0700
Subject: [PATCH 1/2] [mlir][acc] Add GPU workgroup memory operation for
planning
Introduce acc.gpu_shared_memory to represent GPU workgroup
memory slots in a compute region - used for planning before
eventually turned into a `memref.view` of a dynamic slot
within the workgroup allocation.
---
.../mlir/Dialect/OpenACC/OpenACCCGOps.td | 100 ++++++++++++++++++
mlir/lib/Dialect/OpenACC/IR/OpenACCCG.cpp | 33 ++++++
mlir/test/Dialect/OpenACC/cse.mlir | 15 +++
mlir/test/Dialect/OpenACC/invalid-cg.mlir | 32 ++++++
mlir/test/Dialect/OpenACC/ops-cg.mlir | 39 +++++++
5 files changed, 219 insertions(+)
diff --git a/mlir/include/mlir/Dialect/OpenACC/OpenACCCGOps.td b/mlir/include/mlir/Dialect/OpenACC/OpenACCCGOps.td
index 70f16d423d1f8..4d47bcf2f3454 100644
--- a/mlir/include/mlir/Dialect/OpenACC/OpenACCCGOps.td
+++ b/mlir/include/mlir/Dialect/OpenACC/OpenACCCGOps.td
@@ -406,6 +406,106 @@ def OpenACC_ParWidthOp
}];
}
+//===----------------------------------------------------------------------===//
+// acc.gpu_shared_memory
+//===----------------------------------------------------------------------===//
+
+def OpenACC_GPUSharedMemoryOp : OpenACC_Op<"gpu_shared_memory"> {
+ let summary = "GPU workgroup (shared) memory allocation in a compute region";
+ let description = [{
+ Represents a GPU workgroup-memory allocation in a compute region.
+ The result is a typed `memref` view into a byte slab which is later
+ replaced by `memref.view` into a dynamic shared-memory blob at the
+ byte offset within the workgroup allocation.
+
+ Each operation occupies a distinct slot in that collective allocation.
+ `static_upper_bound_bytes` is the conservative byte-size upper bound for
+ the slot. `dynamic_sizes` supply values for dynamic memref result
+ dimensions. The in-kernel layout of the slot is given by
+ `static_upper_bound_bytes` and `dynamic_sizes`.
+
+ Optional `dynamic_shared_memory_scaling_bytes` and
+ `dynamic_shared_memory_fixed_bytes` parameterize
+ `dynamic_shared_memory_size` when the slot footprint depends on launch
+ geometry. They must be specified together. When present:
+
+ dynamic_shared_memory_size =
+ dynamic_shared_memory_scaling_bytes * W
+ + dynamic_shared_memory_fixed_bytes
+
+ where `W` is the launch width that scales the allocation.
+
+ This linear model arises for `acc.cache` regions with dynamic bounds: one
+ cache dimension grows with the thread-parallel launch width, while
+ `dynamic_shared_memory_fixed_bytes` covers bytes that do not scale (for
+ example overlap cells at cache tile boundaries so threads can read
+ neighboring source elements without extra global memory traffic).
+ For a 1D dynamic cache with stencil extent `E`:
+
+ dynamic_shared_memory_fixed_bytes =
+ (E - 1) * dynamic_shared_memory_scaling_bytes
+
+ ```
+ Global array: ... | a | b | c | d | e | f | ...
+ [---- cached tile ----]
+ Thread 0 primary: a (reads neighbor b)
+ Thread 1 primary: b (reads neighbors a, c)
+ ...
+ Scaling portion: dynamic_shared_memory_scaling_bytes * W
+ Fixed portion: dynamic_shared_memory_fixed_bytes ((E - 1) cells)
+ ```
+
+ The scaling attributes affect only `dynamic_shared_memory_size`, not
+ the slot layout. For purely static allocations both are omitted and
+ `dynamic_shared_memory_size` is the sum of aligned
+ `static_upper_bound_bytes` across all slots.
+
+ Example:
+
+ ```mlir
+ %sz = arith.constant 128 : index
+ %cache = acc.gpu_shared_memory(%sz)
+ {num_copies = 1 : i64,
+ static_upper_bound_bytes = 1560 : i64,
+ dynamic_shared_memory_scaling_bytes = 12 : i64,
+ dynamic_shared_memory_fixed_bytes = 24 : i64}
+ : (index) -> memref<?xf32, #gpu.address_space<workgroup>>
+ ```
+ }];
+ let arguments = (ins
+ I64Attr:$num_copies,
+ I64Attr:$static_upper_bound_bytes,
+ Variadic<Index>:$dynamic_sizes,
+ OptionalAttr<I64Attr>:$dynamic_shared_memory_scaling_bytes,
+ OptionalAttr<I64Attr>:$dynamic_shared_memory_fixed_bytes
+ );
+ let results = (outs Res<AnyType, "", [MemAlloc<DefaultResource, 0,
+ FullEffect>]>:$result);
+ let assemblyFormat = [{
+ (`(` $dynamic_sizes^ `)`)? attr-dict `:` functional-type(operands, results)
+ }];
+ let builders = [
+ OpBuilder<(ins "::mlir::Type":$resultType,
+ "::mlir::IntegerAttr":$numCopies,
+ "::mlir::IntegerAttr":$staticUpperBoundBytes,
+ CArg<"::mlir::ValueRange", "{}">:$dynamicSizes,
+ CArg<"::mlir::IntegerAttr", "{}">:$dynamicSharedMemoryScalingBytes,
+ CArg<"::mlir::IntegerAttr", "{}">:$dynamicSharedMemoryFixedBytes), [{
+ $_state.addAttribute("num_copies", numCopies);
+ $_state.addAttribute("static_upper_bound_bytes", staticUpperBoundBytes);
+ $_state.addOperands(dynamicSizes);
+ if (dynamicSharedMemoryScalingBytes)
+ $_state.addAttribute("dynamic_shared_memory_scaling_bytes",
+ dynamicSharedMemoryScalingBytes);
+ if (dynamicSharedMemoryFixedBytes)
+ $_state.addAttribute("dynamic_shared_memory_fixed_bytes",
+ dynamicSharedMemoryFixedBytes);
+ $_state.addTypes(resultType);
+ }]>
+ ];
+ let hasVerifier = 1;
+}
+
//===----------------------------------------------------------------------===//
// acc.compute_region
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/OpenACC/IR/OpenACCCG.cpp b/mlir/lib/Dialect/OpenACC/IR/OpenACCCG.cpp
index 953acb7b7826e..186783b23774d 100644
--- a/mlir/lib/Dialect/OpenACC/IR/OpenACCCG.cpp
+++ b/mlir/lib/Dialect/OpenACC/IR/OpenACCCG.cpp
@@ -792,6 +792,39 @@ ParseResult ComputeRegionOp::parse(OpAsmParser &parser,
return success();
}
+//===----------------------------------------------------------------------===//
+// GPUSharedMemoryOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult GPUSharedMemoryOp::verify() {
+ if (getNumCopies() <= 0)
+ return emitOpError("num_copies must be positive");
+ if (getStaticUpperBoundBytes() <= 0)
+ return emitOpError("static_upper_bound_bytes must be positive");
+
+ bool hasScaling =
+ static_cast<bool>(getDynamicSharedMemoryScalingBytes());
+ bool hasFixed = static_cast<bool>(getDynamicSharedMemoryFixedBytes());
+ if (hasScaling != hasFixed)
+ return emitOpError(
+ "dynamic_shared_memory_scaling_bytes and "
+ "dynamic_shared_memory_fixed_bytes must both be present or both be "
+ "absent");
+
+ if (auto resultTy = dyn_cast<MemRefType>(getResult().getType())) {
+ auto addrSpace = dyn_cast_if_present<gpu::AddressSpaceAttr>(
+ resultTy.getMemorySpace());
+ if (!addrSpace ||
+ addrSpace.getValue() != gpu::GPUDialect::getWorkgroupAddressSpace())
+ return emitOpError(
+ "result memref must use #gpu.address_space<workgroup>");
+ } else {
+ return emitOpError("result must be a memref type");
+ }
+
+ return success();
+}
+
//===----------------------------------------------------------------------===//
// PredicateRegionOp
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/OpenACC/cse.mlir b/mlir/test/Dialect/OpenACC/cse.mlir
index 7f3a11eb21e09..1d617a77eede8 100644
--- a/mlir/test/Dialect/OpenACC/cse.mlir
+++ b/mlir/test/Dialect/OpenACC/cse.mlir
@@ -16,3 +16,18 @@ func.func @cse_across_acc_set(%a: memref<10xf32>, %i: index) -> (f32, f32) {
// CHECK: return %[[V]], %[[V]] : f32, f32
return %v1, %v2 : f32, f32
}
+
+// -----
+
+// Two identical acc.gpu_shared_memory ops must not be CSE'd: each reserves a
+// distinct workgroup-memory slot.
+// CHECK-LABEL: @cse_gpu_shared_memory_not_merged
+func.func @cse_gpu_shared_memory_not_merged() -> (memref<8xf32, #gpu.address_space<workgroup>>, memref<8xf32, #gpu.address_space<workgroup>>) {
+ %0 = acc.gpu_shared_memory {num_copies = 1 : i64, static_upper_bound_bytes = 256 : i64}
+ : () -> memref<8xf32, #gpu.address_space<workgroup>>
+ %1 = acc.gpu_shared_memory {num_copies = 1 : i64, static_upper_bound_bytes = 256 : i64}
+ : () -> memref<8xf32, #gpu.address_space<workgroup>>
+ // CHECK: acc.gpu_shared_memory
+ // CHECK: acc.gpu_shared_memory
+ return %0, %1 : memref<8xf32, #gpu.address_space<workgroup>>, memref<8xf32, #gpu.address_space<workgroup>>
+}
diff --git a/mlir/test/Dialect/OpenACC/invalid-cg.mlir b/mlir/test/Dialect/OpenACC/invalid-cg.mlir
index b37d2da406ee2..ec21263f31840 100644
--- a/mlir/test/Dialect/OpenACC/invalid-cg.mlir
+++ b/mlir/test/Dialect/OpenACC/invalid-cg.mlir
@@ -127,3 +127,35 @@ func.func @predicate_region_outside_compute_region() {
}
return
}
+
+// -----
+
+func.func @gpu_shared_memory_mismatched_runtime_attrs() {
+ // expected-error at +1 {{dynamic_shared_memory_scaling_bytes and dynamic_shared_memory_fixed_bytes must both be present or both be absent}}
+ %sm = acc.gpu_shared_memory()
+ {num_copies = 1 : i64,
+ static_upper_bound_bytes = 512 : i64,
+ dynamic_shared_memory_scaling_bytes = 4 : i64}
+ : () -> memref<8xf32, #gpu.address_space<workgroup>>
+ return
+}
+
+// -----
+
+func.func @gpu_shared_memory_non_workgroup_memref() {
+ // expected-error at +1 {{result memref must use #gpu.address_space<workgroup>}}
+ %sm = acc.gpu_shared_memory()
+ {num_copies = 1 : i64, static_upper_bound_bytes = 512 : i64}
+ : () -> memref<8xf32>
+ return
+}
+
+// -----
+
+func.func @gpu_shared_memory_zero_num_copies() {
+ // expected-error at +1 {{num_copies must be positive}}
+ %sm = acc.gpu_shared_memory()
+ {num_copies = 0 : i64, static_upper_bound_bytes = 512 : i64}
+ : () -> memref<8xf32, #gpu.address_space<workgroup>>
+ return
+}
diff --git a/mlir/test/Dialect/OpenACC/ops-cg.mlir b/mlir/test/Dialect/OpenACC/ops-cg.mlir
index 6e5582a709bfe..765552e8577a7 100644
--- a/mlir/test/Dialect/OpenACC/ops-cg.mlir
+++ b/mlir/test/Dialect/OpenACC/ops-cg.mlir
@@ -428,3 +428,42 @@ func.func @predicate_region_gang_redundant_setup(%idx: memref<i32>, %table: memr
// CHECK: acc.predicate_region {
// CHECK: memref.store
// CHECK: }
+
+// -----
+
+// CHECK-LABEL: func @gpu_shared_memory_static
+func.func @gpu_shared_memory_static() {
+ %c1024 = arith.constant 1024 : index
+ %sm = acc.gpu_shared_memory(%c1024)
+ {num_copies = 1 : i64, static_upper_bound_bytes = 4096 : i64}
+ : (index) -> memref<?xf32, #gpu.address_space<workgroup>>
+ return
+}
+// CHECK: acc.gpu_shared_memory(%{{.*}}) {num_copies = 1 : i64, static_upper_bound_bytes = 4096 : i64}
+// CHECK-SAME: : (index) -> memref<?xf32, #gpu.address_space<workgroup>>
+
+// -----
+
+// CHECK-LABEL: func @gpu_shared_memory_runtime_sized
+func.func @gpu_shared_memory_runtime_sized() {
+ %c128 = arith.constant 128 : index
+ %sm = acc.gpu_shared_memory(%c128)
+ {num_copies = 1 : i64,
+ static_upper_bound_bytes = 1560 : i64,
+ dynamic_shared_memory_scaling_bytes = 12 : i64,
+ dynamic_shared_memory_fixed_bytes = 24 : i64}
+ : (index) -> memref<?xf32, #gpu.address_space<workgroup>>
+ return
+}
+// CHECK: acc.gpu_shared_memory(%{{.*}}) {dynamic_shared_memory_fixed_bytes = 24 : i64, dynamic_shared_memory_scaling_bytes = 12 : i64, num_copies = 1 : i64, static_upper_bound_bytes = 1560 : i64}
+
+// -----
+
+// CHECK-LABEL: func @gpu_shared_memory_worker_copies
+func.func @gpu_shared_memory_worker_copies() {
+ %sm = acc.gpu_shared_memory()
+ {num_copies = 4 : i64, static_upper_bound_bytes = 256 : i64}
+ : () -> memref<8xf32, #gpu.address_space<workgroup>>
+ return
+}
+// CHECK: acc.gpu_shared_memory {num_copies = 4 : i64, static_upper_bound_bytes = 256 : i64}
>From 58d70a54c257e8d308017e902da851e5755d1f8a Mon Sep 17 00:00:00 2001
From: Razvan Lupusoru <rlupusoru at nvidia.com>
Date: Fri, 26 Jun 2026 08:22:13 -0700
Subject: [PATCH 2/2] Fix formatting
---
mlir/lib/Dialect/OpenACC/IR/OpenACCCG.cpp | 7 +++----
1 file changed, 3 insertions(+), 4 deletions(-)
diff --git a/mlir/lib/Dialect/OpenACC/IR/OpenACCCG.cpp b/mlir/lib/Dialect/OpenACC/IR/OpenACCCG.cpp
index 186783b23774d..ed5d2e3e9277b 100644
--- a/mlir/lib/Dialect/OpenACC/IR/OpenACCCG.cpp
+++ b/mlir/lib/Dialect/OpenACC/IR/OpenACCCG.cpp
@@ -802,8 +802,7 @@ LogicalResult GPUSharedMemoryOp::verify() {
if (getStaticUpperBoundBytes() <= 0)
return emitOpError("static_upper_bound_bytes must be positive");
- bool hasScaling =
- static_cast<bool>(getDynamicSharedMemoryScalingBytes());
+ bool hasScaling = static_cast<bool>(getDynamicSharedMemoryScalingBytes());
bool hasFixed = static_cast<bool>(getDynamicSharedMemoryFixedBytes());
if (hasScaling != hasFixed)
return emitOpError(
@@ -812,8 +811,8 @@ LogicalResult GPUSharedMemoryOp::verify() {
"absent");
if (auto resultTy = dyn_cast<MemRefType>(getResult().getType())) {
- auto addrSpace = dyn_cast_if_present<gpu::AddressSpaceAttr>(
- resultTy.getMemorySpace());
+ auto addrSpace =
+ dyn_cast_if_present<gpu::AddressSpaceAttr>(resultTy.getMemorySpace());
if (!addrSpace ||
addrSpace.getValue() != gpu::GPUDialect::getWorkgroupAddressSpace())
return emitOpError(
More information about the Mlir-commits
mailing list