xref: /llvm-project/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp (revision 662133a278f4f3553f061f7999759bae4e842820)
165f66d2cSIvan R. Ivanov //===- MapInfoFinalization.cpp -----------------------------------------===//
265f66d2cSIvan R. Ivanov //
365f66d2cSIvan R. Ivanov // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
465f66d2cSIvan R. Ivanov // See https://llvm.org/LICENSE.txt for license information.
565f66d2cSIvan R. Ivanov // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
665f66d2cSIvan R. Ivanov //
765f66d2cSIvan R. Ivanov //===----------------------------------------------------------------------===//
865f66d2cSIvan R. Ivanov 
965f66d2cSIvan R. Ivanov //===----------------------------------------------------------------------===//
1065f66d2cSIvan R. Ivanov /// \file
1165f66d2cSIvan R. Ivanov /// An OpenMP dialect related pass for FIR/HLFIR which performs some
1265f66d2cSIvan R. Ivanov /// pre-processing of MapInfoOp's after the module has been lowered to
1365f66d2cSIvan R. Ivanov /// finalize them.
1465f66d2cSIvan R. Ivanov ///
1565f66d2cSIvan R. Ivanov /// For example, it expands MapInfoOp's containing descriptor related
1665f66d2cSIvan R. Ivanov /// types (fir::BoxType's) into multiple MapInfoOp's containing the parent
1765f66d2cSIvan R. Ivanov /// descriptor and pointer member components for individual mapping,
1865f66d2cSIvan R. Ivanov /// treating the descriptor type as a record type for later lowering in the
1965f66d2cSIvan R. Ivanov /// OpenMP dialect.
2065f66d2cSIvan R. Ivanov ///
2165f66d2cSIvan R. Ivanov /// The pass also adds MapInfoOp's that are members of a parent object but are
2265f66d2cSIvan R. Ivanov /// not directly used in the body of a target region to its BlockArgument list
2365f66d2cSIvan R. Ivanov /// to maintain consistency across all MapInfoOp's tied to a region directly or
2465f66d2cSIvan R. Ivanov /// indirectly via a parent object.
2565f66d2cSIvan R. Ivanov //===----------------------------------------------------------------------===//
2665f66d2cSIvan R. Ivanov 
27*662133a2SjeanPerier #include "flang/Optimizer/Builder/DirectivesCommon.h"
2865f66d2cSIvan R. Ivanov #include "flang/Optimizer/Builder/FIRBuilder.h"
29e532241bSKareem Ergawy #include "flang/Optimizer/Builder/HLFIRTools.h"
3065f66d2cSIvan R. Ivanov #include "flang/Optimizer/Dialect/FIRType.h"
3165f66d2cSIvan R. Ivanov #include "flang/Optimizer/Dialect/Support/KindMapping.h"
32e532241bSKareem Ergawy #include "flang/Optimizer/HLFIR/HLFIROps.h"
3365f66d2cSIvan R. Ivanov #include "flang/Optimizer/OpenMP/Passes.h"
34e532241bSKareem Ergawy #include "mlir/Analysis/SliceAnalysis.h"
3565f66d2cSIvan R. Ivanov #include "mlir/Dialect/Func/IR/FuncOps.h"
3665f66d2cSIvan R. Ivanov #include "mlir/Dialect/OpenMP/OpenMPDialect.h"
3765f66d2cSIvan R. Ivanov #include "mlir/IR/BuiltinDialect.h"
3865f66d2cSIvan R. Ivanov #include "mlir/IR/BuiltinOps.h"
3965f66d2cSIvan R. Ivanov #include "mlir/IR/Operation.h"
4065f66d2cSIvan R. Ivanov #include "mlir/IR/SymbolTable.h"
4165f66d2cSIvan R. Ivanov #include "mlir/Pass/Pass.h"
4265f66d2cSIvan R. Ivanov #include "mlir/Support/LLVM.h"
4365f66d2cSIvan R. Ivanov #include "llvm/ADT/SmallPtrSet.h"
4465f66d2cSIvan R. Ivanov #include "llvm/Frontend/OpenMP/OMPConstants.h"
45e508baccSagozillon #include <algorithm>
46e508baccSagozillon #include <cstddef>
4765f66d2cSIvan R. Ivanov #include <iterator>
48e508baccSagozillon #include <numeric>
4965f66d2cSIvan R. Ivanov 
5065f66d2cSIvan R. Ivanov namespace flangomp {
5165f66d2cSIvan R. Ivanov #define GEN_PASS_DEF_MAPINFOFINALIZATIONPASS
5265f66d2cSIvan R. Ivanov #include "flang/Optimizer/OpenMP/Passes.h.inc"
5365f66d2cSIvan R. Ivanov } // namespace flangomp
5465f66d2cSIvan R. Ivanov 
5565f66d2cSIvan R. Ivanov namespace {
5665f66d2cSIvan R. Ivanov class MapInfoFinalizationPass
5765f66d2cSIvan R. Ivanov     : public flangomp::impl::MapInfoFinalizationPassBase<
5865f66d2cSIvan R. Ivanov           MapInfoFinalizationPass> {
59e508baccSagozillon   /// Helper class tracking a members parent and its
60e508baccSagozillon   /// placement in the parents member list
61e508baccSagozillon   struct ParentAndPlacement {
62e508baccSagozillon     mlir::omp::MapInfoOp parent;
63e508baccSagozillon     size_t index;
64e508baccSagozillon   };
6565f66d2cSIvan R. Ivanov 
66f4cf93fbSagozillon   /// Tracks any intermediate function/subroutine local allocations we
67f4cf93fbSagozillon   /// generate for the descriptors of box type dummy arguments, so that
68f4cf93fbSagozillon   /// we can retrieve it for subsequent reuses within the functions
69d291e459Smacurtis-amd   /// scope.
70d291e459Smacurtis-amd   ///
71d291e459Smacurtis-amd   ///      descriptor defining op
72d291e459Smacurtis-amd   ///      |                  corresponding local alloca
73d291e459Smacurtis-amd   ///      |                  |
74d291e459Smacurtis-amd   std::map<mlir::Operation *, mlir::Value> localBoxAllocas;
75f4cf93fbSagozillon 
76e508baccSagozillon   /// getMemberUserList gathers all users of a particular MapInfoOp that are
77e508baccSagozillon   /// other MapInfoOp's and places them into the mapMemberUsers list, which
78e508baccSagozillon   /// records the map that the current argument MapInfoOp "op" is part of
79e508baccSagozillon   /// alongside the placement of "op" in the recorded users members list. The
80e508baccSagozillon   /// intent of the generated list is to find all MapInfoOp's that may be
81e508baccSagozillon   /// considered parents of the passed in "op" and in which it shows up in the
82e508baccSagozillon   /// member list, alongside collecting the placement information of "op" in its
83e508baccSagozillon   /// parents member list.
84e508baccSagozillon   void
85e508baccSagozillon   getMemberUserList(mlir::omp::MapInfoOp op,
86e508baccSagozillon                     llvm::SmallVectorImpl<ParentAndPlacement> &mapMemberUsers) {
87e508baccSagozillon     for (auto *user : op->getUsers())
88e508baccSagozillon       if (auto map = mlir::dyn_cast_if_present<mlir::omp::MapInfoOp>(user))
89e508baccSagozillon         for (auto [i, mapMember] : llvm::enumerate(map.getMembers()))
90e508baccSagozillon           if (mapMember.getDefiningOp() == op)
91e508baccSagozillon             mapMemberUsers.push_back({map, i});
92e508baccSagozillon   }
9365f66d2cSIvan R. Ivanov 
94e508baccSagozillon   void getAsIntegers(llvm::ArrayRef<mlir::Attribute> values,
95e508baccSagozillon                      llvm::SmallVectorImpl<int64_t> &ints) {
96e508baccSagozillon     ints.reserve(values.size());
97e508baccSagozillon     llvm::transform(values, std::back_inserter(ints),
98e508baccSagozillon                     [](mlir::Attribute value) {
99e508baccSagozillon                       return mlir::cast<mlir::IntegerAttr>(value).getInt();
100e508baccSagozillon                     });
101e508baccSagozillon   }
102e508baccSagozillon 
103e508baccSagozillon   /// This function will expand a MapInfoOp's member indices back into a vector
104e508baccSagozillon   /// so that they can be trivially modified as unfortunately the attribute type
105e508baccSagozillon   /// that's used does not have modifiable fields at the moment (generally
106e508baccSagozillon   /// awkward to work with)
107e508baccSagozillon   void getMemberIndicesAsVectors(
108e508baccSagozillon       mlir::omp::MapInfoOp mapInfo,
109e508baccSagozillon       llvm::SmallVectorImpl<llvm::SmallVector<int64_t>> &indices) {
110e508baccSagozillon     indices.reserve(mapInfo.getMembersIndexAttr().getValue().size());
111e508baccSagozillon     llvm::transform(mapInfo.getMembersIndexAttr().getValue(),
112e508baccSagozillon                     std::back_inserter(indices), [this](mlir::Attribute value) {
113e508baccSagozillon                       auto memberIndex = mlir::cast<mlir::ArrayAttr>(value);
114e508baccSagozillon                       llvm::SmallVector<int64_t> indexes;
115e508baccSagozillon                       getAsIntegers(memberIndex.getValue(), indexes);
116e508baccSagozillon                       return indexes;
117e508baccSagozillon                     });
118e508baccSagozillon   }
119e508baccSagozillon 
120e508baccSagozillon   /// When provided a MapInfoOp containing a descriptor type that
121e508baccSagozillon   /// we must expand into multiple maps this function will extract
122e508baccSagozillon   /// the value from it and return it, in certain cases we must
123e508baccSagozillon   /// generate a new allocation to store into so that the
124e508baccSagozillon   /// fir::BoxOffsetOp we utilise to access the descriptor datas
125e508baccSagozillon   /// base address can be utilised.
126e508baccSagozillon   mlir::Value getDescriptorFromBoxMap(mlir::omp::MapInfoOp boxMap,
127e508baccSagozillon                                       fir::FirOpBuilder &builder) {
128e508baccSagozillon     mlir::Value descriptor = boxMap.getVarPtr();
129e508baccSagozillon     if (!fir::isTypeWithDescriptor(boxMap.getVarType()))
13065f66d2cSIvan R. Ivanov       if (auto addrOp = mlir::dyn_cast_if_present<fir::BoxAddrOp>(
131e508baccSagozillon               boxMap.getVarPtr().getDefiningOp()))
13265f66d2cSIvan R. Ivanov         descriptor = addrOp.getVal();
133e508baccSagozillon 
134e508baccSagozillon     if (!mlir::isa<fir::BaseBoxType>(descriptor.getType()))
135e508baccSagozillon       return descriptor;
13665f66d2cSIvan R. Ivanov 
137d291e459Smacurtis-amd     mlir::Value &slot = localBoxAllocas[descriptor.getDefiningOp()];
138d291e459Smacurtis-amd     if (slot) {
139d291e459Smacurtis-amd       return slot;
140d291e459Smacurtis-amd     }
141d291e459Smacurtis-amd 
14265f66d2cSIvan R. Ivanov     // The fir::BoxOffsetOp only works with !fir.ref<!fir.box<...>> types, as
14365f66d2cSIvan R. Ivanov     // allowing it to access non-reference box operations can cause some
14465f66d2cSIvan R. Ivanov     // problematic SSA IR. However, in the case of assumed shape's the type
14565f66d2cSIvan R. Ivanov     // is not a !fir.ref, in these cases to retrieve the appropriate
14665f66d2cSIvan R. Ivanov     // !fir.ref<!fir.box<...>> to access the data we need to map we must
14765f66d2cSIvan R. Ivanov     // perform an alloca and then store to it and retrieve the data from the new
14865f66d2cSIvan R. Ivanov     // alloca.
14965f66d2cSIvan R. Ivanov     mlir::OpBuilder::InsertPoint insPt = builder.saveInsertionPoint();
15065f66d2cSIvan R. Ivanov     mlir::Block *allocaBlock = builder.getAllocaBlock();
151e508baccSagozillon     mlir::Location loc = boxMap->getLoc();
15265f66d2cSIvan R. Ivanov     assert(allocaBlock && "No alloca block found for this top level op");
15365f66d2cSIvan R. Ivanov     builder.setInsertionPointToStart(allocaBlock);
15465f66d2cSIvan R. Ivanov     auto alloca = builder.create<fir::AllocaOp>(loc, descriptor.getType());
15565f66d2cSIvan R. Ivanov     builder.restoreInsertionPoint(insPt);
15665f66d2cSIvan R. Ivanov     builder.create<fir::StoreOp>(loc, descriptor, alloca);
157d291e459Smacurtis-amd     return slot = alloca;
158f4cf93fbSagozillon   }
15965f66d2cSIvan R. Ivanov 
160e508baccSagozillon   /// Function that generates a FIR operation accessing the descriptor's
161e508baccSagozillon   /// base address (BoxOffsetOp) and a MapInfoOp for it. The most
162e508baccSagozillon   /// important thing to note is that we normally move the bounds from
163e508baccSagozillon   /// the descriptor map onto the base address map.
164e508baccSagozillon   mlir::omp::MapInfoOp genBaseAddrMap(mlir::Value descriptor,
165e508baccSagozillon                                       mlir::OperandRange bounds,
166e508baccSagozillon                                       int64_t mapType,
167e508baccSagozillon                                       fir::FirOpBuilder &builder) {
168e508baccSagozillon     mlir::Location loc = descriptor.getLoc();
16965f66d2cSIvan R. Ivanov     mlir::Value baseAddrAddr = builder.create<fir::BoxOffsetOp>(
17065f66d2cSIvan R. Ivanov         loc, descriptor, fir::BoxFieldAttr::base_addr);
17165f66d2cSIvan R. Ivanov 
1725137c209Sagozillon     mlir::Type underlyingVarType =
1735137c209Sagozillon         llvm::cast<mlir::omp::PointerLikeType>(
1745137c209Sagozillon             fir::unwrapRefType(baseAddrAddr.getType()))
1755137c209Sagozillon             .getElementType();
1765137c209Sagozillon     if (auto seqType = llvm::dyn_cast<fir::SequenceType>(underlyingVarType))
1775137c209Sagozillon       if (seqType.hasDynamicExtents())
1785137c209Sagozillon         underlyingVarType = seqType.getEleTy();
1795137c209Sagozillon 
18065f66d2cSIvan R. Ivanov     // Member of the descriptor pointing at the allocated data
181e508baccSagozillon     return builder.create<mlir::omp::MapInfoOp>(
18265f66d2cSIvan R. Ivanov         loc, baseAddrAddr.getType(), descriptor,
1835137c209Sagozillon         mlir::TypeAttr::get(underlyingVarType), baseAddrAddr,
1845137c209Sagozillon         /*members=*/mlir::SmallVector<mlir::Value>{},
185e508baccSagozillon         /*membersIndex=*/mlir::ArrayAttr{}, bounds,
186e508baccSagozillon         builder.getIntegerAttr(builder.getIntegerType(64, false), mapType),
18765f66d2cSIvan R. Ivanov         builder.getAttr<mlir::omp::VariableCaptureKindAttr>(
18865f66d2cSIvan R. Ivanov             mlir::omp::VariableCaptureKind::ByRef),
18965f66d2cSIvan R. Ivanov         /*name=*/builder.getStringAttr(""),
19065f66d2cSIvan R. Ivanov         /*partial_map=*/builder.getBoolAttr(false));
191e508baccSagozillon   }
192e508baccSagozillon 
193e508baccSagozillon   /// This function adjusts the member indices vector to include a new
194e508baccSagozillon   /// base address member. We take the position of the descriptor in
195e508baccSagozillon   /// the member indices list, which is the index data that the base
196e508baccSagozillon   /// addresses index will be based off of, as the base address is
197e508baccSagozillon   /// a member of the descriptor. We must also alter other members
198e508baccSagozillon   /// that are members of this descriptor to account for the addition
199e508baccSagozillon   /// of the base address index.
200e508baccSagozillon   void adjustMemberIndices(
201e508baccSagozillon       llvm::SmallVectorImpl<llvm::SmallVector<int64_t>> &memberIndices,
202e508baccSagozillon       size_t memberIndex) {
203e508baccSagozillon     llvm::SmallVector<int64_t> baseAddrIndex = memberIndices[memberIndex];
204e508baccSagozillon 
205e508baccSagozillon     // If we find another member that is "derived/a member of" the descriptor
206e508baccSagozillon     // that is not the descriptor itself, we must insert a 0 for the new base
207e508baccSagozillon     // address we have just added for the descriptor into the list at the
208e508baccSagozillon     // appropriate position to maintain correctness of the positional/index data
209e508baccSagozillon     // for that member.
210e508baccSagozillon     for (llvm::SmallVector<int64_t> &member : memberIndices)
211e508baccSagozillon       if (member.size() > baseAddrIndex.size() &&
212e508baccSagozillon           std::equal(baseAddrIndex.begin(), baseAddrIndex.end(),
213e508baccSagozillon                      member.begin()))
214e508baccSagozillon         member.insert(std::next(member.begin(), baseAddrIndex.size()), 0);
215e508baccSagozillon 
216e508baccSagozillon     // Add the base address index to the main base address member data
217e508baccSagozillon     baseAddrIndex.push_back(0);
218e508baccSagozillon 
219e508baccSagozillon     // Insert our newly created baseAddrIndex into the larger list of indices at
220e508baccSagozillon     // the correct location.
221e508baccSagozillon     memberIndices.insert(std::next(memberIndices.begin(), memberIndex + 1),
222e508baccSagozillon                          baseAddrIndex);
223e508baccSagozillon   }
224e508baccSagozillon 
225e508baccSagozillon   /// Adjusts the descriptor's map type. The main alteration that is done
226e508baccSagozillon   /// currently is transforming the map type to `OMP_MAP_TO` where possible.
227e508baccSagozillon   /// This is because we will always need to map the descriptor to device
228e508baccSagozillon   /// (or at the very least it seems to be the case currently with the
229e508baccSagozillon   /// current lowered kernel IR), as without the appropriate descriptor
230e508baccSagozillon   /// information on the device there is a risk of the kernel IR
231e508baccSagozillon   /// requesting for various data that will not have been copied to
232e508baccSagozillon   /// perform things like indexing. This can cause segfaults and
233e508baccSagozillon   /// memory access errors. However, we do not need this data mapped
234e508baccSagozillon   /// back to the host from the device, as per the OpenMP spec we cannot alter
235e508baccSagozillon   /// the data via resizing or deletion on the device. Discarding any
236e508baccSagozillon   /// descriptor alterations via no map back is reasonable (and required
237e508baccSagozillon   /// for certain segments of descriptor data like the type descriptor that are
238e508baccSagozillon   /// global constants). This alteration is only inapplicable to `target exit`
239e508baccSagozillon   /// and `target update` currently, and that's due to `target exit` not
240e508baccSagozillon   /// allowing `to` mappings, and `target update` not allowing both `to` and
241e508baccSagozillon   /// `from` simultaneously. We currently try to maintain the `implicit` flag
242e508baccSagozillon   /// where necessary, although it does not seem strictly required.
243e508baccSagozillon   unsigned long getDescriptorMapType(unsigned long mapTypeFlag,
244e508baccSagozillon                                      mlir::Operation *target) {
245e508baccSagozillon     if (llvm::isa_and_nonnull<mlir::omp::TargetExitDataOp,
246e508baccSagozillon                               mlir::omp::TargetUpdateOp>(target))
247e508baccSagozillon       return mapTypeFlag;
248e508baccSagozillon 
249e508baccSagozillon     bool hasImplicitMap =
250e508baccSagozillon         (llvm::omp::OpenMPOffloadMappingFlags(mapTypeFlag) &
251e508baccSagozillon          llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_IMPLICIT) ==
252e508baccSagozillon         llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_IMPLICIT;
253e508baccSagozillon 
254e508baccSagozillon     return llvm::to_underlying(
255e508baccSagozillon         hasImplicitMap
256e508baccSagozillon             ? llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_TO |
257e508baccSagozillon                   llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_IMPLICIT
258e508baccSagozillon             : llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_TO);
259e508baccSagozillon   }
260e508baccSagozillon 
261e508baccSagozillon   mlir::omp::MapInfoOp genDescriptorMemberMaps(mlir::omp::MapInfoOp op,
262e508baccSagozillon                                                fir::FirOpBuilder &builder,
263e508baccSagozillon                                                mlir::Operation *target) {
264e508baccSagozillon     llvm::SmallVector<ParentAndPlacement> mapMemberUsers;
265e508baccSagozillon     getMemberUserList(op, mapMemberUsers);
26665f66d2cSIvan R. Ivanov 
26765f66d2cSIvan R. Ivanov     // TODO: map the addendum segment of the descriptor, similarly to the
268e508baccSagozillon     // base address/data pointer member.
269e508baccSagozillon     mlir::Value descriptor = getDescriptorFromBoxMap(op, builder);
270e508baccSagozillon     auto baseAddr = genBaseAddrMap(descriptor, op.getBounds(),
271e508baccSagozillon                                    op.getMapType().value_or(0), builder);
272e508baccSagozillon     mlir::ArrayAttr newMembersAttr;
273e508baccSagozillon     mlir::SmallVector<mlir::Value> newMembers;
274e508baccSagozillon     llvm::SmallVector<llvm::SmallVector<int64_t>> memberIndices;
275e508baccSagozillon 
276e508baccSagozillon     if (!mapMemberUsers.empty() || !op.getMembers().empty())
277e508baccSagozillon       getMemberIndicesAsVectors(
278e508baccSagozillon           !mapMemberUsers.empty() ? mapMemberUsers[0].parent : op,
279e508baccSagozillon           memberIndices);
280e508baccSagozillon 
281e508baccSagozillon     // If the operation that we are expanding with a descriptor has a user
282e508baccSagozillon     // (parent), then we have to expand the parent's member indices to reflect
283e508baccSagozillon     // the adjusted member indices for the base address insertion. However, if
284e508baccSagozillon     // it does not then we are expanding a MapInfoOp without any pre-existing
285e508baccSagozillon     // member information to now have one new member for the base address, or
286e508baccSagozillon     // we are expanding a parent that is a descriptor and we have to adjust
287e508baccSagozillon     // all of its members to reflect the insertion of the base address.
288e508baccSagozillon     if (!mapMemberUsers.empty()) {
289e508baccSagozillon       // Currently, there should only be one user per map when this pass
290e508baccSagozillon       // is executed. Either a parent map, holding the current map in its
291e508baccSagozillon       // member list, or a target operation that holds a map clause. This
292e508baccSagozillon       // may change in the future if we aim to refactor the MLIR for map
293e508baccSagozillon       // clauses to allow sharing of duplicate maps across target
294e508baccSagozillon       // operations.
295e508baccSagozillon       assert(mapMemberUsers.size() == 1 &&
296e508baccSagozillon              "OMPMapInfoFinalization currently only supports single users of a "
297e508baccSagozillon              "MapInfoOp");
298e508baccSagozillon       ParentAndPlacement mapUser = mapMemberUsers[0];
299e508baccSagozillon       adjustMemberIndices(memberIndices, mapUser.index);
300e508baccSagozillon       llvm::SmallVector<mlir::Value> newMemberOps;
301e508baccSagozillon       for (auto v : mapUser.parent.getMembers()) {
302e508baccSagozillon         newMemberOps.push_back(v);
303e508baccSagozillon         if (v == op)
304e508baccSagozillon           newMemberOps.push_back(baseAddr);
305e508baccSagozillon       }
306e508baccSagozillon       mapUser.parent.getMembersMutable().assign(newMemberOps);
307e508baccSagozillon       mapUser.parent.setMembersIndexAttr(
308e508baccSagozillon           builder.create2DI64ArrayAttr(memberIndices));
309e508baccSagozillon     } else {
310e508baccSagozillon       newMembers.push_back(baseAddr);
311e508baccSagozillon       if (!op.getMembers().empty()) {
312e508baccSagozillon         for (auto &indices : memberIndices)
313e508baccSagozillon           indices.insert(indices.begin(), 0);
314e508baccSagozillon         memberIndices.insert(memberIndices.begin(), {0});
315e508baccSagozillon         newMembersAttr = builder.create2DI64ArrayAttr(memberIndices);
316e508baccSagozillon         newMembers.append(op.getMembers().begin(), op.getMembers().end());
317e508baccSagozillon       } else {
318e508baccSagozillon         llvm::SmallVector<llvm::SmallVector<int64_t>> memberIdx = {{0}};
319e508baccSagozillon         newMembersAttr = builder.create2DI64ArrayAttr(memberIdx);
320e508baccSagozillon       }
321e508baccSagozillon     }
32265f66d2cSIvan R. Ivanov 
323d84d0cafSagozillon     mlir::omp::MapInfoOp newDescParentMapOp =
324d84d0cafSagozillon         builder.create<mlir::omp::MapInfoOp>(
32565f66d2cSIvan R. Ivanov             op->getLoc(), op.getResult().getType(), descriptor,
32665f66d2cSIvan R. Ivanov             mlir::TypeAttr::get(fir::unwrapRefType(descriptor.getType())),
327e508baccSagozillon             /*varPtrPtr=*/mlir::Value{}, newMembers, newMembersAttr,
32865f66d2cSIvan R. Ivanov             /*bounds=*/mlir::SmallVector<mlir::Value>{},
329e508baccSagozillon             builder.getIntegerAttr(
330e508baccSagozillon                 builder.getIntegerType(64, false),
331e508baccSagozillon                 getDescriptorMapType(op.getMapType().value_or(0), target)),
332d84d0cafSagozillon             op.getMapCaptureTypeAttr(), op.getNameAttr(),
333e508baccSagozillon             /*partial_map=*/builder.getBoolAttr(false));
334d84d0cafSagozillon     op.replaceAllUsesWith(newDescParentMapOp.getResult());
33565f66d2cSIvan R. Ivanov     op->erase();
336e508baccSagozillon     return newDescParentMapOp;
33765f66d2cSIvan R. Ivanov   }
33865f66d2cSIvan R. Ivanov 
33965f66d2cSIvan R. Ivanov   // We add all mapped record members not directly used in the target region
34065f66d2cSIvan R. Ivanov   // to the block arguments in front of their parent and we place them into
34165f66d2cSIvan R. Ivanov   // the map operands list for consistency.
34265f66d2cSIvan R. Ivanov   //
34365f66d2cSIvan R. Ivanov   // These indirect uses (via accesses to their parent) will still be
34465f66d2cSIvan R. Ivanov   // mapped individually in most cases, and a parent mapping doesn't
34565f66d2cSIvan R. Ivanov   // guarantee the parent will be mapped in its totality, partial
34665f66d2cSIvan R. Ivanov   // mapping is common.
34765f66d2cSIvan R. Ivanov   //
34865f66d2cSIvan R. Ivanov   // For example:
34965f66d2cSIvan R. Ivanov   //    map(tofrom: x%y)
35065f66d2cSIvan R. Ivanov   //
35165f66d2cSIvan R. Ivanov   // Will generate a mapping for "x" (the parent) and "y" (the member).
35265f66d2cSIvan R. Ivanov   // The parent "x" will not be mapped, but the member "y" will.
35365f66d2cSIvan R. Ivanov   // However, we must have the parent as a BlockArg and MapOperand
35465f66d2cSIvan R. Ivanov   // in these cases, to maintain the correct uses within the region and
35565f66d2cSIvan R. Ivanov   // to help tracking that the member is part of a larger object.
35665f66d2cSIvan R. Ivanov   //
35765f66d2cSIvan R. Ivanov   // In the case of:
35865f66d2cSIvan R. Ivanov   //    map(tofrom: x%y, x%z)
35965f66d2cSIvan R. Ivanov   //
36065f66d2cSIvan R. Ivanov   // The parent member becomes more critical, as we perform a partial
36165f66d2cSIvan R. Ivanov   // structure mapping where we link the mapping of the members y
36265f66d2cSIvan R. Ivanov   // and z together via the parent x. We do this at a kernel argument
36365f66d2cSIvan R. Ivanov   // level in LLVM IR and not just MLIR, which is important to maintain
36465f66d2cSIvan R. Ivanov   // similarity to Clang and for the runtime to do the correct thing.
36565f66d2cSIvan R. Ivanov   // However, we still do not map the structure in its totality but
36665f66d2cSIvan R. Ivanov   // rather we generate an un-sized "binding" map entry for it.
36765f66d2cSIvan R. Ivanov   //
36865f66d2cSIvan R. Ivanov   // In the case of:
36965f66d2cSIvan R. Ivanov   //    map(tofrom: x, x%y, x%z)
37065f66d2cSIvan R. Ivanov   //
37165f66d2cSIvan R. Ivanov   // We do actually map the entirety of "x", so the explicit mapping of
37265f66d2cSIvan R. Ivanov   // x%y, x%z becomes unnecessary. It is redundant to write this from a
37365f66d2cSIvan R. Ivanov   // Fortran OpenMP perspective (although it is legal), as even if the
37465f66d2cSIvan R. Ivanov   // members were allocatables or pointers, we are mandated by the
37565f66d2cSIvan R. Ivanov   // specification to map these (and any recursive components) in their
37665f66d2cSIvan R. Ivanov   // entirety, which is different to the C++ equivalent, which requires
37765f66d2cSIvan R. Ivanov   // explicit mapping of these segments.
37865f66d2cSIvan R. Ivanov   void addImplicitMembersToTarget(mlir::omp::MapInfoOp op,
37965f66d2cSIvan R. Ivanov                                   fir::FirOpBuilder &builder,
38065f66d2cSIvan R. Ivanov                                   mlir::Operation *target) {
38165f66d2cSIvan R. Ivanov     auto mapClauseOwner =
382e508baccSagozillon         llvm::dyn_cast_if_present<mlir::omp::MapClauseOwningOpInterface>(
383e508baccSagozillon             target);
384e508baccSagozillon     // TargetDataOp is technically a MapClauseOwningOpInterface, so we
385e508baccSagozillon     // do not need to explicitly check for the extra cases here for use_device
386e508baccSagozillon     // addr/ptr
38765f66d2cSIvan R. Ivanov     if (!mapClauseOwner)
38865f66d2cSIvan R. Ivanov       return;
38965f66d2cSIvan R. Ivanov 
390e508baccSagozillon     auto addOperands = [&](mlir::MutableOperandRange &mutableOpRange,
391e508baccSagozillon                            mlir::Operation *directiveOp,
392e508baccSagozillon                            unsigned blockArgInsertIndex = 0) {
393e508baccSagozillon       if (!llvm::is_contained(mutableOpRange.getAsOperandRange(),
394e508baccSagozillon                               op.getResult()))
395e508baccSagozillon         return;
39665f66d2cSIvan R. Ivanov 
397e508baccSagozillon       // There doesn't appear to be a simple way to convert MutableOperandRange
398e508baccSagozillon       // to a vector currently, so we instead use a for_each to populate our
399e508baccSagozillon       // vector.
400e508baccSagozillon       llvm::SmallVector<mlir::Value> newMapOps;
401e508baccSagozillon       newMapOps.reserve(mutableOpRange.size());
402e508baccSagozillon       llvm::for_each(
403e508baccSagozillon           mutableOpRange.getAsOperandRange(),
404e508baccSagozillon           [&newMapOps](mlir::Value oper) { newMapOps.push_back(oper); });
405e508baccSagozillon 
406e508baccSagozillon       for (auto mapMember : op.getMembers()) {
407e508baccSagozillon         if (llvm::is_contained(mutableOpRange.getAsOperandRange(), mapMember))
408e508baccSagozillon           continue;
40965f66d2cSIvan R. Ivanov         newMapOps.push_back(mapMember);
410e508baccSagozillon         if (directiveOp) {
411e508baccSagozillon           directiveOp->getRegion(0).insertArgument(
412e508baccSagozillon               blockArgInsertIndex, mapMember.getType(), mapMember.getLoc());
413e508baccSagozillon           blockArgInsertIndex++;
41465f66d2cSIvan R. Ivanov         }
41565f66d2cSIvan R. Ivanov       }
416e508baccSagozillon 
417e508baccSagozillon       mutableOpRange.assign(newMapOps);
418e508baccSagozillon     };
419e508baccSagozillon 
420e508baccSagozillon     auto argIface =
421e508baccSagozillon         llvm::dyn_cast<mlir::omp::BlockArgOpenMPOpInterface>(target);
422e508baccSagozillon 
423e508baccSagozillon     if (auto mapClauseOwner =
424e508baccSagozillon             llvm::dyn_cast<mlir::omp::MapClauseOwningOpInterface>(target)) {
425e508baccSagozillon       mlir::MutableOperandRange mapMutableOpRange =
426e508baccSagozillon           mapClauseOwner.getMapVarsMutable();
427e508baccSagozillon       unsigned blockArgInsertIndex =
428e508baccSagozillon           argIface
429e508baccSagozillon               ? argIface.getMapBlockArgsStart() + argIface.numMapBlockArgs()
430e508baccSagozillon               : 0;
431e532241bSKareem Ergawy       addOperands(mapMutableOpRange,
432e532241bSKareem Ergawy                   llvm::dyn_cast_if_present<mlir::omp::TargetOp>(
433e532241bSKareem Ergawy                       argIface.getOperation()),
434e508baccSagozillon                   blockArgInsertIndex);
43565f66d2cSIvan R. Ivanov     }
436e508baccSagozillon 
437e508baccSagozillon     if (auto targetDataOp = llvm::dyn_cast<mlir::omp::TargetDataOp>(target)) {
438e508baccSagozillon       mlir::MutableOperandRange useDevAddrMutableOpRange =
439e508baccSagozillon           targetDataOp.getUseDeviceAddrVarsMutable();
440e508baccSagozillon       addOperands(useDevAddrMutableOpRange, target,
441e508baccSagozillon                   argIface.getUseDeviceAddrBlockArgsStart() +
442e508baccSagozillon                       argIface.numUseDeviceAddrBlockArgs());
443e508baccSagozillon 
444e508baccSagozillon       mlir::MutableOperandRange useDevPtrMutableOpRange =
445e508baccSagozillon           targetDataOp.getUseDevicePtrVarsMutable();
446e508baccSagozillon       addOperands(useDevPtrMutableOpRange, target,
447e508baccSagozillon                   argIface.getUseDevicePtrBlockArgsStart() +
448e508baccSagozillon                       argIface.numUseDevicePtrBlockArgs());
44965f66d2cSIvan R. Ivanov     }
450e508baccSagozillon   }
451e508baccSagozillon 
452e508baccSagozillon   // We retrieve the first user that is a Target operation, of which
453e508baccSagozillon   // there should only be one currently. Every MapInfoOp can be tied to
454e508baccSagozillon   // at most one Target operation and at the minimum no operations.
455e508baccSagozillon   // This may change in the future with IR cleanups/modifications,
456e508baccSagozillon   // in which case this pass will need updating to support cases
457e508baccSagozillon   // where a map can have more than one user and more than one of
458e508baccSagozillon   // those users can be a Target operation. For now, we simply
459e508baccSagozillon   // return the first target operation encountered, which may
460e508baccSagozillon   // be on the parent MapInfoOp in the case of a member mapping.
461e508baccSagozillon   // In that case, we traverse the MapInfoOp chain until we
462e508baccSagozillon   // find the first TargetOp user.
463e508baccSagozillon   mlir::Operation *getFirstTargetUser(mlir::omp::MapInfoOp mapOp) {
464e508baccSagozillon     for (auto *user : mapOp->getUsers()) {
465e508baccSagozillon       if (llvm::isa<mlir::omp::TargetOp, mlir::omp::TargetDataOp,
466e508baccSagozillon                     mlir::omp::TargetUpdateOp, mlir::omp::TargetExitDataOp,
467e508baccSagozillon                     mlir::omp::TargetEnterDataOp>(user))
468e508baccSagozillon         return user;
469e508baccSagozillon 
470e508baccSagozillon       if (auto mapUser = llvm::dyn_cast<mlir::omp::MapInfoOp>(user))
471e508baccSagozillon         return getFirstTargetUser(mapUser);
472e508baccSagozillon     }
473e508baccSagozillon 
474e508baccSagozillon     return nullptr;
47565f66d2cSIvan R. Ivanov   }
47665f66d2cSIvan R. Ivanov 
47765f66d2cSIvan R. Ivanov   // This pass executes on omp::MapInfoOp's containing descriptor based types
47865f66d2cSIvan R. Ivanov   // (allocatables, pointers, assumed shape etc.) and expanding them into
47965f66d2cSIvan R. Ivanov   // multiple omp::MapInfoOp's for each pointer member contained within the
48065f66d2cSIvan R. Ivanov   // descriptor.
48165f66d2cSIvan R. Ivanov   //
48265f66d2cSIvan R. Ivanov   // From the perspective of the MLIR pass manager this runs on the top level
48365f66d2cSIvan R. Ivanov   // operation (usually function) containing the MapInfoOp because this pass
48465f66d2cSIvan R. Ivanov   // will mutate siblings of MapInfoOp.
48565f66d2cSIvan R. Ivanov   void runOnOperation() override {
486e532241bSKareem Ergawy     mlir::ModuleOp module = getOperation();
48765f66d2cSIvan R. Ivanov     if (!module)
48865f66d2cSIvan R. Ivanov       module = getOperation()->getParentOfType<mlir::ModuleOp>();
48965f66d2cSIvan R. Ivanov     fir::KindMapping kindMap = fir::getKindMapping(module);
49065f66d2cSIvan R. Ivanov     fir::FirOpBuilder builder{module, std::move(kindMap)};
49165f66d2cSIvan R. Ivanov 
492f4cf93fbSagozillon     // We wish to maintain some function level scope (currently
493f4cf93fbSagozillon     // just local function scope variables used to load and store box
494f4cf93fbSagozillon     // variables into so we can access their base address, an
495f4cf93fbSagozillon     // quirk of box_offset requires us to have an in memory box, but Fortran
496f4cf93fbSagozillon     // in certain cases does not provide this) whilst not subjecting
497f4cf93fbSagozillon     // ourselves to the possibility of race conditions while this pass
498f4cf93fbSagozillon     // undergoes frequent re-iteration for the near future. So we loop
499f4cf93fbSagozillon     // over function in the module and then map.info inside of those.
500f4cf93fbSagozillon     getOperation()->walk([&](mlir::func::FuncOp func) {
501f4cf93fbSagozillon       // clear all local allocations we made for any boxes in any prior
502f4cf93fbSagozillon       // iterations from previous function scopes.
503f4cf93fbSagozillon       localBoxAllocas.clear();
504f4cf93fbSagozillon 
505e532241bSKareem Ergawy       // First, walk `omp.map.info` ops to see if any record members should be
506e532241bSKareem Ergawy       // implicitly mapped.
507e532241bSKareem Ergawy       func->walk([&](mlir::omp::MapInfoOp op) {
508e532241bSKareem Ergawy         mlir::Type underlyingType =
509e532241bSKareem Ergawy             fir::unwrapRefType(op.getVarPtr().getType());
510e532241bSKareem Ergawy 
511e532241bSKareem Ergawy         // TODO Test with and support more complicated cases; like arrays for
512e532241bSKareem Ergawy         // records, for example.
513e532241bSKareem Ergawy         if (!fir::isRecordWithAllocatableMember(underlyingType))
514e532241bSKareem Ergawy           return mlir::WalkResult::advance();
515e532241bSKareem Ergawy 
516e532241bSKareem Ergawy         // TODO For now, only consider `omp.target` ops. Other ops that support
517e532241bSKareem Ergawy         // `map` clauses will follow later.
518e532241bSKareem Ergawy         mlir::omp::TargetOp target =
519e532241bSKareem Ergawy             mlir::dyn_cast_if_present<mlir::omp::TargetOp>(
520e532241bSKareem Ergawy                 getFirstTargetUser(op));
521e532241bSKareem Ergawy 
522e532241bSKareem Ergawy         if (!target)
523e532241bSKareem Ergawy           return mlir::WalkResult::advance();
524e532241bSKareem Ergawy 
525e532241bSKareem Ergawy         auto mapClauseOwner =
526e532241bSKareem Ergawy             llvm::dyn_cast<mlir::omp::MapClauseOwningOpInterface>(*target);
527e532241bSKareem Ergawy 
528e532241bSKareem Ergawy         int64_t mapVarIdx = mapClauseOwner.getOperandIndexForMap(op);
529e532241bSKareem Ergawy         assert(mapVarIdx >= 0 &&
530e532241bSKareem Ergawy                mapVarIdx <
531e532241bSKareem Ergawy                    static_cast<int64_t>(mapClauseOwner.getMapVars().size()));
532e532241bSKareem Ergawy 
533e532241bSKareem Ergawy         auto argIface =
534e532241bSKareem Ergawy             llvm::dyn_cast<mlir::omp::BlockArgOpenMPOpInterface>(*target);
535e532241bSKareem Ergawy         // TODO How should `map` block argument that correspond to: `private`,
536e532241bSKareem Ergawy         // `use_device_addr`, `use_device_ptr`, be handled?
537e532241bSKareem Ergawy         mlir::BlockArgument opBlockArg = argIface.getMapBlockArgs()[mapVarIdx];
538e532241bSKareem Ergawy         llvm::SetVector<mlir::Operation *> mapVarForwardSlice;
539e532241bSKareem Ergawy         mlir::getForwardSlice(opBlockArg, &mapVarForwardSlice);
540e532241bSKareem Ergawy 
541e532241bSKareem Ergawy         mapVarForwardSlice.remove_if([&](mlir::Operation *sliceOp) {
542e532241bSKareem Ergawy           // TODO Support coordinate_of ops.
543e532241bSKareem Ergawy           //
544e532241bSKareem Ergawy           // TODO Support call ops by recursively examining the forward slice of
545e532241bSKareem Ergawy           // the corresponding parameter to the field in the called function.
546e532241bSKareem Ergawy           return !mlir::isa<hlfir::DesignateOp>(sliceOp);
547e532241bSKareem Ergawy         });
548e532241bSKareem Ergawy 
549e532241bSKareem Ergawy         auto recordType = mlir::cast<fir::RecordType>(underlyingType);
550e532241bSKareem Ergawy         llvm::SmallVector<mlir::Value> newMapOpsForFields;
551e532241bSKareem Ergawy         llvm::SmallVector<int64_t> fieldIndicies;
552e532241bSKareem Ergawy 
553e532241bSKareem Ergawy         for (auto fieldMemTyPair : recordType.getTypeList()) {
554e532241bSKareem Ergawy           auto &field = fieldMemTyPair.first;
555e532241bSKareem Ergawy           auto memTy = fieldMemTyPair.second;
556e532241bSKareem Ergawy 
557e532241bSKareem Ergawy           bool shouldMapField =
558e532241bSKareem Ergawy               llvm::find_if(mapVarForwardSlice, [&](mlir::Operation *sliceOp) {
559e532241bSKareem Ergawy                 if (!fir::isAllocatableType(memTy))
560e532241bSKareem Ergawy                   return false;
561e532241bSKareem Ergawy 
562e532241bSKareem Ergawy                 auto designateOp = mlir::dyn_cast<hlfir::DesignateOp>(sliceOp);
563e532241bSKareem Ergawy                 if (!designateOp)
564e532241bSKareem Ergawy                   return false;
565e532241bSKareem Ergawy 
566e532241bSKareem Ergawy                 return designateOp.getComponent() &&
567e532241bSKareem Ergawy                        designateOp.getComponent()->strref() == field;
568e532241bSKareem Ergawy               }) != mapVarForwardSlice.end();
569e532241bSKareem Ergawy 
570e532241bSKareem Ergawy           // TODO Handle recursive record types. Adapting
571e532241bSKareem Ergawy           // `createParentSymAndGenIntermediateMaps` to work direclty on MLIR
572e532241bSKareem Ergawy           // entities might be helpful here.
573e532241bSKareem Ergawy 
574e532241bSKareem Ergawy           if (!shouldMapField)
575e532241bSKareem Ergawy             continue;
576e532241bSKareem Ergawy 
577e532241bSKareem Ergawy           int64_t fieldIdx = recordType.getFieldIndex(field);
578e532241bSKareem Ergawy           bool alreadyMapped = [&]() {
579e532241bSKareem Ergawy             if (op.getMembersIndexAttr())
580e532241bSKareem Ergawy               for (auto indexList : op.getMembersIndexAttr()) {
581e532241bSKareem Ergawy                 auto indexListAttr = mlir::cast<mlir::ArrayAttr>(indexList);
582e532241bSKareem Ergawy                 if (indexListAttr.size() == 1 &&
583e532241bSKareem Ergawy                     mlir::cast<mlir::IntegerAttr>(indexListAttr[0]).getInt() ==
584e532241bSKareem Ergawy                         fieldIdx)
585e532241bSKareem Ergawy                   return true;
586e532241bSKareem Ergawy               }
587e532241bSKareem Ergawy 
588e532241bSKareem Ergawy             return false;
589e532241bSKareem Ergawy           }();
590e532241bSKareem Ergawy 
591e532241bSKareem Ergawy           if (alreadyMapped)
592e532241bSKareem Ergawy             continue;
593e532241bSKareem Ergawy 
594e532241bSKareem Ergawy           builder.setInsertionPoint(op);
595e532241bSKareem Ergawy           mlir::Value fieldIdxVal = builder.createIntegerConstant(
596e532241bSKareem Ergawy               op.getLoc(), mlir::IndexType::get(builder.getContext()),
597e532241bSKareem Ergawy               fieldIdx);
598e532241bSKareem Ergawy           auto fieldCoord = builder.create<fir::CoordinateOp>(
599e532241bSKareem Ergawy               op.getLoc(), builder.getRefType(memTy), op.getVarPtr(),
600e532241bSKareem Ergawy               fieldIdxVal);
601*662133a2SjeanPerier           fir::factory::AddrAndBoundsInfo info =
602*662133a2SjeanPerier               fir::factory::getDataOperandBaseAddr(
603e532241bSKareem Ergawy                   builder, fieldCoord, /*isOptional=*/false, op.getLoc());
604e532241bSKareem Ergawy           llvm::SmallVector<mlir::Value> bounds =
605*662133a2SjeanPerier               fir::factory::genImplicitBoundsOps<mlir::omp::MapBoundsOp,
606e532241bSKareem Ergawy                                                  mlir::omp::MapBoundsType>(
607e532241bSKareem Ergawy                   builder, info,
608e532241bSKareem Ergawy                   hlfir::translateToExtendedValue(op.getLoc(), builder,
609e532241bSKareem Ergawy                                                   hlfir::Entity{fieldCoord})
610e532241bSKareem Ergawy                       .first,
611e532241bSKareem Ergawy                   /*dataExvIsAssumedSize=*/false, op.getLoc());
612e532241bSKareem Ergawy 
613e532241bSKareem Ergawy           mlir::omp::MapInfoOp fieldMapOp =
614e532241bSKareem Ergawy               builder.create<mlir::omp::MapInfoOp>(
615e532241bSKareem Ergawy                   op.getLoc(), fieldCoord.getResult().getType(),
616e532241bSKareem Ergawy                   fieldCoord.getResult(),
617e532241bSKareem Ergawy                   mlir::TypeAttr::get(
618e532241bSKareem Ergawy                       fir::unwrapRefType(fieldCoord.getResult().getType())),
619e532241bSKareem Ergawy                   /*varPtrPtr=*/mlir::Value{},
620e532241bSKareem Ergawy                   /*members=*/mlir::ValueRange{},
621e532241bSKareem Ergawy                   /*members_index=*/mlir::ArrayAttr{},
622e532241bSKareem Ergawy                   /*bounds=*/bounds, op.getMapTypeAttr(),
623e532241bSKareem Ergawy                   builder.getAttr<mlir::omp::VariableCaptureKindAttr>(
624e532241bSKareem Ergawy                       mlir::omp::VariableCaptureKind::ByRef),
625e532241bSKareem Ergawy                   builder.getStringAttr(op.getNameAttr().strref() + "." +
626e532241bSKareem Ergawy                                         field + ".implicit_map"),
627e532241bSKareem Ergawy                   /*partial_map=*/builder.getBoolAttr(false));
628e532241bSKareem Ergawy           newMapOpsForFields.emplace_back(fieldMapOp);
629e532241bSKareem Ergawy           fieldIndicies.emplace_back(fieldIdx);
630e532241bSKareem Ergawy         }
631e532241bSKareem Ergawy 
632e532241bSKareem Ergawy         if (newMapOpsForFields.empty())
633e532241bSKareem Ergawy           return mlir::WalkResult::advance();
634e532241bSKareem Ergawy 
635e532241bSKareem Ergawy         op.getMembersMutable().append(newMapOpsForFields);
636e532241bSKareem Ergawy         llvm::SmallVector<llvm::SmallVector<int64_t>> newMemberIndices;
637e532241bSKareem Ergawy         mlir::ArrayAttr oldMembersIdxAttr = op.getMembersIndexAttr();
638e532241bSKareem Ergawy 
639e532241bSKareem Ergawy         if (oldMembersIdxAttr)
640e532241bSKareem Ergawy           for (mlir::Attribute indexList : oldMembersIdxAttr) {
641e532241bSKareem Ergawy             llvm::SmallVector<int64_t> listVec;
642e532241bSKareem Ergawy 
643e532241bSKareem Ergawy             for (mlir::Attribute index : mlir::cast<mlir::ArrayAttr>(indexList))
644e532241bSKareem Ergawy               listVec.push_back(mlir::cast<mlir::IntegerAttr>(index).getInt());
645e532241bSKareem Ergawy 
646e532241bSKareem Ergawy             newMemberIndices.emplace_back(std::move(listVec));
647e532241bSKareem Ergawy           }
648e532241bSKareem Ergawy 
649e532241bSKareem Ergawy         for (int64_t newFieldIdx : fieldIndicies)
650e532241bSKareem Ergawy           newMemberIndices.emplace_back(
651e532241bSKareem Ergawy               llvm::SmallVector<int64_t>(1, newFieldIdx));
652e532241bSKareem Ergawy 
653e532241bSKareem Ergawy         op.setMembersIndexAttr(builder.create2DI64ArrayAttr(newMemberIndices));
654e532241bSKareem Ergawy         op.setPartialMap(true);
655e532241bSKareem Ergawy 
656e532241bSKareem Ergawy         return mlir::WalkResult::advance();
657e532241bSKareem Ergawy       });
658e532241bSKareem Ergawy 
659f4cf93fbSagozillon       func->walk([&](mlir::omp::MapInfoOp op) {
660e508baccSagozillon         // TODO: Currently only supports a single user for the MapInfoOp. This
661e508baccSagozillon         // is fine for the moment, as the Fortran frontend will generate a
662e508baccSagozillon         // new MapInfoOp with at most one user currently. In the case of
663e508baccSagozillon         // members of other objects, like derived types, the user would be the
664e508baccSagozillon         // parent. In cases where it's a regular non-member map, the user would
665e508baccSagozillon         // be the target operation it is being mapped by.
666e508baccSagozillon         //
667e508baccSagozillon         // However, when/if we optimise/cleanup the IR we will have to extend
668e508baccSagozillon         // this pass to support multiple users, as we may wish to have a map
669e508baccSagozillon         // be re-used by multiple users (e.g. across multiple targets that map
670e508baccSagozillon         // the variable and have identical map properties).
67165f66d2cSIvan R. Ivanov         assert(llvm::hasSingleElement(op->getUsers()) &&
672f4cf93fbSagozillon                "OMPMapInfoFinalization currently only supports single users "
67365f66d2cSIvan R. Ivanov                "of a MapInfoOp");
67465f66d2cSIvan R. Ivanov 
675e508baccSagozillon         if (fir::isTypeWithDescriptor(op.getVarType()) ||
67665f66d2cSIvan R. Ivanov             mlir::isa_and_present<fir::BoxAddrOp>(
67765f66d2cSIvan R. Ivanov                 op.getVarPtr().getDefiningOp())) {
67865f66d2cSIvan R. Ivanov           builder.setInsertionPoint(op);
679e508baccSagozillon           mlir::Operation *targetUser = getFirstTargetUser(op);
680e508baccSagozillon           assert(targetUser && "expected user of map operation was not found");
681e508baccSagozillon           genDescriptorMemberMaps(op, builder, targetUser);
68265f66d2cSIvan R. Ivanov         }
68365f66d2cSIvan R. Ivanov       });
684e508baccSagozillon 
685e508baccSagozillon       // Wait until after we have generated all of our maps to add them onto
686e508baccSagozillon       // the target's block arguments, simplifying the process as there would be
687e508baccSagozillon       // no need to avoid accidental duplicate additions.
688e508baccSagozillon       func->walk([&](mlir::omp::MapInfoOp op) {
689e508baccSagozillon         mlir::Operation *targetUser = getFirstTargetUser(op);
690e508baccSagozillon         assert(targetUser && "expected user of map operation was not found");
691e508baccSagozillon         addImplicitMembersToTarget(op, builder, targetUser);
692e508baccSagozillon       });
693f4cf93fbSagozillon     });
69465f66d2cSIvan R. Ivanov   }
69565f66d2cSIvan R. Ivanov };
69665f66d2cSIvan R. Ivanov 
69765f66d2cSIvan R. Ivanov } // namespace
698