[Flang][MLIR] Insert descriptor implicit members into map operands and BlockArgs

agozillon · agozillon · commit 3e4d168e7476 · 2024-01-04T12:11:30.000-06:00
diff --git a/flang/lib/Optimizer/Transforms/OMPDescriptorMapInfoGen.cpp b/flang/lib/Optimizer/Transforms/OMPDescriptorMapInfoGen.cpp
@@ -12,6 +12,7 @@
 #include "mlir/Pass/Pass.h"
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/SmallPtrSet.h"
+#include <iterator>
 
 namespace fir {
 #define GEN_PASS_DEF_OMPDESCRIPTORMAPINFOGENPASS
@@ -48,7 +49,8 @@ class OMPDescriptorMapInfoGenPass
   }
 
   void genDescriptorMemberMaps(mlir::omp::MapInfoOp op,
-                               fir::FirOpBuilder &builder) {
+                               fir::FirOpBuilder &builder,
+                               mlir::Operation *target) {
     llvm::SmallVector<mlir::Value> descriptorBaseAddrMembers;
     mlir::Location loc = builder.getUnknownLoc();
     mlir::Value descriptor = op.getVarPtr();
@@ -92,8 +94,36 @@ class OMPDescriptorMapInfoGenPass
 
     op.getVarPtrMutable().assign(descriptor);
     op.setVarType(fir::unwrapRefType(descriptor.getType()));
-    op.getMembersMutable().assign(descriptorBaseAddrMembers);
+    op.getMembersMutable().append(descriptorBaseAddrMembers);
     op.getBoundsMutable().assign(llvm::SmallVector<mlir::Value>{});
+
+    // could use a template to generalise to other TargetOps
+    if (auto mapClauseOwner =
+            llvm::dyn_cast<mlir::omp::MapClauseOwningOpInterface>(target)) {
+      llvm::SmallVector<mlir::Value> newMapOps;
+      for (size_t i = 0; i < mapClauseOwner.getMapOperands().size(); ++i) {
+        if (mapClauseOwner.getMapOperands()[i] == op) {
+          // Push new implicit maps generated for the descriptor.
+          newMapOps.push_back(descriptorBaseAddrMembers[0]);
+
+          // for TargetOp's which have IsolatedFromAbove we must align the
+          // new additional map operand with an appropriate BlockArgument,
+          // as the printing and later processing currently requires a 1:1
+          // mapping of BlockArgs to MapInfoOp's at the same placement in
+          // each array (BlockArgs and MapOperands).
+          if (auto targetOp = llvm::dyn_cast<mlir::omp::TargetOp>(target)) {
+            targetOp.getRegion().insertArgument(
+                i, descriptorBaseAddrMembers[0].getType(), loc);
+          }
+
+          newMapOps.push_back(mapClauseOwner.getMapOperands()[i]);
+        } else {
+          newMapOps.push_back(mapClauseOwner.getMapOperands()[i]);
+        }
+      }
+
+      mapClauseOwner.getMapOperandsMutable().assign(newMapOps);
+    }
   }
 
   // This pass executes on mlir::ModuleOp's finding omp::MapInfoOp's containing
