[llvm-branch-commits] [flang] [Flang][OpenMP] Derived type explicit allocatable member mapping (PR #96266)

via llvm-branch-commits llvm-branch-commits at lists.llvm.org
Wed Jul 17 18:46:21 PDT 2024


================
@@ -85,67 +135,227 @@ class OMPMapInfoFinalizationPass
       descriptor = alloca;
     }
 
+    return descriptor;
+  }
+
+  /// Simple function that will generate a FIR operation accessing
+  /// the descriptors base address (BoxOffsetOp) and then generate a
+  /// MapInfoOp for it, the most important thing to note is that
+  /// we normally move the bounds from the descriptor map onto the
+  /// base address map.
+  mlir::omp::MapInfoOp getBaseAddrMap(mlir::Value descriptor,
+                                      mlir::OperandRange bounds,
+                                      int64_t mapType,
+                                      fir::FirOpBuilder &builder) {
+    mlir::Location loc = descriptor.getLoc();
     mlir::Value baseAddrAddr = builder.create<fir::BoxOffsetOp>(
         loc, descriptor, fir::BoxFieldAttr::base_addr);
 
     // Member of the descriptor pointing at the allocated data
-    mlir::Value baseAddr = builder.create<mlir::omp::MapInfoOp>(
+    return builder.create<mlir::omp::MapInfoOp>(
         loc, baseAddrAddr.getType(), descriptor,
         mlir::TypeAttr::get(llvm::cast<mlir::omp::PointerLikeType>(
                                 fir::unwrapRefType(baseAddrAddr.getType()))
                                 .getElementType()),
-        baseAddrAddr, /*members=*/mlir::SmallVector<mlir::Value>{},
-        /*member_index=*/mlir::DenseIntElementsAttr{}, op.getBounds(),
-        builder.getIntegerAttr(builder.getIntegerType(64, false),
-                               op.getMapType().value()),
+        baseAddrAddr, mlir::SmallVector<mlir::Value>{},
+        mlir::DenseIntElementsAttr{}, bounds,
+        builder.getIntegerAttr(builder.getIntegerType(64, false), mapType),
         builder.getAttr<mlir::omp::VariableCaptureKindAttr>(
             mlir::omp::VariableCaptureKind::ByRef),
-        /*name=*/builder.getStringAttr(""),
-        /*partial_map=*/builder.getBoolAttr(false));
+        builder.getStringAttr("") /*name*/,
+        builder.getBoolAttr(false) /*partial_map*/);
+  }
 
