952 lines · cpp
1//===- DoConcurrentConversion.cpp -- map `DO CONCURRENT` to OpenMP loops --===//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#include "flang/Optimizer/Builder/DirectivesCommon.h"10#include "flang/Optimizer/Builder/FIRBuilder.h"11#include "flang/Optimizer/Builder/HLFIRTools.h"12#include "flang/Optimizer/Builder/Todo.h"13#include "flang/Optimizer/Dialect/FIROps.h"14#include "flang/Optimizer/HLFIR/HLFIROps.h"15#include "flang/Optimizer/OpenMP/Passes.h"16#include "flang/Optimizer/OpenMP/Utils.h"17#include "flang/Support/OpenMP-utils.h"18#include "flang/Utils/OpenMP.h"19#include "mlir/Analysis/SliceAnalysis.h"20#include "mlir/Dialect/OpenMP/OpenMPDialect.h"21#include "mlir/IR/IRMapping.h"22#include "mlir/Transforms/DialectConversion.h"23#include "mlir/Transforms/RegionUtils.h"24#include "llvm/ADT/SmallPtrSet.h"25 26namespace flangomp {27#define GEN_PASS_DEF_DOCONCURRENTCONVERSIONPASS28#include "flang/Optimizer/OpenMP/Passes.h.inc"29} // namespace flangomp30 31#define DEBUG_TYPE "do-concurrent-conversion"32#define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE << "]: ")33 34namespace {35namespace looputils {36/// Stores info needed about the induction/iteration variable for each `do37/// concurrent` in a loop nest.38struct InductionVariableInfo {39 InductionVariableInfo(fir::DoConcurrentLoopOp loop,40 mlir::Value inductionVar) {41 populateInfo(loop, inductionVar);42 }43 /// The operation allocating memory for iteration variable.44 mlir::Operation *iterVarMemDef;45 /// the operation(s) updating the iteration variable with the current46 /// iteration number.47 llvm::SmallVector<mlir::Operation *, 2> indVarUpdateOps;48 49private:50 /// For the \p doLoop parameter, find the following:51 ///52 /// 1. The operation that declares its iteration variable or allocates memory53 /// for it. For example, give the following loop:54 /// ```55 /// ...56 /// %i:2 = hlfir.declare %0 {uniq_name = "_QFEi"} : ...57 /// ...58 /// fir.do_concurrent.loop (%ind_var) = (%lb) to (%ub) step (%s) {59 /// %ind_var_conv = fir.convert %ind_var : (index) -> i3260 /// fir.store %ind_var_conv to %i#1 : !fir.ref<i32>61 /// ...62 /// }63 /// ```64 ///65 /// This function sets the `iterVarMemDef` member to the `hlfir.declare` op66 /// for `%i`.67 ///68 /// 2. The operation(s) that update the loop's iteration variable from its69 /// induction variable. For the above example, the `indVarUpdateOps` is70 /// populated with the first 2 ops in the loop's body.71 ///72 /// Note: The current implementation is dependent on how flang emits loop73 /// bodies; which is sufficient for the current simple test/use cases. If this74 /// proves to be insufficient, this should be made more generic.75 void populateInfo(fir::DoConcurrentLoopOp loop, mlir::Value inductionVar) {76 mlir::Value result = nullptr;77 78 // Checks if a StoreOp is updating the memref of the loop's iteration79 // variable.80 auto isStoringIV = [&](fir::StoreOp storeOp) {81 // Direct store into the IV memref.82 if (storeOp.getValue() == inductionVar) {83 indVarUpdateOps.push_back(storeOp);84 return true;85 }86 87 // Indirect store into the IV memref.88 if (auto convertOp = mlir::dyn_cast<fir::ConvertOp>(89 storeOp.getValue().getDefiningOp())) {90 if (convertOp.getOperand() == inductionVar) {91 indVarUpdateOps.push_back(convertOp);92 indVarUpdateOps.push_back(storeOp);93 return true;94 }95 }96 97 return false;98 };99 100 for (mlir::Operation &op : loop) {101 if (auto storeOp = mlir::dyn_cast<fir::StoreOp>(op))102 if (isStoringIV(storeOp)) {103 result = storeOp.getMemref();104 break;105 }106 }107 108 assert(result != nullptr && result.getDefiningOp() != nullptr);109 iterVarMemDef = result.getDefiningOp();110 }111};112 113using InductionVariableInfos = llvm::SmallVector<InductionVariableInfo>;114 115/// Collect the list of values used inside the loop but defined outside of it.116void collectLoopLiveIns(fir::DoConcurrentLoopOp loop,117 llvm::SmallVectorImpl<mlir::Value> &liveIns) {118 llvm::SmallDenseSet<mlir::Value> seenValues;119 llvm::SmallPtrSet<mlir::Operation *, 8> seenOps;120 121 for (auto [lb, ub, st] : llvm::zip_equal(122 loop.getLowerBound(), loop.getUpperBound(), loop.getStep())) {123 liveIns.push_back(lb);124 liveIns.push_back(ub);125 liveIns.push_back(st);126 }127 128 mlir::visitUsedValuesDefinedAbove(129 loop.getRegion(), [&](mlir::OpOperand *operand) {130 if (!seenValues.insert(operand->get()).second)131 return;132 133 mlir::Operation *definingOp = operand->get().getDefiningOp();134 // We want to collect ops corresponding to live-ins only once.135 if (definingOp && !seenOps.insert(definingOp).second)136 return;137 138 liveIns.push_back(operand->get());139 });140 141 for (mlir::Value local : loop.getLocalVars())142 liveIns.push_back(local);143 144 for (mlir::Value reduce : loop.getReduceVars())145 liveIns.push_back(reduce);146}147 148/// Collects values that are local to a loop: "loop-local values". A loop-local149/// value is one that is used exclusively inside the loop but allocated outside150/// of it. This usually corresponds to temporary values that are used inside the151/// loop body for initialzing other variables for example.152///153/// See `flang/test/Transforms/DoConcurrent/locally_destroyed_temp.f90` for an154/// example of why we need this.155///156/// \param [in] doLoop - the loop within which the function searches for values157/// used exclusively inside.158///159/// \param [out] locals - the list of loop-local values detected for \p doLoop.160void collectLoopLocalValues(fir::DoConcurrentLoopOp loop,161 llvm::SetVector<mlir::Value> &locals) {162 loop.walk([&](mlir::Operation *op) {163 for (mlir::Value operand : op->getOperands()) {164 if (locals.contains(operand))165 continue;166 167 bool isLocal = true;168 169 if (!mlir::isa_and_present<fir::AllocaOp>(operand.getDefiningOp()))170 continue;171 172 // Values defined inside the loop are not interesting since they do not173 // need to be localized.174 if (loop->isAncestor(operand.getDefiningOp()))175 continue;176 177 for (auto *user : operand.getUsers()) {178 if (!loop->isAncestor(user)) {179 isLocal = false;180 break;181 }182 }183 184 if (isLocal)185 locals.insert(operand);186 }187 });188}189 190/// For a "loop-local" value \p local within a loop's scope, localizes that191/// value within the scope of the parallel region the loop maps to. Towards that192/// end, this function moves the allocation of \p local within \p allocRegion.193///194/// \param local - the value used exclusively within a loop's scope (see195/// collectLoopLocalValues).196///197/// \param allocRegion - the parallel region where \p local's allocation will be198/// privatized.199///200/// \param rewriter - builder used for updating \p allocRegion.201static void localizeLoopLocalValue(mlir::Value local, mlir::Region &allocRegion,202 mlir::ConversionPatternRewriter &rewriter) {203 rewriter.moveOpBefore(local.getDefiningOp(), &allocRegion.front().front());204}205} // namespace looputils206 207class DoConcurrentConversion208 : public mlir::OpConversionPattern<fir::DoConcurrentOp> {209private:210 struct TargetDeclareShapeCreationInfo {211 // Note: We use `std::vector` (rather than `llvm::SmallVector` as usual) to212 // interface more easily `ShapeShiftOp::getOrigins()` which returns213 // `std::vector`.214 std::vector<mlir::Value> startIndices;215 std::vector<mlir::Value> extents;216 217 TargetDeclareShapeCreationInfo(mlir::Value liveIn) {218 mlir::Value shape = nullptr;219 mlir::Operation *liveInDefiningOp = liveIn.getDefiningOp();220 auto declareOp =221 mlir::dyn_cast_if_present<hlfir::DeclareOp>(liveInDefiningOp);222 223 if (declareOp != nullptr)224 shape = declareOp.getShape();225 226 if (!shape)227 return;228 229 auto shapeOp =230 mlir::dyn_cast_if_present<fir::ShapeOp>(shape.getDefiningOp());231 auto shapeShiftOp =232 mlir::dyn_cast_if_present<fir::ShapeShiftOp>(shape.getDefiningOp());233 234 if (!shapeOp && !shapeShiftOp)235 TODO(liveIn.getLoc(),236 "Shapes not defined by `fir.shape` or `fir.shape_shift` op's are"237 "not supported yet.");238 239 if (shapeShiftOp != nullptr)240 startIndices = shapeShiftOp.getOrigins();241 242 extents = shapeOp != nullptr243 ? std::vector<mlir::Value>(shapeOp.getExtents().begin(),244 shapeOp.getExtents().end())245 : shapeShiftOp.getExtents();246 }247 248 bool isShapedValue() const { return !extents.empty(); }249 bool isShapeShiftedValue() const { return !startIndices.empty(); }250 };251 252 using LiveInShapeInfoMap =253 llvm::DenseMap<mlir::Value, TargetDeclareShapeCreationInfo>;254 255public:256 using mlir::OpConversionPattern<fir::DoConcurrentOp>::OpConversionPattern;257 258 DoConcurrentConversion(259 mlir::MLIRContext *context, bool mapToDevice,260 llvm::DenseSet<fir::DoConcurrentOp> &concurrentLoopsToSkip,261 mlir::SymbolTable &moduleSymbolTable)262 : OpConversionPattern(context), mapToDevice(mapToDevice),263 concurrentLoopsToSkip(concurrentLoopsToSkip),264 moduleSymbolTable(moduleSymbolTable) {}265 266 mlir::LogicalResult267 matchAndRewrite(fir::DoConcurrentOp doLoop, OpAdaptor adaptor,268 mlir::ConversionPatternRewriter &rewriter) const override {269 looputils::InductionVariableInfos ivInfos;270 auto loop = mlir::cast<fir::DoConcurrentLoopOp>(271 doLoop.getRegion().back().getTerminator());272 273 auto indVars = loop.getLoopInductionVars();274 assert(indVars.has_value());275 276 for (mlir::Value indVar : *indVars)277 ivInfos.emplace_back(loop, indVar);278 279 llvm::SmallVector<mlir::Value> loopNestLiveIns;280 looputils::collectLoopLiveIns(loop, loopNestLiveIns);281 assert(!loopNestLiveIns.empty());282 283 llvm::SetVector<mlir::Value> locals;284 looputils::collectLoopLocalValues(loop, locals);285 286 // We do not want to map "loop-local" values to the device through287 // `omp.map.info` ops. Therefore, we remove them from the list of live-ins.288 loopNestLiveIns.erase(llvm::remove_if(loopNestLiveIns,289 [&](mlir::Value liveIn) {290 return locals.contains(liveIn);291 }),292 loopNestLiveIns.end());293 294 mlir::omp::TargetOp targetOp;295 mlir::omp::LoopNestOperands loopNestClauseOps;296 297 mlir::IRMapping mapper;298 299 if (mapToDevice) {300 mlir::ModuleOp module = doLoop->getParentOfType<mlir::ModuleOp>();301 bool isTargetDevice =302 llvm::cast<mlir::omp::OffloadModuleInterface>(*module)303 .getIsTargetDevice();304 305 mlir::omp::TargetOperands targetClauseOps;306 genLoopNestClauseOps(doLoop.getLoc(), rewriter, loop, loopNestClauseOps,307 isTargetDevice ? nullptr : &targetClauseOps);308 309 LiveInShapeInfoMap liveInShapeInfoMap;310 fir::FirOpBuilder builder(311 rewriter,312 fir::getKindMapping(doLoop->getParentOfType<mlir::ModuleOp>()));313 314 for (mlir::Value liveIn : loopNestLiveIns) {315 targetClauseOps.mapVars.push_back(316 genMapInfoOpForLiveIn(builder, liveIn));317 liveInShapeInfoMap.insert(318 {liveIn, TargetDeclareShapeCreationInfo(liveIn)});319 }320 321 targetOp =322 genTargetOp(doLoop.getLoc(), rewriter, mapper, loopNestLiveIns,323 targetClauseOps, loopNestClauseOps, liveInShapeInfoMap);324 genTeamsOp(rewriter, loop, mapper);325 }326 327 mlir::omp::ParallelOp parallelOp =328 genParallelOp(rewriter, loop, ivInfos, mapper);329 330 // Only set as composite when part of `distribute parallel do`.331 parallelOp.setComposite(mapToDevice);332 333 if (!mapToDevice)334 genLoopNestClauseOps(doLoop.getLoc(), rewriter, loop, loopNestClauseOps);335 336 for (mlir::Value local : locals)337 looputils::localizeLoopLocalValue(local, parallelOp.getRegion(),338 rewriter);339 340 if (mapToDevice)341 genDistributeOp(doLoop.getLoc(), rewriter).setComposite(/*val=*/true);342 343 auto [loopNestOp, wsLoopOp] =344 genWsLoopOp(rewriter, loop, mapper, loopNestClauseOps,345 /*isComposite=*/mapToDevice);346 347 // `local` region arguments are transferred/cloned from the `do concurrent`348 // loop to the loopnest op when the region is cloned above. Instead, these349 // region arguments should be on the workshare loop's region.350 if (mapToDevice) {351 for (auto [parallelArg, loopNestArg] : llvm::zip_equal(352 parallelOp.getRegion().getArguments(),353 loopNestOp.getRegion().getArguments().slice(354 loop.getLocalOperandsStart(), loop.getNumLocalOperands())))355 rewriter.replaceAllUsesWith(loopNestArg, parallelArg);356 357 for (auto [wsloopArg, loopNestArg] : llvm::zip_equal(358 wsLoopOp.getRegion().getArguments(),359 loopNestOp.getRegion().getArguments().slice(360 loop.getReduceOperandsStart(), loop.getNumReduceOperands())))361 rewriter.replaceAllUsesWith(loopNestArg, wsloopArg);362 } else {363 for (auto [wsloopArg, loopNestArg] :364 llvm::zip_equal(wsLoopOp.getRegion().getArguments(),365 loopNestOp.getRegion().getArguments().drop_front(366 loopNestClauseOps.loopLowerBounds.size())))367 rewriter.replaceAllUsesWith(loopNestArg, wsloopArg);368 }369 370 for (unsigned i = 0;371 i < loop.getLocalVars().size() + loop.getReduceVars().size(); ++i)372 loopNestOp.getRegion().eraseArgument(373 loopNestClauseOps.loopLowerBounds.size());374 375 rewriter.setInsertionPoint(doLoop);376 fir::FirOpBuilder builder(377 rewriter,378 fir::getKindMapping(doLoop->getParentOfType<mlir::ModuleOp>()));379 380 // Collect iteration variable(s) allocations so that we can move them381 // outside the `fir.do_concurrent` wrapper (before erasing it).382 llvm::SmallVector<mlir::Operation *> opsToMove;383 for (mlir::Operation &op : llvm::drop_end(doLoop))384 opsToMove.push_back(&op);385 386 mlir::Block *allocBlock = builder.getAllocaBlock();387 388 for (mlir::Operation *op : llvm::reverse(opsToMove)) {389 rewriter.moveOpBefore(op, allocBlock, allocBlock->begin());390 }391 392 // Mark `unordered` loops that are not perfectly nested to be skipped from393 // the legality check of the `ConversionTarget` since we are not interested394 // in mapping them to OpenMP.395 loopNestOp->walk([&](fir::DoConcurrentOp doLoop) {396 concurrentLoopsToSkip.insert(doLoop);397 });398 399 rewriter.eraseOp(doLoop);400 401 return mlir::success();402 }403 404private:405 mlir::omp::ParallelOp406 genParallelOp(mlir::ConversionPatternRewriter &rewriter,407 fir::DoConcurrentLoopOp loop,408 looputils::InductionVariableInfos &ivInfos,409 mlir::IRMapping &mapper) const {410 mlir::omp::ParallelOperands parallelOps;411 412 if (mapToDevice)413 genPrivatizers(rewriter, mapper, loop, parallelOps);414 415 mlir::Location loc = loop.getLoc();416 auto parallelOp = mlir::omp::ParallelOp::create(rewriter, loc, parallelOps);417 Fortran::common::openmp::EntryBlockArgs parallelArgs;418 parallelArgs.priv.vars = parallelOps.privateVars;419 Fortran::common::openmp::genEntryBlock(rewriter, parallelArgs,420 parallelOp.getRegion());421 rewriter.setInsertionPoint(mlir::omp::TerminatorOp::create(rewriter, loc));422 423 genLoopNestIndVarAllocs(rewriter, ivInfos, mapper);424 return parallelOp;425 }426 427 void genLoopNestIndVarAllocs(mlir::ConversionPatternRewriter &rewriter,428 looputils::InductionVariableInfos &ivInfos,429 mlir::IRMapping &mapper) const {430 431 for (auto &indVarInfo : ivInfos)432 genInductionVariableAlloc(rewriter, indVarInfo.iterVarMemDef, mapper);433 }434 435 mlir::Operation *436 genInductionVariableAlloc(mlir::ConversionPatternRewriter &rewriter,437 mlir::Operation *indVarMemDef,438 mlir::IRMapping &mapper) const {439 assert(440 indVarMemDef != nullptr &&441 "Induction variable memdef is expected to have a defining operation.");442 443 llvm::SmallSetVector<mlir::Operation *, 2> indVarDeclareAndAlloc;444 for (auto operand : indVarMemDef->getOperands())445 indVarDeclareAndAlloc.insert(operand.getDefiningOp());446 indVarDeclareAndAlloc.insert(indVarMemDef);447 448 mlir::Operation *result;449 for (mlir::Operation *opToClone : indVarDeclareAndAlloc)450 result = rewriter.clone(*opToClone, mapper);451 452 return result;453 }454 455 void genLoopNestClauseOps(456 mlir::Location loc, mlir::ConversionPatternRewriter &rewriter,457 fir::DoConcurrentLoopOp loop,458 mlir::omp::LoopNestOperands &loopNestClauseOps,459 mlir::omp::TargetOperands *targetClauseOps = nullptr) const {460 assert(loopNestClauseOps.loopLowerBounds.empty() &&461 "Loop nest bounds were already emitted!");462 463 auto populateBounds = [](mlir::Value var,464 llvm::SmallVectorImpl<mlir::Value> &bounds) {465 bounds.push_back(var.getDefiningOp()->getResult(0));466 };467 468 auto hostEvalCapture = [&](mlir::Value var,469 llvm::SmallVectorImpl<mlir::Value> &bounds) {470 populateBounds(var, bounds);471 472 // Ensure that loop-nest bounds are evaluated in the host and forwarded to473 // the nested omp constructs when we map to the device.474 if (targetClauseOps)475 targetClauseOps->hostEvalVars.push_back(var);476 };477 478 for (auto [lb, ub, st] : llvm::zip_equal(479 loop.getLowerBound(), loop.getUpperBound(), loop.getStep())) {480 hostEvalCapture(lb, loopNestClauseOps.loopLowerBounds);481 hostEvalCapture(ub, loopNestClauseOps.loopUpperBounds);482 hostEvalCapture(st, loopNestClauseOps.loopSteps);483 }484 485 loopNestClauseOps.loopInclusive = rewriter.getUnitAttr();486 }487 488 std::pair<mlir::omp::LoopNestOp, mlir::omp::WsloopOp>489 genWsLoopOp(mlir::ConversionPatternRewriter &rewriter,490 fir::DoConcurrentLoopOp loop, mlir::IRMapping &mapper,491 const mlir::omp::LoopNestOperands &clauseOps,492 bool isComposite) const {493 mlir::omp::WsloopOperands wsloopClauseOps;494 if (!mapToDevice)495 genPrivatizers(rewriter, mapper, loop, wsloopClauseOps);496 497 genReductions(rewriter, mapper, loop, wsloopClauseOps);498 499 auto wsloopOp =500 mlir::omp::WsloopOp::create(rewriter, loop.getLoc(), wsloopClauseOps);501 wsloopOp.setComposite(isComposite);502 503 Fortran::common::openmp::EntryBlockArgs wsloopArgs;504 wsloopArgs.priv.vars = wsloopClauseOps.privateVars;505 wsloopArgs.reduction.vars = wsloopClauseOps.reductionVars;506 Fortran::common::openmp::genEntryBlock(rewriter, wsloopArgs,507 wsloopOp.getRegion());508 509 auto loopNestOp =510 mlir::omp::LoopNestOp::create(rewriter, loop.getLoc(), clauseOps);511 512 // Clone the loop's body inside the loop nest construct using the513 // mapped values.514 rewriter.cloneRegionBefore(loop.getRegion(), loopNestOp.getRegion(),515 loopNestOp.getRegion().begin(), mapper);516 517 rewriter.setInsertionPointToEnd(&loopNestOp.getRegion().back());518 mlir::omp::YieldOp::create(rewriter, loop->getLoc());519 520 return {loopNestOp, wsloopOp};521 }522 523 void genBoundsOps(fir::FirOpBuilder &builder, mlir::Value liveIn,524 mlir::Value rawAddr,525 llvm::SmallVectorImpl<mlir::Value> &boundsOps) const {526 fir::ExtendedValue extVal =527 hlfir::translateToExtendedValue(rawAddr.getLoc(), builder,528 hlfir::Entity{liveIn},529 /*contiguousHint=*/530 true)531 .first;532 fir::factory::AddrAndBoundsInfo info = fir::factory::getDataOperandBaseAddr(533 builder, rawAddr, /*isOptional=*/false, rawAddr.getLoc());534 boundsOps = fir::factory::genImplicitBoundsOps<mlir::omp::MapBoundsOp,535 mlir::omp::MapBoundsType>(536 builder, info, extVal,537 /*dataExvIsAssumedSize=*/false, rawAddr.getLoc());538 }539 540 mlir::omp::MapInfoOp genMapInfoOpForLiveIn(fir::FirOpBuilder &builder,541 mlir::Value liveIn) const {542 mlir::Value rawAddr = liveIn;543 llvm::StringRef name;544 545 mlir::Operation *liveInDefiningOp = liveIn.getDefiningOp();546 auto declareOp =547 mlir::dyn_cast_if_present<hlfir::DeclareOp>(liveInDefiningOp);548 549 if (declareOp != nullptr) {550 // Use the raw address to avoid unboxing `fir.box` values whenever551 // possible. Put differently, if we have access to the direct value memory552 // reference/address, we use it.553 rawAddr = declareOp.getOriginalBase();554 name = declareOp.getUniqName();555 }556 557 if (!llvm::isa<mlir::omp::PointerLikeType>(rawAddr.getType())) {558 mlir::OpBuilder::InsertionGuard guard(builder);559 builder.setInsertionPointAfter(liveInDefiningOp);560 auto copyVal = builder.createTemporary(liveIn.getLoc(), liveIn.getType());561 builder.createStoreWithConvert(copyVal.getLoc(), liveIn, copyVal);562 rawAddr = copyVal;563 }564 565 mlir::Type liveInType = liveIn.getType();566 mlir::Type eleType = liveInType;567 if (auto refType = mlir::dyn_cast<fir::ReferenceType>(liveInType))568 eleType = refType.getElementType();569 570 mlir::omp::ClauseMapFlags mapFlag = mlir::omp::ClauseMapFlags::implicit;571 mlir::omp::VariableCaptureKind captureKind =572 mlir::omp::VariableCaptureKind::ByRef;573 574 if (fir::isa_trivial(eleType) || fir::isa_char(eleType)) {575 captureKind = mlir::omp::VariableCaptureKind::ByCopy;576 } else if (!fir::isa_builtin_cptr_type(eleType)) {577 mapFlag |= mlir::omp::ClauseMapFlags::to;578 mapFlag |= mlir::omp::ClauseMapFlags::from;579 }580 581 llvm::SmallVector<mlir::Value> boundsOps;582 genBoundsOps(builder, liveIn, rawAddr, boundsOps);583 584 return Fortran::utils::openmp::createMapInfoOp(585 builder, liveIn.getLoc(), rawAddr,586 /*varPtrPtr=*/{}, name.str(), boundsOps,587 /*members=*/{},588 /*membersIndex=*/mlir::ArrayAttr{}, mapFlag, captureKind,589 rawAddr.getType());590 }591 592 mlir::omp::TargetOp593 genTargetOp(mlir::Location loc, mlir::ConversionPatternRewriter &rewriter,594 mlir::IRMapping &mapper, llvm::ArrayRef<mlir::Value> mappedVars,595 mlir::omp::TargetOperands &clauseOps,596 mlir::omp::LoopNestOperands &loopNestClauseOps,597 const LiveInShapeInfoMap &liveInShapeInfoMap) const {598 auto targetOp = mlir::omp::TargetOp::create(rewriter, loc, clauseOps);599 auto argIface = llvm::cast<mlir::omp::BlockArgOpenMPOpInterface>(*targetOp);600 601 mlir::Region ®ion = targetOp.getRegion();602 603 llvm::SmallVector<mlir::Type> regionArgTypes;604 llvm::SmallVector<mlir::Location> regionArgLocs;605 606 for (auto var : llvm::concat<const mlir::Value>(clauseOps.hostEvalVars,607 clauseOps.mapVars)) {608 regionArgTypes.push_back(var.getType());609 regionArgLocs.push_back(var.getLoc());610 }611 612 rewriter.createBlock(®ion, {}, regionArgTypes, regionArgLocs);613 fir::FirOpBuilder builder(614 rewriter,615 fir::getKindMapping(targetOp->getParentOfType<mlir::ModuleOp>()));616 617 // Within the loop, it is possible that we discover other values that need618 // to be mapped to the target region (the shape info values for arrays, for619 // example). Therefore, the map block args might be extended and resized.620 // Hence, we invoke `argIface.getMapBlockArgs()` every iteration to make621 // sure we access the proper vector of data.622 int idx = 0;623 for (auto [mapInfoOp, mappedVar] :624 llvm::zip_equal(clauseOps.mapVars, mappedVars)) {625 auto miOp = mlir::cast<mlir::omp::MapInfoOp>(mapInfoOp.getDefiningOp());626 hlfir::DeclareOp liveInDeclare =627 genLiveInDeclare(builder, targetOp, argIface.getMapBlockArgs()[idx],628 miOp, liveInShapeInfoMap.at(mappedVar));629 ++idx;630 631 // If `mappedVar.getDefiningOp()` is a `fir::BoxAddrOp`, we probably632 // need to "unpack" the box by getting the defining op of it's value.633 // However, we did not hit this case in reality yet so leaving it as a634 // todo for now.635 if (mlir::isa<fir::BoxAddrOp>(mappedVar.getDefiningOp()))636 TODO(mappedVar.getLoc(),637 "Mapped variabled defined by `BoxAddrOp` are not supported yet");638 639 auto mapHostValueToDevice = [&](mlir::Value hostValue,640 mlir::Value deviceValue) {641 if (!llvm::isa<mlir::omp::PointerLikeType>(hostValue.getType()))642 mapper.map(hostValue,643 builder.loadIfRef(hostValue.getLoc(), deviceValue));644 else645 mapper.map(hostValue, deviceValue);646 };647 648 mapHostValueToDevice(mappedVar, liveInDeclare.getOriginalBase());649 650 if (auto origDeclareOp = mlir::dyn_cast_if_present<hlfir::DeclareOp>(651 mappedVar.getDefiningOp()))652 mapHostValueToDevice(origDeclareOp.getBase(), liveInDeclare.getBase());653 }654 655 for (auto [arg, hostEval] : llvm::zip_equal(argIface.getHostEvalBlockArgs(),656 clauseOps.hostEvalVars))657 mapper.map(hostEval, arg);658 659 for (unsigned i = 0; i < loopNestClauseOps.loopLowerBounds.size(); ++i) {660 loopNestClauseOps.loopLowerBounds[i] =661 mapper.lookup(loopNestClauseOps.loopLowerBounds[i]);662 loopNestClauseOps.loopUpperBounds[i] =663 mapper.lookup(loopNestClauseOps.loopUpperBounds[i]);664 loopNestClauseOps.loopSteps[i] =665 mapper.lookup(loopNestClauseOps.loopSteps[i]);666 }667 668 // Check if cloning the bounds introduced any dependency on the outer669 // region. If so, then either clone them as well if they are670 // MemoryEffectFree, or else copy them to a new temporary and add them to671 // the map and block_argument lists and replace their uses with the new672 // temporary.673 Fortran::utils::openmp::cloneOrMapRegionOutsiders(builder, targetOp);674 rewriter.setInsertionPoint(675 mlir::omp::TerminatorOp::create(rewriter, targetOp.getLoc()));676 677 return targetOp;678 }679 680 hlfir::DeclareOp genLiveInDeclare(681 fir::FirOpBuilder &builder, mlir::omp::TargetOp targetOp,682 mlir::Value liveInArg, mlir::omp::MapInfoOp liveInMapInfoOp,683 const TargetDeclareShapeCreationInfo &targetShapeCreationInfo) const {684 mlir::Type liveInType = liveInArg.getType();685 std::string liveInName = liveInMapInfoOp.getName().has_value()686 ? liveInMapInfoOp.getName().value().str()687 : std::string("");688 if (fir::isa_ref_type(liveInType))689 liveInType = fir::unwrapRefType(liveInType);690 691 mlir::Value shape = [&]() -> mlir::Value {692 if (!targetShapeCreationInfo.isShapedValue())693 return {};694 695 if (targetShapeCreationInfo.isShapeShiftedValue()) {696 llvm::SmallVector<mlir::Value> shapeShiftOperands;697 698 size_t shapeIdx = 0;699 for (auto [startIndex, extent] :700 llvm::zip_equal(targetShapeCreationInfo.startIndices,701 targetShapeCreationInfo.extents)) {702 shapeShiftOperands.push_back(703 Fortran::utils::openmp::mapTemporaryValue(704 builder, targetOp, startIndex,705 liveInName + ".start_idx.dim" + std::to_string(shapeIdx)));706 shapeShiftOperands.push_back(707 Fortran::utils::openmp::mapTemporaryValue(708 builder, targetOp, extent,709 liveInName + ".extent.dim" + std::to_string(shapeIdx)));710 ++shapeIdx;711 }712 713 auto shapeShiftType = fir::ShapeShiftType::get(714 builder.getContext(), shapeShiftOperands.size() / 2);715 return fir::ShapeShiftOp::create(builder, liveInArg.getLoc(),716 shapeShiftType, shapeShiftOperands);717 }718 719 llvm::SmallVector<mlir::Value> shapeOperands;720 size_t shapeIdx = 0;721 for (auto extent : targetShapeCreationInfo.extents) {722 shapeOperands.push_back(Fortran::utils::openmp::mapTemporaryValue(723 builder, targetOp, extent,724 liveInName + ".extent.dim" + std::to_string(shapeIdx)));725 ++shapeIdx;726 }727 728 return fir::ShapeOp::create(builder, liveInArg.getLoc(), shapeOperands);729 }();730 731 return hlfir::DeclareOp::create(builder, liveInArg.getLoc(), liveInArg,732 liveInName, shape);733 }734 735 mlir::omp::TeamsOp genTeamsOp(mlir::ConversionPatternRewriter &rewriter,736 fir::DoConcurrentLoopOp loop,737 mlir::IRMapping &mapper) const {738 mlir::omp::TeamsOperands teamsOps;739 genReductions(rewriter, mapper, loop, teamsOps);740 741 mlir::Location loc = loop.getLoc();742 auto teamsOp = mlir::omp::TeamsOp::create(rewriter, loc, teamsOps);743 Fortran::common::openmp::EntryBlockArgs teamsArgs;744 teamsArgs.reduction.vars = teamsOps.reductionVars;745 Fortran::common::openmp::genEntryBlock(rewriter, teamsArgs,746 teamsOp.getRegion());747 748 rewriter.setInsertionPoint(mlir::omp::TerminatorOp::create(rewriter, loc));749 750 for (auto [loopVar, teamsArg] : llvm::zip_equal(751 loop.getReduceVars(), teamsOp.getRegion().getArguments())) {752 mapper.map(loopVar, teamsArg);753 }754 755 return teamsOp;756 }757 758 mlir::omp::DistributeOp759 genDistributeOp(mlir::Location loc,760 mlir::ConversionPatternRewriter &rewriter) const {761 auto distOp = mlir::omp::DistributeOp::create(762 rewriter, loc, /*clauses=*/mlir::omp::DistributeOperands{});763 764 rewriter.createBlock(&distOp.getRegion());765 return distOp;766 }767 768 void cloneFIRRegionToOMP(mlir::ConversionPatternRewriter &rewriter,769 mlir::Region &firRegion,770 mlir::Region &ompRegion) const {771 if (!firRegion.empty()) {772 rewriter.cloneRegionBefore(firRegion, ompRegion, ompRegion.begin());773 auto firYield =774 mlir::cast<fir::YieldOp>(ompRegion.back().getTerminator());775 rewriter.setInsertionPoint(firYield);776 mlir::omp::YieldOp::create(rewriter, firYield.getLoc(),777 firYield.getOperands());778 rewriter.eraseOp(firYield);779 }780 }781 782 /// Generate bodies of OpenMP privatizers by cloning the bodies of FIR783 /// privatizers.784 ///785 /// \param [in] rewriter - used to driver IR generation for privatizers.786 /// \param [in] mapper - value mapping from FIR to OpenMP constructs.787 /// \param [in] loop - FIR loop to convert its localizers.788 ///789 /// \param [out] privateClauseOps - OpenMP privatizers to gen their bodies.790 void genPrivatizers(mlir::ConversionPatternRewriter &rewriter,791 mlir::IRMapping &mapper, fir::DoConcurrentLoopOp loop,792 mlir::omp::PrivateClauseOps &privateClauseOps) const {793 // For `local` (and `local_init`) operands, emit corresponding `private`794 // clauses and attach these clauses to the workshare loop.795 if (!loop.getLocalVars().empty())796 for (auto [var, sym, arg] : llvm::zip_equal(797 loop.getLocalVars(),798 loop.getLocalSymsAttr().getAsRange<mlir::SymbolRefAttr>(),799 loop.getRegionLocalArgs())) {800 auto localizer = moduleSymbolTable.lookup<fir::LocalitySpecifierOp>(801 sym.getLeafReference());802 if (localizer.getLocalitySpecifierType() ==803 fir::LocalitySpecifierType::LocalInit)804 TODO(localizer.getLoc(),805 "local_init conversion is not supported yet");806 807 mlir::OpBuilder::InsertionGuard guard(rewriter);808 rewriter.setInsertionPointAfter(localizer);809 810 auto privatizer = mlir::omp::PrivateClauseOp::create(811 rewriter, localizer.getLoc(), sym.getLeafReference().str() + ".omp",812 localizer.getTypeAttr().getValue(),813 mlir::omp::DataSharingClauseType::Private);814 815 cloneFIRRegionToOMP(rewriter, localizer.getInitRegion(),816 privatizer.getInitRegion());817 cloneFIRRegionToOMP(rewriter, localizer.getDeallocRegion(),818 privatizer.getDeallocRegion());819 820 moduleSymbolTable.insert(privatizer);821 822 privateClauseOps.privateVars.push_back(mapToDevice ? mapper.lookup(var)823 : var);824 privateClauseOps.privateSyms.push_back(825 mlir::SymbolRefAttr::get(privatizer));826 }827 }828 829 void genReductions(mlir::ConversionPatternRewriter &rewriter,830 mlir::IRMapping &mapper, fir::DoConcurrentLoopOp loop,831 mlir::omp::ReductionClauseOps &reductionClauseOps) const {832 if (!loop.getReduceVars().empty()) {833 for (auto [var, byRef, sym, arg] : llvm::zip_equal(834 loop.getReduceVars(), loop.getReduceByrefAttr().asArrayRef(),835 loop.getReduceSymsAttr().getAsRange<mlir::SymbolRefAttr>(),836 loop.getRegionReduceArgs())) {837 auto firReducer = moduleSymbolTable.lookup<fir::DeclareReductionOp>(838 sym.getLeafReference());839 840 mlir::OpBuilder::InsertionGuard guard(rewriter);841 rewriter.setInsertionPointAfter(firReducer);842 std::string ompReducerName = sym.getLeafReference().str() + ".omp";843 844 auto ompReducer =845 moduleSymbolTable.lookup<mlir::omp::DeclareReductionOp>(846 rewriter.getStringAttr(ompReducerName));847 848 if (!ompReducer) {849 ompReducer = mlir::omp::DeclareReductionOp::create(850 rewriter, firReducer.getLoc(), ompReducerName,851 firReducer.getTypeAttr().getValue(),852 firReducer.getByrefElementTypeAttr());853 854 cloneFIRRegionToOMP(rewriter, firReducer.getAllocRegion(),855 ompReducer.getAllocRegion());856 cloneFIRRegionToOMP(rewriter, firReducer.getInitializerRegion(),857 ompReducer.getInitializerRegion());858 cloneFIRRegionToOMP(rewriter, firReducer.getReductionRegion(),859 ompReducer.getReductionRegion());860 cloneFIRRegionToOMP(rewriter, firReducer.getAtomicReductionRegion(),861 ompReducer.getAtomicReductionRegion());862 cloneFIRRegionToOMP(rewriter, firReducer.getCleanupRegion(),863 ompReducer.getCleanupRegion());864 moduleSymbolTable.insert(ompReducer);865 }866 867 reductionClauseOps.reductionVars.push_back(868 mapToDevice ? mapper.lookup(var) : var);869 reductionClauseOps.reductionByref.push_back(byRef);870 reductionClauseOps.reductionSyms.push_back(871 mlir::SymbolRefAttr::get(ompReducer));872 }873 }874 }875 876 bool mapToDevice;877 llvm::DenseSet<fir::DoConcurrentOp> &concurrentLoopsToSkip;878 mlir::SymbolTable &moduleSymbolTable;879};880 881/// A listener that forwards notifyOperationErased to the given callback.882struct CallbackListener : public mlir::RewriterBase::Listener {883 CallbackListener(std::function<void(mlir::Operation *op)> onOperationErased)884 : onOperationErased(onOperationErased) {}885 886 void notifyOperationErased(mlir::Operation *op) override {887 onOperationErased(op);888 }889 890 std::function<void(mlir::Operation *op)> onOperationErased;891};892 893class DoConcurrentConversionPass894 : public flangomp::impl::DoConcurrentConversionPassBase<895 DoConcurrentConversionPass> {896public:897 DoConcurrentConversionPass() = default;898 899 DoConcurrentConversionPass(900 const flangomp::DoConcurrentConversionPassOptions &options)901 : DoConcurrentConversionPassBase(options) {}902 903 void runOnOperation() override {904 mlir::ModuleOp module = getOperation();905 mlir::MLIRContext *context = &getContext();906 mlir::SymbolTable moduleSymbolTable(module);907 908 if (mapTo != flangomp::DoConcurrentMappingKind::DCMK_Host &&909 mapTo != flangomp::DoConcurrentMappingKind::DCMK_Device) {910 mlir::emitWarning(mlir::UnknownLoc::get(context),911 "DoConcurrentConversionPass: invalid `map-to` value. "912 "Valid values are: `host` or `device`");913 return;914 }915 916 llvm::DenseSet<fir::DoConcurrentOp> concurrentLoopsToSkip;917 CallbackListener callbackListener([&](mlir::Operation *op) {918 if (auto loop = mlir::dyn_cast<fir::DoConcurrentOp>(op))919 concurrentLoopsToSkip.erase(loop);920 });921 mlir::RewritePatternSet patterns(context);922 patterns.insert<DoConcurrentConversion>(923 context, mapTo == flangomp::DoConcurrentMappingKind::DCMK_Device,924 concurrentLoopsToSkip, moduleSymbolTable);925 mlir::ConversionTarget target(*context);926 target.addDynamicallyLegalOp<fir::DoConcurrentOp>(927 [&](fir::DoConcurrentOp op) {928 return concurrentLoopsToSkip.contains(op);929 });930 target.markUnknownOpDynamicallyLegal(931 [](mlir::Operation *) { return true; });932 933 mlir::ConversionConfig config;934 config.allowPatternRollback = false;935 config.listener = &callbackListener;936 if (mlir::failed(mlir::applyFullConversion(module, target,937 std::move(patterns), config))) {938 signalPassFailure();939 }940 }941};942} // namespace943 944std::unique_ptr<mlir::Pass>945flangomp::createDoConcurrentConversionPass(bool mapToDevice) {946 DoConcurrentConversionPassOptions options;947 options.mapTo = mapToDevice ? flangomp::DoConcurrentMappingKind::DCMK_Device948 : flangomp::DoConcurrentMappingKind::DCMK_Host;949 950 return std::make_unique<DoConcurrentConversionPass>(options);951}952