[MLIR][OpenMP] Added target data, exit data, and enter data operation definition...

This includes a basic implementation for the OpenMP 5.1 Target Data, Target Exit Data and Target Enter Data constructs
operation.

TODO:
  - Depend clause support for Target Enter and Exit Data.
  - Mapper and Iterator value support for Map Type Modifiers.
  - Verifier for the operations.

Co-authored-by: abidmalikwaterloo <amalik@bnl.gov>
Co-authored-by: raghavendra <Raghu.Maddhipatla@amd.com>
Differential Revision: https://reviews.llvm.org/D131915

diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index af02e79..cb5571a 100644 (file)
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -817,6 +817,158 @@ def FlushOp : OpenMP_Op<"flush"> {
     }
   }];
 }
+
+//===---------------------------------------------------------------------===//
+// 2.14.2 target data Construct
+//===---------------------------------------------------------------------===//
+
+def Target_DataOp: OpenMP_Op<"target_data", [AttrSizedOperandSegments]>{
+  let summary = "target data construct";
+  let description = [{
+    Map variables to a device data environment for the extent of the region.
+
+    The omp target data directive maps variables to a device data
+    environment, and defines the lexical scope of the data environment
+    that is created. The omp target data directive can reduce data copies
+    to and from the offloading device when multiple target regions are using
+    the same data.
+
+    The optional $if_expr parameter specifies a boolean result of a
+    conditional check. If this value is 1 or is not provided then the target
+    region runs on a device, if it is 0 then the target region is executed
+    on the host device.
+
+    The optional $device parameter specifies the device number for the target
+    region.
+
+    The optional $use_device_ptr specifies the device pointers to the
+    corresponding list items in the device data environment.
+
+    The optional $use_device_addr specifies the address of the objects in the
+    device data enviornment.
+
+    The $map_operands specifies the locator-list operands of the map clause.
+
+    The $map_types specifies the types and modifiers for the map clause.
+
+    TODO:  depend clause and map_type_modifier values iterator and mapper.
+  }];
+
+  let arguments = (ins Optional<I1>:$if_expr,
+                       Optional<AnyInteger>:$device,
+                       Variadic<AnyType>:$use_device_ptr,
+                       Variadic<AnyType>:$use_device_addr,
+                       Variadic<AnyType>:$map_operands,
+                       I64ArrayAttr:$map_types);
+
+  let regions = (region AnyRegion:$region);
+
+  let assemblyFormat = [{
+    oilist(`if` `(` $if_expr `:` type($if_expr) `)`
+    | `device` `(` $device `:` type($device) `)`
+    | `use_device_ptr` `(` $use_device_ptr `:` type($use_device_ptr) `)`
+    | `use_device_addr` `(` $use_device_addr `:` type($use_device_addr) `)`)
+    `map` `(` custom<MapClause>($map_operands, type($map_operands), $map_types) `)`
+    $region attr-dict
+  }];
+
+  let hasVerifier = 1;
+}
+
+//===---------------------------------------------------------------------===//
+// 2.14.3 target enter data Construct
+//===---------------------------------------------------------------------===//
+
+def Target_EnterDataOp: OpenMP_Op<"target_enter_data",
+                                                 [AttrSizedOperandSegments]>{
+  let  summary = "target enter data construct";
+  let description = [{
+    The target enter data directive specifies that variables are mapped to
+    a device data environment. The target enter data directive is a
+    stand-alone directive.
+
+    The optional $if_expr parameter specifies a boolean result of a
+    conditional check. If this value is 1 or is not provided then the target
+    region runs on a device, if it is 0 then the target region is executed on
+    the host device.
+
+    The optional $device parameter specifies the device number for the
+    target region.
+
+    The optional $nowait eliminates the implicit barrier so the parent task
+    can make progress even if the target task is not yet completed.
+
+    The $map_operands specifies the locator-list operands of the map clause.
+
+    The $map_types specifies the types and modifiers for the map clause.
+
+    TODO:  depend clause and map_type_modifier values iterator and mapper.
+  }];
+
+  let arguments = (ins Optional<I1>:$if_expr,
+                       Optional<AnyInteger>:$device,
+                       UnitAttr:$nowait,
+                       Variadic<AnyType>:$map_operands,
+                       I64ArrayAttr:$map_types);
+
+  let assemblyFormat = [{
+    oilist(`if` `(` $if_expr `:` type($if_expr) `)`
+    | `device` `(` $device `:` type($device) `)`
+    | `nowait` $nowait)
+    `map` `(` custom<MapClause>($map_operands, type($map_operands), $map_types) `)`
+    attr-dict
+   }];
+
+  let hasVerifier = 1;
+}
+
+//===---------------------------------------------------------------------===//
+// 2.14.4 target exit data Construct
+//===---------------------------------------------------------------------===//
+
+def Target_ExitDataOp: OpenMP_Op<"target_exit_data",
+                                                 [AttrSizedOperandSegments]>{
+  let  summary = "target exit data construct";
+  let description = [{
+    The target exit data directive specifies that variables are mapped to a
+    device data environment. The target exit data directive is
+    a stand-alone directive.
+
+    The optional $if_expr parameter specifies a boolean result of a
+    conditional check. If this value is 1 or is not provided then the target
+    region runs on a device, if it is 0 then the target region is executed
+    on the host device.
+
+    The optional $device parameter specifies the device number for the
+    target region.
+
+    The optional $nowait eliminates the implicit barrier so the parent
+    task can make progress even if the target task is not yet completed.
+
+    The $map_operands specifies the locator-list operands of the map clause.
+
+    The $map_types specifies the types and modifiers for the map clause.
+
+    TODO:  depend clause and map_type_modifier values iterator and mapper.
+  }];
+
+  let arguments = (ins Optional<I1>:$if_expr,
+                       Optional<AnyInteger>:$device,
+                       UnitAttr:$nowait,
+                       Variadic<AnyType>:$map_operands,
+                       I64ArrayAttr:$map_types);
+
+  let assemblyFormat = [{
+    oilist(`if` `(` $if_expr `:` type($if_expr) `)`
+    | `device` `(` $device `:` type($device) `)`
+    | `nowait` $nowait)
+    `map` `(` custom<MapClause>($map_operands, type($map_operands), $map_types) `)`
+    attr-dict
+   }];
+
+  let hasVerifier = 1;
+}
+
 //===----------------------------------------------------------------------===//
 // 2.14.5 target construct
 //===----------------------------------------------------------------------===//

diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index daa7ff8..6bf4da1 100644 (file)
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -22,6 +22,7 @@
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/TypeSwitch.h"
+#include "llvm/Frontend/OpenMP/OMPConstants.h"
 #include <cstddef>
 
 #include "mlir/Dialect/OpenMP/OpenMPOpsDialect.cpp.inc"
@@ -553,6 +554,191 @@ static LogicalResult verifySynchronizationHint(Operation *op, uint64_t hint) {
 }
 
 //===----------------------------------------------------------------------===//
+// Parser, printer and verifier for Target Data
+//===----------------------------------------------------------------------===//
+/// 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
+parseMapClause(OpAsmParser &parser,
+               SmallVectorImpl<OpAsmParser::UnresolvedOperand> &map_operands,
+               SmallVectorImpl<Type> &map_operand_types, ArrayAttr &map_types) {
+  StringRef mapTypeMod;
+  OpAsmParser::UnresolvedOperand arg1;
+  Type arg1Type;
+  IntegerAttr arg2;
+  SmallVector<IntegerAttr> mapTypesVec;
+  llvm::omp::OpenMPOffloadMappingFlags mapTypeBits;
+
+  auto parseTypeAndMod = [&]() -> ParseResult {
+    if (parser.parseKeyword(&mapTypeMod))
+      return failure();
+
+    if (mapTypeMod == "always")
+      mapTypeBits |= llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_ALWAYS;
+    if (mapTypeMod == "close")
+      mapTypeBits |= llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_CLOSE;
+    if (mapTypeMod == "present")
+      mapTypeBits |= llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_PRESENT;
+
+    if (mapTypeMod == "to")
+      mapTypeBits |= llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_TO;
+    if (mapTypeMod == "from")
+      mapTypeBits |= llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_FROM;
+    if (mapTypeMod == "tofrom")
+      mapTypeBits |= llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_TO |
+                     llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_FROM;
+    if (mapTypeMod == "delete")
+      mapTypeBits |= llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_DELETE;
+    return success();
+  };
+
+  auto parseMap = [&]() -> ParseResult {
+    mapTypeBits = llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_NONE;
+
+    if (parser.parseLParen() ||
+        parser.parseCommaSeparatedList(parseTypeAndMod) ||
+        parser.parseArrow() || parser.parseOperand(arg1) ||
+        parser.parseColon() || parser.parseType(arg1Type) ||
+        parser.parseRParen())
+      return failure();
+    map_operands.push_back(arg1);
+    map_operand_types.push_back(arg1Type);
+    arg2 = parser.getBuilder().getIntegerAttr(
+        parser.getBuilder().getI64Type(),
+        static_cast<
+            std::underlying_type_t<llvm::omp::OpenMPOffloadMappingFlags>>(
+            mapTypeBits));
+    mapTypesVec.push_back(arg2);
+    return success();
+  };
+
+  if (parser.parseCommaSeparatedList(parseMap))
+    return failure();
+
+  SmallVector<Attribute> mapTypesAttr(mapTypesVec.begin(), mapTypesVec.end());
+  map_types = ArrayAttr::get(parser.getContext(), mapTypesAttr);
+  return success();
+}
+
+static void printMapClause(OpAsmPrinter &p, Operation *op,
+                           OperandRange map_operands,
+                           TypeRange map_operand_types, 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);
+  };
+
+  assert(map_operands.size() == map_types.size());
+
+  for (unsigned i = 0, e = map_operands.size(); i < e; i++) {
+    int64_t mapTypeBits = 0x00;
+    Value mapOp = map_operands[i];
+    Attribute mapTypeOp = map_types[i];
+
+    assert(mapTypeOp.isa<mlir::IntegerAttr>());
+    mapTypeBits = mapTypeOp.cast<mlir::IntegerAttr>().getInt();
+
+    bool always = bitAnd(mapTypeBits,
+                         llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_ALWAYS);
+    bool close = bitAnd(mapTypeBits,
+                        llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_CLOSE);
+    bool present = bitAnd(
+        mapTypeBits, llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_PRESENT);
+
+    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);
+
+    std::string typeModStr, typeStr;
+    llvm::raw_string_ostream typeMod(typeModStr), type(typeStr);
+
+    if (always)
+      typeMod << "always, ";
+    if (close)
+      typeMod << "close, ";
+    if (present)
+      typeMod << "present, ";
+
+    if (to)
+      type << "to";
+    if (from)
+      type << "from";
+    if (del)
+      type << "delete";
+    if (type.str().empty())
+      type << (isa<ExitDataOp>(op) ? "release" : "alloc");
+
+    p << '(' << typeMod.str() << type.str() << " -> " << mapOp << " : "
+      << mapOp.getType() << ')';
+    if (i + 1 < e)
+      p << ", ";
+  }
+}
+
+static LogicalResult verifyMapClause(Operation *op, OperandRange map_operands,
+                                     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())
+    return failure();
+
+  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<EnterDataOp>(op) && (from || del))
+      return failure();
+    if (isa<ExitDataOp>(op) && to)
+      return failure();
+  }
+
+  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());
+}
+
+//===----------------------------------------------------------------------===//
 // ParallelOp
 //===----------------------------------------------------------------------===//
 

