[Mlir-commits] [mlir] 31e43e2 - [MLIR] [OpenMP] Modify definition of ALLOCATOR clause to support allocator type defined in user program. (#157399)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Sep 18 07:19:24 PDT 2025
Author: Raghu Maddhipatla
Date: 2025-09-18T09:19:19-05:00
New Revision: 31e43e2fb8634f35a70699f636c49c9d2451e81c
URL: https://github.com/llvm/llvm-project/commit/31e43e2fb8634f35a70699f636c49c9d2451e81c
DIFF: https://github.com/llvm/llvm-project/commit/31e43e2fb8634f35a70699f636c49c9d2451e81c.diff
LOG: [MLIR] [OpenMP] Modify definition of ALLOCATOR clause to support allocator type defined in user program. (#157399)
Earlier only predefined allocator types mentioned in OpenMP spec were allowed. This patch addresses support for user defined allocator type in allocator clause increasing the scope of allocator clause.
Added:
Modified:
mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
mlir/include/mlir/Dialect/OpenMP/OpenMPEnums.td
mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
mlir/test/Dialect/OpenMP/invalid.mlir
mlir/test/Dialect/OpenMP/ops.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
index 5f40abe62a0f6..1eda5e4bc1618 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
@@ -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 = [{
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPEnums.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPEnums.td
index c080c3fac87d4..9dbe6897a3304 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPEnums.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPEnums.td
@@ -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
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index 830b36f440098..5c77e215467e4 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -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";
diff --git a/mlir/test/Dialect/OpenMP/invalid.mlir b/mlir/test/Dialect/OpenMP/invalid.mlir
index 763f41c5420b8..af24d969064ab 100644
--- a/mlir/test/Dialect/OpenMP/invalid.mlir
+++ b/mlir/test/Dialect/OpenMP/invalid.mlir
@@ -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 {
diff --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir
index 60b1f61135ac2..cbd863f88fd1f 100644
--- a/mlir/test/Dialect/OpenMP/ops.mlir
+++ b/mlir/test/Dialect/OpenMP/ops.mlir
@@ -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>) {
@@ -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
}
More information about the Mlir-commits
mailing list