================
@@ -225,32 +359,79 @@ class MapInfoFinalizationPass
fir::FirOpBuilder &builder,
mlir::Operation *target) {
auto mapClauseOwner =
- llvm::dyn_cast<mlir::omp::MapClauseOwningOpInterface>(target);
+ llvm::dyn_cast_if_present<mlir::omp::MapClauseOwningOpInterface>(
+ target);
+ // TargetDataOp is technically a MapClauseOwningOpInterface, so we
+ // do not need to explicitly check for the extra cases here for use_device
+ // addr/ptr
if (!mapClauseOwner)
return;
- llvm::SmallVector<mlir::Value> newMapOps;
- mlir::OperandRange mapVarsArr = mapClauseOwner.getMapVars();
- auto targetOp = llvm::dyn_cast<mlir::omp::TargetOp>(target);
-
- for (size_t i = 0; i < mapVarsArr.size(); ++i) {
- if (mapVarsArr[i] == op) {
- for (auto [j, mapMember] : llvm::enumerate(op.getMembers())) {
- newMapOps.push_back(mapMember);
- // 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 MapVars).
- if (targetOp) {
- targetOp.getRegion().insertArgument(i + j, mapMember.getType(),
- targetOp->getLoc());
+ auto addOperands = [&](mlir::OperandRange &mapVarsArr,
+ mlir::MutableOperandRange &mutableOpRange,
+ auto directiveOp) {
+ llvm::SmallVector<mlir::Value> newMapOps;
+ for (size_t i = 0; i < mapVarsArr.size(); ++i) {
+ if (mapVarsArr[i] == op) {
+ for (auto [j, mapMember] : llvm::enumerate(op.getMembers())) {
+ newMapOps.push_back(mapMember);
+ // 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 MapVars).
+ if (directiveOp) {
+ directiveOp.getRegion().insertArgument(i + j,
mapMember.getType(),
+ directiveOp->getLoc());
+ }
}
}
+ newMapOps.push_back(mapVarsArr[i]);
}
- newMapOps.push_back(mapVarsArr[i]);
+ mutableOpRange.assign(newMapOps);
+ };
+
+ if (auto mapClauseOwner =
+ llvm::dyn_cast<mlir::omp::MapClauseOwningOpInterface>(target)) {
+ mlir::OperandRange mapVarsArr = mapClauseOwner.getMapVars();
+ mlir::MutableOperandRange mapMutableOpRange =
+ mapClauseOwner.getMapVarsMutable();
+ mlir::omp::TargetOp targetOp =
+ llvm::dyn_cast<mlir::omp::TargetOp>(target);
+ addOperands(mapVarsArr, 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);
}
- mapClauseOwner.getMapVarsMutable().assign(newMapOps);
+ }
+
+ // We retrieve the first user that is a Target operation, of which
+ // there should only be one currently. Every MapInfoOp can be tied to
+ // at most one Target operation and at the minimum no operations.
+ // This may change in the future with IR cleanups/modifications,
+ // in which case this pass will need updating to support cases
+ // where a map can have more than one user and more than one of
+ // those users can be a Target operation. For now, we simply
+ // return the first target operation encountered, which may
+ // be on the parent MapInfoOp in the case of a member mapping.
+ // In that case, we traverse the MapInfoOp chain until we
+ // find the first TargetOp user.
+ mlir::Operation *getFirstTargetUser(mlir::omp::MapInfoOp mapOp) {
+ for (auto *user : mapOp->getUsers()) {
+ if (llvm::isa<mlir::omp::TargetOp, mlir::omp::TargetDataOp,
+ mlir::omp::TargetUpdateOp, mlir::omp::TargetExitDataOp,
+ mlir::omp::TargetEnterDataOp>(user))
+ return user;
+
+ if (auto mapUser = llvm::dyn_cast_if_present<mlir::omp::MapInfoOp>(user))
+ return getFirstTargetUser(mapUser);
----------------
agozillon wrote:
Going to opt for some assert checks, not a huge fan of llvm_unreachable outside
of certain cases!
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
