[OpenMP][MLIR] Refactor and extend current map support by adding MapInfoOp and DataBoundsOp operations to the OpenMP Dialect

This patch adds two new operations:

The first is the DataBoundsOp, which is based on OpenACC's DataBoundsOp,
which holds stride, index, extent, lower bound and upper bounds
which will be used in future follow up patches to perform initial
array sectioning of mapped arrays, and Fortran pointer and
allocatable mapping. Similarly to OpenACC, this new OpenMP
DataBoundsOp also comes with a new OpenMP type, which
helps to restrict operations to accepting only
DataBoundsOp as an input or output where necessary
(or other related operations that implement this type as
a return).

The patch also adds the MapInfoOp which rolls up some of
the old map information stored in target
operations into this new operation, and adds new
information that will be utilised in the lowering of mapped
variables, e.g. the aforementioned DataBoundsOp, but also a
new ByCapture OpenMP MLIR attribute, and isImplicit boolean
attribute. Both the ByCapture and isImplicit arguments will
affect the lowering from the OpenMP dialect to LLVM-IR in
minor but important ways, such as shifting the final maptype
or generating different load/store combinations to maintain
semantics with the OpenMP standard and alignment with the
current Clang OpenMP output as best as possible.

This MapInfoOp operation is slightly based on OpenACC's
DataEntryOp, the main difference other than some slightly
different fields (e,g, isImplicit/MapType/ByCapture) is that
OpenACC's data operations "inherit" (the MLIR ODS
equivalent) from this operation, whereas in OpenMP operations
that utilise MapInfoOp's are composed of/contain them.

A series of these MapInfoOp (one per map clause list item) is
now held by target operations that represent OpenMP
directives that utilise map clauses, e.g. TargetOp. MapInfoOp's
do not have their own specialised lowering to LLVM-IR, instead
the lowering is dependent on the particular container of the
MapInfoOp's, e.g. TargetOp has its own lowering to LLVM-IR
which utilised the information stored inside of MapInfoOp's to
affect it's lowering and the end result of the LLVM-IR generated,
which in turn can differ for host and device.

This patch contains these operations, minor changes to the
printing and parsing to support them, changes to tests (only
those relevant to this segment of the patch, other test
additions and changes are in other dependent
patches in this series) and some alterations to the OpenMPToLLVM
rewriter to support the new OpenMP type and operations.

This patch is one in a series that are dependent on each
other:

https://reviews.llvm.org/D158734
https://reviews.llvm.org/D158735
https://reviews.llvm.org/D158737

Reviewers: kiranchandramohan, TIFitis, razvanlupusoru

Differential Revision: https://reviews.llvm.org/D158732

Author: agozillon

mlir/include/mlir/Dialect/OpenMP/CMakeLists.txt

@@ -7,6 +7,8 @@ mlir_tablegen(OpenMPOpsDialect.h.inc -gen-dialect-decls -dialect=omp)
 mlir_tablegen(OpenMPOpsDialect.cpp.inc -gen-dialect-defs -dialect=omp)
 mlir_tablegen(OpenMPOps.h.inc -gen-op-decls)
 mlir_tablegen(OpenMPOps.cpp.inc -gen-op-defs)
+mlir_tablegen(OpenMPOpsTypes.h.inc -gen-typedef-decls -typedefs-dialect=omp)
+mlir_tablegen(OpenMPOpsTypes.cpp.inc -gen-typedef-defs -typedefs-dialect=omp)
 mlir_tablegen(OpenMPOpsEnums.h.inc -gen-enum-decls)
 mlir_tablegen(OpenMPOpsEnums.cpp.inc -gen-enum-defs)
 mlir_tablegen(OpenMPOpsAttributes.h.inc -gen-attrdef-decls -attrdefs-dialect=omp)

mlir/include/mlir/Dialect/OpenMP/OpenMPDialect.h

@@ -22,6 +22,9 @@
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
+#define GET_TYPEDEF_CLASSES
+#include "mlir/Dialect/OpenMP/OpenMPOpsTypes.h.inc"
+
 #include "mlir/Dialect/OpenMP/OpenMPOpsDialect.h.inc"
 #include "mlir/Dialect/OpenMP/OpenMPOpsEnums.h.inc"
 #include "mlir/Dialect/OpenMP/OpenMPTypeInterfaces.h.inc"

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