diff --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir
index 44b489a..56d898e 100644 (file)
--- a/mlir/test/Dialect/OpenMP/ops.mlir
+++ b/mlir/test/Dialect/OpenMP/ops.mlir
@@ -495,6 +495,26 @@ func.func @omp_target(%if_cond : i1, %device : si32,  %num_threads : i32) -> ()
     return
 }
 
+// CHECK-LABEL: omp_target_data
+func.func @omp_target_data (%if_cond : i1, %device : si32, %device_ptr: memref<i32>, %device_addr: memref<?xi32>, %map1: memref<?xi32>, %map2: memref<?xi32>) -> () {
+    // CHECK: omp.target_data if(%[[VAL_0:.*]] : i1) device(%[[VAL_1:.*]] : si32) map((always, from -> %[[VAL_2:.*]] : memref<?xi32>))
+    omp.target_data if(%if_cond : i1) device(%device : si32) map((always, from -> %map1 : memref<?xi32>)){}
+
+    // CHECK: omp.target_data use_device_ptr(%[[VAL_3:.*]] : memref<i32>) use_device_addr(%[[VAL_4:.*]] : memref<?xi32>) map((close, present, to -> %[[VAL_2:.*]] : memref<?xi32>))
+    omp.target_data use_device_ptr(%device_ptr : memref<i32>) use_device_addr(%device_addr : memref<?xi32>) map((close, present, to -> %map1 : memref<?xi32>)){}
+
+    // CHECK: omp.target_data map((tofrom -> %[[VAL_2]] : memref<?xi32>), (alloc -> %[[VAL_5:.*]] : memref<?xi32>))
+    omp.target_data map((tofrom -> %map1 : memref<?xi32>), (alloc -> %map2 : memref<?xi32>)){}
+
+    // CHECK: omp.target_enter_data if(%[[VAL_0]] : i1) device(%[[VAL_1]] : si32) nowait map((alloc -> %[[VAL_2]] : memref<?xi32>))
+    omp.target_enter_data if(%if_cond : i1) device(%device : si32) nowait map((alloc -> %map1 : memref<?xi32>))
+
+    // CHECK: omp.target_exit_data if(%[[VAL_0]] : i1) device(%[[VAL_1]] : si32) nowait map((release -> %[[VAL_5]] : memref<?xi32>))
+    omp.target_exit_data if(%if_cond : i1) device(%device : si32) nowait map((release -> %map2 : memref<?xi32>))
+
+    return
+}
+
 // CHECK-LABEL: omp_target_pretty
 func.func @omp_target_pretty(%if_cond : i1, %device : si32,  %num_threads : i32) -> () {
     // CHECK: omp.target if({{.*}}) device({{.*}})