1 | //===- MapsForPrivatizedSymbols.cpp ---------------------------------------===// |
2 | // |
3 | // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. |
4 | // See https://llvm.org/LICENSE.txt for license information. |
5 | // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
6 | // |
7 | //===----------------------------------------------------------------------===// |
8 | |
9 | //===----------------------------------------------------------------------===// |
10 | /// \file |
11 | /// An OpenMP dialect related pass for FIR/HLFIR which creates MapInfoOp |
12 | /// instances for certain privatized symbols. |
13 | /// For example, if an allocatable variable is used in a private clause attached |
14 | /// to a omp.target op, then the allocatable variable's descriptor will be |
15 | /// needed on the device (e.g. GPU). This descriptor needs to be separately |
16 | /// mapped onto the device. This pass creates the necessary omp.map.info ops for |
17 | /// this. |
18 | //===----------------------------------------------------------------------===// |
19 | // TODO: |
20 | // 1. Before adding omp.map.info, check if we already have an omp.map.info for |
21 | // the variable in question. |
22 | // 2. Generalize this for more than just omp.target ops. |
23 | //===----------------------------------------------------------------------===// |
24 | |
25 | #include "flang/Optimizer/Builder/DirectivesCommon.h" |
26 | #include "flang/Optimizer/Builder/FIRBuilder.h" |
27 | #include "flang/Optimizer/Builder/HLFIRTools.h" |
28 | #include "flang/Optimizer/Dialect/FIRType.h" |
29 | #include "flang/Optimizer/Dialect/Support/KindMapping.h" |
30 | #include "flang/Optimizer/HLFIR/HLFIROps.h" |
31 | #include "flang/Optimizer/OpenMP/Passes.h" |
32 | |
33 | #include "mlir/Dialect/Func/IR/FuncOps.h" |
34 | #include "mlir/Dialect/OpenMP/OpenMPDialect.h" |
35 | #include "mlir/IR/BuiltinAttributes.h" |
36 | #include "mlir/IR/SymbolTable.h" |
37 | #include "mlir/Pass/Pass.h" |
38 | #include "llvm/Frontend/OpenMP/OMPConstants.h" |
39 | #include "llvm/Support/Debug.h" |
40 | #include <type_traits> |
41 | |
42 | #define DEBUG_TYPE "omp-maps-for-privatized-symbols" |
43 | #define PDBGS() (llvm::dbgs() << "[" << DEBUG_TYPE << "]: ") |
44 | namespace flangomp { |
45 | #define GEN_PASS_DEF_MAPSFORPRIVATIZEDSYMBOLSPASS |
46 | #include "flang/Optimizer/OpenMP/Passes.h.inc" |
47 | } // namespace flangomp |
48 | |
49 | using namespace mlir; |
50 | |
51 | namespace { |
52 | class MapsForPrivatizedSymbolsPass |
53 | : public flangomp::impl::MapsForPrivatizedSymbolsPassBase< |
54 | MapsForPrivatizedSymbolsPass> { |
55 | |
56 | omp::MapInfoOp createMapInfo(Location loc, Value var, |
57 | fir::FirOpBuilder &builder) { |
58 | uint64_t mapTypeTo = static_cast< |
59 | std::underlying_type_t<llvm::omp::OpenMPOffloadMappingFlags>>( |
60 | llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_TO); |
61 | Operation *definingOp = var.getDefiningOp(); |
62 | |
63 | Value varPtr = var; |
64 | // We want the first result of the hlfir.declare op because our goal |
65 | // is to map the descriptor (fir.box or fir.boxchar) and the first |
66 | // result for hlfir.declare is the descriptor if a the symbol being |
67 | // declared needs a descriptor. |
68 | // Some types are boxed immediately before privatization. These have other |
69 | // operations in between the privatization and the declaration. It is safe |
70 | // to use var directly here because they will be boxed anyway. |
71 | if (auto declOp = llvm::dyn_cast_if_present<hlfir::DeclareOp>(definingOp)) |
72 | varPtr = declOp.getBase(); |
73 | |
74 | // If we do not have a reference to a descriptor but the descriptor itself, |
75 | // then we need to store that on the stack so that we can map the |
76 | // address of the descriptor. |
77 | if (mlir::isa<fir::BaseBoxType>(varPtr.getType()) || |
78 | mlir::isa<fir::BoxCharType>(varPtr.getType())) { |
79 | OpBuilder::InsertPoint savedInsPoint = builder.saveInsertionPoint(); |
80 | mlir::Block *allocaBlock = builder.getAllocaBlock(); |
81 | assert(allocaBlock && "No allocablock found for a funcOp" ); |
82 | builder.setInsertionPointToStart(allocaBlock); |
83 | auto alloca = builder.create<fir::AllocaOp>(loc, varPtr.getType()); |
84 | builder.restoreInsertionPoint(savedInsPoint); |
85 | builder.create<fir::StoreOp>(loc, varPtr, alloca); |
86 | varPtr = alloca; |
87 | } |
88 | assert(mlir::isa<omp::PointerLikeType>(varPtr.getType()) && |
89 | "Dealing with a varPtr that is not a PointerLikeType" ); |
90 | |
91 | // Figure out the bounds because knowing the bounds will help the subsequent |
92 | // MapInfoFinalizationPass map the underlying data of the descriptor. |
93 | llvm::SmallVector<mlir::Value> boundsOps; |
94 | if (needsBoundsOps(varPtr)) |
95 | genBoundsOps(builder, varPtr, boundsOps); |
96 | |
97 | return builder.create<omp::MapInfoOp>( |
98 | loc, varPtr.getType(), varPtr, |
99 | TypeAttr::get(llvm::cast<omp::PointerLikeType>(varPtr.getType()) |
100 | .getElementType()), |
101 | builder.getIntegerAttr(builder.getIntegerType(64, /*isSigned=*/false), |
102 | mapTypeTo), |
103 | builder.getAttr<omp::VariableCaptureKindAttr>( |
104 | omp::VariableCaptureKind::ByRef), |
105 | /*varPtrPtr=*/Value{}, |
106 | /*members=*/SmallVector<Value>{}, |
107 | /*member_index=*/mlir::ArrayAttr{}, |
108 | /*bounds=*/boundsOps, |
109 | /*mapperId=*/mlir::FlatSymbolRefAttr(), /*name=*/StringAttr(), |
110 | builder.getBoolAttr(false)); |
111 | } |
112 | void addMapInfoOp(omp::TargetOp targetOp, omp::MapInfoOp mapInfoOp) { |
113 | auto argIface = llvm::cast<omp::BlockArgOpenMPOpInterface>(*targetOp); |
114 | unsigned insertIndex = |
115 | argIface.getMapBlockArgsStart() + argIface.numMapBlockArgs(); |
116 | targetOp.getMapVarsMutable().append(ValueRange{mapInfoOp}); |
117 | targetOp.getRegion().insertArgument(insertIndex, mapInfoOp.getType(), |
118 | mapInfoOp.getLoc()); |
119 | } |
120 | void addMapInfoOps(omp::TargetOp targetOp, |
121 | llvm::SmallVectorImpl<omp::MapInfoOp> &mapInfoOps) { |
122 | for (auto mapInfoOp : mapInfoOps) |
123 | addMapInfoOp(targetOp, mapInfoOp); |
124 | } |
125 | void runOnOperation() override { |
126 | ModuleOp module = getOperation()->getParentOfType<ModuleOp>(); |
127 | fir::KindMapping kindMap = fir::getKindMapping(module); |
128 | fir::FirOpBuilder builder{module, std::move(kindMap)}; |
129 | llvm::DenseMap<Operation *, llvm::SmallVector<omp::MapInfoOp, 4>> |
130 | mapInfoOpsForTarget; |
131 | |
132 | getOperation()->walk([&](omp::TargetOp targetOp) { |
133 | if (targetOp.getPrivateVars().empty()) |
134 | return; |
135 | OperandRange privVars = targetOp.getPrivateVars(); |
136 | llvm::SmallVector<int64_t> privVarMapIdx; |
137 | |
138 | std::optional<ArrayAttr> privSyms = targetOp.getPrivateSyms(); |
139 | SmallVector<omp::MapInfoOp, 4> mapInfoOps; |
140 | for (auto [privVar, privSym] : llvm::zip_equal(privVars, *privSyms)) { |
141 | |
142 | SymbolRefAttr privatizerName = llvm::cast<SymbolRefAttr>(privSym); |
143 | omp::PrivateClauseOp privatizer = |
144 | SymbolTable::lookupNearestSymbolFrom<omp::PrivateClauseOp>( |
145 | targetOp, privatizerName); |
146 | if (!privatizer.needsMap()) { |
147 | privVarMapIdx.push_back(-1); |
148 | continue; |
149 | } |
150 | |
151 | privVarMapIdx.push_back(targetOp.getMapVars().size() + |
152 | mapInfoOps.size()); |
153 | |
154 | builder.setInsertionPoint(targetOp); |
155 | Location loc = targetOp.getLoc(); |
156 | omp::MapInfoOp mapInfoOp = createMapInfo(loc, privVar, builder); |
157 | mapInfoOps.push_back(mapInfoOp); |
158 | |
159 | LLVM_DEBUG(PDBGS() << "MapsForPrivatizedSymbolsPass created ->\n" |
160 | << mapInfoOp << "\n" ); |
161 | } |
162 | if (!mapInfoOps.empty()) { |
163 | mapInfoOpsForTarget.insert({targetOp.getOperation(), mapInfoOps}); |
164 | targetOp.setPrivateMapsAttr( |
165 | mlir::DenseI64ArrayAttr::get(targetOp.getContext(), privVarMapIdx)); |
166 | } |
167 | }); |
168 | if (!mapInfoOpsForTarget.empty()) { |
169 | for (auto &[targetOp, mapInfoOps] : mapInfoOpsForTarget) { |
170 | addMapInfoOps(static_cast<omp::TargetOp>(targetOp), mapInfoOps); |
171 | } |
172 | } |
173 | } |
174 | // As the name suggests, this function examines var to determine if |
175 | // it has dynamic size. If true, this pass'll have to extract these |
176 | // bounds from descriptor of var and add the bounds to the resultant |
177 | // MapInfoOp. |
178 | bool needsBoundsOps(mlir::Value var) { |
179 | assert(mlir::isa<omp::PointerLikeType>(var.getType()) && |
180 | "needsBoundsOps can deal only with pointer types" ); |
181 | mlir::Type t = fir::unwrapRefType(var.getType()); |
182 | // t could be a box, so look inside the box |
183 | auto innerType = fir::dyn_cast_ptrOrBoxEleTy(t); |
184 | if (innerType) |
185 | return fir::hasDynamicSize(innerType); |
186 | return fir::hasDynamicSize(t); |
187 | } |
188 | |
189 | void genBoundsOps(fir::FirOpBuilder &builder, mlir::Value var, |
190 | llvm::SmallVector<mlir::Value> &boundsOps) { |
191 | mlir::Location loc = var.getLoc(); |
192 | fir::factory::AddrAndBoundsInfo info = |
193 | fir::factory::getDataOperandBaseAddr(builder, var, |
194 | /*isOptional=*/false, loc); |
195 | fir::ExtendedValue extendedValue = |
196 | hlfir::translateToExtendedValue(loc, builder, hlfir::Entity{info.addr}, |
197 | /*continguousHint=*/true) |
198 | .first; |
199 | llvm::SmallVector<mlir::Value> boundsOpsVec = |
200 | fir::factory::genImplicitBoundsOps<mlir::omp::MapBoundsOp, |
201 | mlir::omp::MapBoundsType>( |
202 | builder, info, extendedValue, |
203 | /*dataExvIsAssumedSize=*/false, loc); |
204 | for (auto bounds : boundsOpsVec) |
205 | boundsOps.push_back(bounds); |
206 | } |
207 | }; |
208 | } // namespace |
209 | |