[MLIR] [OpenMP] Modify definition of ALLOCATOR clause to support allocator type defined in user program. #157399
Conversation
raghavendhra
requested review from
DominikAdamski,
NimishMishra,
TIFitis,
abidh,
agozillon,
anchuraj,
bhandarkar-pranav,
dpalermo,
jsjodin,
kiranchandramohan,
kiranktp,
kparzysz,
skatrak,
skc7 and
tblah
|
@llvm/pr-subscribers-mlir @llvm/pr-subscribers-flang-openmp Author: Raghu Maddhipatla (raghavendhra) ChangesEarlier 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. Full diff: https://github.com/llvm/llvm-project/pull/157399.diff 4 Files Affected:
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
index 311c57fb4446c..2fbde606056c0 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
+ DefaultValuedOptionalAttr<I64Attr, "0">:$allocator
);
let optAssemblyFormat = [{
- `allocator` `(` custom<ClauseAttr>($allocator) `)`
+ `allocator` `(` custom<AllocatorHandle>($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/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index 6e43f28e8d93d..dfef05d242465 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -1245,6 +1245,93 @@ verifyReductionVarList(Operation *op, std::optional<ArrayAttr> reductionSyms,
return success();
}
+//===----------------------------------------------------------------------===//
+// Parser, printer and verifier for Allocator (Section 8.4 in OpenMP 6.0)
+//===----------------------------------------------------------------------===//
+
+/// Parses a allocator clause. The value of allocator handle is an integer
+/// which is a combination of different allocator handles from
+/// `omp_allocator_handle_t`.
+///
+/// allocator-clause = `allocator` `(` allocator-value `)`
+static ParseResult parseAllocatorHandle(OpAsmParser &parser,
+ IntegerAttr &allocatorHandleAttr) {
+ StringRef allocatorKeyword;
+ int64_t allocator = 0;
+ if (succeeded(parser.parseOptionalKeyword("none"))) {
+ allocatorHandleAttr = IntegerAttr::get(parser.getBuilder().getI64Type(), 0);
+ return success();
+ }
+ auto parseKeyword = [&]() -> ParseResult {
+ if (failed(parser.parseKeyword(&allocatorKeyword)))
+ return failure();
+ if (allocatorKeyword == "omp_null_allocator")
+ allocator = 0;
+ else if (allocatorKeyword == "omp_default_mem_alloc")
+ allocator = 1;
+ else if (allocatorKeyword == "omp_large_cap_mem_alloc")
+ allocator = 2;
+ else if (allocatorKeyword == "omp_const_mem_alloc")
+ allocator = 3;
+ else if (allocatorKeyword == "omp_high_bw_mem_alloc")
+ allocator = 4;
+ else if (allocatorKeyword == "omp_low_lat_mem_alloc")
+ allocator = 5;
+ else if (allocatorKeyword == "omp_cgroup_mem_alloc")
+ allocator = 6;
+ else if (allocatorKeyword == "omp_pteam_mem_alloc")
+ allocator = 7;
+ else if (allocatorKeyword == "omp_thread_mem_alloc")
+ allocator = 8;
+ else
+ return parser.emitError(parser.getCurrentLocation())
+ << allocatorKeyword << " is not a valid allocator";
+ return success();
+ };
+ if (parser.parseCommaSeparatedList(parseKeyword))
+ return failure();
+ allocatorHandleAttr =
+ IntegerAttr::get(parser.getBuilder().getI64Type(), allocator);
+ return success();
+}
+
+/// Prints a allocator clause
+static void printAllocatorHandle(OpAsmPrinter &p, Operation *op,
+ IntegerAttr allocatorHandleAttr) {
+ int64_t allocator = allocatorHandleAttr.getInt();
+ StringRef allocatorHandle;
+ switch (allocator) {
+ case 0:
+ allocatorHandle = "omp_null_allocator";
+ break;
+ case 1:
+ allocatorHandle = "omp_default_mem_alloc";
+ break;
+ case 2:
+ allocatorHandle = "omp_large_cap_mem_alloc";
+ break;
+ case 3:
+ allocatorHandle = "omp_const_mem_alloc";
+ break;
+ case 4:
+ allocatorHandle = "omp_high_bw_mem_alloc";
+ break;
+ case 5:
+ allocatorHandle = "omp_low_lat_mem_alloc";
+ break;
+ case 6:
+ allocatorHandle = "omp_cgroup_mem_alloc";
+ break;
+ case 7:
+ allocatorHandle = "omp_pteam_mem_alloc";
+ break;
+ case 8:
+ allocatorHandle = "omp_thread_mem_alloc";
+ break;
+ }
+ p << allocatorHandle;
+}
+
//===----------------------------------------------------------------------===//
// Parser, printer and verifier for Copyprivate
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/OpenMP/invalid.mlir b/mlir/test/Dialect/OpenMP/invalid.mlir
index 986c3844d0bb9..4e8266813a47e 100644
--- a/mlir/test/Dialect/OpenMP/invalid.mlir
+++ b/mlir/test/Dialect/OpenMP/invalid.mlir
@@ -3010,14 +3010,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 {
|
tblah
left a comment
tblah
left a comment
There was a problem hiding this comment.
It looks like the intention of this patch is to allow specifying other numbers for the allocator.
If this representation is to be used, please could you add a parse/unparse test demonstrating that this works.
But I'm unsure if this is the right approach. As I understand it, the custom allocator would be a handle returned by omp_init_allocator. That would be an mlir (runtime) value rather than a (compile-time) attribute.
I have tests already in ops.mlir https://github.com/llvm/llvm-project/blob/3f3f7d1fd99eea891ddd643201617d22c634bbfb/mlir/test/Dialect/OpenMP/ops.mlir#L3226C3-L3236C102 Is there any other location you would like me to add tests? |
|
Please could you add one showing how you intend a non-default allocator to be represented. |
raghavendhra
force-pushed
the
allocate_directive_allocator_clause
branch
from
a113f11 to
2e248fd
Compare
tblah
left a comment
tblah
left a comment
There was a problem hiding this comment.
Thank you for adding the test. Why did you decide to go with an integer attribute for the user defined allocator handle?
For predefined allocator_handle_kinds, they are defined as integer kind as per OpenMP Spec and in openmp/runtime https://github.com/llvm/llvm-project/blob/b6be44ad0deeb86e920873de87875d2eaa6c2d8b/openmp/runtime/src/include/omp_lib.F90.var#L180C9-L188C96 And, So used integer attribute to be consistent. |
|
My concern is that the integer returned from Previously for the built-in allocators an attribute made sense because the mapping from the keyword to the integer is always known. This isn't true for user-defined allocators because the handle is returned from a runtime call. I'm worried this will make lowering code difficult. |
Thanks for the explanation @tblah :) Can you please suggest type which I can use in Op definition which covers both and predefined and user-defined allocators? |
You could stick with your existing approach using Integers. But make it a value argument instead of an attribute. For the pre-defined allocators lowering should insert |
allocator type defined in user program.
…t instead of IntegerAttr.
raghavendhra
force-pushed
the
allocate_directive_allocator_clause
branch
from
46f7189 to
c89d0ff
Compare
tblah
left a comment
tblah
left a comment
There was a problem hiding this comment.
LGTM. Thanks for the updates
3 participants
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.