Add test-case for user-defined allocator value

Author: raghavendhra

mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

@@ -1267,34 +1267,32 @@ static ParseResult parseAllocatorHandle(OpAsmParser &parser,
     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";
+  OptionalParseResult parsedInteger = parser.parseOptionalInteger(allocator);
+  if (parsedInteger.has_value()) {
+    allocatorHandleAttr =
+        IntegerAttr::get(parser.getBuilder().getI64Type(), allocator);
     return success();
-  };
-  if (parser.parseCommaSeparatedList(parseKeyword))
+  }
+  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();
@@ -1333,6 +1331,9 @@ static void printAllocatorHandle(OpAsmPrinter &p, Operation *op,
   case 8:
     allocatorHandle = "omp_thread_mem_alloc";
     break;
+  default:
+    p << Twine(allocator).str();
+    return;
   }
   p << allocatorHandle;
 }

mlir/test/Dialect/OpenMP/ops.mlir

@@ -3279,7 +3279,7 @@ func.func @omp_allocate_dir(%arg0 : memref<i32>, %arg1 : memref<i32>) -> () {
 
   // 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)
+  omp.allocate_dir (%arg0 : memref<i32>) allocator(7)
 
   // Test with one data var, align clause and allocator clause
   // CHECK: omp.allocate_dir(%[[ARG0]] : memref<i32>) align(2) allocator(omp_thread_mem_alloc)
@@ -3289,6 +3289,10 @@ func.func @omp_allocate_dir(%arg0 : memref<i32>, %arg1 : memref<i32>) -> () {
   // 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)
 
+  // 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)
+
   return
 }