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

[Sergio Afonso via llvm-branch-commits]

================
@@ -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))
----------------
skatrak wrote:

Nit: The check above already expects that `user` is not null, and it seems like it would be unexpected that users would be null.
```suggestion
      if (auto mapUser = llvm::dyn_cast<mlir::omp::MapInfoOp>(user))
```

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