@@ -29,6 +29,7 @@ def OpenMP_Dialect : Dialect {
   let cppNamespace = "::mlir::omp";
   let dependentDialects = ["::mlir::LLVM::LLVMDialect, ::mlir::func::FuncDialect"];
   let useDefaultAttributePrinterParser = 1;
+  let useDefaultTypePrinterParser = 1;
 }
 
 // OmpCommon requires definition of OpenACC_Dialect.
@@ -89,6 +90,10 @@ def IntLikeType : AnyTypeOf<[AnyInteger, Index]>;
 def OpenMP_PointerLikeType : TypeAlias<OpenMP_PointerLikeTypeInterface,
 	"OpenMP-compatible variable type">;
 
+class OpenMP_Type<string name, string typeMnemonic> : TypeDef<OpenMP_Dialect, name> {
+  let mnemonic = typeMnemonic;
+}
+
 //===----------------------------------------------------------------------===//
 //  2.12.7 Declare Target Directive
 //===----------------------------------------------------------------------===//
@@ -1004,6 +1009,220 @@ def FlushOp : OpenMP_Op<"flush"> {
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// Map related constructs
+//===----------------------------------------------------------------------===//
+
+def CaptureThis : I32EnumAttrCase<"This", 0>;
+def CaptureByRef : I32EnumAttrCase<"ByRef", 1>;
+def CaptureByCopy : I32EnumAttrCase<"ByCopy", 2>;
+def CaptureVLAType : I32EnumAttrCase<"VLAType", 3>;
+
+def VariableCaptureKind : I32EnumAttr<
+    "VariableCaptureKind",
+    "variable capture kind",
+    [CaptureThis, CaptureByRef, CaptureByCopy, CaptureVLAType]> {
+  let genSpecializedAttr = 0;
+  let cppNamespace = "::mlir::omp";
+}
+
+def VariableCaptureKindAttr : EnumAttr<OpenMP_Dialect, VariableCaptureKind,
+                                       "variable_capture_kind"> {
+  let assemblyFormat = "`(` $value `)`";
+}
+
+def DataBoundsType : OpenMP_Type<"DataBounds", "data_bounds_ty"> {
+  let summary = "Type for representing omp data clause bounds information";
+}
+
+def DataBoundsOp : OpenMP_Op<"bounds",
+    [AttrSizedOperandSegments, NoMemoryEffect]> {
+  let summary = "Represents normalized bounds information for map clauses.";
+
+  let description = [{
+    This operation is a variation on the OpenACC dialects DataBoundsOp. Within
+    the OpenMP dialect it stores the bounds/range of data to be mapped to a
+    device specified by map clauses on target directives. Within,
+    the OpenMP dialect the DataBoundsOp is associated with MapInfoOp,
+    helping to store bounds information for the mapped variable.
+
+    It is used to support OpenMP array sectioning, Fortran pointer and
+    allocatable mapping and pointer/allocatable member of derived types.
+    In all cases the DataBoundsOp holds information on the section of
+    data to be mapped. Such as the upper bound and lower bound of the
+    section of data to be mapped. This information is currently
+    utilised by the LLVM-IR lowering to help generate instructions to
+    copy data to and from the device when processing target operations.
+
+    The example below copys a section of a 10-element array; all except the
+    first element, utilising OpenMP array sectioning syntax where array
+    subscripts are provided to specify the bounds to be mapped to device.
+    To simplify the examples, the constants are used directly, in reality
+    they will be MLIR SSA values.
+
+    C++:
+    ```
+    int array[10];
+    #pragma target map(array[1:9])
+    ```
+    =>
+    ```mlir
+    omp.bounds lower_bound(1) upper_bound(9) extent(9) start_idx(0)
+    ```
+
+    Fortran:
+    ```
+    integer :: array(1:10)
+    !$target map(array(2:10))
+    ```
+    =>
+    ```mlir
+    omp.bounds lower_bound(1) upper_bound(9) extent(9) start_idx(1)
+    ```
+
+    For Fortran pointers and allocatables (as well as those that are
+    members of derived types) the bounds information is provided by
+    the Fortran compiler and runtime through descriptor information.
+
+    A basic pointer example can be found below (constants again
+    provided for simplicity, where in reality SSA values will be
+    used, in this case that point to data yielded by Fortran's
+    descriptors):
+
+    Fortran:
+    ```
+    integer, pointer :: ptr(:)
+    allocate(ptr(10))
+    !$target map(ptr)
+    ```
+    =>
+    ```mlir
+    omp.bounds lower_bound(0) upper_bound(9) extent(10) start_idx(1)
+    ```
+
+    This operation records the bounds information in a normalized fashion
+    (zero-based). This works well with the `PointerLikeType`
+    requirement in data clauses - since a `lower_bound` of 0 means looking
+    at data at the zero offset from pointer.
+
+    This operation must have an `upper_bound` or `extent` (or both are allowed -
+    but not checked for consistency). When the source language's arrays are
+    not zero-based, the `start_idx` must specify the zero-position index.
+  }];
+
+  let arguments = (ins Optional<IntLikeType>:$lower_bound,
+                       Optional<IntLikeType>:$upper_bound,
+                       Optional<IntLikeType>:$extent,
+                       Optional<IntLikeType>:$stride,
+                       DefaultValuedAttr<BoolAttr, "false">:$stride_in_bytes,
+                       Optional<IntLikeType>:$start_idx);
+  let results = (outs DataBoundsType:$result);
+
+  let assemblyFormat = [{
+    oilist(
+        `lower_bound` `(` $lower_bound `:` type($lower_bound) `)`
+      | `upper_bound` `(` $upper_bound `:` type($upper_bound) `)`
+      | `extent` `(` $extent `:` type($extent) `)`
+      | `stride` `(` $stride `:` type($stride) `)`
+      | `start_idx` `(` $start_idx `:` type($start_idx) `)`
+    ) attr-dict
+  }];
+
+  let extraClassDeclaration = [{
+    /// The number of variable operands.
+    unsigned getNumVariableOperands() {
+      return getNumOperands();
+    }
+
+    /// The i-th variable operand passed.
+    Value getVariableOperand(unsigned i) {
+      return getOperands()[i];
+    }
+  }];
+
+  let hasVerifier = 1;
+}
+
+def MapInfoOp : OpenMP_Op<"map_info", [AttrSizedOperandSegments]> {
+  let arguments = (ins OpenMP_PointerLikeType:$var_ptr,
+                       Optional<OpenMP_PointerLikeType>:$var_ptr_ptr,
+                       Variadic<DataBoundsType>:$bounds, /* rank-0 to rank-{n-1} */
+                       OptionalAttr<UI64Attr>:$map_type,
+                       OptionalAttr<VariableCaptureKindAttr>:$map_capture_type,
+                       DefaultValuedAttr<BoolAttr, "false">:$implicit,
+                       OptionalAttr<StrAttr>:$name);
+  let results = (outs OpenMP_PointerLikeType:$omp_ptr);
+
+  let description = [{
+    The MapInfoOp captures information relating to individual OpenMP map clauses
+    that are applied to certain OpenMP directives such as Target and Target Data.
+
+    For example, the map type modifier; such as from, tofrom and to, the variable
+    being captured or the bounds of an array section being mapped.
+
+    It can be used to capture both implicit and explicit map information, where
+    explicit is an argument directly specified to an OpenMP map clause or implicit
+    where a variable is utilised in a target region but is defined externally to
+    the target region.
+
+    This map information is later used to aid the lowering of the target operations
+    they are attached to providing argument input and output context for kernels
+    generated or the target data mapping environment.
+
+    Example (Fortran):
+
+    ```
+    integer :: index
+    !$target map(to: index)
+    ```
+    =>
+    ```mlir
+    omp.map_info var_ptr(%index_ssa) map_type(to) map_capture_type(ByRef) implicit(false)
+      name(index)
+    ```
+
+    Description of arguments:
+    - `var_ptr`: The address of variable to copy.
+    - `var_ptr_ptr`: Used when the variable copied is a member of a class, structure
+      or derived type and refers to the originating struct.
+    - `bounds`: Used when copying slices of array's, pointers or pointer members of
+      objects (e.g. derived types or classes), indicates the bounds to be copied
+      of the variable. When it's an array slice it is in rank order where rank 0
+      is the inner-most dimension.
+    - `implicit`: indicates where the map item has been specified explicitly in a
+      map clause or captured implicitly by being used in a target region with no
+      map or other data mapping construct.
+    - 'map_clauses': OpenMP map type for this map capture, for example: from, to and
+      always. It's a bitfield composed of the OpenMP runtime flags stored in
+      OpenMPOffloadMappingFlags.
+    - 'map_capture_type': Capture type for the variable e.g. this, byref, byvalue, byvla
+       this can affect how the variable is lowered.
+    - `name`: Holds the name of variable as specified in user clause (including bounds).
+  }];
+
+  let assemblyFormat = [{
+    `var_ptr` `(` $var_ptr `:` type($var_ptr) `)`
+    oilist(
+        `var_ptr_ptr` `(` $var_ptr_ptr `:` type($var_ptr_ptr) `)`
+      | `map_clauses` `(` custom<MapClause>($map_type) `)`
+      | `capture` `(` custom<CaptureType>($map_capture_type) `)`
+      | `bounds` `(` $bounds `)`
+    ) `->` type($omp_ptr) attr-dict
+  }];
+
+  let extraClassDeclaration = [{
+    /// The number of variable operands.
+    unsigned getNumVariableOperands() {
+      return getNumOperands();
+    }
+
+    /// The i-th variable operand passed.
+    Value getVariableOperand(unsigned i) {
+      return getOperands()[i];
+    }
+  }];
+}
+
 //===---------------------------------------------------------------------===//
 // 2.14.2 target data Construct
 //===---------------------------------------------------------------------===//
@@ -1044,16 +1263,14 @@ def Target_DataOp: OpenMP_Op<"target_data", [AttrSizedOperandSegments]>{
                        Optional<AnyInteger>:$device,
                        Variadic<OpenMP_PointerLikeType>:$use_device_ptr,
                        Variadic<OpenMP_PointerLikeType>:$use_device_addr,
-                       Variadic<OpenMP_PointerLikeType>:$map_operands,
-                       OptionalAttr<I64ArrayAttr>:$map_types);
+                       Variadic<OpenMP_PointerLikeType>:$map_operands);
 
   let regions = (region AnyRegion:$region);
 
   let assemblyFormat = [{
     oilist(`if` `(` $if_expr `:` type($if_expr) `)`
     | `device` `(` $device `:` type($device) `)`
-    | `map`
-      `(` custom<MapClause>($map_operands, type($map_operands), $map_types) `)`
+    | `map_entries` `(` $map_operands `:` type($map_operands) `)`
     | `use_device_ptr` `(` $use_device_ptr `:` type($use_device_ptr) `)`
     | `use_device_addr` `(` $use_device_addr `:` type($use_device_addr) `)`)
     $region attr-dict
@@ -1095,15 +1312,14 @@ def Target_EnterDataOp: OpenMP_Op<"target_enter_data",
   let arguments = (ins Optional<I1>:$if_expr,
                        Optional<AnyInteger>:$device,
                        UnitAttr:$nowait,
-                       Variadic<OpenMP_PointerLikeType>:$map_operands,
-                       I64ArrayAttr:$map_types);
+                       Variadic<OpenMP_PointerLikeType>:$map_operands);
 
   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
+    | `nowait` $nowait
+    | `map_entries` `(` $map_operands `:` type($map_operands) `)`
+    ) attr-dict
    }];
 
   let hasVerifier = 1;
