https://github.com/skatrak created https://github.com/llvm/llvm-project/pull/127816
This patch adds the `OpenMPIRBuilder::createDistribute()` function and updates `OpenMPIRBuilder::applyStaticWorkshareLoop()` in preparation for adding `distribute` support to flang. >From a79b7a2d6a443ef26bf4beaf73ec3c8042d968d1 Mon Sep 17 00:00:00 2001 From: Dominik Adamski <dominik.adam...@amd.com> Date: Mon, 17 Feb 2025 14:25:40 +0000 Subject: [PATCH] [OpenMPIRBuilder] Add support for distribute constructs This patch adds the `OpenMPIRBuilder::createDistribute()` function and updates `OpenMPIRBuilder::applyStaticWorkshareLoop()` in preparation for adding `distribute` support to flang. Co-authored-by: Sergio Afonso <safon...@amd.com> --- .../llvm/Frontend/OpenMP/OMPIRBuilder.h | 17 ++++-- llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp | 59 ++++++++++++++++--- 2 files changed, 64 insertions(+), 12 deletions(-) diff --git a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h b/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h index d25077cae63e4..9ad85413acd34 100644 --- a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h +++ b/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h @@ -1004,12 +1004,12 @@ class OpenMPIRBuilder { /// preheader of the loop. /// \param NeedsBarrier Indicates whether a barrier must be inserted after /// the loop. + /// \param LoopType Type of workshare loop. /// /// \returns Point where to insert code after the workshare construct. - InsertPointOrErrorTy applyStaticWorkshareLoop(DebugLoc DL, - CanonicalLoopInfo *CLI, - InsertPointTy AllocaIP, - bool NeedsBarrier); + InsertPointOrErrorTy applyStaticWorkshareLoop( + DebugLoc DL, CanonicalLoopInfo *CLI, InsertPointTy AllocaIP, + omp::WorksharingLoopType LoopType, bool NeedsBarrier); /// Modifies the canonical loop a statically-scheduled workshare loop with a /// user-specified chunk size. @@ -2660,6 +2660,15 @@ class OpenMPIRBuilder { Value *NumTeamsLower = nullptr, Value *NumTeamsUpper = nullptr, Value *ThreadLimit = nullptr, Value *IfExpr = nullptr); + /// Generator for `#omp distribute` + /// + /// \param Loc The location where the distribute construct was encountered. + /// \param AllocaIP The insertion points to be used for alloca instructions. + /// \param BodyGenCB Callback that will generate the region code. + InsertPointOrErrorTy createDistribute(const LocationDescription &Loc, + InsertPointTy AllocaIP, + BodyGenCallbackTy BodyGenCB); + /// Generate conditional branch and relevant BasicBlocks through which private /// threads copy the 'copyin' variables from Master copy to threadprivate /// copies. diff --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp index 04acab1e5765e..9e380bf2d3dbe 100644 --- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp +++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp @@ -2295,7 +2295,8 @@ OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::createSections( return LoopInfo.takeError(); InsertPointOrErrorTy WsloopIP = - applyStaticWorkshareLoop(Loc.DL, *LoopInfo, AllocaIP, !IsNowait); + applyStaticWorkshareLoop(Loc.DL, *LoopInfo, AllocaIP, + WorksharingLoopType::ForStaticLoop, !IsNowait); if (!WsloopIP) return WsloopIP.takeError(); InsertPointTy AfterIP = *WsloopIP; @@ -4145,10 +4146,9 @@ static FunctionCallee getKmpcForStaticInitForType(Type *Ty, Module &M, llvm_unreachable("unknown OpenMP loop iterator bitwidth"); } -OpenMPIRBuilder::InsertPointOrErrorTy -OpenMPIRBuilder::applyStaticWorkshareLoop(DebugLoc DL, CanonicalLoopInfo *CLI, - InsertPointTy AllocaIP, - bool NeedsBarrier) { +OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::applyStaticWorkshareLoop( + DebugLoc DL, CanonicalLoopInfo *CLI, InsertPointTy AllocaIP, + WorksharingLoopType LoopType, bool NeedsBarrier) { assert(CLI->isValid() && "Requires a valid canonical loop"); assert(!isConflictIP(AllocaIP, CLI->getPreheaderIP()) && "Require dedicated allocate IP"); @@ -4191,8 +4191,12 @@ OpenMPIRBuilder::applyStaticWorkshareLoop(DebugLoc DL, CanonicalLoopInfo *CLI, Value *ThreadNum = getOrCreateThreadID(SrcLoc); - Constant *SchedulingType = ConstantInt::get( - I32Type, static_cast<int>(OMPScheduleType::UnorderedStatic)); + OMPScheduleType SchedType = + (LoopType == WorksharingLoopType::DistributeStaticLoop) + ? OMPScheduleType::OrderedDistribute + : OMPScheduleType::UnorderedStatic; + Constant *SchedulingType = + ConstantInt::get(I32Type, static_cast<int>(SchedType)); // Call the "init" function and update the trip count of the loop with the // value it produced. @@ -4452,6 +4456,7 @@ static void createTargetLoopWorkshareCall( RealArgs.push_back(TripCount); if (LoopType == WorksharingLoopType::DistributeStaticLoop) { RealArgs.push_back(ConstantInt::get(TripCountTy, 0)); + Builder.restoreIP({InsertBlock, std::prev(InsertBlock->end())}); Builder.CreateCall(RTLFn, RealArgs); return; } @@ -4645,7 +4650,7 @@ OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::applyWorkshareLoop( return applyDynamicWorkshareLoop(DL, CLI, AllocaIP, EffectiveScheduleType, NeedsBarrier, ChunkSize); // FIXME: Monotonicity ignored? - return applyStaticWorkshareLoop(DL, CLI, AllocaIP, NeedsBarrier); + return applyStaticWorkshareLoop(DL, CLI, AllocaIP, LoopType, NeedsBarrier); case OMPScheduleType::BaseStaticChunked: if (IsOrdered) @@ -9245,6 +9250,44 @@ OpenMPIRBuilder::createTeams(const LocationDescription &Loc, return Builder.saveIP(); } +OpenMPIRBuilder::InsertPointOrErrorTy +OpenMPIRBuilder::createDistribute(const LocationDescription &Loc, + InsertPointTy OuterAllocaIP, + BodyGenCallbackTy BodyGenCB) { + if (!updateToLocation(Loc)) + return InsertPointTy(); + + BasicBlock *OuterAllocaBB = OuterAllocaIP.getBlock(); + + if (OuterAllocaBB == Builder.GetInsertBlock()) { + BasicBlock *BodyBB = + splitBB(Builder, /*CreateBranch=*/true, "distribute.entry"); + Builder.SetInsertPoint(BodyBB, BodyBB->begin()); + } + BasicBlock *ExitBB = + splitBB(Builder, /*CreateBranch=*/true, "distribute.exit"); + BasicBlock *BodyBB = + splitBB(Builder, /*CreateBranch=*/true, "distribute.body"); + BasicBlock *AllocaBB = + splitBB(Builder, /*CreateBranch=*/true, "distribute.alloca"); + + // Generate the body of distribute clause + InsertPointTy AllocaIP(AllocaBB, AllocaBB->begin()); + InsertPointTy CodeGenIP(BodyBB, BodyBB->begin()); + if (Error Err = BodyGenCB(AllocaIP, CodeGenIP)) + return Err; + + OutlineInfo OI; + OI.OuterAllocaBB = OuterAllocaIP.getBlock(); + OI.EntryBB = AllocaBB; + OI.ExitBB = ExitBB; + + addOutlineInfo(std::move(OI)); + Builder.SetInsertPoint(ExitBB, ExitBB->begin()); + + return Builder.saveIP(); +} + GlobalVariable * OpenMPIRBuilder::createOffloadMapnames(SmallVectorImpl<llvm::Constant *> &Names, std::string VarName) { _______________________________________________ llvm-branch-commits mailing list llvm-branch-commits@lists.llvm.org https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-branch-commits