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

Sergio Afonso via llvm-branch-commits llvm-branch-commits at lists.llvm.org
Mon Jul 29 05:49:07 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);
----------------
skatrak wrote:

The `map_type` argument is defined as optional in MLIR, so a valid operation could have it not set, and there are no prior checks for that before trying to access it. This pass could be triggered manually on some handwritten MLIR, as far as we know, we can't just rely on it coming from the output of Flang lowering, even if that execution path always sets this argument.

Since there's a well-defined default value we could use (`llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_NONE = 0`), maybe we could change the definition of this argument in OpenMPOps.td to `DefaultValuedAttr<UI64Attr, "0">:$map_type`. If that's not an option, maybe we should either emit an error or use `op.getMapType().value_or(...)` to avoid crashing the compiler in this case.

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


More information about the llvm-branch-commits mailing list