================ @@ -82,104 +132,188 @@ class MapInfoFinalizationPass // perform an alloca and then store to it and retrieve the data from the new // alloca. if (mlir::isa<fir::BaseBoxType>(descriptor.getType())) { - // If we have already created a local allocation for this BoxType, - // we must be sure to re-use it so that we end up with the same - // allocations being utilised for the same descriptor across all map uses, - // this prevents runtime issues such as not appropriately releasing or - // deleting all mapped data. - auto find = localBoxAllocas.find(descriptor.getAsOpaquePointer()); - if (find != localBoxAllocas.end()) { - builder.create<fir::StoreOp>(loc, descriptor, find->second); - descriptor = find->second; - } else { - mlir::OpBuilder::InsertPoint insPt = builder.saveInsertionPoint(); - mlir::Block *allocaBlock = builder.getAllocaBlock(); - assert(allocaBlock && "No alloca block found for this top level op"); - builder.setInsertionPointToStart(allocaBlock); - auto alloca = builder.create<fir::AllocaOp>(loc, descriptor.getType()); - builder.restoreInsertionPoint(insPt); - builder.create<fir::StoreOp>(loc, descriptor, alloca); - localBoxAllocas[descriptor.getAsOpaquePointer()] = alloca; - descriptor = alloca; - } + mlir::OpBuilder::InsertPoint insPt = builder.saveInsertionPoint(); + mlir::Block *allocaBlock = builder.getAllocaBlock(); + mlir::Location loc = boxMap->getLoc(); + assert(allocaBlock && "No alloca block found for this top level op"); + builder.setInsertionPointToStart(allocaBlock); + auto alloca = builder.create<fir::AllocaOp>(loc, descriptor.getType()); + builder.restoreInsertionPoint(insPt); + builder.create<fir::StoreOp>(loc, descriptor, alloca); + descriptor = alloca; } + return descriptor; + } + + /// Function that generates a FIR operation accessing the descriptor's + /// base address (BoxOffsetOp) and 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()), + /*membersIndex=*/mlir::ArrayAttr{}, 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)); + } - // 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 member's + void adjustMemberIndices( + llvm::SmallVector<llvm::SmallVector<int64_t>> &memberIndices, + size_t memberIndex) { + llvm::SmallVector<int64_t> baseAddrIndex = memberIndices[memberIndex]; + baseAddrIndex.push_back(0); - auto addOperands = [&](mlir::OperandRange &operandsArr, - mlir::MutableOperandRange &mutableOpRange, - auto directiveOp) { - llvm::SmallVector<mlir::Value> newMapOps; - for (size_t i = 0; i < operandsArr.size(); ++i) { - if (operandsArr[i] == op) { - // Push new implicit maps generated for the descriptor. - newMapOps.push_back(baseAddr); - - // 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 (directiveOp) { - directiveOp.getRegion().insertArgument(i, baseAddr.getType(), loc); - } - } - newMapOps.push_back(operandsArr[i]); + // 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(), std::prev(baseAddrIndex.end())); + for (llvm::SmallVector<long> member : memberIndices) { + if (member.size() > insertPosition && + std::equal(baseAddrIndex.begin(), std::prev(baseAddrIndex.end()), + member.begin())) { + member.insert(std::next(member.begin(), insertPosition), 0); } - mutableOpRange.assign(newMapOps); - }; - if (auto mapClauseOwner = - llvm::dyn_cast<mlir::omp::MapClauseOwningOpInterface>(target)) { - mlir::OperandRange mapOperandsArr = mapClauseOwner.getMapVars(); - mlir::MutableOperandRange mapMutableOpRange = - mapClauseOwner.getMapVarsMutable(); - mlir::omp::TargetOp targetOp = - llvm::dyn_cast<mlir::omp::TargetOp>(target); - addOperands(mapOperandsArr, mapMutableOpRange, targetOp); } - if (auto targetDataOp = llvm::dyn_cast<mlir::omp::TargetDataOp>(target)) { - mlir::OperandRange useDevAddrArr = targetDataOp.getUseDeviceAddrVars(); - mlir::MutableOperandRange useDevAddrMutableOpRange = - targetDataOp.getUseDeviceAddrVarsMutable(); - addOperands(useDevAddrArr, useDevAddrMutableOpRange, targetDataOp); + + // 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 descriptor's map type. The main alteration that is done + /// currently is transforming the map type to `OMP_MAP_TO` where possible. + /// 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 kernel 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. However, we do not need this data mapped + /// back to the host from the device, as per the OpenMP spec we cannot alter + /// the data via resizing or deletion on the device. Discarding any + /// descriptor alterations via no map back is reasonable (and required + /// for certain segments of descriptor data like the type descriptor that are + /// global constants). This alteration is only inapplicable to `target exit` + /// and `target update` currently, and that's due to `target exit` not + /// allowing `to` mappings, and `target update` not allowing both `to` and + /// `from` simultaneously. We currently try to maintain the `implicit` flag + /// where necessary, although it does not seem strictly required. + unsigned long getDescriptorMapType(unsigned long mapTypeFlag, + mlir::Operation *target) { + if (llvm::isa_and_nonnull<mlir::omp::TargetExitDataOp, + mlir::omp::TargetUpdateOp>(target)) + return mapTypeFlag; + + bool hasImplicitMap = + (llvm::omp::OpenMPOffloadMappingFlags(mapTypeFlag) & + llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_IMPLICIT) == + llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_IMPLICIT; + + return llvm::to_underlying( + hasImplicitMap + ? llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_TO | + llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_IMPLICIT + : llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_TO); + } + + 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_or(0), builder); + mlir::ArrayAttr newMembersAttr; + mlir::SmallVector<mlir::Value> newMembers; + llvm::SmallVector<llvm::SmallVector<int64_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 parent's 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 its members to reflect the insertion of the base address. + if (!mapMemberUsers.empty()) { + // Currently, there should only be one user per map when this pass ---------------- agozillon wrote:
There's an equivalent on entry to the pass, but happy to add another fail safe. https://github.com/llvm/llvm-project/pull/111192 _______________________________________________ llvm-branch-commits mailing list llvm-branch-commits@lists.llvm.org https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-branch-commits