Changed allocator clause definition to use Integer type value argument instead of

IntegerAttr.

Author: raghavendhra

mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td

@@ -120,11 +120,11 @@ class OpenMP_AllocatorClauseSkip<
                     extraClassDeclaration> {
 
   let arguments = (ins
-    DefaultValuedOptionalAttr<I64Attr, "0">:$allocator
+    Optional<I64>:$allocator
   );
 
   let optAssemblyFormat = [{
-    `allocator` `(` custom<AllocatorHandle>($allocator) `)`
+    `allocator` `(` $allocator `)`
   }];
 
   let description = [{

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/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

@@ -1250,94 +1250,6 @@ 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();
-  }
-  OptionalParseResult parsedInteger = parser.parseOptionalInteger(allocator);
-  if (parsedInteger.has_value()) {
-    allocatorHandleAttr =
-        IntegerAttr::get(parser.getBuilder().getI64Type(), allocator);
-    return success();
-  }
-  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;
-  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;
-  default:
-    p << Twine(allocator).str();
-    return;
-  }
-  p << allocatorHandle;
-}
-
 //===----------------------------------------------------------------------===//
 // Parser, printer and verifier for Copyprivate
 //===----------------------------------------------------------------------===//

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,20 +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(7)
+  // 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: omp.allocate_dir(%[[ARG0]] : memref<i32>) allocator(9)
-  omp.allocate_dir (%arg0 : memref<i32>) allocator(9)
+  // 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
 }