[llvm] [OpenMPIRBuilder] Add support for distribute constructs (PR #127816)
Sergio Afonso via llvm-commits
llvm-commits at lists.llvm.org
Mon Feb 24 05:40:18 PST 2025
https://github.com/skatrak updated https://github.com/llvm/llvm-project/pull/127816
>From 00cc8b6421c910240aed89732e5297036f034a81 Mon Sep 17 00:00:00 2001
From: Dominik Adamski <dominik.adamski at 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 <safonsof at 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 33b3d7bad4a71..5fc2bd6b78fc5 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.
@@ -2666,6 +2666,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 18bc82fc827f7..1e6a19539b85c 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)
@@ -9275,6 +9280,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) {
More information about the llvm-commits
mailing list