[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:10 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,
+ memberIndices);
+
+ // If the operation that we are expanding with a descriptor has a user
+ // (parent), then we have to expand the parents member indices to reflect
+ // the adjusted member indices for the base address insertion. However, if
+ // it does not then we are expanding a MapInfoOp without any pre-existing
+ // member information to now have one new member for the base address or we
+ // are expanding a parent that is a descriptor and we have to adjust all of
+ // it's members to reflect the insertion of the base address.
+ if (!mapMemberUsers.empty()) {
+ auto baseAddrIndex = memberIndices[mapMemberUsers[0].index];
+ adjustMemberIndices(memberIndices, mapMemberUsers[0].index);
+ llvm::SmallVector<mlir::Value> newMemberOps;
+ mlir::OperandRange membersArr = mapMemberUsers[0].parent.getMembers();
+ for (size_t i = 0; i < membersArr.size(); ++i) {
+ newMemberOps.push_back(membersArr[i]);
+ if (membersArr[i] == op)
+ newMemberOps.push_back(baseAddr);
+ }
+ mapMemberUsers[0].parent.getMembersMutable().assign(newMemberOps);
+ Fortran::lower::omp::fillMemberIndices(memberIndices);
+ mapMemberUsers[0].parent.setMembersIndexAttr(
+ Fortran::lower::omp::createDenseElementsAttrFromIndices(memberIndices,
+ builder));
+ } else {
+ newMembers.push_back(baseAddr);
+ if (!op.getMembers().empty()) {
+ for (auto &indices : memberIndices)
+ indices.insert(indices.begin(), 0);
+ llvm::SmallVector<int> baseAddrIndex;
+ baseAddrIndex.resize(memberIndices[0].size());
+ std::fill(baseAddrIndex.begin(), baseAddrIndex.end(), -1);
+ baseAddrIndex[0] = 0;
+ memberIndices.insert(memberIndices.begin(), baseAddrIndex);
+ Fortran::lower::omp::fillMemberIndices(memberIndices);
+ newMembersAttr =
+ Fortran::lower::omp::createDenseElementsAttrFromIndices(
+ memberIndices, builder);
+ newMembers.append(op.getMembers().begin(), op.getMembers().end());
+ } else {
+ newMembersAttr = mlir::DenseIntElementsAttr::get(
mlir::VectorType::get(
llvm::ArrayRef<int64_t>({1, 1}),
mlir::IntegerType::get(builder.getContext(), 32)),
- llvm::ArrayRef<int32_t>({0})),
- /*bounds=*/mlir::SmallVector<mlir::Value>{},
- builder.getIntegerAttr(builder.getIntegerType(64, false),
- op.getMapType().value()),
- op.getMapCaptureTypeAttr(), op.getNameAttr(), op.getPartialMapAttr());
- op.replaceAllUsesWith(newDescParentMapOp);
+ llvm::ArrayRef<int32_t>({0}));
+ }
+ }
+
+ mlir::omp::MapInfoOp newDescParentMapOp =
+ builder.create<mlir::omp::MapInfoOp>(
+ op->getLoc(), op.getResult().getType(), descriptor,
+ mlir::TypeAttr::get(fir::unwrapRefType(descriptor.getType())),
----------------
skatrak wrote:
Nit: Document argument names directly initialized in the call, using the `/*name=*/...` format.
https://github.com/llvm/llvm-project/pull/96266
More information about the llvm-branch-commits
mailing list