1848 lines · cpp
1//===- LowerWorkdistribute.cpp2//-------------------------------------------------===//3//4// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.5// See https://llvm.org/LICENSE.txt for license information.6// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception7//8//===----------------------------------------------------------------------===//9//10// This file implements the lowering and optimisations of omp.workdistribute.11//12// Fortran array statements are lowered to fir as fir.do_loop unordered.13// lower-workdistribute pass works mainly on identifying fir.do_loop unordered14// that is nested in target{teams{workdistribute{fir.do_loop unordered}}} and15// lowers it to target{teams{parallel{distribute{wsloop{loop_nest}}}}}.16// It hoists all the other ops outside target region.17// Relaces heap allocation on target with omp.target_allocmem and18// deallocation with omp.target_freemem from host. Also replaces19// runtime function "Assign" with omp_target_memcpy.20//21//===----------------------------------------------------------------------===//22 23#include "flang/Optimizer/Builder/FIRBuilder.h"24#include "flang/Optimizer/Dialect/FIRDialect.h"25#include "flang/Optimizer/Dialect/FIROps.h"26#include "flang/Optimizer/Dialect/FIRType.h"27#include "flang/Optimizer/HLFIR/Passes.h"28#include "flang/Optimizer/OpenMP/Utils.h"29#include "flang/Optimizer/Transforms/Passes.h"30#include "mlir/Analysis/SliceAnalysis.h"31#include "mlir/Dialect/OpenMP/OpenMPDialect.h"32#include "mlir/IR/Builders.h"33#include "mlir/IR/Value.h"34#include "mlir/Transforms/DialectConversion.h"35#include "mlir/Transforms/GreedyPatternRewriteDriver.h"36#include "mlir/Transforms/RegionUtils.h"37#include "llvm/Frontend/OpenMP/OMPConstants.h"38#include <mlir/Dialect/Arith/IR/Arith.h>39#include <mlir/Dialect/LLVMIR/LLVMTypes.h>40#include <mlir/Dialect/Utils/IndexingUtils.h>41#include <mlir/IR/BlockSupport.h>42#include <mlir/IR/BuiltinOps.h>43#include <mlir/IR/Diagnostics.h>44#include <mlir/IR/IRMapping.h>45#include <mlir/IR/PatternMatch.h>46#include <mlir/Interfaces/SideEffectInterfaces.h>47#include <mlir/Support/LLVM.h>48#include <optional>49#include <variant>50 51namespace flangomp {52#define GEN_PASS_DEF_LOWERWORKDISTRIBUTE53#include "flang/Optimizer/OpenMP/Passes.h.inc"54} // namespace flangomp55 56#define DEBUG_TYPE "lower-workdistribute"57 58using namespace mlir;59 60namespace {61 62/// This string is used to identify the Fortran-specific runtime FortranAAssign.63static constexpr llvm::StringRef FortranAssignStr = "_FortranAAssign";64 65/// The isRuntimeCall function is a utility designed to determine66/// if a given operation is a call to a Fortran-specific runtime function.67static bool isRuntimeCall(Operation *op) {68 if (auto callOp = dyn_cast<fir::CallOp>(op)) {69 auto callee = callOp.getCallee();70 if (!callee)71 return false;72 auto *func = op->getParentOfType<ModuleOp>().lookupSymbol(*callee);73 if (func->getAttr(fir::FIROpsDialect::getFirRuntimeAttrName()))74 return true;75 }76 return false;77}78 79/// This is the single source of truth about whether we should parallelize an80/// operation nested in an omp.workdistribute region.81/// Parallelize here refers to dividing into units of work.82static bool shouldParallelize(Operation *op) {83 // True if the op is a runtime call to Assign84 if (isRuntimeCall(op)) {85 fir::CallOp runtimeCall = cast<fir::CallOp>(op);86 auto funcName = runtimeCall.getCallee()->getRootReference().getValue();87 if (funcName == FortranAssignStr) {88 return true;89 }90 }91 // We cannot parallelize ops with side effects.92 // Parallelizable operations should not produce93 // values that other operations depend on94 if (llvm::any_of(op->getResults(),95 [](OpResult v) -> bool { return !v.use_empty(); }))96 return false;97 // We will parallelize unordered loops - these come from array syntax98 if (auto loop = dyn_cast<fir::DoLoopOp>(op)) {99 auto unordered = loop.getUnordered();100 if (!unordered)101 return false;102 return *unordered;103 }104 // We cannot parallelize anything else.105 return false;106}107 108/// The getPerfectlyNested function is a generic utility for finding109/// a single, "perfectly nested" operation within a parent operation.110template <typename T>111static T getPerfectlyNested(Operation *op) {112 if (op->getNumRegions() != 1)113 return nullptr;114 auto ®ion = op->getRegion(0);115 if (region.getBlocks().size() != 1)116 return nullptr;117 auto *block = ®ion.front();118 auto *firstOp = &block->front();119 if (auto nested = dyn_cast<T>(firstOp))120 if (firstOp->getNextNode() == block->getTerminator())121 return nested;122 return nullptr;123}124 125/// verifyTargetTeamsWorkdistribute method verifies that126/// omp.target { teams { workdistribute { ... } } } is well formed127/// and fails for function calls that don't have lowering implemented yet.128static LogicalResult129verifyTargetTeamsWorkdistribute(omp::WorkdistributeOp workdistribute) {130 OpBuilder rewriter(workdistribute);131 auto loc = workdistribute->getLoc();132 auto teams = dyn_cast<omp::TeamsOp>(workdistribute->getParentOp());133 if (!teams) {134 emitError(loc, "workdistribute not nested in teams\n");135 return failure();136 }137 if (workdistribute.getRegion().getBlocks().size() != 1) {138 emitError(loc, "workdistribute with multiple blocks\n");139 return failure();140 }141 if (teams.getRegion().getBlocks().size() != 1) {142 emitError(loc, "teams with multiple blocks\n");143 return failure();144 }145 146 bool foundWorkdistribute = false;147 for (auto &op : teams.getOps()) {148 if (isa<omp::WorkdistributeOp>(op)) {149 if (foundWorkdistribute) {150 emitError(loc, "teams has multiple workdistribute ops.\n");151 return failure();152 }153 foundWorkdistribute = true;154 continue;155 }156 // Identify any omp dialect ops present before/after workdistribute.157 if (op.getDialect() && isa<omp::OpenMPDialect>(op.getDialect()) &&158 !isa<omp::TerminatorOp>(op)) {159 emitError(loc, "teams has omp ops other than workdistribute. Lowering "160 "not implemented yet.\n");161 return failure();162 }163 }164 165 omp::TargetOp targetOp = dyn_cast<omp::TargetOp>(teams->getParentOp());166 // return if not omp.target167 if (!targetOp)168 return success();169 170 for (auto &op : workdistribute.getOps()) {171 if (auto callOp = dyn_cast<fir::CallOp>(op)) {172 if (isRuntimeCall(&op)) {173 auto funcName = (*callOp.getCallee()).getRootReference().getValue();174 // _FortranAAssign is handled. Other runtime calls are not supported175 // in omp.workdistribute yet.176 if (funcName == FortranAssignStr)177 continue;178 else {179 emitError(loc, "Runtime call " + funcName +180 " lowering not supported for workdistribute yet.");181 return failure();182 }183 }184 }185 }186 return success();187}188 189/// fissionWorkdistribute method finds the parallelizable ops190/// within teams {workdistribute} region and moves them to their191/// own teams{workdistribute} region.192///193/// If B() and D() are parallelizable,194///195/// omp.teams {196/// omp.workdistribute {197/// A()198/// B()199/// C()200/// D()201/// E()202/// }203/// }204///205/// becomes206///207/// A()208/// omp.teams {209/// omp.workdistribute {210/// B()211/// }212/// }213/// C()214/// omp.teams {215/// omp.workdistribute {216/// D()217/// }218/// }219/// E()220static FailureOr<bool>221fissionWorkdistribute(omp::WorkdistributeOp workdistribute) {222 OpBuilder rewriter(workdistribute);223 auto loc = workdistribute->getLoc();224 auto teams = dyn_cast<omp::TeamsOp>(workdistribute->getParentOp());225 auto *teamsBlock = &teams.getRegion().front();226 bool changed = false;227 // Move the ops inside teams and before workdistribute outside.228 IRMapping irMapping;229 llvm::SmallVector<Operation *> teamsHoisted;230 for (auto &op : teams.getOps()) {231 if (&op == workdistribute) {232 break;233 }234 if (shouldParallelize(&op)) {235 emitError(loc, "teams has parallelize ops before first workdistribute\n");236 return failure();237 } else {238 rewriter.setInsertionPoint(teams);239 rewriter.clone(op, irMapping);240 teamsHoisted.push_back(&op);241 changed = true;242 }243 }244 for (auto *op : llvm::reverse(teamsHoisted)) {245 op->replaceAllUsesWith(irMapping.lookup(op));246 op->erase();247 }248 249 // While we have unhandled operations in the original workdistribute250 auto *workdistributeBlock = &workdistribute.getRegion().front();251 auto *terminator = workdistributeBlock->getTerminator();252 while (&workdistributeBlock->front() != terminator) {253 rewriter.setInsertionPoint(teams);254 IRMapping mapping;255 llvm::SmallVector<Operation *> hoisted;256 Operation *parallelize = nullptr;257 for (auto &op : workdistribute.getOps()) {258 if (&op == terminator) {259 break;260 }261 if (shouldParallelize(&op)) {262 parallelize = &op;263 break;264 } else {265 rewriter.clone(op, mapping);266 hoisted.push_back(&op);267 changed = true;268 }269 }270 271 for (auto *op : llvm::reverse(hoisted)) {272 op->replaceAllUsesWith(mapping.lookup(op));273 op->erase();274 }275 276 if (parallelize && hoisted.empty() &&277 parallelize->getNextNode() == terminator)278 break;279 if (parallelize) {280 auto newTeams = rewriter.cloneWithoutRegions(teams);281 auto *newTeamsBlock = rewriter.createBlock(282 &newTeams.getRegion(), newTeams.getRegion().begin(), {}, {});283 for (auto arg : teamsBlock->getArguments())284 newTeamsBlock->addArgument(arg.getType(), arg.getLoc());285 auto newWorkdistribute = omp::WorkdistributeOp::create(rewriter, loc);286 omp::TerminatorOp::create(rewriter, loc);287 rewriter.createBlock(&newWorkdistribute.getRegion(),288 newWorkdistribute.getRegion().begin(), {}, {});289 auto *cloned = rewriter.clone(*parallelize);290 parallelize->replaceAllUsesWith(cloned);291 parallelize->erase();292 omp::TerminatorOp::create(rewriter, loc);293 changed = true;294 }295 }296 return changed;297}298 299/// Generate omp.parallel operation with an empty region.300static void genParallelOp(Location loc, OpBuilder &rewriter, bool composite) {301 auto parallelOp = mlir::omp::ParallelOp::create(rewriter, loc);302 parallelOp.setComposite(composite);303 rewriter.createBlock(¶llelOp.getRegion());304 rewriter.setInsertionPoint(mlir::omp::TerminatorOp::create(rewriter, loc));305 return;306}307 308/// Generate omp.distribute operation with an empty region.309static void genDistributeOp(Location loc, OpBuilder &rewriter, bool composite) {310 mlir::omp::DistributeOperands distributeClauseOps;311 auto distributeOp =312 mlir::omp::DistributeOp::create(rewriter, loc, distributeClauseOps);313 distributeOp.setComposite(composite);314 auto distributeBlock = rewriter.createBlock(&distributeOp.getRegion());315 rewriter.setInsertionPointToStart(distributeBlock);316 return;317}318 319/// Generate loop nest clause operands from fir.do_loop operation.320static void321genLoopNestClauseOps(OpBuilder &rewriter, fir::DoLoopOp loop,322 mlir::omp::LoopNestOperands &loopNestClauseOps) {323 assert(loopNestClauseOps.loopLowerBounds.empty() &&324 "Loop nest bounds were already emitted!");325 loopNestClauseOps.loopLowerBounds.push_back(loop.getLowerBound());326 loopNestClauseOps.loopUpperBounds.push_back(loop.getUpperBound());327 loopNestClauseOps.loopSteps.push_back(loop.getStep());328 loopNestClauseOps.loopInclusive = rewriter.getUnitAttr();329}330 331/// Generate omp.wsloop operation with an empty region and332/// clone the body of fir.do_loop operation inside the loop nest region.333static void genWsLoopOp(mlir::OpBuilder &rewriter, fir::DoLoopOp doLoop,334 const mlir::omp::LoopNestOperands &clauseOps,335 bool composite) {336 337 auto wsloopOp = mlir::omp::WsloopOp::create(rewriter, doLoop.getLoc());338 wsloopOp.setComposite(composite);339 rewriter.createBlock(&wsloopOp.getRegion());340 341 auto loopNestOp =342 mlir::omp::LoopNestOp::create(rewriter, doLoop.getLoc(), clauseOps);343 344 // Clone the loop's body inside the loop nest construct using the345 // mapped values.346 rewriter.cloneRegionBefore(doLoop.getRegion(), loopNestOp.getRegion(),347 loopNestOp.getRegion().begin());348 Block *clonedBlock = &loopNestOp.getRegion().back();349 mlir::Operation *terminatorOp = clonedBlock->getTerminator();350 351 // Erase fir.result op of do loop and create yield op.352 if (auto resultOp = dyn_cast<fir::ResultOp>(terminatorOp)) {353 rewriter.setInsertionPoint(terminatorOp);354 mlir::omp::YieldOp::create(rewriter, doLoop->getLoc());355 terminatorOp->erase();356 }357}358 359/// workdistributeDoLower method finds the fir.do_loop unoredered360/// nested in teams {workdistribute{fir.do_loop unoredered}} and361/// lowers it to teams {parallel { distribute {wsloop {loop_nest}}}}.362///363/// If fir.do_loop is present inside teams workdistribute364///365/// omp.teams {366/// omp.workdistribute {367/// fir.do_loop unoredered {368/// ...369/// }370/// }371/// }372///373/// Then, its lowered to374///375/// omp.teams {376/// omp.parallel {377/// omp.distribute {378/// omp.wsloop {379/// omp.loop_nest380/// ...381/// }382/// }383/// }384/// }385/// }386static bool387workdistributeDoLower(omp::WorkdistributeOp workdistribute,388 SetVector<omp::TargetOp> &targetOpsToProcess) {389 OpBuilder rewriter(workdistribute);390 auto doLoop = getPerfectlyNested<fir::DoLoopOp>(workdistribute);391 auto wdLoc = workdistribute->getLoc();392 if (doLoop && shouldParallelize(doLoop)) {393 assert(doLoop.getReduceOperands().empty());394 395 // Record the target ops to process later396 if (auto teamsOp = dyn_cast<omp::TeamsOp>(workdistribute->getParentOp())) {397 auto targetOp = dyn_cast<omp::TargetOp>(teamsOp->getParentOp());398 if (targetOp) {399 targetOpsToProcess.insert(targetOp);400 }401 }402 // Generate the nested parallel, distribute, wsloop and loop_nest ops.403 genParallelOp(wdLoc, rewriter, true);404 genDistributeOp(wdLoc, rewriter, true);405 mlir::omp::LoopNestOperands loopNestClauseOps;406 genLoopNestClauseOps(rewriter, doLoop, loopNestClauseOps);407 genWsLoopOp(rewriter, doLoop, loopNestClauseOps, true);408 workdistribute.erase();409 return true;410 }411 return false;412}413 414/// Check if the enclosed type in fir.ref is fir.box and fir.box encloses array415static bool isEnclosedTypeRefToBoxArray(Type type) {416 // Check if it's a reference type417 if (auto refType = dyn_cast<fir::ReferenceType>(type)) {418 // Get the referenced type (should be fir.box)419 auto referencedType = refType.getEleTy();420 // Check if referenced type is a box421 if (auto boxType = dyn_cast<fir::BoxType>(referencedType)) {422 // Get the boxed type and check if it's an array423 auto boxedType = boxType.getEleTy();424 // Check if boxed type is a sequence (array)425 return isa<fir::SequenceType>(boxedType);426 }427 }428 return false;429}430 431/// Check if the enclosed type in fir.box is scalar (not array)432static bool isEnclosedTypeBoxScalar(Type type) {433 // Check if it's a box type434 if (auto boxType = dyn_cast<fir::BoxType>(type)) {435 // Get the boxed type436 auto boxedType = boxType.getEleTy();437 // Check if boxed type is NOT a sequence (array)438 return !isa<fir::SequenceType>(boxedType);439 }440 return false;441}442 443/// Check if the FortranAAssign call has src as scalar and dest as array444static bool isFortranAssignSrcScalarAndDestArray(fir::CallOp callOp) {445 if (callOp.getNumOperands() < 2)446 return false;447 auto srcArg = callOp.getOperand(1);448 auto destArg = callOp.getOperand(0);449 // Both operands should be fir.convert ops450 auto srcConvert = srcArg.getDefiningOp<fir::ConvertOp>();451 auto destConvert = destArg.getDefiningOp<fir::ConvertOp>();452 if (!srcConvert || !destConvert) {453 emitError(callOp->getLoc(),454 "Unimplemented: FortranAssign to OpenMP lowering\n");455 return false;456 }457 // Get the original types before conversion458 auto srcOrigType = srcConvert.getValue().getType();459 auto destOrigType = destConvert.getValue().getType();460 461 // Check if src is scalar and dest is array462 bool srcIsScalar = isEnclosedTypeBoxScalar(srcOrigType);463 bool destIsArray = isEnclosedTypeRefToBoxArray(destOrigType);464 return srcIsScalar && destIsArray;465}466 467/// Convert a flat index to multi-dimensional indices for an array box468/// Example: 2D array with shape (2,4)469/// Col 1 Col 2 Col 3 Col 4470/// Row 1: (1,1) (1,2) (1,3) (1,4)471/// Row 2: (2,1) (2,2) (2,3) (2,4)472///473/// extents: (2,4)474///475/// flatIdx: 0 1 2 3 4 5 6 7476/// Indices: (1,1) (1,2) (1,3) (1,4) (2,1) (2,2) (2,3) (2,4)477static SmallVector<Value> convertFlatToMultiDim(OpBuilder &builder,478 Location loc, Value flatIdx,479 Value arrayBox) {480 // Get array type and rank481 auto boxType = cast<fir::BoxType>(arrayBox.getType());482 auto seqType = cast<fir::SequenceType>(boxType.getEleTy());483 int rank = seqType.getDimension();484 485 // Get all extents486 SmallVector<Value> extents;487 // Get extents for each dimension488 for (int i = 0; i < rank; ++i) {489 auto dimIdx = arith::ConstantIndexOp::create(builder, loc, i);490 auto boxDims = fir::BoxDimsOp::create(builder, loc, arrayBox, dimIdx);491 extents.push_back(boxDims.getResult(1));492 }493 494 // Convert flat index to multi-dimensional indices495 SmallVector<Value> indices(rank);496 Value temp = flatIdx;497 auto c1 = arith::ConstantIndexOp::create(builder, loc, 1);498 499 // Work backwards through dimensions (row-major order)500 for (int i = rank - 1; i >= 0; --i) {501 Value zeroBasedIdx = arith::RemSIOp::create(builder, loc, temp, extents[i]);502 // Convert to one-based index503 indices[i] = arith::AddIOp::create(builder, loc, zeroBasedIdx, c1);504 if (i > 0) {505 temp = arith::DivSIOp::create(builder, loc, temp, extents[i]);506 }507 }508 509 return indices;510}511 512/// Calculate the total number of elements in the array box513/// (totalElems = extent(1) * extent(2) * ... * extent(n))514static Value CalculateTotalElements(OpBuilder &builder, Location loc,515 Value arrayBox) {516 auto boxType = cast<fir::BoxType>(arrayBox.getType());517 auto seqType = cast<fir::SequenceType>(boxType.getEleTy());518 int rank = seqType.getDimension();519 520 Value totalElems = nullptr;521 for (int i = 0; i < rank; ++i) {522 auto dimIdx = arith::ConstantIndexOp::create(builder, loc, i);523 auto boxDims = fir::BoxDimsOp::create(builder, loc, arrayBox, dimIdx);524 Value extent = boxDims.getResult(1);525 if (i == 0) {526 totalElems = extent;527 } else {528 totalElems = arith::MulIOp::create(builder, loc, totalElems, extent);529 }530 }531 return totalElems;532}533 534/// Replace the FortranAAssign runtime call with an unordered do loop535static void replaceWithUnorderedDoLoop(OpBuilder &builder, Location loc,536 omp::TeamsOp teamsOp,537 omp::WorkdistributeOp workdistribute,538 fir::CallOp callOp) {539 auto destConvert = callOp.getOperand(0).getDefiningOp<fir::ConvertOp>();540 auto srcConvert = callOp.getOperand(1).getDefiningOp<fir::ConvertOp>();541 542 Value destBox = destConvert.getValue();543 Value srcBox = srcConvert.getValue();544 545 // get defining alloca op of destBox and srcBox546 auto destAlloca = destBox.getDefiningOp<fir::AllocaOp>();547 548 if (!destAlloca) {549 emitError(loc, "Unimplemented: FortranAssign to OpenMP lowering\n");550 return;551 }552 553 // get the store op that stores to the alloca554 for (auto user : destAlloca->getUsers()) {555 if (auto storeOp = dyn_cast<fir::StoreOp>(user)) {556 destBox = storeOp.getValue();557 break;558 }559 }560 561 builder.setInsertionPoint(teamsOp);562 // Load destination array box (if it's a reference)563 Value arrayBox = destBox;564 if (isa<fir::ReferenceType>(destBox.getType()))565 arrayBox = fir::LoadOp::create(builder, loc, destBox);566 567 auto scalarValue = fir::BoxAddrOp::create(builder, loc, srcBox);568 Value scalar = fir::LoadOp::create(builder, loc, scalarValue);569 570 // Calculate total number of elements (flattened)571 auto c0 = arith::ConstantIndexOp::create(builder, loc, 0);572 auto c1 = arith::ConstantIndexOp::create(builder, loc, 1);573 Value totalElems = CalculateTotalElements(builder, loc, arrayBox);574 575 auto *workdistributeBlock = &workdistribute.getRegion().front();576 builder.setInsertionPointToStart(workdistributeBlock);577 // Create single unordered loop for flattened array578 auto doLoop = fir::DoLoopOp::create(builder, loc, c0, totalElems, c1, true);579 Block *loopBlock = &doLoop.getRegion().front();580 builder.setInsertionPointToStart(doLoop.getBody());581 582 auto flatIdx = loopBlock->getArgument(0);583 SmallVector<Value> indices =584 convertFlatToMultiDim(builder, loc, flatIdx, arrayBox);585 // Use fir.array_coor for linear addressing586 auto elemPtr = fir::ArrayCoorOp::create(587 builder, loc, fir::ReferenceType::get(scalar.getType()), arrayBox,588 nullptr, nullptr, ValueRange{indices}, ValueRange{});589 590 fir::StoreOp::create(builder, loc, scalar, elemPtr);591}592 593/// workdistributeRuntimeCallLower method finds the runtime calls594/// nested in teams {workdistribute{}} and595/// lowers FortranAAssign to unordered do loop if src is scalar and dest is596/// array. Other runtime calls are not handled currently.597static FailureOr<bool>598workdistributeRuntimeCallLower(omp::WorkdistributeOp workdistribute,599 SetVector<omp::TargetOp> &targetOpsToProcess) {600 OpBuilder rewriter(workdistribute);601 auto loc = workdistribute->getLoc();602 auto teams = dyn_cast<omp::TeamsOp>(workdistribute->getParentOp());603 if (!teams) {604 emitError(loc, "workdistribute not nested in teams\n");605 return failure();606 }607 if (workdistribute.getRegion().getBlocks().size() != 1) {608 emitError(loc, "workdistribute with multiple blocks\n");609 return failure();610 }611 if (teams.getRegion().getBlocks().size() != 1) {612 emitError(loc, "teams with multiple blocks\n");613 return failure();614 }615 bool changed = false;616 // Get the target op parent of teams617 omp::TargetOp targetOp = dyn_cast<omp::TargetOp>(teams->getParentOp());618 SmallVector<Operation *> opsToErase;619 for (auto &op : workdistribute.getOps()) {620 if (isRuntimeCall(&op)) {621 rewriter.setInsertionPoint(&op);622 fir::CallOp runtimeCall = cast<fir::CallOp>(op);623 auto funcName = runtimeCall.getCallee()->getRootReference().getValue();624 if (funcName == FortranAssignStr) {625 if (isFortranAssignSrcScalarAndDestArray(runtimeCall) && targetOp) {626 // Record the target ops to process later627 targetOpsToProcess.insert(targetOp);628 replaceWithUnorderedDoLoop(rewriter, loc, teams, workdistribute,629 runtimeCall);630 opsToErase.push_back(&op);631 changed = true;632 }633 }634 }635 }636 // Erase the runtime calls that have been replaced.637 for (auto *op : opsToErase) {638 op->erase();639 }640 return changed;641}642 643/// teamsWorkdistributeToSingleOp method hoists all the ops inside644/// teams {workdistribute{}} before teams op.645///646/// If A() and B () are present inside teams workdistribute647///648/// omp.teams {649/// omp.workdistribute {650/// A()651/// B()652/// }653/// }654///655/// Then, its lowered to656///657/// A()658/// B()659///660/// If only the terminator remains in teams after hoisting, we erase teams op.661static bool662teamsWorkdistributeToSingleOp(omp::TeamsOp teamsOp,663 SetVector<omp::TargetOp> &targetOpsToProcess) {664 auto workdistributeOp = getPerfectlyNested<omp::WorkdistributeOp>(teamsOp);665 if (!workdistributeOp)666 return false;667 // Get the block containing teamsOp (the parent block).668 Block *parentBlock = teamsOp->getBlock();669 Block &workdistributeBlock = *workdistributeOp.getRegion().begin();670 // Record the target ops to process later671 for (auto &op : workdistributeBlock.getOperations()) {672 if (shouldParallelize(&op)) {673 auto targetOp = dyn_cast<omp::TargetOp>(teamsOp->getParentOp());674 if (targetOp) {675 targetOpsToProcess.insert(targetOp);676 }677 }678 }679 auto insertPoint = Block::iterator(teamsOp);680 // Get the range of operations to move (excluding the terminator).681 auto workdistributeBegin = workdistributeBlock.begin();682 auto workdistributeEnd = workdistributeBlock.getTerminator()->getIterator();683 // Move the operations from workdistribute block to before teamsOp.684 parentBlock->getOperations().splice(insertPoint,685 workdistributeBlock.getOperations(),686 workdistributeBegin, workdistributeEnd);687 // Erase the now-empty workdistributeOp.688 workdistributeOp.erase();689 Block &teamsBlock = *teamsOp.getRegion().begin();690 // Check if only the terminator remains and erase teams op.691 if (teamsBlock.getOperations().size() == 1 &&692 teamsBlock.getTerminator() != nullptr) {693 teamsOp.erase();694 }695 return true;696}697 698/// If multiple workdistribute are nested in a target regions, we will need to699/// split the target region, but we want to preserve the data semantics of the700/// original data region and avoid unnecessary data movement at each of the701/// subkernels - we split the target region into a target_data{target}702/// nest where only the outer one moves the data703FailureOr<omp::TargetOp> splitTargetData(omp::TargetOp targetOp,704 RewriterBase &rewriter) {705 auto loc = targetOp->getLoc();706 if (targetOp.getMapVars().empty()) {707 emitError(loc, "Target region has no data maps\n");708 return failure();709 }710 // Collect all the mapinfo ops711 SmallVector<omp::MapInfoOp> mapInfos;712 for (auto opr : targetOp.getMapVars()) {713 auto mapInfo = cast<omp::MapInfoOp>(opr.getDefiningOp());714 mapInfos.push_back(mapInfo);715 }716 717 rewriter.setInsertionPoint(targetOp);718 SmallVector<Value> innerMapInfos;719 SmallVector<Value> outerMapInfos;720 // Create new mapinfo ops for the inner target region721 for (auto mapInfo : mapInfos) {722 mlir::omp::ClauseMapFlags originalMapType = mapInfo.getMapType();723 auto originalCaptureType = mapInfo.getMapCaptureType();724 mlir::omp::ClauseMapFlags newMapType;725 mlir::omp::VariableCaptureKind newCaptureType;726 // For bycopy, we keep the same map type and capture type727 // For byref, we change the map type to none and keep the capture type728 if (originalCaptureType == mlir::omp::VariableCaptureKind::ByCopy) {729 newMapType = originalMapType;730 newCaptureType = originalCaptureType;731 } else if (originalCaptureType == mlir::omp::VariableCaptureKind::ByRef) {732 newMapType = mlir::omp::ClauseMapFlags::storage;733 newCaptureType = originalCaptureType;734 outerMapInfos.push_back(mapInfo);735 } else {736 emitError(targetOp->getLoc(), "Unhandled case");737 return failure();738 }739 auto innerMapInfo = cast<omp::MapInfoOp>(rewriter.clone(*mapInfo));740 innerMapInfo.setMapTypeAttr(741 rewriter.getAttr<omp::ClauseMapFlagsAttr>(newMapType));742 innerMapInfo.setMapCaptureType(newCaptureType);743 innerMapInfos.push_back(innerMapInfo.getResult());744 }745 746 rewriter.setInsertionPoint(targetOp);747 auto device = targetOp.getDevice();748 auto ifExpr = targetOp.getIfExpr();749 auto deviceAddrVars = targetOp.getHasDeviceAddrVars();750 auto devicePtrVars = targetOp.getIsDevicePtrVars();751 // Create the target data op752 auto targetDataOp =753 omp::TargetDataOp::create(rewriter, loc, device, ifExpr, outerMapInfos,754 deviceAddrVars, devicePtrVars);755 auto taregtDataBlock = rewriter.createBlock(&targetDataOp.getRegion());756 mlir::omp::TerminatorOp::create(rewriter, loc);757 rewriter.setInsertionPointToStart(taregtDataBlock);758 // Create the inner target op759 auto newTargetOp = omp::TargetOp::create(760 rewriter, targetOp.getLoc(), targetOp.getAllocateVars(),761 targetOp.getAllocatorVars(), targetOp.getBareAttr(),762 targetOp.getDependKindsAttr(), targetOp.getDependVars(),763 targetOp.getDevice(), targetOp.getHasDeviceAddrVars(),764 targetOp.getHostEvalVars(), targetOp.getIfExpr(),765 targetOp.getInReductionVars(), targetOp.getInReductionByrefAttr(),766 targetOp.getInReductionSymsAttr(), targetOp.getIsDevicePtrVars(),767 innerMapInfos, targetOp.getNowaitAttr(), targetOp.getPrivateVars(),768 targetOp.getPrivateSymsAttr(), targetOp.getPrivateNeedsBarrierAttr(),769 targetOp.getThreadLimit(), targetOp.getPrivateMapsAttr());770 rewriter.inlineRegionBefore(targetOp.getRegion(), newTargetOp.getRegion(),771 newTargetOp.getRegion().begin());772 rewriter.replaceOp(targetOp, targetDataOp);773 return newTargetOp;774}775 776/// getNestedOpToIsolate function is designed to identify a specific teams777/// parallel op within the body of an omp::TargetOp that should be "isolated."778/// This returns a tuple of op, if its first op in targetBlock, or if the op is779/// last op in the traget block.780static std::optional<std::tuple<Operation *, bool, bool>>781getNestedOpToIsolate(omp::TargetOp targetOp) {782 if (targetOp.getRegion().empty())783 return std::nullopt;784 auto *targetBlock = &targetOp.getRegion().front();785 for (auto &op : *targetBlock) {786 bool first = &op == &*targetBlock->begin();787 bool last = op.getNextNode() == targetBlock->getTerminator();788 if (first && last)789 return std::nullopt;790 791 if (isa<omp::TeamsOp>(&op))792 return {{&op, first, last}};793 }794 return std::nullopt;795}796 797/// Temporary structure to hold the two mapinfo ops798struct TempOmpVar {799 omp::MapInfoOp from, to;800};801 802/// isPtr checks if the type is a pointer or reference type.803static bool isPtr(Type ty) {804 return isa<fir::ReferenceType>(ty) || isa<LLVM::LLVMPointerType>(ty);805}806 807/// getPtrTypeForOmp returns an LLVM pointer type for the given type.808static Type getPtrTypeForOmp(Type ty) {809 if (isPtr(ty))810 return LLVM::LLVMPointerType::get(ty.getContext());811 else812 return fir::ReferenceType::get(ty);813}814 815/// allocateTempOmpVar allocates a temporary variable for OpenMP mapping816static TempOmpVar allocateTempOmpVar(Location loc, Type ty,817 RewriterBase &rewriter) {818 MLIRContext &ctx = *ty.getContext();819 Value alloc;820 Type allocType;821 auto llvmPtrTy = LLVM::LLVMPointerType::get(&ctx);822 // Get the appropriate type for allocation823 if (isPtr(ty)) {824 Type intTy = rewriter.getI32Type();825 auto one = LLVM::ConstantOp::create(rewriter, loc, intTy, 1);826 allocType = llvmPtrTy;827 alloc = LLVM::AllocaOp::create(rewriter, loc, llvmPtrTy, allocType, one);828 allocType = intTy;829 } else {830 allocType = ty;831 alloc = fir::AllocaOp::create(rewriter, loc, allocType);832 }833 // Lambda to create mapinfo ops834 auto getMapInfo = [&](mlir::omp::ClauseMapFlags mappingFlags,835 const char *name) {836 return omp::MapInfoOp::create(837 rewriter, loc, alloc.getType(), alloc, TypeAttr::get(allocType),838 rewriter.getAttr<omp::ClauseMapFlagsAttr>(mappingFlags),839 rewriter.getAttr<omp::VariableCaptureKindAttr>(840 omp::VariableCaptureKind::ByRef),841 /*varPtrPtr=*/Value{},842 /*members=*/SmallVector<Value>{},843 /*member_index=*/mlir::ArrayAttr{},844 /*bounds=*/ValueRange(),845 /*mapperId=*/mlir::FlatSymbolRefAttr(),846 /*name=*/rewriter.getStringAttr(name), rewriter.getBoolAttr(false));847 };848 // Create mapinfo ops.849 auto mapInfoFrom = getMapInfo(mlir::omp::ClauseMapFlags::from,850 "__flang_workdistribute_from");851 auto mapInfoTo =852 getMapInfo(mlir::omp::ClauseMapFlags::to, "__flang_workdistribute_to");853 return TempOmpVar{mapInfoFrom, mapInfoTo};854}855 856// usedOutsideSplit checks if a value is used outside the split operation.857static bool usedOutsideSplit(Value v, Operation *split) {858 if (!split)859 return false;860 auto targetOp = cast<omp::TargetOp>(split->getParentOp());861 auto *targetBlock = &targetOp.getRegion().front();862 for (auto *user : v.getUsers()) {863 while (user->getBlock() != targetBlock) {864 user = user->getParentOp();865 }866 if (!user->isBeforeInBlock(split))867 return true;868 }869 return false;870}871 872/// isRecomputableAfterFission checks if an operation can be recomputed873static bool isRecomputableAfterFission(Operation *op, Operation *splitBefore) {874 // If the op has side effects, it cannot be recomputed.875 // We consider fir.declare as having no side effects.876 return isa<fir::DeclareOp>(op) || isMemoryEffectFree(op);877}878 879/// collectNonRecomputableDeps collects dependencies that cannot be recomputed880static void collectNonRecomputableDeps(Value &v, omp::TargetOp targetOp,881 SetVector<Operation *> &nonRecomputable,882 SetVector<Operation *> &toCache,883 SetVector<Operation *> &toRecompute) {884 Operation *op = v.getDefiningOp();885 // If v is a block argument, it must be from the targetOp.886 if (!op) {887 assert(cast<BlockArgument>(v).getOwner()->getParentOp() == targetOp);888 return;889 }890 // If the op is in the nonRecomputable set, add it to toCache and return.891 if (nonRecomputable.contains(op)) {892 toCache.insert(op);893 return;894 }895 // Add the op to toRecompute.896 toRecompute.insert(op);897 for (auto opr : op->getOperands())898 collectNonRecomputableDeps(opr, targetOp, nonRecomputable, toCache,899 toRecompute);900}901 902/// createBlockArgsAndMap creates block arguments and maps them903static void createBlockArgsAndMap(Location loc, RewriterBase &rewriter,904 omp::TargetOp &targetOp, Block *targetBlock,905 Block *newTargetBlock,906 SmallVector<Value> &hostEvalVars,907 SmallVector<Value> &mapOperands,908 SmallVector<Value> &allocs,909 IRMapping &irMapping) {910 // FIRST: Map `host_eval_vars` to block arguments911 unsigned originalHostEvalVarsSize = targetOp.getHostEvalVars().size();912 for (unsigned i = 0; i < hostEvalVars.size(); ++i) {913 Value originalValue;914 BlockArgument newArg;915 if (i < originalHostEvalVarsSize) {916 originalValue = targetBlock->getArgument(i); // Host_eval args come first917 newArg = newTargetBlock->addArgument(originalValue.getType(),918 originalValue.getLoc());919 } else {920 originalValue = hostEvalVars[i];921 newArg = newTargetBlock->addArgument(originalValue.getType(),922 originalValue.getLoc());923 }924 irMapping.map(originalValue, newArg);925 }926 927 // SECOND: Map `map_operands` to block arguments928 unsigned originalMapVarsSize = targetOp.getMapVars().size();929 for (unsigned i = 0; i < mapOperands.size(); ++i) {930 Value originalValue;931 BlockArgument newArg;932 // Map the new arguments from the original block.933 if (i < originalMapVarsSize) {934 originalValue = targetBlock->getArgument(originalHostEvalVarsSize +935 i); // Offset by host_eval count936 newArg = newTargetBlock->addArgument(originalValue.getType(),937 originalValue.getLoc());938 }939 // Map the new arguments from the `allocs`.940 else {941 originalValue = allocs[i - originalMapVarsSize];942 newArg = newTargetBlock->addArgument(943 getPtrTypeForOmp(originalValue.getType()), originalValue.getLoc());944 }945 irMapping.map(originalValue, newArg);946 }947 948 // THIRD: Map `private_vars` to block arguments (if any)949 unsigned originalPrivateVarsSize = targetOp.getPrivateVars().size();950 for (unsigned i = 0; i < originalPrivateVarsSize; ++i) {951 auto originalArg = targetBlock->getArgument(originalHostEvalVarsSize +952 originalMapVarsSize + i);953 auto newArg = newTargetBlock->addArgument(originalArg.getType(),954 originalArg.getLoc());955 irMapping.map(originalArg, newArg);956 }957 return;958}959 960/// reloadCacheAndRecompute reloads cached values and recomputes operations961static void reloadCacheAndRecompute(962 Location loc, RewriterBase &rewriter, Operation *splitBefore,963 omp::TargetOp &targetOp, Block *targetBlock, Block *newTargetBlock,964 SmallVector<Value> &hostEvalVars, SmallVector<Value> &mapOperands,965 SmallVector<Value> &allocs, SetVector<Operation *> &toRecompute,966 IRMapping &irMapping) {967 // Handle the load operations for the allocs.968 rewriter.setInsertionPointToStart(newTargetBlock);969 auto llvmPtrTy = LLVM::LLVMPointerType::get(targetOp.getContext());970 971 unsigned originalMapVarsSize = targetOp.getMapVars().size();972 unsigned hostEvalVarsSize = hostEvalVars.size();973 // Create load operations for each allocated variable.974 for (unsigned i = 0; i < allocs.size(); ++i) {975 Value original = allocs[i];976 // Get the new block argument for this specific allocated value.977 Value newArg =978 newTargetBlock->getArgument(hostEvalVarsSize + originalMapVarsSize + i);979 Value restored;980 // If the original value is a pointer or reference, load and convert if981 // necessary.982 if (isPtr(original.getType())) {983 restored = LLVM::LoadOp::create(rewriter, loc, llvmPtrTy, newArg);984 if (!isa<LLVM::LLVMPointerType>(original.getType()))985 restored =986 fir::ConvertOp::create(rewriter, loc, original.getType(), restored);987 } else {988 restored = fir::LoadOp::create(rewriter, loc, newArg);989 }990 irMapping.map(original, restored);991 }992 // Clone the operations if they are in the toRecompute set.993 for (auto it = targetBlock->begin(); it != splitBefore->getIterator(); it++) {994 if (toRecompute.contains(&*it))995 rewriter.clone(*it, irMapping);996 }997}998 999/// Given a teamsOp, navigate down the nested structure to find the1000/// innermost LoopNestOp. The expected nesting is:1001/// teams -> parallel -> distribute -> wsloop -> loop_nest1002static mlir::omp::LoopNestOp getLoopNestFromTeams(mlir::omp::TeamsOp teamsOp) {1003 if (teamsOp.getRegion().empty())1004 return nullptr;1005 // Ensure the teams region has a single block.1006 if (teamsOp.getRegion().getBlocks().size() != 1)1007 return nullptr;1008 // Find parallel op inside teams1009 mlir::omp::ParallelOp parallelOp = nullptr;1010 // Look for the parallel op in the teams region1011 for (auto &op : teamsOp.getRegion().front()) {1012 if (auto parallel = dyn_cast<mlir::omp::ParallelOp>(op)) {1013 parallelOp = parallel;1014 break;1015 }1016 }1017 if (!parallelOp)1018 return nullptr;1019 1020 // Find distribute op inside parallel1021 mlir::omp::DistributeOp distributeOp = nullptr;1022 for (auto &op : parallelOp.getRegion().front()) {1023 if (auto distribute = dyn_cast<mlir::omp::DistributeOp>(op)) {1024 distributeOp = distribute;1025 break;1026 }1027 }1028 if (!distributeOp)1029 return nullptr;1030 1031 // Find wsloop op inside distribute1032 mlir::omp::WsloopOp wsloopOp = nullptr;1033 for (auto &op : distributeOp.getRegion().front()) {1034 if (auto wsloop = dyn_cast<mlir::omp::WsloopOp>(op)) {1035 wsloopOp = wsloop;1036 break;1037 }1038 }1039 if (!wsloopOp)1040 return nullptr;1041 1042 // Find loop_nest op inside wsloop1043 for (auto &op : wsloopOp.getRegion().front()) {1044 if (auto loopNest = dyn_cast<mlir::omp::LoopNestOp>(op)) {1045 return loopNest;1046 }1047 }1048 1049 return nullptr;1050}1051 1052/// Generate LLVM constant operations for i32 and i64 types.1053static mlir::LLVM::ConstantOp1054genI32Constant(mlir::Location loc, mlir::RewriterBase &rewriter, int value) {1055 mlir::Type i32Ty = rewriter.getI32Type();1056 mlir::IntegerAttr attr = rewriter.getI32IntegerAttr(value);1057 return mlir::LLVM::ConstantOp::create(rewriter, loc, i32Ty, attr);1058}1059 1060/// Given a box descriptor, extract the base address of the data it describes.1061/// If the box descriptor is a reference, load it first.1062/// The base address is returned as an i8* pointer.1063static Value genDescriptorGetBaseAddress(fir::FirOpBuilder &builder,1064 Location loc, Value boxDesc) {1065 Value box = boxDesc;1066 if (auto refBox = dyn_cast<fir::ReferenceType>(boxDesc.getType())) {1067 box = fir::LoadOp::create(builder, loc, boxDesc);1068 }1069 assert(isa<fir::BoxType>(box.getType()) &&1070 "Unknown type passed to genDescriptorGetBaseAddress");1071 auto i8Type = builder.getI8Type();1072 auto unknownArrayType =1073 fir::SequenceType::get({fir::SequenceType::getUnknownExtent()}, i8Type);1074 auto i8BoxType = fir::BoxType::get(unknownArrayType);1075 auto typedBox = fir::ConvertOp::create(builder, loc, i8BoxType, box);1076 auto rawAddr = fir::BoxAddrOp::create(builder, loc, typedBox);1077 return rawAddr;1078}1079 1080/// Given a box descriptor, extract the total number of elements in the array it1081/// describes. If the box descriptor is a reference, load it first.1082/// The total number of elements is returned as an i64 value.1083static Value genDescriptorGetTotalElements(fir::FirOpBuilder &builder,1084 Location loc, Value boxDesc) {1085 Value box = boxDesc;1086 if (auto refBox = dyn_cast<fir::ReferenceType>(boxDesc.getType())) {1087 box = fir::LoadOp::create(builder, loc, boxDesc);1088 }1089 assert(isa<fir::BoxType>(box.getType()) &&1090 "Unknown type passed to genDescriptorGetTotalElements");1091 auto i64Type = builder.getI64Type();1092 return fir::BoxTotalElementsOp::create(builder, loc, i64Type, box);1093}1094 1095/// Given a box descriptor, extract the size of each element in the array it1096/// describes. If the box descriptor is a reference, load it first.1097/// The element size is returned as an i64 value.1098static Value genDescriptorGetEleSize(fir::FirOpBuilder &builder, Location loc,1099 Value boxDesc) {1100 Value box = boxDesc;1101 if (auto refBox = dyn_cast<fir::ReferenceType>(boxDesc.getType())) {1102 box = fir::LoadOp::create(builder, loc, boxDesc);1103 }1104 assert(isa<fir::BoxType>(box.getType()) &&1105 "Unknown type passed to genDescriptorGetElementSize");1106 auto i64Type = builder.getI64Type();1107 return fir::BoxEleSizeOp::create(builder, loc, i64Type, box);1108}1109 1110/// Given a box descriptor, compute the total size in bytes of the data it1111/// describes. This is done by multiplying the total number of elements by the1112/// size of each element. If the box descriptor is a reference, load it first.1113/// The total size in bytes is returned as an i64 value.1114static Value genDescriptorGetDataSizeInBytes(fir::FirOpBuilder &builder,1115 Location loc, Value boxDesc) {1116 Value box = boxDesc;1117 if (auto refBox = dyn_cast<fir::ReferenceType>(boxDesc.getType())) {1118 box = fir::LoadOp::create(builder, loc, boxDesc);1119 }1120 assert(isa<fir::BoxType>(box.getType()) &&1121 "Unknown type passed to genDescriptorGetElementSize");1122 Value eleSize = genDescriptorGetEleSize(builder, loc, box);1123 Value totalElements = genDescriptorGetTotalElements(builder, loc, box);1124 return mlir::arith::MulIOp::create(builder, loc, totalElements, eleSize);1125}1126 1127/// Generate a call to the OpenMP runtime function `omp_get_mapped_ptr` to1128/// retrieve the device pointer corresponding to a given host pointer and device1129/// number. If no mapping exists, the original host pointer is returned.1130/// Signature:1131/// void *omp_get_mapped_ptr(void *host_ptr, int device_num);1132static mlir::Value genOmpGetMappedPtrIfPresent(fir::FirOpBuilder &builder,1133 mlir::Location loc,1134 mlir::Value hostPtr,1135 mlir::Value deviceNum,1136 mlir::ModuleOp module) {1137 auto *context = builder.getContext();1138 auto voidPtrType = fir::LLVMPointerType::get(context, builder.getI8Type());1139 auto i32Type = builder.getI32Type();1140 auto funcName = "omp_get_mapped_ptr";1141 auto funcOp = module.lookupSymbol<mlir::func::FuncOp>(funcName);1142 1143 if (!funcOp) {1144 auto funcType =1145 mlir::FunctionType::get(context, {voidPtrType, i32Type}, {voidPtrType});1146 1147 mlir::OpBuilder::InsertionGuard guard(builder);1148 builder.setInsertionPointToStart(module.getBody());1149 1150 funcOp = mlir::func::FuncOp::create(builder, loc, funcName, funcType);1151 funcOp.setPrivate();1152 }1153 1154 llvm::SmallVector<mlir::Value> args;1155 args.push_back(fir::ConvertOp::create(builder, loc, voidPtrType, hostPtr));1156 args.push_back(fir::ConvertOp::create(builder, loc, i32Type, deviceNum));1157 auto callOp = fir::CallOp::create(builder, loc, funcOp, args);1158 auto mappedPtr = callOp.getResult(0);1159 auto isNull = builder.genIsNullAddr(loc, mappedPtr);1160 auto convertedHostPtr =1161 fir::ConvertOp::create(builder, loc, voidPtrType, hostPtr);1162 auto result = arith::SelectOp::create(builder, loc, isNull, convertedHostPtr,1163 mappedPtr);1164 return result;1165}1166 1167/// Generate a call to the OpenMP runtime function `omp_target_memcpy` to1168/// perform memory copy between host and device or between devices.1169/// Signature:1170/// int omp_target_memcpy(void *dst, const void *src, size_t length,1171/// size_t dst_offset, size_t src_offset,1172/// int dst_device, int src_device);1173static void genOmpTargetMemcpyCall(fir::FirOpBuilder &builder,1174 mlir::Location loc, mlir::Value dst,1175 mlir::Value src, mlir::Value length,1176 mlir::Value dstOffset, mlir::Value srcOffset,1177 mlir::Value device, mlir::ModuleOp module) {1178 auto *context = builder.getContext();1179 auto funcName = "omp_target_memcpy";1180 auto voidPtrType = fir::LLVMPointerType::get(context, builder.getI8Type());1181 auto sizeTType = builder.getI64Type(); // assuming size_t is 64-bit1182 auto i32Type = builder.getI32Type();1183 auto funcOp = module.lookupSymbol<mlir::func::FuncOp>(funcName);1184 1185 if (!funcOp) {1186 mlir::OpBuilder::InsertionGuard guard(builder);1187 builder.setInsertionPointToStart(module.getBody());1188 llvm::SmallVector<mlir::Type> argTypes = {1189 voidPtrType, voidPtrType, sizeTType, sizeTType,1190 sizeTType, i32Type, i32Type};1191 auto funcType = mlir::FunctionType::get(context, argTypes, {i32Type});1192 funcOp = mlir::func::FuncOp::create(builder, loc, funcName, funcType);1193 funcOp.setPrivate();1194 }1195 1196 llvm::SmallVector<mlir::Value> args{dst, src, length, dstOffset,1197 srcOffset, device, device};1198 fir::CallOp::create(builder, loc, funcOp, args);1199 return;1200}1201 1202/// Generate code to replace a Fortran array assignment call with OpenMP1203/// runtime calls to perform the equivalent operation on the device.1204/// This involves extracting the source and destination pointers from the1205/// Fortran array descriptors, retrieving their mapped device pointers (if any),1206/// and invoking `omp_target_memcpy` to copy the data on the device.1207static void genFortranAssignOmpReplacement(fir::FirOpBuilder &builder,1208 mlir::Location loc,1209 fir::CallOp callOp,1210 mlir::Value device,1211 mlir::ModuleOp module) {1212 assert(callOp.getNumResults() == 0 &&1213 "Expected _FortranAAssign to have no results");1214 assert(callOp.getNumOperands() >= 2 &&1215 "Expected _FortranAAssign to have at least two operands");1216 1217 // Extract the source and destination pointers from the call operands.1218 mlir::Value dest = callOp.getOperand(0);1219 mlir::Value src = callOp.getOperand(1);1220 1221 // Get the base addresses of the source and destination arrays.1222 mlir::Value srcBase = genDescriptorGetBaseAddress(builder, loc, src);1223 mlir::Value destBase = genDescriptorGetBaseAddress(builder, loc, dest);1224 1225 // Get the total size in bytes of the data to be copied.1226 mlir::Value srcDataSize = genDescriptorGetDataSizeInBytes(builder, loc, src);1227 1228 // Retrieve the mapped device pointers for source and destination.1229 // If no mapping exists, the original host pointer is used.1230 Value destPtr =1231 genOmpGetMappedPtrIfPresent(builder, loc, destBase, device, module);1232 Value srcPtr =1233 genOmpGetMappedPtrIfPresent(builder, loc, srcBase, device, module);1234 Value zero = LLVM::ConstantOp::create(builder, loc, builder.getI64Type(),1235 builder.getI64IntegerAttr(0));1236 1237 // Generate the call to omp_target_memcpy to perform the data copy on the1238 // device.1239 genOmpTargetMemcpyCall(builder, loc, destPtr, srcPtr, srcDataSize, zero, zero,1240 device, module);1241}1242 1243/// Struct to hold the host eval vars corresponding to loop bounds and steps1244struct HostEvalVars {1245 SmallVector<Value> lbs;1246 SmallVector<Value> ubs;1247 SmallVector<Value> steps;1248};1249 1250/// moveToHost method clones all the ops from target region outside of it.1251/// It hoists runtime function "_FortranAAssign" and replaces it with omp1252/// version. Also hoists and replaces fir.allocmem with omp.target_allocmem and1253/// fir.freemem with omp.target_freemem1254static LogicalResult moveToHost(omp::TargetOp targetOp, RewriterBase &rewriter,1255 mlir::ModuleOp module,1256 struct HostEvalVars &hostEvalVars) {1257 OpBuilder::InsertionGuard guard(rewriter);1258 Block *targetBlock = &targetOp.getRegion().front();1259 assert(targetBlock == &targetOp.getRegion().back());1260 IRMapping mapping;1261 1262 // Get the parent target_data op1263 auto targetDataOp = cast<omp::TargetDataOp>(targetOp->getParentOp());1264 if (!targetDataOp) {1265 emitError(targetOp->getLoc(),1266 "Expected target op to be inside target_data op");1267 return failure();1268 }1269 // create mapping for host_eval_vars1270 unsigned hostEvalVarCount = targetOp.getHostEvalVars().size();1271 for (unsigned i = 0; i < targetOp.getHostEvalVars().size(); ++i) {1272 Value hostEvalVar = targetOp.getHostEvalVars()[i];1273 BlockArgument arg = targetBlock->getArguments()[i];1274 mapping.map(arg, hostEvalVar);1275 }1276 // create mapping for map_vars1277 for (unsigned i = 0; i < targetOp.getMapVars().size(); ++i) {1278 Value mapInfo = targetOp.getMapVars()[i];1279 BlockArgument arg = targetBlock->getArguments()[hostEvalVarCount + i];1280 Operation *op = mapInfo.getDefiningOp();1281 assert(op);1282 auto mapInfoOp = cast<omp::MapInfoOp>(op);1283 // map the block argument to the host-side variable pointer1284 mapping.map(arg, mapInfoOp.getVarPtr());1285 }1286 // create mapping for private_vars1287 unsigned mapSize = targetOp.getMapVars().size();1288 for (unsigned i = 0; i < targetOp.getPrivateVars().size(); ++i) {1289 Value privateVar = targetOp.getPrivateVars()[i];1290 // The mapping should link the device-side variable to the host-side one.1291 BlockArgument arg =1292 targetBlock->getArguments()[hostEvalVarCount + mapSize + i];1293 // Map the device-side copy (`arg`) to the host-side value (`privateVar`).1294 mapping.map(arg, privateVar);1295 }1296 1297 rewriter.setInsertionPoint(targetOp);1298 SmallVector<Operation *> opsToReplace;1299 Value device = targetOp.getDevice();1300 1301 // If device is not specified, default to device 0.1302 if (!device) {1303 device = genI32Constant(targetOp.getLoc(), rewriter, 0);1304 }1305 // Clone all operations.1306 for (auto it = targetBlock->begin(), end = std::prev(targetBlock->end());1307 it != end; ++it) {1308 auto *op = &*it;1309 Operation *clonedOp = rewriter.clone(*op, mapping);1310 // Map the results of the original op to the cloned op.1311 for (unsigned i = 0; i < op->getNumResults(); ++i) {1312 mapping.map(op->getResult(i), clonedOp->getResult(i));1313 }1314 // fir.declare changes its type when hoisting it out of omp.target to1315 // omp.target_data Introduce a load, if original declareOp input is not of1316 // reference type, but cloned delcareOp input is reference type.1317 if (fir::DeclareOp clonedDeclareOp = dyn_cast<fir::DeclareOp>(clonedOp)) {1318 auto originalDeclareOp = cast<fir::DeclareOp>(op);1319 Type originalInType = originalDeclareOp.getMemref().getType();1320 Type clonedInType = clonedDeclareOp.getMemref().getType();1321 1322 fir::ReferenceType originalRefType =1323 dyn_cast<fir::ReferenceType>(originalInType);1324 fir::ReferenceType clonedRefType =1325 dyn_cast<fir::ReferenceType>(clonedInType);1326 if (!originalRefType && clonedRefType) {1327 Type clonedEleTy = clonedRefType.getElementType();1328 if (clonedEleTy == originalDeclareOp.getType()) {1329 opsToReplace.push_back(clonedOp);1330 }1331 }1332 }1333 // Collect the ops to be replaced.1334 if (isa<fir::AllocMemOp>(clonedOp) || isa<fir::FreeMemOp>(clonedOp))1335 opsToReplace.push_back(clonedOp);1336 // Check for runtime calls to be replaced.1337 if (isRuntimeCall(clonedOp)) {1338 fir::CallOp runtimeCall = cast<fir::CallOp>(op);1339 auto funcName = runtimeCall.getCallee()->getRootReference().getValue();1340 if (funcName == FortranAssignStr) {1341 opsToReplace.push_back(clonedOp);1342 } else {1343 emitError(runtimeCall->getLoc(), "Unhandled runtime call hoisting.");1344 return failure();1345 }1346 }1347 }1348 // Replace fir.allocmem with omp.target_allocmem.1349 for (Operation *op : opsToReplace) {1350 if (auto allocOp = dyn_cast<fir::AllocMemOp>(op)) {1351 rewriter.setInsertionPoint(allocOp);1352 auto ompAllocmemOp = omp::TargetAllocMemOp::create(1353 rewriter, allocOp.getLoc(), rewriter.getI64Type(), device,1354 allocOp.getInTypeAttr(), allocOp.getUniqNameAttr(),1355 allocOp.getBindcNameAttr(), allocOp.getTypeparams(),1356 allocOp.getShape());1357 auto firConvertOp = fir::ConvertOp::create(rewriter, allocOp.getLoc(),1358 allocOp.getResult().getType(),1359 ompAllocmemOp.getResult());1360 rewriter.replaceOp(allocOp, firConvertOp.getResult());1361 }1362 // Replace fir.freemem with omp.target_freemem.1363 else if (auto freeOp = dyn_cast<fir::FreeMemOp>(op)) {1364 rewriter.setInsertionPoint(freeOp);1365 auto firConvertOp =1366 fir::ConvertOp::create(rewriter, freeOp.getLoc(),1367 rewriter.getI64Type(), freeOp.getHeapref());1368 omp::TargetFreeMemOp::create(rewriter, freeOp.getLoc(), device,1369 firConvertOp.getResult());1370 rewriter.eraseOp(freeOp);1371 }1372 // fir.declare changes its type when hoisting it out of omp.target to1373 // omp.target_data Introduce a load, if original declareOp input is not of1374 // reference type, but cloned delcareOp input is reference type.1375 else if (fir::DeclareOp clonedDeclareOp = dyn_cast<fir::DeclareOp>(op)) {1376 Type clonedInType = clonedDeclareOp.getMemref().getType();1377 fir::ReferenceType clonedRefType =1378 dyn_cast<fir::ReferenceType>(clonedInType);1379 Type clonedEleTy = clonedRefType.getElementType();1380 rewriter.setInsertionPoint(op);1381 Value loadedValue =1382 fir::LoadOp::create(rewriter, clonedDeclareOp.getLoc(), clonedEleTy,1383 clonedDeclareOp.getMemref());1384 clonedDeclareOp.getResult().replaceAllUsesWith(loadedValue);1385 }1386 // Replace runtime calls with omp versions.1387 else if (isRuntimeCall(op)) {1388 fir::CallOp runtimeCall = cast<fir::CallOp>(op);1389 auto funcName = runtimeCall.getCallee()->getRootReference().getValue();1390 if (funcName == FortranAssignStr) {1391 rewriter.setInsertionPoint(op);1392 fir::FirOpBuilder builder{rewriter, op};1393 1394 mlir::Location loc = runtimeCall.getLoc();1395 genFortranAssignOmpReplacement(builder, loc, runtimeCall, device,1396 module);1397 rewriter.eraseOp(op);1398 } else {1399 emitError(runtimeCall->getLoc(), "Unhandled runtime call hoisting.");1400 return failure();1401 }1402 } else {1403 emitError(op->getLoc(), "Unhandled op hoisting.");1404 return failure();1405 }1406 }1407 1408 // Update the host_eval_vars to use the mapped values.1409 for (size_t i = 0; i < hostEvalVars.lbs.size(); ++i) {1410 hostEvalVars.lbs[i] = mapping.lookup(hostEvalVars.lbs[i]);1411 hostEvalVars.ubs[i] = mapping.lookup(hostEvalVars.ubs[i]);1412 hostEvalVars.steps[i] = mapping.lookup(hostEvalVars.steps[i]);1413 }1414 // Finally erase the original targetOp.1415 rewriter.eraseOp(targetOp);1416 return success();1417}1418 1419/// Result of isolateOp method1420struct SplitResult {1421 omp::TargetOp preTargetOp;1422 omp::TargetOp isolatedTargetOp;1423 omp::TargetOp postTargetOp;1424};1425 1426/// computeAllocsCacheRecomputable method computes the allocs needed to cache1427/// the values that are used outside the split point. It also computes the ops1428/// that need to be cached and the ops that can be recomputed after the split.1429static void computeAllocsCacheRecomputable(1430 omp::TargetOp targetOp, Operation *splitBeforeOp, RewriterBase &rewriter,1431 SmallVector<Value> &preMapOperands, SmallVector<Value> &postMapOperands,1432 SmallVector<Value> &allocs, SmallVector<Value> &requiredVals,1433 SetVector<Operation *> &nonRecomputable, SetVector<Operation *> &toCache,1434 SetVector<Operation *> &toRecompute) {1435 auto *targetBlock = &targetOp.getRegion().front();1436 // Find all values that are used outside the split point.1437 for (auto it = targetBlock->begin(); it != splitBeforeOp->getIterator();1438 it++) {1439 // Check if any of the results are used outside the split point.1440 for (auto res : it->getResults()) {1441 if (usedOutsideSplit(res, splitBeforeOp)) {1442 requiredVals.push_back(res);1443 }1444 }1445 // If the op is not recomputable, add it to the nonRecomputable set.1446 if (!isRecomputableAfterFission(&*it, splitBeforeOp)) {1447 nonRecomputable.insert(&*it);1448 }1449 }1450 // For each required value, collect its dependencies.1451 for (auto requiredVal : requiredVals)1452 collectNonRecomputableDeps(requiredVal, targetOp, nonRecomputable, toCache,1453 toRecompute);1454 // For each op in toCache, create an alloc and update the pre and post map1455 // operands.1456 for (Operation *op : toCache) {1457 for (auto res : op->getResults()) {1458 auto alloc =1459 allocateTempOmpVar(targetOp.getLoc(), res.getType(), rewriter);1460 allocs.push_back(res);1461 preMapOperands.push_back(alloc.from);1462 postMapOperands.push_back(alloc.to);1463 }1464 }1465}1466 1467/// genPreTargetOp method generates the preTargetOp that contains all the ops1468/// before the split point. It also creates the block arguments and maps the1469/// values accordingly. It also creates the store operations for the allocs.1470static omp::TargetOp1471genPreTargetOp(omp::TargetOp targetOp, SmallVector<Value> &preMapOperands,1472 SmallVector<Value> &allocs, Operation *splitBeforeOp,1473 RewriterBase &rewriter, struct HostEvalVars &hostEvalVars,1474 bool isTargetDevice) {1475 auto loc = targetOp.getLoc();1476 auto *targetBlock = &targetOp.getRegion().front();1477 SmallVector<Value> preHostEvalVars{targetOp.getHostEvalVars()};1478 // update the hostEvalVars of preTargetOp1479 omp::TargetOp preTargetOp = omp::TargetOp::create(1480 rewriter, targetOp.getLoc(), targetOp.getAllocateVars(),1481 targetOp.getAllocatorVars(), targetOp.getBareAttr(),1482 targetOp.getDependKindsAttr(), targetOp.getDependVars(),1483 targetOp.getDevice(), targetOp.getHasDeviceAddrVars(), preHostEvalVars,1484 targetOp.getIfExpr(), targetOp.getInReductionVars(),1485 targetOp.getInReductionByrefAttr(), targetOp.getInReductionSymsAttr(),1486 targetOp.getIsDevicePtrVars(), preMapOperands, targetOp.getNowaitAttr(),1487 targetOp.getPrivateVars(), targetOp.getPrivateSymsAttr(),1488 targetOp.getPrivateNeedsBarrierAttr(), targetOp.getThreadLimit(),1489 targetOp.getPrivateMapsAttr());1490 auto *preTargetBlock = rewriter.createBlock(1491 &preTargetOp.getRegion(), preTargetOp.getRegion().begin(), {}, {});1492 IRMapping preMapping;1493 // Create block arguments and map the values.1494 createBlockArgsAndMap(loc, rewriter, targetOp, targetBlock, preTargetBlock,1495 preHostEvalVars, preMapOperands, allocs, preMapping);1496 1497 // Handle the store operations for the allocs.1498 rewriter.setInsertionPointToStart(preTargetBlock);1499 auto llvmPtrTy = LLVM::LLVMPointerType::get(targetOp.getContext());1500 1501 // Clone the original operations.1502 for (auto it = targetBlock->begin(); it != splitBeforeOp->getIterator();1503 it++) {1504 rewriter.clone(*it, preMapping);1505 }1506 1507 unsigned originalHostEvalVarsSize = preHostEvalVars.size();1508 unsigned originalMapVarsSize = targetOp.getMapVars().size();1509 // Create Stores for allocs.1510 for (unsigned i = 0; i < allocs.size(); ++i) {1511 Value originalResult = allocs[i];1512 Value toStore = preMapping.lookup(originalResult);1513 // Get the new block argument for this specific allocated value.1514 Value newArg = preTargetBlock->getArgument(originalHostEvalVarsSize +1515 originalMapVarsSize + i);1516 // Create the store operation.1517 if (isPtr(originalResult.getType())) {1518 if (!isa<LLVM::LLVMPointerType>(toStore.getType()))1519 toStore = fir::ConvertOp::create(rewriter, loc, llvmPtrTy, toStore);1520 LLVM::StoreOp::create(rewriter, loc, toStore, newArg);1521 } else {1522 fir::StoreOp::create(rewriter, loc, toStore, newArg);1523 }1524 }1525 omp::TerminatorOp::create(rewriter, loc);1526 1527 // Update hostEvalVars with the mapped values for the loop bounds if we have1528 // a loopNestOp and we are not generating code for the target device.1529 omp::LoopNestOp loopNestOp =1530 getLoopNestFromTeams(cast<omp::TeamsOp>(splitBeforeOp));1531 if (loopNestOp && !isTargetDevice) {1532 for (size_t i = 0; i < loopNestOp.getLoopLowerBounds().size(); ++i) {1533 Value lb = loopNestOp.getLoopLowerBounds()[i];1534 Value ub = loopNestOp.getLoopUpperBounds()[i];1535 Value step = loopNestOp.getLoopSteps()[i];1536 1537 hostEvalVars.lbs.push_back(preMapping.lookup(lb));1538 hostEvalVars.ubs.push_back(preMapping.lookup(ub));1539 hostEvalVars.steps.push_back(preMapping.lookup(step));1540 }1541 }1542 1543 return preTargetOp;1544}1545 1546/// genIsolatedTargetOp method generates the isolatedTargetOp that contains the1547/// ops between the split point. It also creates the block arguments and maps1548/// the values accordingly. It also creates the load operations for the allocs1549/// and recomputes the necessary ops.1550static omp::TargetOp1551genIsolatedTargetOp(omp::TargetOp targetOp, SmallVector<Value> &postMapOperands,1552 Operation *splitBeforeOp, RewriterBase &rewriter,1553 SmallVector<Value> &allocs,1554 SetVector<Operation *> &toRecompute,1555 struct HostEvalVars &hostEvalVars, bool isTargetDevice) {1556 auto loc = targetOp.getLoc();1557 auto *targetBlock = &targetOp.getRegion().front();1558 SmallVector<Value> isolatedHostEvalVars{targetOp.getHostEvalVars()};1559 // update the hostEvalVars of isolatedTargetOp1560 if (!hostEvalVars.lbs.empty() && !isTargetDevice) {1561 isolatedHostEvalVars.append(hostEvalVars.lbs.begin(),1562 hostEvalVars.lbs.end());1563 isolatedHostEvalVars.append(hostEvalVars.ubs.begin(),1564 hostEvalVars.ubs.end());1565 isolatedHostEvalVars.append(hostEvalVars.steps.begin(),1566 hostEvalVars.steps.end());1567 }1568 // Create the isolated target op1569 omp::TargetOp isolatedTargetOp = omp::TargetOp::create(1570 rewriter, targetOp.getLoc(), targetOp.getAllocateVars(),1571 targetOp.getAllocatorVars(), targetOp.getBareAttr(),1572 targetOp.getDependKindsAttr(), targetOp.getDependVars(),1573 targetOp.getDevice(), targetOp.getHasDeviceAddrVars(),1574 isolatedHostEvalVars, targetOp.getIfExpr(), targetOp.getInReductionVars(),1575 targetOp.getInReductionByrefAttr(), targetOp.getInReductionSymsAttr(),1576 targetOp.getIsDevicePtrVars(), postMapOperands, targetOp.getNowaitAttr(),1577 targetOp.getPrivateVars(), targetOp.getPrivateSymsAttr(),1578 targetOp.getPrivateNeedsBarrierAttr(), targetOp.getThreadLimit(),1579 targetOp.getPrivateMapsAttr());1580 auto *isolatedTargetBlock =1581 rewriter.createBlock(&isolatedTargetOp.getRegion(),1582 isolatedTargetOp.getRegion().begin(), {}, {});1583 IRMapping isolatedMapping;1584 // Create block arguments and map the values.1585 createBlockArgsAndMap(loc, rewriter, targetOp, targetBlock,1586 isolatedTargetBlock, isolatedHostEvalVars,1587 postMapOperands, allocs, isolatedMapping);1588 // Handle the load operations for the allocs and recompute ops.1589 reloadCacheAndRecompute(loc, rewriter, splitBeforeOp, targetOp, targetBlock,1590 isolatedTargetBlock, isolatedHostEvalVars,1591 postMapOperands, allocs, toRecompute,1592 isolatedMapping);1593 1594 // Clone the original operations.1595 rewriter.clone(*splitBeforeOp, isolatedMapping);1596 omp::TerminatorOp::create(rewriter, loc);1597 1598 // update the loop bounds in the isolatedTargetOp if we have host_eval vars1599 // and we are not generating code for the target device.1600 if (!hostEvalVars.lbs.empty() && !isTargetDevice) {1601 omp::TeamsOp teamsOp;1602 for (auto &op : *isolatedTargetBlock) {1603 if (isa<omp::TeamsOp>(&op))1604 teamsOp = cast<omp::TeamsOp>(&op);1605 }1606 assert(teamsOp && "No teamsOp found in isolated target region");1607 // Get the loopNestOp inside the teamsOp1608 auto loopNestOp = getLoopNestFromTeams(teamsOp);1609 // Get the BlockArgs related to host_eval vars and update loop_nest bounds1610 // to them1611 unsigned originalHostEvalVarsSize = targetOp.getHostEvalVars().size();1612 unsigned index = originalHostEvalVarsSize;1613 // Replace loop bounds with the block arguments passed down via host_eval1614 SmallVector<Value> lbs, ubs, steps;1615 1616 // Collect new lb/ub/step values from target block args1617 for (size_t i = 0; i < hostEvalVars.lbs.size(); ++i)1618 lbs.push_back(isolatedTargetBlock->getArgument(index++));1619 1620 for (size_t i = 0; i < hostEvalVars.ubs.size(); ++i)1621 ubs.push_back(isolatedTargetBlock->getArgument(index++));1622 1623 for (size_t i = 0; i < hostEvalVars.steps.size(); ++i)1624 steps.push_back(isolatedTargetBlock->getArgument(index++));1625 1626 // Reset the loop bounds1627 loopNestOp.getLoopLowerBoundsMutable().assign(lbs);1628 loopNestOp.getLoopUpperBoundsMutable().assign(ubs);1629 loopNestOp.getLoopStepsMutable().assign(steps);1630 }1631 1632 return isolatedTargetOp;1633}1634 1635/// genPostTargetOp method generates the postTargetOp that contains all the ops1636/// after the split point. It also creates the block arguments and maps the1637/// values accordingly. It also creates the load operations for the allocs1638/// and recomputes the necessary ops.1639static omp::TargetOp genPostTargetOp(omp::TargetOp targetOp,1640 Operation *splitBeforeOp,1641 SmallVector<Value> &postMapOperands,1642 RewriterBase &rewriter,1643 SmallVector<Value> &allocs,1644 SetVector<Operation *> &toRecompute) {1645 auto loc = targetOp.getLoc();1646 auto *targetBlock = &targetOp.getRegion().front();1647 SmallVector<Value> postHostEvalVars{targetOp.getHostEvalVars()};1648 // Create the post target op1649 omp::TargetOp postTargetOp = omp::TargetOp::create(1650 rewriter, targetOp.getLoc(), targetOp.getAllocateVars(),1651 targetOp.getAllocatorVars(), targetOp.getBareAttr(),1652 targetOp.getDependKindsAttr(), targetOp.getDependVars(),1653 targetOp.getDevice(), targetOp.getHasDeviceAddrVars(), postHostEvalVars,1654 targetOp.getIfExpr(), targetOp.getInReductionVars(),1655 targetOp.getInReductionByrefAttr(), targetOp.getInReductionSymsAttr(),1656 targetOp.getIsDevicePtrVars(), postMapOperands, targetOp.getNowaitAttr(),1657 targetOp.getPrivateVars(), targetOp.getPrivateSymsAttr(),1658 targetOp.getPrivateNeedsBarrierAttr(), targetOp.getThreadLimit(),1659 targetOp.getPrivateMapsAttr());1660 // Create the block for postTargetOp1661 auto *postTargetBlock = rewriter.createBlock(1662 &postTargetOp.getRegion(), postTargetOp.getRegion().begin(), {}, {});1663 IRMapping postMapping;1664 // Create block arguments and map the values.1665 createBlockArgsAndMap(loc, rewriter, targetOp, targetBlock, postTargetBlock,1666 postHostEvalVars, postMapOperands, allocs, postMapping);1667 // Handle the load operations for the allocs and recompute ops.1668 reloadCacheAndRecompute(loc, rewriter, splitBeforeOp, targetOp, targetBlock,1669 postTargetBlock, postHostEvalVars, postMapOperands,1670 allocs, toRecompute, postMapping);1671 assert(splitBeforeOp->getNumResults() == 0 ||1672 llvm::all_of(splitBeforeOp->getResults(),1673 [](Value result) { return result.use_empty(); }));1674 // Clone the original operations after the split point.1675 for (auto it = std::next(splitBeforeOp->getIterator());1676 it != targetBlock->end(); it++)1677 rewriter.clone(*it, postMapping);1678 return postTargetOp;1679}1680 1681/// isolateOp method rewrites a omp.target_data { omp.target } in to1682/// omp.target_data {1683/// // preTargetOp region contains ops before splitBeforeOp.1684/// omp.target {}1685/// // isolatedTargetOp region contains splitBeforeOp,1686/// omp.target {}1687/// // postTargetOp region contains ops after splitBeforeOp.1688/// omp.target {}1689/// }1690/// It also handles the mapping of variables and the caching/recomputing1691/// of values as needed.1692static FailureOr<SplitResult> isolateOp(Operation *splitBeforeOp,1693 bool splitAfter, RewriterBase &rewriter,1694 mlir::ModuleOp module,1695 bool isTargetDevice) {1696 auto targetOp = cast<omp::TargetOp>(splitBeforeOp->getParentOp());1697 assert(targetOp);1698 rewriter.setInsertionPoint(targetOp);1699 1700 // Prepare the map operands for preTargetOp and postTargetOp1701 auto preMapOperands = SmallVector<Value>(targetOp.getMapVars());1702 auto postMapOperands = SmallVector<Value>(targetOp.getMapVars());1703 1704 // Vectors to hold analysis results1705 SmallVector<Value> requiredVals;1706 SetVector<Operation *> toCache;1707 SetVector<Operation *> toRecompute;1708 SetVector<Operation *> nonRecomputable;1709 SmallVector<Value> allocs;1710 struct HostEvalVars hostEvalVars;1711 1712 // Analyze the ops in target region to determine which ops need to be1713 // cached and which ops need to be recomputed1714 computeAllocsCacheRecomputable(1715 targetOp, splitBeforeOp, rewriter, preMapOperands, postMapOperands,1716 allocs, requiredVals, nonRecomputable, toCache, toRecompute);1717 1718 rewriter.setInsertionPoint(targetOp);1719 1720 // Generate the preTargetOp that contains all the ops before splitBeforeOp.1721 auto preTargetOp =1722 genPreTargetOp(targetOp, preMapOperands, allocs, splitBeforeOp, rewriter,1723 hostEvalVars, isTargetDevice);1724 1725 // Move the ops of preTarget to host.1726 auto res = moveToHost(preTargetOp, rewriter, module, hostEvalVars);1727 if (failed(res))1728 return failure();1729 rewriter.setInsertionPoint(targetOp);1730 1731 // Generate the isolatedTargetOp1732 omp::TargetOp isolatedTargetOp =1733 genIsolatedTargetOp(targetOp, postMapOperands, splitBeforeOp, rewriter,1734 allocs, toRecompute, hostEvalVars, isTargetDevice);1735 1736 omp::TargetOp postTargetOp = nullptr;1737 // Generate the postTargetOp that contains all the ops after splitBeforeOp.1738 if (splitAfter) {1739 rewriter.setInsertionPoint(targetOp);1740 postTargetOp = genPostTargetOp(targetOp, splitBeforeOp, postMapOperands,1741 rewriter, allocs, toRecompute);1742 }1743 // Finally erase the original targetOp.1744 rewriter.eraseOp(targetOp);1745 return SplitResult{preTargetOp, isolatedTargetOp, postTargetOp};1746}1747 1748/// Recursively fission target ops until no more nested ops can be isolated.1749static LogicalResult fissionTarget(omp::TargetOp targetOp,1750 RewriterBase &rewriter,1751 mlir::ModuleOp module, bool isTargetDevice) {1752 auto tuple = getNestedOpToIsolate(targetOp);1753 if (!tuple) {1754 LLVM_DEBUG(llvm::dbgs() << " No op to isolate\n");1755 struct HostEvalVars hostEvalVars;1756 return moveToHost(targetOp, rewriter, module, hostEvalVars);1757 }1758 Operation *toIsolate = std::get<0>(*tuple);1759 bool splitBefore = !std::get<1>(*tuple);1760 bool splitAfter = !std::get<2>(*tuple);1761 // Recursively isolate the target op.1762 if (splitBefore && splitAfter) {1763 auto res =1764 isolateOp(toIsolate, splitAfter, rewriter, module, isTargetDevice);1765 if (failed(res))1766 return failure();1767 return fissionTarget((*res).postTargetOp, rewriter, module, isTargetDevice);1768 }1769 // Isolate only before the op.1770 if (splitBefore) {1771 auto res =1772 isolateOp(toIsolate, splitAfter, rewriter, module, isTargetDevice);1773 if (failed(res))1774 return failure();1775 } else {1776 emitError(toIsolate->getLoc(), "Unhandled case in fissionTarget");1777 return failure();1778 }1779 return success();1780}1781 1782/// Pass to lower omp.workdistribute ops.1783class LowerWorkdistributePass1784 : public flangomp::impl::LowerWorkdistributeBase<LowerWorkdistributePass> {1785public:1786 void runOnOperation() override {1787 MLIRContext &context = getContext();1788 auto moduleOp = getOperation();1789 bool changed = false;1790 SetVector<omp::TargetOp> targetOpsToProcess;1791 auto verify =1792 moduleOp->walk([&](mlir::omp::WorkdistributeOp workdistribute) {1793 if (failed(verifyTargetTeamsWorkdistribute(workdistribute)))1794 return WalkResult::interrupt();1795 return WalkResult::advance();1796 });1797 if (verify.wasInterrupted())1798 return signalPassFailure();1799 1800 auto fission =1801 moduleOp->walk([&](mlir::omp::WorkdistributeOp workdistribute) {1802 auto res = fissionWorkdistribute(workdistribute);1803 if (failed(res))1804 return WalkResult::interrupt();1805 changed |= *res;1806 return WalkResult::advance();1807 });1808 if (fission.wasInterrupted())1809 return signalPassFailure();1810 1811 auto rtCallLower =1812 moduleOp->walk([&](mlir::omp::WorkdistributeOp workdistribute) {1813 auto res = workdistributeRuntimeCallLower(workdistribute,1814 targetOpsToProcess);1815 if (failed(res))1816 return WalkResult::interrupt();1817 changed |= *res;1818 return WalkResult::advance();1819 });1820 if (rtCallLower.wasInterrupted())1821 return signalPassFailure();1822 1823 moduleOp->walk([&](mlir::omp::WorkdistributeOp workdistribute) {1824 changed |= workdistributeDoLower(workdistribute, targetOpsToProcess);1825 });1826 1827 moduleOp->walk([&](mlir::omp::TeamsOp teams) {1828 changed |= teamsWorkdistributeToSingleOp(teams, targetOpsToProcess);1829 });1830 if (changed) {1831 bool isTargetDevice =1832 llvm::cast<mlir::omp::OffloadModuleInterface>(*moduleOp)1833 .getIsTargetDevice();1834 IRRewriter rewriter(&context);1835 for (auto targetOp : targetOpsToProcess) {1836 auto res = splitTargetData(targetOp, rewriter);1837 if (failed(res))1838 return signalPassFailure();1839 if (*res) {1840 if (failed(fissionTarget(*res, rewriter, moduleOp, isTargetDevice)))1841 return signalPassFailure();1842 }1843 }1844 }1845 }1846};1847} // namespace1848