Skip to content

[MLIR] [OpenMP] Modify definition of ALLOCATOR clause to support allocator type defined in user program. #157399

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 6 commits into from
Sep 18, 2025
Merged

[MLIR] [OpenMP] Modify definition of ALLOCATOR clause to support allocator type defined in user program. #157399

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
a6b5d96
[MLIR] [OpenMP] Modify definition of ALLOCATOR clause to support
raghavendhra Sep 8, 2025
d455dca
Add test-case for user-defined allocator value
raghavendhra Sep 9, 2025
c89d0ff
Changed allocator clause definition to use Integer type value argumen…
raghavendhra Sep 17, 2025
bfce138
Merge branch 'main' into allocate_directive_allocator_clause
raghavendhra Sep 17, 2025
859de37
Merge branch 'main' into allocate_directive_allocator_clause
raghavendhra Sep 17, 2025
3ebc9d1
Merge branch 'main' into allocate_directive_allocator_clause
raghavendhra Sep 18, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -120,11 +120,11 @@ class OpenMP_AllocatorClauseSkip<
extraClassDeclaration> {

let arguments = (ins
OptionalAttr<AllocatorHandleAttr>:$allocator
Optional<I64>:$allocator
);

let optAssemblyFormat = [{
`allocator` `(` custom<ClauseAttr>($allocator) `)`
`allocator` `(` $allocator `)`
}];

let description = [{
Expand Down
30 changes: 0 additions & 30 deletions mlir/include/mlir/Dialect/OpenMP/OpenMPEnums.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -263,34 +263,4 @@ def VariableCaptureKindAttr : OpenMP_EnumAttr<VariableCaptureKind,
let assemblyFormat = "`(` $value `)`";
}


//===----------------------------------------------------------------------===//
// allocator_handle enum.
//===----------------------------------------------------------------------===//

def OpenMP_AllocatorHandleNullAllocator : I32EnumAttrCase<"omp_null_allocator", 0>;
def OpenMP_AllocatorHandleDefaultMemAlloc : I32EnumAttrCase<"omp_default_mem_alloc", 1>;
def OpenMP_AllocatorHandleLargeCapMemAlloc : I32EnumAttrCase<"omp_large_cap_mem_alloc", 2>;
def OpenMP_AllocatorHandleConstMemAlloc : I32EnumAttrCase<"omp_const_mem_alloc", 3>;
def OpenMP_AllocatorHandleHighBwMemAlloc : I32EnumAttrCase<"omp_high_bw_mem_alloc", 4>;
def OpenMP_AllocatorHandleLowLatMemAlloc : I32EnumAttrCase<"omp_low_lat_mem_alloc", 5>;
def OpenMP_AllocatorHandleCgroupMemAlloc : I32EnumAttrCase<"omp_cgroup_mem_alloc", 6>;
def OpenMP_AllocatorHandlePteamMemAlloc : I32EnumAttrCase<"omp_pteam_mem_alloc", 7>;
def OpenMP_AllocatorHandlethreadMemAlloc : I32EnumAttrCase<"omp_thread_mem_alloc", 8>;

def AllocatorHandle : OpenMP_I32EnumAttr<
"AllocatorHandle",
"OpenMP allocator_handle", [
OpenMP_AllocatorHandleNullAllocator,
OpenMP_AllocatorHandleDefaultMemAlloc,
OpenMP_AllocatorHandleLargeCapMemAlloc,
OpenMP_AllocatorHandleConstMemAlloc,
OpenMP_AllocatorHandleHighBwMemAlloc,
OpenMP_AllocatorHandleLowLatMemAlloc,
OpenMP_AllocatorHandleCgroupMemAlloc,
OpenMP_AllocatorHandlePteamMemAlloc,
OpenMP_AllocatorHandlethreadMemAlloc
]>;

def AllocatorHandleAttr : OpenMP_EnumAttr<AllocatorHandle, "allocator_handle">;
#endif // OPENMP_ENUMS
2 changes: 1 addition & 1 deletion mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2100,7 +2100,7 @@ def MaskedOp : OpenMP_Op<"masked", clauses = [
//===----------------------------------------------------------------------===//
// [Spec 5.2] 6.5 allocate Directive
//===----------------------------------------------------------------------===//
def AllocateDirOp : OpenMP_Op<"allocate_dir", clauses = [
def AllocateDirOp : OpenMP_Op<"allocate_dir", [AttrSizedOperandSegments], clauses = [
OpenMP_AlignClause, OpenMP_AllocatorClause
]> {
let summary = "allocate directive";
Expand Down
8 changes: 0 additions & 8 deletions mlir/test/Dialect/OpenMP/invalid.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3033,14 +3033,6 @@ func.func @invalid_allocate_align_2(%arg0 : memref<i32>) -> () {
return
}

// -----
func.func @invalid_allocate_allocator(%arg0 : memref<i32>) -> () {
// expected-error @below {{invalid clause value}}
omp.allocate_dir (%arg0 : memref<i32>) allocator(omp_small_cap_mem_alloc)

return
}

// -----
func.func @invalid_workdistribute_empty_region() -> () {
omp.teams {
Expand Down
29 changes: 23 additions & 6 deletions mlir/test/Dialect/OpenMP/ops.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3260,6 +3260,10 @@ func.func @omp_workshare_loop_wrapper_attrs(%idx : index) {
return
}

func.func @omp_init_allocator(%custom_allocator : i64) -> i64 {
return %custom_allocator : i64
}

// CHECK-LABEL: func.func @omp_allocate_dir(
// CHECK-SAME: %[[ARG0:.*]]: memref<i32>,
// CHECK-SAME: %[[ARG1:.*]]: memref<i32>) {
Expand All @@ -3278,16 +3282,29 @@ func.func @omp_allocate_dir(%arg0 : memref<i32>, %arg1 : memref<i32>) -> () {
omp.allocate_dir (%arg0 : memref<i32>) align(2)

// Test with one data var and allocator clause
// CHECK: omp.allocate_dir(%[[ARG0]] : memref<i32>) allocator(omp_pteam_mem_alloc)
omp.allocate_dir (%arg0 : memref<i32>) allocator(omp_pteam_mem_alloc)
// CHECK: %[[VAL_1:.*]] = arith.constant 1 : i64
%omp_default_mem_alloc = arith.constant 1 : i64
// CHECK: omp.allocate_dir(%[[ARG0]] : memref<i32>) allocator(%[[VAL_1:.*]])
omp.allocate_dir (%arg0 : memref<i32>) allocator(%omp_default_mem_alloc)

// Test with one data var, align clause and allocator clause
// CHECK: omp.allocate_dir(%[[ARG0]] : memref<i32>) align(2) allocator(omp_thread_mem_alloc)
omp.allocate_dir (%arg0 : memref<i32>) align(2) allocator(omp_thread_mem_alloc)
// CHECK: %[[VAL_2:.*]] = arith.constant 7 : i64
%omp_pteam_mem_alloc = arith.constant 7 : i64
// CHECK: omp.allocate_dir(%[[ARG0]] : memref<i32>) align(4) allocator(%[[VAL_2:.*]])
omp.allocate_dir (%arg0 : memref<i32>) align(4) allocator(%omp_pteam_mem_alloc)

// Test with two data vars, align clause and allocator clause
// CHECK: omp.allocate_dir(%[[ARG0]], %[[ARG1]] : memref<i32>, memref<i32>) align(2) allocator(omp_cgroup_mem_alloc)
omp.allocate_dir (%arg0, %arg1 : memref<i32>, memref<i32>) align(2) allocator(omp_cgroup_mem_alloc)
// CHECK: %[[VAL_3:.*]] = arith.constant 6 : i64
%omp_cgroup_mem_alloc = arith.constant 6 : i64
// CHECK: omp.allocate_dir(%[[ARG0]], %[[ARG1]] : memref<i32>, memref<i32>) align(8) allocator(%[[VAL_3:.*]])
omp.allocate_dir (%arg0, %arg1 : memref<i32>, memref<i32>) align(8) allocator(%omp_cgroup_mem_alloc)

// Test with one data var and user defined allocator clause
// CHECK: %[[VAL_4:.*]] = arith.constant 9 : i64
%custom_allocator = arith.constant 9 : i64
%custom_mem_alloc = func.call @omp_init_allocator(%custom_allocator) : (i64) -> (i64)
// CHECK: omp.allocate_dir(%[[ARG0]] : memref<i32>) allocator(%[[VAL_5:.*]])
omp.allocate_dir (%arg0 : memref<i32>) allocator(%custom_mem_alloc)

return
}
Expand Down