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 << "]: ")
44namespace flangomp {
45#define GEN_PASS_DEF_MAPSFORPRIVATIZEDSYMBOLSPASS
46#include "flang/Optimizer/OpenMP/Passes.h.inc"
47} // namespace flangomp
48
49using namespace mlir;
50
51namespace {
52class 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

source code of flang/lib/Optimizer/OpenMP/MapsForPrivatizedSymbols.cpp