[MLIR][OpenMP] Added map clause support for Target

[TIFitis]

Summary

Added map clause support for the OMP Target directive with test.

[raghavendhra]

LGTM

[TIFitis]

ping for reviews. Thanks :)

[kiranchandramohan]

LG. Couple of Nit comments.

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

815

Nit: Add a test for the verifier as well.

mlir/test/Dialect/OpenMP/ops.mlir

492

Nit: You mean map type? Add more variants.

[TIFitis]

Addressed reviewer comments.

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

815

Could you please let me know where to add tests for the verifier?

mlir/test/Dialect/OpenMP/ops.mlir

492

Hi, the map clause itself is optional for the target directive.

I've added more variants.

[kiranchandramohan]

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

815

mlir/test/Dialect/OpenMP/invalid.mlir

[TIFitis]

Changed verifier to emit error instead of silent failure.
Added verifier tests.

[TIFitis]

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

815

I've added the verifier tests.

Let me know if the error msg is alright. Thanks :)

[kiranchandramohan]

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

815

I think values are not capitalised. Please check the other messages and keep it consistent.

[TIFitis]

Changed error message to lwoer case.

[kiranchandramohan]

LG

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

 class OpenMP_Attr<string name, string attrMnemonic,
                 list<Trait> traits = [],
                 string baseCppClass = "::mlir::Attribute">
     : AttrDef<OpenMP_Dialect, name, traits, baseCppClass> {
   let mnemonic = attrMnemonic;
 }
 
 def IsDeviceAttr : OpenMP_Attr<"IsDevice", "isdevice"> {
-  let parameters = (ins 
+  let parameters = (ins
     "bool":$is_device
   );
-  
   let assemblyFormat = "`<` struct(params) `>`";
 }
 
 //===----------------------------------------------------------------------===//
 // Runtime library flag's attribute that holds information for lowering to LLVM
 //===----------------------------------------------------------------------===//
 
 def FlagsAttr : OpenMP_Attr<"Flags", "flags"> {
-  let parameters = (ins 
+  let parameters = (ins
     DefaultValuedParameter<"uint32_t", "0">:$debug_kind,
     DefaultValuedParameter<"bool", "false">:$assume_teams_oversubscription,
     DefaultValuedParameter<"bool", "false">:$assume_threads_oversubscription,
     DefaultValuedParameter<"bool", "false">:$assume_no_thread_state,
     DefaultValuedParameter<"bool", "false">:$assume_no_nested_parallelism
   );
-  
   let assemblyFormat = "`<` struct(params) `>`";
 }
 
 def TargetAttr : OpenMP_Attr<"Target", "target"> {
   let parameters = (ins
     StringRefParameter<>:$target_cpu,
     StringRefParameter<>:$target_features
   );
let description = [{
 
     The optional $device parameter specifies the device number for the target region.
 
     The optional $thread_limit specifies the limit on the number of threads
 
     The optional $nowait elliminates the implicit barrier so the parent task can make progress
     even if the target task is not yet completed.
 
-    TODO:  map, is_device_ptr, depend, defaultmap, in_reduction
+    TODO:  is_device_ptr, depend, defaultmap, in_reduction
 
   }];
 
   let arguments = (ins Optional<I1>:$if_expr,
                        Optional<AnyInteger>:$device,
                        Optional<AnyInteger>:$thread_limit,
-                       UnitAttr:$nowait);
+                       UnitAttr:$nowait,
+                       Variadic<OpenMP_PointerLikeType>:$map_operands,
+                       OptionalAttr<I64ArrayAttr>:$map_types);
 
   let regions = (region AnyRegion:$region);
 
   let assemblyFormat = [{
     oilist( `if` `(` $if_expr `)`
-          | `device` `(` $device `:` type($device) `)`
-          | `thread_limit` `(` $thread_limit `:` type($thread_limit) `)`
-          | `nowait` $nowait
+    | `map` `(` custom<MapClause>($map_operands, type($map_operands), $map_types) `)`
     ) $region attr-dict
   }];
+
+  let hasVerifier = 1;
 }
 
 
 //===----------------------------------------------------------------------===//
 // 2.16 master Construct
 //===----------------------------------------------------------------------===//
 def MasterOp : OpenMP_Op<"master"> {
   let summary = "master construct";

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

return op->emitOpError() << "the hints omp_sync_hint_uncontended and "
                                 "omp_sync_hint_contended cannot be combined";
   if (nonspeculative && speculative)
     return op->emitOpError() << "the hints omp_sync_hint_nonspeculative and "
                                 "omp_sync_hint_speculative cannot be combined.";
   return success();
 }
 
 //===----------------------------------------------------------------------===//
-// Parser, printer and verifier for Target Data
+// Parser, printer and verifier for Target
 //===----------------------------------------------------------------------===//
 /// Parses a Map Clause.
 ///
 /// map-clause = `map (` ( `(` `always, `? `close, `? `present, `? ( `to` |
 /// `from` | `delete` ) ` -> ` symbol-ref ` : ` type(symbol-ref) `)` )+ `)`
 /// Eg: map((release -> %1 : !llvm.ptr<array<1024 x i32>>), (always, close, from
 /// -> %2 : !llvm.ptr<array<1024 x i32>>))
 static ParseResult
for (unsigned i = 0, e = map_operands.size(); i < e; i++) {
     p << '(' << typeMod.str() << type.str() << " -> " << mapOp << " : "
       << mapOp.getType() << ')';
     if (i + 1 < e)
       p << ", ";
   }
 }
 
 static LogicalResult verifyMapClause(Operation *op, OperandRange map_operands,
-                                     ArrayAttr map_types) {
+                                     std::optional<ArrayAttr> map_types) {
   // Helper function to get bitwise AND of `value` and 'flag'
   auto bitAnd = [](int64_t value,
                    llvm::omp::OpenMPOffloadMappingFlags flag) -> bool {
     return value &
            static_cast<
                std::underlying_type_t<llvm::omp::OpenMPOffloadMappingFlags>>(
                flag);
   };
-  if (map_operands.size() != map_types.size())
+  if (!map_types.has_value())
+    return success();
+
+  if (map_operands.size() != map_types->size())
     return failure();
 
-  for (const auto &mapTypeOp : map_types) {
+  for (const auto &mapTypeOp : *map_types) {
     int64_t mapTypeBits = 0x00;
 
     if (!mapTypeOp.isa<mlir::IntegerAttr>())
       return failure();
 
     mapTypeBits = mapTypeOp.cast<mlir::IntegerAttr>().getInt();
 
     bool to =
         bitAnd(mapTypeBits, llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_TO);
     bool from =
         bitAnd(mapTypeBits, llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_FROM);
     bool del = bitAnd(mapTypeBits,
                       llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_DELETE);
 
-    if (isa<DataOp>(op) && del)
-      return failure();
+    if ((isa<DataOp>(op) || isa<TargetOp>(op)) && del)
+      return emitError(op->getLoc(),
+                       "to, from, tofrom and alloc map types are permitted");
     if (isa<EnterDataOp>(op) && (from || del))
-      return failure();
+      return emitError(op->getLoc(), "to and alloc map types are permitted");
     if (isa<ExitDataOp>(op) && to)
-      return failure();
+      return emitError(op->getLoc(),
+                       "from, release and delete map types are permitted");
   }
 
   return success();
 }
 
 LogicalResult DataOp::verify() {
   return verifyMapClause(*this, getMapOperands(), getMapTypes());
 }
 
 LogicalResult EnterDataOp::verify() {
   return verifyMapClause(*this, getMapOperands(), getMapTypes());
 }
 
 LogicalResult ExitDataOp::verify() {
   return verifyMapClause(*this, getMapOperands(), getMapTypes());
 }
 
+LogicalResult TargetOp::verify() {
+  return verifyMapClause(*this, getMapOperands(), getMapTypes());
+}
+
 //===----------------------------------------------------------------------===//
 // ParallelOp
 //===----------------------------------------------------------------------===//
 
 void ParallelOp::build(OpBuilder &builder, OperationState &state,
                        ArrayRef<NamedAttribute> attributes) {
   ParallelOp::build(
       builder, state, /*if_expr_var=*/nullptr, /*num_threads_var=*/nullptr,

mlir/test/Dialect/OpenMP/invalid.mlir

 
 func.func @omp_threadprivate() {
   %1 = llvm.mlir.addressof @_QFsubEx : !llvm.ptr<i32>
   // expected-error @below {{op failed to verify that all of {sym_addr, tls_addr} have same type}}
   %2 = omp.threadprivate %1 : !llvm.ptr<i32> -> memref<i32>
   return
 }
 
+// -----
+
+func.func @omp_target(%map1: memref<?xi32>) {
+  // expected-error @below {{to, from, tofrom and alloc map types are permitted}}
+  omp.target map((delete -> %map1 : memref<?xi32>)){}
+  return
+}
+
+// -----
+
+func.func @omp_target_data(%map1: memref<?xi32>) {
+  // expected-error @below {{to, from, tofrom and alloc map types are permitted}}
+  omp.target_data map((delete -> %map1 : memref<?xi32>)){}
+  return
+}
+
+// -----
+
+func.func @omp_target_enter_data(%map1: memref<?xi32>) {
+  // expected-error @below {{to and alloc map types are permitted}}
+  omp.target_enter_data map((from -> %map1 : memref<?xi32>)){}
+  return
+}
+
+// -----
+
+func.func @omp_target_exit_data(%map1: memref<?xi32>) {
+  // expected-error @below {{from, release and delete map types are permitted}}
+  omp.target_exit_data map((to -> %map1 : memref<?xi32>)){}
+  return
+}
+
 llvm.mlir.global internal @_QFsubEx() : i32

mlir/test/Dialect/OpenMP/ops.mlir

func.func @omp_simdloop_pretty_multiple(%lb1 : index, %ub1 : index, %step1 : index, %lb2 : index, %ub2 : index, %step2 : index) -> () {
   // CHECK: omp.simdloop for (%{{.*}}, %{{.*}}) : index = (%{{.*}}, %{{.*}}) to (%{{.*}}, %{{.*}}) step (%{{.*}}, %{{.*}})
   omp.simdloop for (%iv1, %iv2) : index = (%lb1, %lb2) to (%ub1, %ub2) step (%step1, %step2) {
     omp.yield
   }
   return
 }
 
 // CHECK-LABEL: omp_target
-func.func @omp_target(%if_cond : i1, %device : si32,  %num_threads : i32) -> () {
+func.func @omp_target(%if_cond : i1, %device : si32,  %num_threads : i32, %map1: memref<?xi32>, %map2: memref<?xi32>) -> () {
 
     // Test with optional operands; if_expr, device, thread_limit, private, firstprivate and nowait.
     // CHECK: omp.target if({{.*}}) device({{.*}}) thread_limit({{.*}}) nowait
     "omp.target"(%if_cond, %device, %num_threads) ({
        // CHECK: omp.terminator
        omp.terminator
-    }) {nowait, operand_segment_sizes = array<i32: 1,1,1>} : ( i1, si32, i32 ) -> ()
+    }) {nowait, operand_segment_sizes = array<i32: 1,1,1,0>} : ( i1, si32, i32 ) -> ()
+
+    // Test with optional map clause.
+    // CHECK: omp.target map((tofrom -> %{{.*}} : memref<?xi32>), (alloc -> %{{.*}} : memref<?xi32>)) {
+    omp.target map((tofrom -> %map1 : memref<?xi32>), (alloc -> %map2 : memref<?xi32>)){}
+
+    // CHECK: omp.target map((to -> %{{.*}} : memref<?xi32>), (always, from -> %{{.*}} : memref<?xi32>)) {
+    omp.target map((to -> %map1 : memref<?xi32>), (always, from -> %map2 : memref<?xi32>)){}
 
     // CHECK: omp.barrier
     omp.barrier
 
     return
 }
 
 // CHECK-LABEL: omp_target_data