[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.

Author: raghavendhra

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 = [{

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

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";

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 {

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
 }