[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(),
----------------
skatrak wrote:

This probably needs rewriting based on the `BlockArgOpenMPOpInterface`, since this index-based approach easily breaks if other entry block argument-generating clauses are introduced before `map` (such as `in_reduction`). This may be a helpful reference: https://github.com/llvm/llvm-project/blob/db1a76206902e6926464e30cc47249e217405dab/flang/lib/Lower/OpenMP/OpenMP.cpp#L1010-L1014

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