237 lines · cpp
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-exception6//7//===----------------------------------------------------------------------===//8 9//===----------------------------------------------------------------------===//10/// \file11/// An OpenMP dialect related pass for FIR/HLFIR which creates MapInfoOp12/// instances for certain privatized symbols.13/// For example, if an allocatable variable is used in a private clause attached14/// to a omp.target op, then the allocatable variable's descriptor will be15/// needed on the device (e.g. GPU). This descriptor needs to be separately16/// mapped onto the device. This pass creates the necessary omp.map.info ops for17/// this.18//===----------------------------------------------------------------------===//19// TODO:20// 1. Before adding omp.map.info, check if we already have an omp.map.info for21// 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/Support/Debug.h"39#include <type_traits>40 41#define DEBUG_TYPE "omp-maps-for-privatized-symbols"42#define PDBGS() (llvm::dbgs() << "[" << DEBUG_TYPE << "]: ")43namespace flangomp {44#define GEN_PASS_DEF_MAPSFORPRIVATIZEDSYMBOLSPASS45#include "flang/Optimizer/OpenMP/Passes.h.inc"46} // namespace flangomp47 48using namespace mlir;49 50namespace {51class MapsForPrivatizedSymbolsPass52 : public flangomp::impl::MapsForPrivatizedSymbolsPassBase<53 MapsForPrivatizedSymbolsPass> {54 55 // TODO Use `createMapInfoOp` from `flang/Utils/OpenMP.h`.56 omp::MapInfoOp createMapInfo(Location loc, Value var,57 fir::FirOpBuilder &builder) {58 // Check if a value of type `type` can be passed to the kernel by value.59 // All kernel parameters are of pointer type, so if the value can be60 // represented inside of a pointer, then it can be passed by value.61 auto canPassByValue = [&](mlir::Type type) {62 const mlir::DataLayout &dl = builder.getDataLayout();63 mlir::Type ptrTy = mlir::LLVM::LLVMPointerType::get(builder.getContext());64 uint64_t ptrSize = dl.getTypeSize(ptrTy);65 uint64_t ptrAlign = dl.getTypePreferredAlignment(ptrTy);66 67 auto [size, align] = fir::getTypeSizeAndAlignmentOrCrash(68 loc, type, dl, builder.getKindMap());69 return size <= ptrSize && align <= ptrAlign;70 };71 72 Operation *definingOp = var.getDefiningOp();73 74 Value varPtr = var;75 // We want the first result of the hlfir.declare op because our goal76 // is to map the descriptor (fir.box or fir.boxchar) and the first77 // result for hlfir.declare is the descriptor if a the symbol being78 // declared needs a descriptor.79 // Some types are boxed immediately before privatization. These have other80 // operations in between the privatization and the declaration. It is safe81 // to use var directly here because they will be boxed anyway.82 if (auto declOp = llvm::dyn_cast_if_present<hlfir::DeclareOp>(definingOp))83 varPtr = declOp.getBase();84 85 // If we do not have a reference to a descriptor but the descriptor itself,86 // then we need to store that on the stack so that we can map the87 // address of the descriptor.88 if (mlir::isa<fir::BaseBoxType>(varPtr.getType()) ||89 mlir::isa<fir::BoxCharType>(varPtr.getType())) {90 OpBuilder::InsertPoint savedInsPoint = builder.saveInsertionPoint();91 mlir::Block *allocaBlock = builder.getAllocaBlock();92 assert(allocaBlock && "No allocablock found for a funcOp");93 builder.setInsertionPointToStart(allocaBlock);94 auto alloca = fir::AllocaOp::create(builder, loc, varPtr.getType());95 builder.restoreInsertionPoint(savedInsPoint);96 fir::StoreOp::create(builder, loc, varPtr, alloca);97 varPtr = alloca;98 }99 assert(mlir::isa<omp::PointerLikeType>(varPtr.getType()) &&100 "Dealing with a varPtr that is not a PointerLikeType");101 102 // Figure out the bounds because knowing the bounds will help the subsequent103 // MapInfoFinalizationPass map the underlying data of the descriptor.104 llvm::SmallVector<mlir::Value> boundsOps;105 if (needsBoundsOps(varPtr))106 genBoundsOps(builder, varPtr, boundsOps);107 mlir::Type varType = varPtr.getType();108 109 mlir::omp::VariableCaptureKind captureKind =110 mlir::omp::VariableCaptureKind::ByRef;111 if (fir::isa_trivial(fir::unwrapRefType(varType)) ||112 fir::isa_char(fir::unwrapRefType(varType))) {113 if (canPassByValue(fir::unwrapRefType(varType))) {114 captureKind = mlir::omp::VariableCaptureKind::ByCopy;115 }116 }117 118 // Use tofrom if what we are mapping is not a trivial type. In all119 // likelihood, it is a descriptor120 mlir::omp::ClauseMapFlags mapFlag;121 if (fir::isa_trivial(fir::unwrapRefType(varType)) ||122 fir::isa_char(fir::unwrapRefType(varType)))123 mapFlag = mlir::omp::ClauseMapFlags::to;124 else125 mapFlag = mlir::omp::ClauseMapFlags::to | mlir::omp::ClauseMapFlags::from;126 127 return omp::MapInfoOp::create(128 builder, loc, varType, varPtr,129 TypeAttr::get(130 llvm::cast<omp::PointerLikeType>(varType).getElementType()),131 builder.getAttr<omp::ClauseMapFlagsAttr>(mapFlag),132 builder.getAttr<omp::VariableCaptureKindAttr>(captureKind),133 /*varPtrPtr=*/Value{},134 /*members=*/SmallVector<Value>{},135 /*member_index=*/mlir::ArrayAttr{},136 /*bounds=*/boundsOps,137 /*mapperId=*/mlir::FlatSymbolRefAttr(), /*name=*/StringAttr(),138 builder.getBoolAttr(false));139 }140 void addMapInfoOp(omp::TargetOp targetOp, omp::MapInfoOp mapInfoOp) {141 auto argIface = llvm::cast<omp::BlockArgOpenMPOpInterface>(*targetOp);142 unsigned insertIndex =143 argIface.getMapBlockArgsStart() + argIface.numMapBlockArgs();144 targetOp.getMapVarsMutable().append(ValueRange{mapInfoOp});145 targetOp.getRegion().insertArgument(insertIndex, mapInfoOp.getType(),146 mapInfoOp.getLoc());147 }148 void addMapInfoOps(omp::TargetOp targetOp,149 llvm::SmallVectorImpl<omp::MapInfoOp> &mapInfoOps) {150 for (auto mapInfoOp : mapInfoOps)151 addMapInfoOp(targetOp, mapInfoOp);152 }153 void runOnOperation() override {154 ModuleOp module = getOperation()->getParentOfType<ModuleOp>();155 fir::KindMapping kindMap = fir::getKindMapping(module);156 fir::FirOpBuilder builder{module, std::move(kindMap)};157 llvm::DenseMap<Operation *, llvm::SmallVector<omp::MapInfoOp, 4>>158 mapInfoOpsForTarget;159 160 getOperation()->walk([&](omp::TargetOp targetOp) {161 if (targetOp.getPrivateVars().empty())162 return;163 OperandRange privVars = targetOp.getPrivateVars();164 llvm::SmallVector<int64_t> privVarMapIdx;165 166 std::optional<ArrayAttr> privSyms = targetOp.getPrivateSyms();167 SmallVector<omp::MapInfoOp, 4> mapInfoOps;168 for (auto [privVar, privSym] : llvm::zip_equal(privVars, *privSyms)) {169 170 SymbolRefAttr privatizerName = llvm::cast<SymbolRefAttr>(privSym);171 omp::PrivateClauseOp privatizer =172 SymbolTable::lookupNearestSymbolFrom<omp::PrivateClauseOp>(173 targetOp, privatizerName);174 if (!privatizer.needsMap()) {175 privVarMapIdx.push_back(-1);176 continue;177 }178 179 privVarMapIdx.push_back(targetOp.getMapVars().size() +180 mapInfoOps.size());181 182 builder.setInsertionPoint(targetOp);183 Location loc = targetOp.getLoc();184 omp::MapInfoOp mapInfoOp = createMapInfo(loc, privVar, builder);185 mapInfoOps.push_back(mapInfoOp);186 187 LLVM_DEBUG(PDBGS() << "MapsForPrivatizedSymbolsPass created ->\n"188 << mapInfoOp << "\n");189 }190 if (!mapInfoOps.empty()) {191 mapInfoOpsForTarget.insert({targetOp.getOperation(), mapInfoOps});192 targetOp.setPrivateMapsAttr(193 mlir::DenseI64ArrayAttr::get(targetOp.getContext(), privVarMapIdx));194 }195 });196 if (!mapInfoOpsForTarget.empty()) {197 for (auto &[targetOp, mapInfoOps] : mapInfoOpsForTarget) {198 addMapInfoOps(static_cast<omp::TargetOp>(targetOp), mapInfoOps);199 }200 }201 }202 // As the name suggests, this function examines var to determine if203 // it has dynamic size. If true, this pass'll have to extract these204 // bounds from descriptor of var and add the bounds to the resultant205 // MapInfoOp.206 bool needsBoundsOps(mlir::Value var) {207 assert(mlir::isa<omp::PointerLikeType>(var.getType()) &&208 "needsBoundsOps can deal only with pointer types");209 mlir::Type t = fir::unwrapRefType(var.getType());210 // t could be a box, so look inside the box211 auto innerType = fir::dyn_cast_ptrOrBoxEleTy(t);212 if (innerType)213 return fir::hasDynamicSize(innerType);214 return fir::hasDynamicSize(t);215 }216 217 void genBoundsOps(fir::FirOpBuilder &builder, mlir::Value var,218 llvm::SmallVector<mlir::Value> &boundsOps) {219 mlir::Location loc = var.getLoc();220 fir::factory::AddrAndBoundsInfo info =221 fir::factory::getDataOperandBaseAddr(builder, var,222 /*isOptional=*/false, loc);223 fir::ExtendedValue extendedValue =224 hlfir::translateToExtendedValue(loc, builder, hlfir::Entity{info.addr},225 /*continguousHint=*/true)226 .first;227 llvm::SmallVector<mlir::Value> boundsOpsVec =228 fir::factory::genImplicitBoundsOps<mlir::omp::MapBoundsOp,229 mlir::omp::MapBoundsType>(230 builder, info, extendedValue,231 /*dataExvIsAssumedSize=*/false, loc);232 for (auto bounds : boundsOpsVec)233 boundsOps.push_back(bounds);234 }235};236} // namespace237