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

[Sergio Afonso via llvm-branch-commits]

================
@@ -81,105 +135,192 @@ class MapInfoFinalizationPass
     // !fir.ref<!fir.box<...>> to access the data we need to map we must
     // 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);
+    return alloca;
+  }
 
+  /// 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 genBaseAddrMap(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 members
+  /// that are members of this descriptor to account for the addition
+  /// of the base address index.
+  void adjustMemberIndices(
+      llvm::SmallVectorImpl<llvm::SmallVector<int64_t>> &memberIndices,
+      size_t memberIndex) {
+    llvm::SmallVector<int64_t> baseAddrIndex = memberIndices[memberIndex];
 
-    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(), baseAddrIndex.end());
+    for (llvm::SmallVector<int64_t> member : memberIndices) {
+      if (member.size() > insertPosition &&
+          std::equal(baseAddrIndex.begin(), baseAddrIndex.end(),
+                     member.begin())) {
+        member.insert(std::next(member.begin(), insertPosition), 0);
----------------
skatrak wrote:

If I understand it correctly, we're inserting into a local copy here, so this statement wouldn't serve any purpose (unless there's some corner inside of C++ or the `llvm::SmallVector` class I'm missing). Shouldn't `member` be declared as a reference above for this to work?

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