@@ -1142,15 +1358,14 @@ def Target_ExitDataOp: OpenMP_Op<"target_exit_data",
   let arguments = (ins Optional<I1>:$if_expr,
                        Optional<AnyInteger>:$device,
                        UnitAttr:$nowait,
-                       Variadic<OpenMP_PointerLikeType>:$map_operands,
-                       I64ArrayAttr:$map_types);
+                       Variadic<OpenMP_PointerLikeType>:$map_operands);
 
   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
+    | `nowait` $nowait
+    | `map_entries` `(` $map_operands `:` type($map_operands) `)`
+    ) attr-dict
    }];
 
   let hasVerifier = 1;
@@ -1186,8 +1401,7 @@ def TargetOp : OpenMP_Op<"target",[AttrSizedOperandSegments]> {
                        Optional<AnyInteger>:$device,
                        Optional<AnyInteger>:$thread_limit,
                        UnitAttr:$nowait,
-                       Variadic<OpenMP_PointerLikeType>:$map_operands,
-                       OptionalAttr<I64ArrayAttr>:$map_types);
+                       Variadic<OpenMP_PointerLikeType>:$map_operands);
 
   let regions = (region AnyRegion:$region);
 
@@ -1196,7 +1410,7 @@ def TargetOp : OpenMP_Op<"target",[AttrSizedOperandSegments]> {
     | `device` `(` $device `:` type($device) `)`
     | `thread_limit` `(` $thread_limit `:` type($thread_limit) `)`
     | `nowait` $nowait
-    | `map` `(` custom<MapClause>($map_operands, type($map_operands), $map_types) `)`
+    | `map_entries` `(` $map_operands `:` type($map_operands) `)`
     ) $region attr-dict
   }];
 

mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp

@@ -207,10 +207,10 @@ void mlir::configureOpenMPToLLVMConversionLegality(
            typeConverter.isLegal(op->getOperandTypes()) &&
            typeConverter.isLegal(op->getResultTypes());
   });
-  target.addDynamicallyLegalOp<mlir::omp::AtomicReadOp,
-                               mlir::omp::AtomicWriteOp, mlir::omp::FlushOp,
-                               mlir::omp::ThreadprivateOp, mlir::omp::YieldOp,
-                               mlir::omp::EnterDataOp, mlir::omp::ExitDataOp>(
+  target.addDynamicallyLegalOp<
+      mlir::omp::AtomicReadOp, mlir::omp::AtomicWriteOp, mlir::omp::FlushOp,
+      mlir::omp::ThreadprivateOp, mlir::omp::YieldOp, mlir::omp::EnterDataOp,
+      mlir::omp::ExitDataOp, mlir::omp::DataBoundsOp, mlir::omp::MapInfoOp>(
       [&](Operation *op) {
         return typeConverter.isLegal(op->getOperandTypes()) &&
                typeConverter.isLegal(op->getResultTypes());
@@ -230,6 +230,12 @@ void mlir::configureOpenMPToLLVMConversionLegality(
 
 void mlir::populateOpenMPToLLVMConversionPatterns(LLVMTypeConverter &converter,
                                                   RewritePatternSet &patterns) {
+  // This type is allowed when converting OpenMP to LLVM Dialect, it carries
+  // bounds information for map clauses and the operation and type are
+  // discarded on lowering to LLVM-IR from the OpenMP dialect.
+  converter.addConversion(
+      [&](omp::DataBoundsType type) -> Type { return type; });
+
   patterns.add<
       AtomicReadOpConversion, ReductionOpConversion,
       ReductionDeclareOpConversion, RegionOpConversion<omp::CriticalOp>,
@@ -246,7 +252,9 @@ void mlir::populateOpenMPToLLVMConversionPatterns(LLVMTypeConverter &converter,
       RegionLessOpWithVarOperandsConversion<omp::ThreadprivateOp>,
       RegionLessOpConversion<omp::YieldOp>,
       RegionLessOpConversion<omp::EnterDataOp>,
-      RegionLessOpConversion<omp::ExitDataOp>>(converter);
+      RegionLessOpConversion<omp::ExitDataOp>,
+      RegionLessOpWithVarOperandsConversion<omp::DataBoundsOp>,
+      RegionLessOpWithVarOperandsConversion<omp::MapInfoOp>>(converter);
 }
 
 namespace {