@@ -108,7 +138,9 @@ class OMPDescriptorMapInfoGenPass
       if (fir::isTypeWithDescriptor(op.getVarType()) ||
           mlir::isa<fir::BoxAddrOp>(op.getVarPtr().getDefiningOp())) {
         builder.setInsertionPoint(op);
-        genDescriptorMemberMaps(op, builder);
+        // Currently a MapInfoOp argument can only show up on a single target
+        // user so we can retrieve and use the first user.
+        genDescriptorMemberMaps(op, builder, *op->getUsers().begin());
       }
     });
   }
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -1232,7 +1232,8 @@ def MapInfoOp : OpenMP_Op<"map_info", [AttrSizedOperandSegments]> {
 // 2.14.2 target data Construct
 //===---------------------------------------------------------------------===//
 
-def Target_DataOp: OpenMP_Op<"target_data", [AttrSizedOperandSegments]>{
+def Target_DataOp: OpenMP_Op<"target_data", [AttrSizedOperandSegments, 
+                                             MapClauseOwningOpInterface]>{
   let summary = "target data construct";
   let description = [{
     Map variables to a device data environment for the extent of the region.
@@ -1289,7 +1290,8 @@ def Target_DataOp: OpenMP_Op<"target_data", [AttrSizedOperandSegments]>{
 //===---------------------------------------------------------------------===//
 
 def Target_EnterDataOp: OpenMP_Op<"target_enter_data",
-                                                 [AttrSizedOperandSegments]>{
+                                                 [AttrSizedOperandSegments, 
+                                                  MapClauseOwningOpInterface]>{
   let  summary = "target enter data construct";
   let description = [{
     The target enter data directive specifies that variables are mapped to
@@ -1335,7 +1337,8 @@ def Target_EnterDataOp: OpenMP_Op<"target_enter_data",
 //===---------------------------------------------------------------------===//
 
 def Target_ExitDataOp: OpenMP_Op<"target_exit_data",
-                                                 [AttrSizedOperandSegments]>{
+                                                 [AttrSizedOperandSegments,
+                                                  MapClauseOwningOpInterface]>{
   let  summary = "target exit data construct";
   let description = [{
     The target exit data directive specifies that variables are mapped to a
@@ -1381,7 +1384,8 @@ def Target_ExitDataOp: OpenMP_Op<"target_exit_data",
 //===---------------------------------------------------------------------===//
 
 def Target_UpdateDataOp: OpenMP_Op<"target_update_data",
-                                                 [AttrSizedOperandSegments]>{
+                                                 [AttrSizedOperandSegments,
+                                                  MapClauseOwningOpInterface]>{
   let  summary = "target update data construct";
   let description = [{
     The target update directive makes the corresponding list items in the device
@@ -1413,13 +1417,13 @@ def Target_UpdateDataOp: OpenMP_Op<"target_update_data",
   let arguments = (ins Optional<I1>:$if_expr,
                        Optional<AnyInteger>:$device,
                        UnitAttr:$nowait,
-                       Variadic<OpenMP_PointerLikeType>:$motion_operands);
+                       Variadic<OpenMP_PointerLikeType>:$map_operands);
 
   let assemblyFormat = [{
     oilist(`if` `(` $if_expr `:` type($if_expr) `)`
     | `device` `(` $device `:` type($device) `)`
     | `nowait` $nowait
-    | `motion_entries` `(` $motion_operands `:` type($motion_operands) `)`
+    | `motion_entries` `(` $map_operands `:` type($map_operands) `)`
     ) attr-dict
    }];
 
@@ -1430,7 +1434,8 @@ def Target_UpdateDataOp: OpenMP_Op<"target_update_data",
 // 2.14.5 target construct
 //===----------------------------------------------------------------------===//
 
-def TargetOp : OpenMP_Op<"target",[IsolatedFromAbove, OutlineableOpenMPOpInterface, AttrSizedOperandSegments]> {
+def TargetOp : OpenMP_Op<"target",[IsolatedFromAbove, MapClauseOwningOpInterface,
+                                   OutlineableOpenMPOpInterface, AttrSizedOperandSegments]> {
   let summary = "target construct";
   let description = [{
     The target construct includes a region of code which is to be executed
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOpsInterfaces.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOpsInterfaces.td
@@ -18,7 +18,7 @@ include "mlir/IR/OpBase.td"
 def OutlineableOpenMPOpInterface : OpInterface<"OutlineableOpenMPOpInterface"> {
   let description = [{
     OpenMP operations whose region will be outlined will implement this
-    interface. These operations will
+    interface.
   }];
 
   let cppNamespace = "::mlir::omp";
@@ -31,6 +31,28 @@ def OutlineableOpenMPOpInterface : OpInterface<"OutlineableOpenMPOpInterface"> {
   ];
 }
 
+def MapClauseOwningOpInterface : OpInterface<"MapClauseOwningOpInterface"> {
+  let description = [{
+    OpenMP operations which own a list of omp::MapInfoOp's implement this interface
+    to allow generic access to deal with map operands to more easily manipulate 
+    this class of operations. 
+  }];
+
+  let cppNamespace = "::mlir::omp";
+
+  let methods = [
+    InterfaceMethod<"Get map operands", "::mlir::OperandRange", "getMapOperands",
+      (ins), [{
+        return $_op.getMapOperands();
+      }]>,
+      InterfaceMethod<"Get mutable map operands", "::mlir::MutableOperandRange", 
+                      "getMapOperandsMutable",
+      (ins), [{
+        return $_op.getMapOperandsMutable();
+      }]>,
+  ];
+}
+
 def ReductionClauseInterface : OpInterface<"ReductionClauseInterface"> {
   let description = [{
     OpenMP operations that support reduction clause have this interface.
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -980,7 +980,7 @@ LogicalResult ExitDataOp::verify() {
 }
 
 LogicalResult UpdateDataOp::verify() {
-  return verifyMapClause(*this, getMotionOperands());
+  return verifyMapClause(*this, getMapOperands());
 }
 
 LogicalResult TargetOp::verify() {
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -1640,6 +1640,7 @@ getRefPtrIfDeclareTarget(mlir::Value value,
 // value) more than neccessary.
 struct MapInfoData : llvm::OpenMPIRBuilder::MapInfosTy {
   llvm::SmallVector<bool, 4> IsDeclareTarget;
+  llvm::SmallVector<bool, 4> IsAMember;
   llvm::SmallVector<mlir::Operation *, 4> MapClause;
   llvm::SmallVector<llvm::Value *, 4> OriginalValue;
   // Stripped off array/pointer to get the underlying
@@ -1728,13 +1729,11 @@ void collectMapDataFromMapOperands(MapInfoData &mapData,
                                    LLVM::ModuleTranslation &moduleTranslation,
                                    DataLayout &dl,
                                    llvm::IRBuilderBase &builder) {
-  auto mapFill = [&](mlir::Value mapValue) {
-    assert(mlir::isa<mlir::omp::MapInfoOp>(mapValue.getDefiningOp()) &&
-           "missing map info operation or incorrect map info operation type");
+  for (mlir::Value mapValue : mapOperands) {
     if (auto mapOp = mlir::dyn_cast_if_present<mlir::omp::MapInfoOp>(
             mapValue.getDefiningOp())) {
-      mapData.OriginalValue.push_back(moduleTranslation.lookupValue(
-          mapOp.getVarPtrPtr() ? mapOp.getVarPtrPtr() : mapOp.getVarPtr()));
+      mapData.OriginalValue.push_back(
+          moduleTranslation.lookupValue(mapOp.getVarPtr()));
       mapData.Pointers.push_back(mapData.OriginalValue.back());
 
       if (llvm::Value *refPtr =
@@ -1759,41 +1758,27 @@ void collectMapDataFromMapOperands(MapInfoData &mapData,
           mapOp.getLoc(), *moduleTranslation.getOpenMPBuilder()));
       mapData.DevicePointers.push_back(
           llvm::OpenMPIRBuilder::DeviceInfoTy::None);
-    }
-  };
 
-  // In the case of Fortran descriptors some members get added implicitly
-  // after the target region has been generated during CodeGen lowering
-  // which prevents them from being added trivially to the target region
-  // as map arguments, we must handle this case here by generating
-  // MapInfoData for them.
-  // TODO: Revisit this when implementing derived types explicit member
-  // mapping, we likely want to represent these identically to simplify
-  // the overall lowering
-  SmallVector<Value> mapMemberOperands;
-  for (size_t i = 0; i < mapOperands.size(); ++i) {
-    auto mapInfoOp =
-        mlir::dyn_cast<mlir::omp::MapInfoOp>(mapOperands[i].getDefiningOp());
-    for (auto members : mapInfoOp.getMembers()) {
-      if (!std::any_of(mapOperands.begin(), mapOperands.end(),
-                       [&](auto mapOp) {
-                         return mapOp.getDefiningOp() ==
-                                members.getDefiningOp();
-                       }) &&
-          !std::any_of(mapMemberOperands.begin(), mapMemberOperands.end(),
-                       [&](auto mapOp) {
-                         return mapOp.getDefiningOp() ==
-                                members.getDefiningOp();
-                       }))
-        mapMemberOperands.push_back(members);
+      // Check if this is a member mapping and correctly assign that it is, if
+      // it is a member of a larger object.
+      // TODO: Need better handling of members, and distinguishing of members
+      // that are implicitly allocated on device vs explicitly passed in as
+      // arguments.
+      // TODO: May require some further additions to support nested record
+      // types, i.e. member maps that can have member maps.
+      mapData.IsAMember.push_back(false);
+      for (mlir::Value mapValue : mapOperands) {
+        if (auto map = mlir::dyn_cast_if_present<mlir::omp::MapInfoOp>(
+                mapValue.getDefiningOp())) {
+          for (auto member : map.getMembers()) {
+            if (member == mapOp) {
+              mapData.IsAMember.back() = true;
+            }
+          }
+        }
+      }
     }
   }
-
-  for (mlir::Value mapValue : mapOperands)
-    mapFill(mapValue);
-
-  for (mlir::Value mapValue : mapMemberOperands)
-    mapFill(mapValue);
 }
 
 static void processMapWithMembersOf(
@@ -1960,35 +1945,22 @@ static void genMapInfos(llvm::IRBuilderBase &builder,
     combinedInfo.Names.clear();
   };
 
-  llvm::SmallVector<size_t, 4> primaryMapIdx;
-  for (size_t i = 0; i < mapData.MapClause.size(); ++i) {
-    primaryMapIdx.push_back(i);
-  }
-
-  // TODO: Handle nested MembersOf, currently only cares about the first level
-  // of nesting (all that was relevant for Fortran descriptors), but a slight
-  // refactoring of mapInfoData to hold nestings or membersOf may be a better
-  // approach to simplify things.
-  for (size_t i = 0; i < mapData.MapClause.size(); ++i) {
-    auto mapInfoOp = mlir::dyn_cast<mlir::omp::MapInfoOp>(mapData.MapClause[i]);
-    for (auto member : mapInfoOp.getMembers()) {
-      for (size_t j = 0; j < primaryMapIdx.size(); j++) {
-        if (member.getDefiningOp() == mapData.MapClause[primaryMapIdx[j]]) {
-          primaryMapIdx.erase(&primaryMapIdx[j]);
-          j--;
-        }
-      }
-    }
-  }
-
   // We operate under the assumption that all vectors that are
   // required in MapInfoData are of equal lengths (either filled with
   // default constructed data or appropiate information) so we can
   // utilise the size from any component of MapInfoData, if we can't
   // something is missing from the initial MapInfoData construction.
-  for (unsigned long i : primaryMapIdx) {
-    auto mapInfoOp = mlir::dyn_cast<mlir::omp::MapInfoOp>(mapData.MapClause[i]);
+  for (size_t i = 0; i < mapData.MapClause.size(); ++i) {
+    // NOTE/TODO: We currently do not handle member mapping seperately from it's
+    // parent or explicit mapping of a parent and member in the same operation,
+    // this will need to change in the near future, for now we primarily handle
+    // descriptor mapping from fortran, generalised as mapping record types
+    // with implicit member maps. This lowering needs further generalisation to
+    // fully support fortran derived types, and C/C++ structures and classes.
+    if (mapData.IsAMember[i])
+      continue;
 
+    auto mapInfoOp = mlir::dyn_cast<mlir::omp::MapInfoOp>(mapData.MapClause[i]);
     if (!mapInfoOp.getMembers().empty()) {
       processMapWithMembersOf(moduleTranslation, builder, *ompBuilder, dl,
                               combinedInfo, mapData, i, isTargetParams);
@@ -2136,7 +2108,7 @@ convertOmpTargetData(Operation *op, llvm::IRBuilderBase &builder,
                   deviceID = intAttr.getInt();
 
             RTLFn = llvm::omp::OMPRTL___tgt_target_data_update_mapper;
-            mapOperands = updateDataOp.getMotionOperands();
+            mapOperands = updateDataOp.getMapOperands();
             return success();
           })
           .Default([&](Operation *op) {
@@ -2633,7 +2605,12 @@ convertOmpTarget(Operation &opInst, llvm::IRBuilderBase &builder,
   llvm::SmallVector<llvm::Value *, 4> kernelInput;
   for (size_t i = 0; i < mapOperands.size(); ++i) {
     // declare target arguments are not passed to kernels as arguments
-    if (!mapData.IsDeclareTarget[i])
+    // TODO: We currently do not handle cases where a member is explicitly
+    // passed in as an argument, this will likley need to be handled in
+    // the near future, rather than using IsAMember, it may be better to
+    // test if the relevant BlockArg is used within the target region and
+    // then use that as a basis for exclusion in the kernel inputs.
+    if (!mapData.IsDeclareTarget[i] && !mapData.IsAMember[i])
       kernelInput.push_back(mapData.OriginalValue[i]);
   }
 

Original file line number	Diff line number	Diff line change
`@@ -980,7 +980,7 @@ LogicalResult ExitDataOp::verify() {`
`980`	`980`	`}`
`981`	`981`
`982`	`982`	`LogicalResult UpdateDataOp::verify() {`
`983`		`- return verifyMapClause(*this, getMotionOperands());`
	`983`	`+ return verifyMapClause(*this, getMapOperands());`
`984`	`984`	`}`
`985`	`985`
`986`	`986`	`LogicalResult TargetOp::verify() {`