-    // TODO: map the addendum segment of the descriptor, similarly to the
-    // above base address/data pointer member.
+  /// This function adjusts the member indices vector to include a new
+  /// base address member, we take the position of the descriptor in
+  /// the member indices list, which is the index data that the base
+  /// addresses index will be based off of, as the base address is
+  /// a member of the descriptor, we must also alter other members
+  /// indices in the list to account for this new addition. This
+  /// requires extending all members with -1's if the addition of
+  /// the new base address has increased the member vector past the
+  /// original size, as we must make sure all member indices are of
+  /// the same length (think rectangle matrix) due to DenseIntElementsAttr
+  /// requiring this. We also need to be aware that we are inserting
+  /// into the middle of a member index vector in some cases (i.e.
+  /// we could be accessing the member of a descriptor type with a
+  /// subsequent map, so we must be sure to adjust any of these cases
+  /// with the addition of the new base address index value).
+  void adjustMemberIndices(
+      llvm::SmallVector<llvm::SmallVector<int32_t>> &memberIndices,
+      size_t memberIndex) {
+    // Find if the descriptor member we are basing our new base address index
+    // off of has a -1 somewhere, indicating an empty index already exists (due
+    // to a larger sized member position elsewhere) which allows us to simplify
+    // later steps a little
+    auto baseAddrIndex = memberIndices[memberIndex];
+    auto *iterPos = std::find(baseAddrIndex.begin(), baseAddrIndex.end(), -1);
 
-    if (auto mapClauseOwner =
-            llvm::dyn_cast<mlir::omp::MapClauseOwningOpInterface>(target)) {
-      llvm::SmallVector<mlir::Value> newMapOps;
-      mlir::OperandRange mapOperandsArr = mapClauseOwner.getMapOperands();
+    // If we aren't at the end, as we found a -1, we can simply modify the -1
+    // to the base addresses index in the descriptor (which will always be the
+    // first member in the descriptor, so 0). If not, then we're extending the
+    // index list and have to push on a 0 and adjust the position to the new
+    // end.
+    if (iterPos != baseAddrIndex.end()) {
+      *iterPos = 0;
+    } else {
+      baseAddrIndex.push_back(0);
+      iterPos = baseAddrIndex.end();
+    }
 
-      for (size_t i = 0; i < mapOperandsArr.size(); ++i) {
-        if (mapOperandsArr[i] == op) {
-          // Push new implicit maps generated for the descriptor.
-          newMapOps.push_back(baseAddr);
+    auto isEqual = [](auto first1, auto last1, auto first2, auto last2) {
+      int v1, v2;
+      for (; first1 != last1; ++first1, ++first2) {
+        v1 = (first1 == last1) ? -1 : *first1;
+        v2 = (first2 == last2) ? -1 : *first2;
 
-          // 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, baseAddr.getType(), loc);
-        }
-        newMapOps.push_back(mapOperandsArr[i]);
+        if (!(v1 == v2))
+          return false;
+      }
+      return true;
+    };
+
+    // If we find another member that is "derived/a member of" the descriptor
+    // that is not the descriptor itself, we must insert a 0 for the new base
+    // address we have just added for the descriptor into the list at the
+    // appropriate position to maintain correctness of the positional/index data
+    // for that member.
+    size_t insertPosition = std::distance(baseAddrIndex.begin(), iterPos);
+    for (size_t i = 0; i < memberIndices.size(); ++i) {
+      if (isEqual(baseAddrIndex.begin(), iterPos, memberIndices[i].begin(),
+                  memberIndices[i].end())) {
+        if (i == memberIndex)
+          continue;
+
+        memberIndices[i].insert(
+            std::next(memberIndices[i].begin(), insertPosition), 0);
       }
-      mapClauseOwner.getMapOperandsMutable().assign(newMapOps);
     }
 
-    mlir::Value newDescParentMapOp = builder.create<mlir::omp::MapInfoOp>(
-        op->getLoc(), op.getResult().getType(), descriptor,
-        mlir::TypeAttr::get(fir::unwrapRefType(descriptor.getType())),
-        /*varPtrPtr=*/mlir::Value{},
-        /*members=*/mlir::SmallVector<mlir::Value>{baseAddr},
-        /*members_index=*/
-        mlir::DenseIntElementsAttr::get(
+    // Insert our newly created baseAddrIndex into the larger list of indices at
+    // the correct location.
+    memberIndices.insert(std::next(memberIndices.begin(), memberIndex + 1),
+                         baseAddrIndex);
+  }
+
+  // Adjusts the descriptors map type based on the map type of the original
+  // map type which will apply to the base address, the main alteration that
+  // is done currently is appending OMP_MAP_TO in cases where we only have
+  // OMP_MAP_FROM or an ALLOC (lack of flag). This is because we will
+  // always need to map the descriptor to device (or at the very least it
+  // seems to be the case currently with the current lowered IR), as without
+  // the appropriate descriptor information on the device there is a risk
+  // of the kernel IR requesting for various data that will not have been
+  // copied to perform things like indexing, this can cause segfaults and
+  // memory access errors. These alterations are only unapplicable to
+  // target exit currently.
+  unsigned long getDescriptorMapType(unsigned long mapTypeFlag,
+                                     mlir::Operation *target) {
+    auto newDescFlag = llvm::omp::OpenMPOffloadMappingFlags(mapTypeFlag);
+
+    if ((llvm::isa_and_nonnull<mlir::omp::TargetDataOp>(target) ||
+         llvm::isa_and_nonnull<mlir::omp::TargetOp>(target)) &&
+        static_cast<
+            std::underlying_type_t<llvm::omp::OpenMPOffloadMappingFlags>>(
+            (newDescFlag &
+             llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_FROM)) &&
+        static_cast<
+            std::underlying_type_t<llvm::omp::OpenMPOffloadMappingFlags>>(
+            (newDescFlag & llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_TO) !=
+            llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_TO))
+      return static_cast<
+          std::underlying_type_t<llvm::omp::OpenMPOffloadMappingFlags>>(
+          newDescFlag | llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_TO);
+
+    if ((llvm::isa_and_nonnull<mlir::omp::TargetDataOp>(target) ||
+         llvm::isa_and_nonnull<mlir::omp::TargetEnterDataOp>(target)) &&
+        mapTypeFlag == 0)
+      return static_cast<
+          std::underlying_type_t<llvm::omp::OpenMPOffloadMappingFlags>>(
+          llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_TO);
+
+    return mapTypeFlag;
+  }
+
+  mlir::omp::MapInfoOp genDescriptorMemberMaps(mlir::omp::MapInfoOp op,
+                                               fir::FirOpBuilder &builder,
+                                               mlir::Operation *target) {
+    llvm::SmallVector<ParentAndPlacement> mapMemberUsers;
+    getMemberUserList(op, mapMemberUsers);
+
+    // TODO: map the addendum segment of the descriptor, similarly to the
+    // base address/data pointer member.
+    mlir::Value descriptor = getDescriptorFromBoxMap(op, builder);
+    auto baseAddr = getBaseAddrMap(descriptor, op.getBounds(),
+                                   op.getMapType().value(), builder);
+    mlir::DenseIntElementsAttr newMembersAttr;
+    mlir::SmallVector<mlir::Value> newMembers;
+    llvm::SmallVector<llvm::SmallVector<int32_t>> memberIndices;
+    if (!mapMemberUsers.empty() || !op.getMembers().empty())
+      getMemberIndicesAsVectors(
+          !mapMemberUsers.empty() ? mapMemberUsers[0].parent : op,
----------------
agozillon wrote:

it's intentional for the moment, there should never be more than one mapMemberUsers currently when we reach this pass, as we do not generate MLIR that allows the re-use of MapInfoOp's across operations, so technically the only cases where a MapInfoOp should be used, is a parent map member operation (in which case we add an entry to mapMemberUsers) or when it's tied to an existing MapInfoOp using operation (TargetOp/TargetExitData etc.).  

So, perhaps being a bit pre-emptive and adding some initial support for this case, for when we may end up trying to generating tidier MLIR where we re-use map info operations allowing them to have multiple users. As that would be ideal in the future!

Previously an assert checking that we had no more than 1 user was above but i think it is covered by the prior assert at line 488 so it felt a little redundant.

https://github.com/llvm/llvm-project/pull/96266


More information about the llvm-branch-commits mailing list