================
@@ -391,29 +414,237 @@ class MapInfoFinalizationPass
   /// of the base address index.
   void adjustMemberIndices(
       llvm::SmallVectorImpl<llvm::SmallVector<int64_t>> &memberIndices,
-      size_t memberIndex) {
-    llvm::SmallVector<int64_t> baseAddrIndex = memberIndices[memberIndex];
+      ParentAndPlacement parentAndPlacement) {
+    llvm::SmallVector<int64_t> baseAddrIndex =
+        memberIndices[parentAndPlacement.index];
+    auto &expansionIndexes = expandedBaseAddr[parentAndPlacement.parent];
 
     // 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.
-    for (llvm::SmallVector<int64_t> &member : memberIndices)
+    for (auto [i, member] : llvm::enumerate(memberIndices)) {
+      if (std::find(expansionIndexes.begin(), expansionIndexes.end(), i) !=
+          expansionIndexes.end())
+        if (member.size() == baseAddrIndex.size() + 1 &&
+            member[baseAddrIndex.size()] == 0)
+          continue;
+
       if (member.size() > baseAddrIndex.size() &&
           std::equal(baseAddrIndex.begin(), baseAddrIndex.end(),
                      member.begin()))
         member.insert(std::next(member.begin(), baseAddrIndex.size()), 0);
+    }
 
     // Add the base address index to the main base address member data
     baseAddrIndex.push_back(0);
 
-    // Insert our newly created baseAddrIndex into the larger list of indices 
at
-    // the correct location.
-    memberIndices.insert(std::next(memberIndices.begin(), memberIndex + 1),
+    uint64_t newIdxInsert = parentAndPlacement.index + 1;
+    expansionIndexes.push_back(newIdxInsert);
+
+    // Insert our newly created baseAddrIndex into the larger list of
+    // indices at the correct location.
+    memberIndices.insert(std::next(memberIndices.begin(), newIdxInsert),
                          baseAddrIndex);
   }
 
+  void
+  insertIntoMapClauseInterface(mlir::Operation *target,
+                               std::function<void(mlir::MutableOperandRange &,
+                                                  mlir::Operation *, unsigned)>
+                                   addOperands) {
+    auto argIface =
+        llvm::dyn_cast<mlir::omp::BlockArgOpenMPOpInterface>(target);
+
+    if (auto mapClauseOwner =
+            llvm::dyn_cast<mlir::omp::MapClauseOwningOpInterface>(target)) {
+      mlir::MutableOperandRange mapVarsArr = 
mapClauseOwner.getMapVarsMutable();
+      unsigned blockArgInsertIndex =
+          argIface
+              ? argIface.getMapBlockArgsStart() + argIface.numMapBlockArgs()
+              : 0;
+      addOperands(mapVarsArr,
+                  llvm::dyn_cast_if_present<mlir::omp::TargetOp>(
+                      argIface.getOperation()),
+                  blockArgInsertIndex);
+    }
+
+    if (auto targetDataOp = llvm::dyn_cast<mlir::omp::TargetDataOp>(target)) {
+      mlir::MutableOperandRange useDevAddrMutableOpRange =
+          targetDataOp.getUseDeviceAddrVarsMutable();
+      addOperands(useDevAddrMutableOpRange, target,
+                  argIface.getUseDeviceAddrBlockArgsStart() +
+                      argIface.numUseDeviceAddrBlockArgs());
+
+      mlir::MutableOperandRange useDevPtrMutableOpRange =
+          targetDataOp.getUseDevicePtrVarsMutable();
+      addOperands(useDevPtrMutableOpRange, target,
+                  argIface.getUseDevicePtrBlockArgsStart() +
+                      argIface.numUseDevicePtrBlockArgs());
+    } else if (auto targetOp = llvm::dyn_cast<mlir::omp::TargetOp>(target)) {
+      mlir::MutableOperandRange hasDevAddrMutableOpRange =
+          targetOp.getHasDeviceAddrVarsMutable();
+      addOperands(hasDevAddrMutableOpRange, target,
+                  argIface.getHasDeviceAddrBlockArgsStart() +
+                      argIface.numHasDeviceAddrBlockArgs());
+    }
+  }
+
+  // This functions aims to insert new maps derived from existing maps into the
+  // corresponding clause list, interlinking it correctly with block arguments
+  // where required .
+  void addDerivedMemberToTarget(
+      mlir::omp::MapInfoOp owner, mlir::omp::MapInfoOp derived,
+      llvm::SmallVectorImpl<ParentAndPlacement> &mapMemberUsers,
+      fir::FirOpBuilder &builder, mlir::Operation *target) {
+    auto addOperands = [&](mlir::MutableOperandRange &mapVarsArr,
+                           mlir::Operation *directiveOp,
+                           unsigned blockArgInsertIndex = 0) {
+      // Check we're inserting into the correct MapInfoOp list
+      if (!llvm::is_contained(mapVarsArr.getAsOperandRange(),
+                              mapMemberUsers.empty()
+                                  ? owner.getResult()
+                                  : mapMemberUsers[0].parent.getResult()))
+        return;
+
+      // Check we're not inserting a duplicate map.
+      if (llvm::is_contained(mapVarsArr.getAsOperandRange(),
+                             derived.getResult()))
+        return;
+
+      // There doesn't appear to be a simple way to convert MutableOperandRange
+      // to a vector currently, so we instead use a for_each to populate our
+      // vector.
+      llvm::SmallVector<mlir::Value> newMapOps;
+      newMapOps.reserve(mapVarsArr.size());
+      llvm::for_each(
+          mapVarsArr.getAsOperandRange(),
+          [&newMapOps](mlir::Value oper) { newMapOps.push_back(oper); });
+
+      newMapOps.push_back(derived);
+      if (directiveOp) {
+        directiveOp->getRegion(0).insertArgument(
+            blockArgInsertIndex, derived.getType(), derived.getLoc());
+        blockArgInsertIndex++;
+      }
+
+      mapVarsArr.assign(newMapOps);
+    };
+
+    insertIntoMapClauseInterface(target, addOperands);
+  }
+
+  // We add all mapped record members not directly used in the target region
+  // to the block arguments in front of their parent and we place them into
+  // the map operands list for consistency.
+  //
+  // These indirect uses (via accesses to their parent) will still be
+  // mapped individually in most cases, and a parent mapping doesn't
+  // guarantee the parent will be mapped in its totality, partial
+  // mapping is common.
+  //
+  // For example:
+  //    map(tofrom: x%y)
+  //
+  // Will generate a mapping for "x" (the parent) and "y" (the member).
+  // The parent "x" will not be mapped, but the member "y" will.
+  // However, we must have the parent as a BlockArg and MapOperand
+  // in these cases, to maintain the correct uses within the region and
+  // to help tracking that the member is part of a larger object.
+  //
+  // In the case of:
+  //    map(tofrom: x%y, x%z)
+  //
+  // The parent member becomes more critical, as we perform a partial
+  // structure mapping where we link the mapping of the members y
+  // and z together via the parent x. We do this at a kernel argument
+  // level in LLVM IR and not just MLIR, which is important to maintain
+  // similarity to Clang and for the runtime to do the correct thing.
+  // However, we still do not map the structure in its totality but
+  // rather we generate an un-sized "binding" map entry for it.
+  //
+  // In the case of:
+  //    map(tofrom: x, x%y, x%z)
+  //
+  // We do actually map the entirety of "x", so the explicit mapping of
+  // x%y, x%z becomes unnecessary. It is redundant to write this from a
+  // Fortran OpenMP perspective (although it is legal), as even if the
+  // members were allocatables or pointers, we are mandated by the
+  // specification to map these (and any recursive components) in their
+  // entirety, which is different to the C++ equivalent, which requires
+  // explicit mapping of these segments.
+  void addImplicitMembersToTarget(mlir::omp::MapInfoOp op,
+                                  fir::FirOpBuilder &builder,
+                                  mlir::Operation *target) {
+    auto mapClauseOwner =
+        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;
+
+    auto addOperands = [&](mlir::MutableOperandRange &mapVarsArr,
+                           mlir::Operation *directiveOp,
+                           unsigned blockArgInsertIndex = 0) {
+      if (!llvm::is_contained(mapVarsArr.getAsOperandRange(), op.getResult()))
+        return;
+
+      // There doesn't appear to be a simple way to convert MutableOperandRange
+      // to a vector currently, so we instead use a for_each to populate our
+      // vector.
+      llvm::SmallVector<mlir::Value> newMapOps;
+      newMapOps.reserve(mapVarsArr.size());
+      llvm::for_each(mapVarsArr.getAsOperandRange(),
+                     [&newMapOps](mlir::Value oper) {
+                       if (oper)
+                         newMapOps.push_back(oper);
+                     });
+
+      for (auto mapMember : op.getMembers()) {
+        if (llvm::is_contained(mapVarsArr.getAsOperandRange(), mapMember))
+          continue;
+        newMapOps.push_back(mapMember);
+        if (directiveOp) {
+          directiveOp->getRegion(0).insertArgument(
+              blockArgInsertIndex, mapMember.getType(), mapMember.getLoc());
+          blockArgInsertIndex++;
+        }
+      }
+      mapVarsArr.assign(newMapOps);
+    };
+
+    insertIntoMapClauseInterface(target, addOperands);
+  }
+
+  // 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.
----------------
bhandarkar-pranav wrote:

I agree that an assert here is the right way to go. Short of making any changes 
to the dialect itself, at least, it makes explicit our assumptions here in 
code. 

https://github.com/llvm/llvm-project/pull/177715
_______________________________________________
llvm-branch-commits mailing list
[email protected]
https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-branch-commits

Reply via email to