[Mlir-commits] [mlir] [mlir][acc] Introduce privatization operations for codegen (PR #195273)
Razvan Lupusoru
llvmlistbot at llvm.org
Fri May 1 07:54:43 PDT 2026
https://github.com/razvanlupusoru created https://github.com/llvm/llvm-project/pull/195273
This change adds codegen-oriented operations for representing private-variable storage and materializing the storage that a particular parallel execution actually uses.
The two operations are meant to be used together:
- acc.privatize introduces an abstract handle for the privatized storage,
including the parallel levels that determine the ultimate size of the storage needed. Which parallel levels apply can be stated when that structure is known, or omitted so the same representation can be refined later as launch and loop parallelism are decided.
- acc.private_local takes that handle and yields the concrete storage for the current execution context(for example the slice that corresponds to this gang or worker).
>From fa21e0ba51a49a23178b4087ff081badbda131d4 Mon Sep 17 00:00:00 2001
From: Razvan Lupusoru <rlupusoru at nvidia.com>
Date: Fri, 1 May 2026 07:53:08 -0700
Subject: [PATCH] [mlir][acc] Introduce privatization operations for codegen
This change adds codegen-oriented operations for representing
private-variable storage and materializing the storage that a
particular parallel execution actually uses.
The two operations are meant to be used together:
- acc.privatize introduces an abstract handle for the privatized
storage,
including the parallel levels that determine the ultimate size of the
storage needed. Which parallel levels apply can be stated when that
structure is known, or omitted so the same representation can be
refined later as launch and loop parallelism are decided.
- acc.private_local takes that handle and yields the concrete storage
for the current execution context(for example the slice that
corresponds to this gang or worker).
---
.../mlir/Dialect/OpenACC/OpenACCCGOps.td | 65 ++++++++++
.../mlir/Dialect/OpenACC/OpenACCOpsTypes.td | 11 ++
.../Dialect/OpenACC/ops-cg-privatization.mlir | 116 ++++++++++++++++++
3 files changed, 192 insertions(+)
create mode 100644 mlir/test/Dialect/OpenACC/ops-cg-privatization.mlir
diff --git a/mlir/include/mlir/Dialect/OpenACC/OpenACCCGOps.td b/mlir/include/mlir/Dialect/OpenACC/OpenACCCGOps.td
index 69848101a5e4d..1e4618f5b6ca7 100644
--- a/mlir/include/mlir/Dialect/OpenACC/OpenACCCGOps.td
+++ b/mlir/include/mlir/Dialect/OpenACC/OpenACCCGOps.td
@@ -221,6 +221,71 @@ def OpenACC_FirstprivateMapInitialOp
let extraClassDeclaration = extraClassDeclarationBase;
}
+//===----------------------------------------------------------------------===//
+// acc.privatize
+//===----------------------------------------------------------------------===//
+
+def OpenACC_PrivatizeOp : OpenACC_Op<"privatize", []> {
+ let summary = "Create a handle for privatized storage along parallel dimensions";
+ let description = [{
+ Introduces a privatization handle for storage that varies across the active
+ parallel dimensions (for example OpenACC `private` / `firstprivate` after
+ recipe materialization). The handle type is `acc.private_type<T>` where
+ `T` is the logical storage type (commonly a `memref`).
+
+ Optional `index` operands supply dynamic sizes when the privatized shape
+ depends on SSA values (for example `memref<?xi32>` row lengths).
+
+ The optional `acc.par_dims` attribute records which GPU parallel dimensions
+ participate in the privatization.
+ }];
+ let arguments = (ins Variadic<Index>:$dynamicSizes,
+ OptionalAttr<OpenACC_GPUParallelDimsAttr>:$par_dims);
+ let results = (outs OpenACC_PrivateType:$result);
+ let assemblyFormat = [{
+ (`(` $dynamicSizes^ `)`)? (`[` qualified($par_dims)^ `]`)? attr-dict `:` functional-type(operands, results)
+ }];
+ let builders = [
+ OpBuilder<(ins "::mlir::Type":$resultType,
+ CArg<"::mlir::ValueRange", "{}">:$dynamicSizes), [{
+ $_state.addTypes(resultType);
+ $_state.addOperands(dynamicSizes);
+ }]>
+ ];
+}
+
+//===----------------------------------------------------------------------===//
+// acc.private_local
+//===----------------------------------------------------------------------===//
+
+def OpenACC_PrivateLocalOp
+ : OpenACC_Op<"private_local", [NoMemoryEffect, AlwaysSpeculatable]> {
+ let summary = "Materialize privatized storage for the current parallelism context";
+ let description = [{
+ Given a value of type `acc.private_type<T>`, materializes the underlying
+ storage (typically a `memref`) for the current thread / iteration. Which
+ slice of the privatized allocation is selected is determined by surrounding
+ parallelism assigned in the context.
+
+ The result type is usually `T` (often a `memref` that matches the logical
+ storage type). The result may instead use a different surface type that
+ still aliases the same underlying storage as `T`; in that case passes that
+ consume this operation must treat the handle type and the result type as
+ describing the same storage layout.
+ }];
+ let arguments = (ins OpenACC_PrivateType:$privatized);
+ let results = (outs AnyType:$output);
+ let assemblyFormat = [{
+ $privatized attr-dict `:` functional-type(operands, results)
+ }];
+ let builders = [
+ OpBuilder<(ins "::mlir::Type":$resultType, "::mlir::Value":$privatized), [{
+ $_state.addOperands(privatized);
+ $_state.addTypes(resultType);
+ }]>
+ ];
+}
+
//===----------------------------------------------------------------------===//
// acc.par_width
//===----------------------------------------------------------------------===//
diff --git a/mlir/include/mlir/Dialect/OpenACC/OpenACCOpsTypes.td b/mlir/include/mlir/Dialect/OpenACC/OpenACCOpsTypes.td
index 117272693d626..e942a49be1fb9 100644
--- a/mlir/include/mlir/Dialect/OpenACC/OpenACCOpsTypes.td
+++ b/mlir/include/mlir/Dialect/OpenACC/OpenACCOpsTypes.td
@@ -33,4 +33,15 @@ def OpenACC_DeclareTokenType : OpenACC_Type<"DeclareToken", "declare_token"> {
}];
}
+def OpenACC_PrivateType : OpenACC_Type<"Private", "private_type"> {
+ let summary = "Handle type for privatized storage across parallel dimensions";
+ let description = [{
+ A handle type for privatized memory used in OpenACC codegen. It is
+ produced by `acc.privatize` and accessed through `acc.private_local`, which
+ materializes the underlying storage for the current parallel execution context.
+ }];
+ let parameters = (ins "mlir::Type":$baseTy);
+ let assemblyFormat = "`<` $baseTy `>`";
+}
+
#endif // OPENACC_OPS_TYPES
diff --git a/mlir/test/Dialect/OpenACC/ops-cg-privatization.mlir b/mlir/test/Dialect/OpenACC/ops-cg-privatization.mlir
new file mode 100644
index 0000000000000..206466e94abd7
--- /dev/null
+++ b/mlir/test/Dialect/OpenACC/ops-cg-privatization.mlir
@@ -0,0 +1,116 @@
+// RUN: mlir-opt -split-input-file %s | FileCheck %s --check-prefixes=CHECK
+// RUN: mlir-opt -split-input-file %s | mlir-opt -split-input-file | FileCheck %s --check-prefixes=CHECK
+// RUN: mlir-opt -split-input-file -mlir-print-op-generic %s | mlir-opt -split-input-file | FileCheck %s --check-prefixes=CHECK
+
+// -----
+
+// CHECK-LABEL: func @privatize_static_scalar
+func.func @privatize_static_scalar() {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c128 = arith.constant 128 : index
+ %h = acc.privatize : () -> !acc.private_type<memref<f32>>
+ scf.parallel (%tx) = (%c0) to (%c128) step (%c1) {
+ %loc = acc.private_local %h : (!acc.private_type<memref<f32>>) -> memref<f32>
+ %z = arith.constant 0.0 : f32
+ memref.store %z, %loc[] : memref<f32>
+ scf.reduce
+ } {acc.par_dims = #acc<par_dims[thread_x]>}
+ return
+}
+// CHECK-DAG: %[[H:.*]] = acc.privatize
+// CHECK-SAME: () -> !acc.private_type<memref<f32>>
+// CHECK: scf.parallel
+// CHECK: %{{.*}} = acc.private_local %[[H]] : (!acc.private_type<memref<f32>>) -> memref<f32>
+// CHECK: memref.store
+// CHECK: } {acc.par_dims = #acc<par_dims[thread_x]>}
+
+// -----
+
+// CHECK-LABEL: func @privatize_with_par_dims_attr
+func.func @privatize_with_par_dims_attr() {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c8 = arith.constant 8 : index
+ %c128 = arith.constant 128 : index
+ %h = acc.privatize [#acc<par_dims[block_x, thread_x]>] : () -> !acc.private_type<memref<i32>>
+ scf.parallel (%bx, %tx) = (%c0, %c0) to (%c8, %c128) step (%c1, %c1) {
+ %loc = acc.private_local %h : (!acc.private_type<memref<i32>>) -> memref<i32>
+ %z = arith.constant 0 : i32
+ memref.store %z, %loc[] : memref<i32>
+ scf.reduce
+ } {acc.par_dims = #acc<par_dims[block_x, thread_x]>}
+ return
+}
+// CHECK-DAG: %[[H2:.*]] = acc.privatize [#acc<par_dims[block_x, thread_x]>] : () -> !acc.private_type<memref<i32>>
+// CHECK: %{{.*}} = acc.private_local %[[H2]] : (!acc.private_type<memref<i32>>) -> memref<i32>
+
+// -----
+
+// CHECK-LABEL: func @privatize_dynamic_row
+func.func @privatize_dynamic_row(%n : index) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c32 = arith.constant 32 : index
+ %h = acc.privatize(%n) : (index) -> !acc.private_type<memref<?xf32>>
+ scf.parallel (%tx) = (%c0) to (%c32) step (%c1) {
+ %row = acc.private_local %h : (!acc.private_type<memref<?xf32>>) -> memref<?xf32>
+ %c0_f32 = arith.constant 0.0 : f32
+ memref.store %c0_f32, %row[%c0] : memref<?xf32>
+ scf.reduce
+ } {acc.par_dims = #acc<par_dims[thread_x]>}
+ return
+}
+// CHECK: acc.privatize(%{{.*}}) : (index) -> !acc.private_type<memref<?xf32>>
+// CHECK: %{{.*}} = acc.private_local %{{.*}} : (!acc.private_type<memref<?xf32>>) -> memref<?xf32>
+
+// -----
+
+// CHECK-LABEL: func @privatize_dynamic_and_par_dims
+func.func @privatize_dynamic_and_par_dims(%rows : index, %cols : index) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c4 = arith.constant 4 : index
+ %c16 = arith.constant 16 : index
+ %h = acc.privatize(%rows, %cols) [#acc<par_dims[block_x, thread_x]>] : (index, index) -> !acc.private_type<memref<?x?xi32>>
+ scf.parallel (%bx, %tx) = (%c0, %c0) to (%c4, %c16) step (%c1, %c1) {
+ %tile = acc.private_local %h : (!acc.private_type<memref<?x?xi32>>) -> memref<?x?xi32>
+ %z = arith.constant 0 : i32
+ memref.store %z, %tile[%c0, %c0] : memref<?x?xi32>
+ scf.reduce
+ } {acc.par_dims = #acc<par_dims[block_x, thread_x]>}
+ return
+}
+// CHECK: acc.privatize(%{{.*}}, %{{.*}}) [#acc<par_dims[block_x, thread_x]>] : (index, index) -> !acc.private_type<memref<?x?xi32>>
+
+// -----
+
+// CHECK-LABEL: func @privatize_inside_compute_region
+func.func @privatize_inside_compute_region(%data : memref<64xf32>) {
+ %copy = acc.copyin varPtr(%data : memref<64xf32>) -> memref<64xf32>
+ acc.kernel_environment dataOperands(%copy : memref<64xf32>) {
+ %c64_kw = arith.constant 64 : index
+ %w = acc.par_width %c64_kw {par_dim = #acc.par_dim<thread_x>}
+ acc.compute_region launch(%lw = %w) ins(%d = %copy) : (memref<64xf32>) {
+ %priv = acc.privatize [#acc<par_dims[thread_x]>] : () -> !acc.private_type<memref<f32>>
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %ub = arith.constant 64 : index
+ scf.parallel (%iv) = (%c0) to (%ub) step (%c1) {
+ %acc = acc.private_local %priv : (!acc.private_type<memref<f32>>) -> memref<f32>
+ %v = memref.load %d[%iv] : memref<64xf32>
+ memref.store %v, %acc[] : memref<f32>
+ scf.reduce
+ } {acc.par_dims = #acc<par_dims[thread_x]>}
+ acc.yield
+ } {origin = "acc.parallel"}
+ }
+ acc.delete accPtr(%copy : memref<64xf32>)
+ return
+}
+// CHECK: acc.kernel_environment
+// CHECK: acc.compute_region
+// CHECK: acc.privatize [#acc<par_dims[thread_x]>] : () -> !acc.private_type<memref<f32>>
+// CHECK: acc.private_local
+// CHECK: memref.load %{{.*}}[%{{.*}}] : memref<64xf32>
+// CHECK: } {origin = "acc.parallel"}
More information about the Mlir-commits
mailing list