diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td --- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td +++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td @@ -786,6 +786,157 @@ let assemblyFormat = "attr-dict"; } +//===---------------------------------------------------------------------===// +// 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:$if_expr, + Optional:$device, + Variadic:$use_device_ptr, + Variadic:$use_device_addr, + Variadic:$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($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:$if_expr, + Optional:$device, + UnitAttr:$nowait, + Variadic:$map_operands, + I64ArrayAttr:$map_types); + + let assemblyFormat = [{ + oilist(`if` `(` $if_expr `:` type($if_expr) `)` + | `device` `(` $device `:` type($device) `)` + | `nowait` $nowait) + `map` `(` custom($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:$if_expr, + Optional:$device, + UnitAttr:$nowait, + Variadic:$map_operands, + I64ArrayAttr:$map_types); + + let assemblyFormat = [{ + oilist(`if` `(` $if_expr `:` type($if_expr) `)` + | `device` `(` $device `:` type($device) `)` + | `nowait` $nowait) + `map` `(` custom($map_operands, type($map_operands), $map_types) `)` + attr-dict + }]; + + let hasVerifier = 1; +} + //===----------------------------------------------------------------------===// // 2.13.7 flush Construct //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp --- 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 #include "mlir/Dialect/OpenMP/OpenMPOpsDialect.cpp.inc" @@ -536,6 +537,191 @@ return success(); } +//===----------------------------------------------------------------------===// +// 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>), (always, close, from +/// -> %2 : !llvm.ptr>)) +static ParseResult +parseMapClause(OpAsmParser &parser, + SmallVectorImpl &map_operands, + SmallVectorImpl &map_operand_types, ArrayAttr &map_types) { + StringRef mapTypeMod; + OpAsmParser::UnresolvedOperand arg1; + Type arg1Type; + IntegerAttr arg2; + SmallVector 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>( + mapTypeBits)); + mapTypesVec.push_back(arg2); + return success(); + }; + + if (parser.parseCommaSeparatedList(parseMap)) + return failure(); + + SmallVector 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>( + flag); + }; + + assert(map_operands.size() == map_types.size()); + + for (unsigned i = 0, e = map_operands.size(); i < e; i++) { + int64_t mapTypeBits = 0x00; + auto mapOp = map_operands[i]; + auto mapTypeOp = map_types[i]; + + assert(mapTypeOp.isa()); + mapTypeBits = mapTypeOp.cast().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(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>( + flag); + }; + if (map_operands.size() != map_types.size()) + return failure(); + + for (const auto &mapTypeOp : map_types) { + int64_t mapTypeBits = 0x00; + + if (!mapTypeOp.isa()) + return failure(); + + mapTypeBits = mapTypeOp.cast().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(op) && del) + return failure(); + if (isa(op) && (from || del)) + return failure(); + if (isa(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 --- a/mlir/test/Dialect/OpenMP/ops.mlir +++ b/mlir/test/Dialect/OpenMP/ops.mlir @@ -451,6 +451,26 @@ return } +// CHECK-LABEL: omp_target_data +func.func @omp_target_data (%if_cond : i1, %device : si32, %device_ptr: memref, %device_addr: memref, %map1: memref, %map2: memref) -> () { + // CHECK: omp.target_data if(%[[VAL_0:.*]] : i1) device(%[[VAL_1:.*]] : si32) map((always, from -> %[[VAL_2:.*]] : memref)) + omp.target_data if(%if_cond : i1) device(%device : si32) map((always, from -> %map1 : memref)){} + + // CHECK: omp.target_data use_device_ptr(%[[VAL_3:.*]] : memref) use_device_addr(%[[VAL_4:.*]] : memref) map((close, present, to -> %[[VAL_2:.*]] : memref)) + omp.target_data use_device_ptr(%device_ptr : memref) use_device_addr(%device_addr : memref) map((close, present, to -> %map1 : memref)){} + + // CHECK: omp.target_data map((tofrom -> %[[VAL_2]] : memref), (alloc -> %[[VAL_5:.*]] : memref)) + omp.target_data map((tofrom -> %map1 : memref), (alloc -> %map2 : memref)){} + + // CHECK: omp.target_enter_data if(%[[VAL_0]] : i1) device(%[[VAL_1]] : si32) nowait map((alloc -> %[[VAL_2]] : memref)) + omp.target_enter_data if(%if_cond : i1) device(%device : si32) nowait map((alloc -> %map1 : memref)) + + // CHECK: omp.target_exit_data if(%[[VAL_0]] : i1) device(%[[VAL_1]] : si32) nowait map((release -> %[[VAL_5]] : memref)) + omp.target_exit_data if(%if_cond : i1) device(%device : si32) nowait map((release -> %map2 : memref)) + + return +} + // CHECK-LABEL: omp_target_pretty func.func @omp_target_pretty(%if_cond : i1, %device : si32, %num_threads : i32) -> () { // CHECK: omp.target if({{.*}}) device({{.*}})