[llvm] e41eaf4 - [OpenMPIRBuilder] Add ThreadLimit and NumTeams clauses to teams construct (#68364)
via llvm-commits
llvm-commits at lists.llvm.org
Wed Oct 11 08:36:09 PDT 2023
Author: Shraiysh
Date: 2023-10-11T10:36:03-05:00
New Revision: e41eaf4896d96b5e9c81959130cc61d453f471fa
URL: https://github.com/llvm/llvm-project/commit/e41eaf4896d96b5e9c81959130cc61d453f471fa
DIFF: https://github.com/llvm/llvm-project/commit/e41eaf4896d96b5e9c81959130cc61d453f471fa.diff
LOG: [OpenMPIRBuilder] Add ThreadLimit and NumTeams clauses to teams construct (#68364)
This patch adds support for `thread_limit` and bounds on `num_teams`
clause for the teams construct in OpenMP.
Added testcases for the same.
Added:
Modified:
llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
llvm/include/llvm/Frontend/OpenMP/OMPKinds.def
llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp
Removed:
################################################################################
diff --git a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h b/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
index b505daae7e75f80..9d2adf229b78654 100644
--- a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
+++ b/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
@@ -1917,8 +1917,17 @@ class OpenMPIRBuilder {
///
/// \param Loc The location where the teams construct was encountered.
/// \param BodyGenCB Callback that will generate the region code.
+ /// \param NumTeamsLower Lower bound on number of teams. If this is nullptr,
+ /// it is as if lower bound is specified as equal to upperbound. If
+ /// this is non-null, then upperbound must also be non-null.
+ /// \param NumTeamsUpper Upper bound on the number of teams.
+ /// \param ThreadLimit on the number of threads that may participate in a
+ /// contention group created by each team.
InsertPointTy createTeams(const LocationDescription &Loc,
- BodyGenCallbackTy BodyGenCB);
+ BodyGenCallbackTy BodyGenCB,
+ Value *NumTeamsLower = nullptr,
+ Value *NumTeamsUpper = nullptr,
+ Value *ThreadLimit = nullptr);
/// Generate conditional branch and relevant BasicBlocks through which private
/// threads copy the 'copyin' variables from Master copy to threadprivate
diff --git a/llvm/include/llvm/Frontend/OpenMP/OMPKinds.def b/llvm/include/llvm/Frontend/OpenMP/OMPKinds.def
index 176b883fe68f7ad..4823c4cc6b833ec 100644
--- a/llvm/include/llvm/Frontend/OpenMP/OMPKinds.def
+++ b/llvm/include/llvm/Frontend/OpenMP/OMPKinds.def
@@ -387,6 +387,7 @@ __OMP_RTL(__kmpc_cancellationpoint, false, Int32, IdentPtr, Int32, Int32)
__OMP_RTL(__kmpc_fork_teams, true, Void, IdentPtr, Int32, ParallelTaskPtr)
__OMP_RTL(__kmpc_push_num_teams, false, Void, IdentPtr, Int32, Int32, Int32)
+__OMP_RTL(__kmpc_push_num_teams_51, false, Void, IdentPtr, Int32, Int32, Int32, Int32)
__OMP_RTL(__kmpc_set_thread_limit, false, Void, IdentPtr, Int32, Int32)
__OMP_RTL(__kmpc_copyprivate, false, Void, IdentPtr, Int32, SizeTy, VoidPtr,
diff --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
index 8e240ba96050082..a658990f2d45355 100644
--- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
+++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
@@ -5733,7 +5733,8 @@ OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::createAtomicCompare(
OpenMPIRBuilder::InsertPointTy
OpenMPIRBuilder::createTeams(const LocationDescription &Loc,
- BodyGenCallbackTy BodyGenCB) {
+ BodyGenCallbackTy BodyGenCB, Value *NumTeamsLower,
+ Value *NumTeamsUpper, Value *ThreadLimit) {
if (!updateToLocation(Loc))
return InsertPointTy();
@@ -5771,6 +5772,26 @@ OpenMPIRBuilder::createTeams(const LocationDescription &Loc,
BasicBlock *AllocaBB =
splitBB(Builder, /*CreateBranch=*/true, "teams.alloca");
+ // Push num_teams
+ if (NumTeamsLower || NumTeamsUpper || ThreadLimit) {
+ assert((NumTeamsLower == nullptr || NumTeamsUpper != nullptr) &&
+ "if lowerbound is non-null, then upperbound must also be non-null "
+ "for bounds on num_teams");
+
+ if (NumTeamsUpper == nullptr)
+ NumTeamsUpper = Builder.getInt32(0);
+
+ if (NumTeamsLower == nullptr)
+ NumTeamsLower = NumTeamsUpper;
+
+ if (ThreadLimit == nullptr)
+ ThreadLimit = Builder.getInt32(0);
+
+ Value *ThreadNum = getOrCreateThreadID(Ident);
+ Builder.CreateCall(
+ getOrCreateRuntimeFunctionPtr(OMPRTL___kmpc_push_num_teams_51),
+ {Ident, ThreadNum, NumTeamsLower, NumTeamsUpper, ThreadLimit});
+ }
// Generate the body of teams.
InsertPointTy AllocaIP(AllocaBB, AllocaBB->begin());
InsertPointTy CodeGenIP(BodyBB, BodyBB->begin());
diff --git a/llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp b/llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp
index c60aca259d1859f..d770facc1730252 100644
--- a/llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp
+++ b/llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp
@@ -4074,6 +4074,216 @@ TEST_F(OpenMPIRBuilderTest, CreateTeams) {
[](Instruction &inst) { return isa<ICmpInst>(&inst); }));
}
+TEST_F(OpenMPIRBuilderTest, CreateTeamsWithThreadLimit) {
+ using InsertPointTy = OpenMPIRBuilder::InsertPointTy;
+ OpenMPIRBuilder OMPBuilder(*M);
+ OMPBuilder.initialize();
+ F->setName("func");
+ IRBuilder<> &Builder = OMPBuilder.Builder;
+ Builder.SetInsertPoint(BB);
+
+ Function *FakeFunction =
+ Function::Create(FunctionType::get(Builder.getVoidTy(), false),
+ GlobalValue::ExternalLinkage, "fakeFunction", M.get());
+
+ auto BodyGenCB = [&](InsertPointTy AllocaIP, InsertPointTy CodeGenIP) {
+ Builder.restoreIP(CodeGenIP);
+ Builder.CreateCall(FakeFunction, {});
+ };
+
+ // `F` has an argument - an integer, so we use that as the thread limit.
+ Builder.restoreIP(OMPBuilder.createTeams(/*=*/Builder, BodyGenCB,
+ /*NumTeamsLower=*/nullptr,
+ /*NumTeamsUpper=*/nullptr,
+ /*ThreadLimit=*/F->arg_begin()));
+
+ Builder.CreateRetVoid();
+ OMPBuilder.finalize();
+
+ ASSERT_FALSE(verifyModule(*M));
+
+ CallInst *PushNumTeamsCallInst =
+ findSingleCall(F, OMPRTL___kmpc_push_num_teams_51, OMPBuilder);
+ ASSERT_NE(PushNumTeamsCallInst, nullptr);
+
+ EXPECT_EQ(PushNumTeamsCallInst->getArgOperand(2), Builder.getInt32(0));
+ EXPECT_EQ(PushNumTeamsCallInst->getArgOperand(3), Builder.getInt32(0));
+ EXPECT_EQ(PushNumTeamsCallInst->getArgOperand(4), &*F->arg_begin());
+
+ // Verifying that the next instruction to execute is kmpc_fork_teams
+ BranchInst *BrInst =
+ dyn_cast<BranchInst>(PushNumTeamsCallInst->getNextNonDebugInstruction());
+ ASSERT_NE(BrInst, nullptr);
+ ASSERT_EQ(BrInst->getNumSuccessors(), 1U);
+ Instruction *NextInstruction =
+ BrInst->getSuccessor(0)->getFirstNonPHIOrDbgOrLifetime();
+ CallInst *ForkTeamsCI = dyn_cast_if_present<CallInst>(NextInstruction);
+ ASSERT_NE(ForkTeamsCI, nullptr);
+ EXPECT_EQ(ForkTeamsCI->getCalledFunction(),
+ OMPBuilder.getOrCreateRuntimeFunctionPtr(OMPRTL___kmpc_fork_teams));
+}
+
+TEST_F(OpenMPIRBuilderTest, CreateTeamsWithNumTeamsUpper) {
+ using InsertPointTy = OpenMPIRBuilder::InsertPointTy;
+ OpenMPIRBuilder OMPBuilder(*M);
+ OMPBuilder.initialize();
+ F->setName("func");
+ IRBuilder<> &Builder = OMPBuilder.Builder;
+ Builder.SetInsertPoint(BB);
+
+ Function *FakeFunction =
+ Function::Create(FunctionType::get(Builder.getVoidTy(), false),
+ GlobalValue::ExternalLinkage, "fakeFunction", M.get());
+
+ auto BodyGenCB = [&](InsertPointTy AllocaIP, InsertPointTy CodeGenIP) {
+ Builder.restoreIP(CodeGenIP);
+ Builder.CreateCall(FakeFunction, {});
+ };
+
+ // `F` already has an integer argument, so we use that as upper bound to
+ // `num_teams`
+ Builder.restoreIP(OMPBuilder.createTeams(Builder, BodyGenCB,
+ /*NumTeamsLower=*/nullptr,
+ /*NumTeamsUpper=*/F->arg_begin()));
+
+ Builder.CreateRetVoid();
+ OMPBuilder.finalize();
+
+ ASSERT_FALSE(verifyModule(*M));
+
+ CallInst *PushNumTeamsCallInst =
+ findSingleCall(F, OMPRTL___kmpc_push_num_teams_51, OMPBuilder);
+ ASSERT_NE(PushNumTeamsCallInst, nullptr);
+
+ EXPECT_EQ(PushNumTeamsCallInst->getArgOperand(2), &*F->arg_begin());
+ EXPECT_EQ(PushNumTeamsCallInst->getArgOperand(3), &*F->arg_begin());
+ EXPECT_EQ(PushNumTeamsCallInst->getArgOperand(4), Builder.getInt32(0));
+
+ // Verifying that the next instruction to execute is kmpc_fork_teams
+ BranchInst *BrInst =
+ dyn_cast<BranchInst>(PushNumTeamsCallInst->getNextNonDebugInstruction());
+ ASSERT_NE(BrInst, nullptr);
+ ASSERT_EQ(BrInst->getNumSuccessors(), 1U);
+ Instruction *NextInstruction =
+ BrInst->getSuccessor(0)->getFirstNonPHIOrDbgOrLifetime();
+ CallInst *ForkTeamsCI = dyn_cast_if_present<CallInst>(NextInstruction);
+ ASSERT_NE(ForkTeamsCI, nullptr);
+ EXPECT_EQ(ForkTeamsCI->getCalledFunction(),
+ OMPBuilder.getOrCreateRuntimeFunctionPtr(OMPRTL___kmpc_fork_teams));
+}
+
+TEST_F(OpenMPIRBuilderTest, CreateTeamsWithNumTeamsBoth) {
+ using InsertPointTy = OpenMPIRBuilder::InsertPointTy;
+ OpenMPIRBuilder OMPBuilder(*M);
+ OMPBuilder.initialize();
+ F->setName("func");
+ IRBuilder<> &Builder = OMPBuilder.Builder;
+ Builder.SetInsertPoint(BB);
+
+ Function *FakeFunction =
+ Function::Create(FunctionType::get(Builder.getVoidTy(), false),
+ GlobalValue::ExternalLinkage, "fakeFunction", M.get());
+
+ Value *NumTeamsLower =
+ Builder.CreateAdd(F->arg_begin(), Builder.getInt32(5), "numTeamsLower");
+ Value *NumTeamsUpper =
+ Builder.CreateAdd(F->arg_begin(), Builder.getInt32(10), "numTeamsUpper");
+
+ auto BodyGenCB = [&](InsertPointTy AllocaIP, InsertPointTy CodeGenIP) {
+ Builder.restoreIP(CodeGenIP);
+ Builder.CreateCall(FakeFunction, {});
+ };
+
+ // `F` already has an integer argument, so we use that as upper bound to
+ // `num_teams`
+ Builder.restoreIP(
+ OMPBuilder.createTeams(Builder, BodyGenCB, NumTeamsLower, NumTeamsUpper));
+
+ Builder.CreateRetVoid();
+ OMPBuilder.finalize();
+
+ ASSERT_FALSE(verifyModule(*M));
+
+ CallInst *PushNumTeamsCallInst =
+ findSingleCall(F, OMPRTL___kmpc_push_num_teams_51, OMPBuilder);
+ ASSERT_NE(PushNumTeamsCallInst, nullptr);
+
+ EXPECT_EQ(PushNumTeamsCallInst->getArgOperand(2), NumTeamsLower);
+ EXPECT_EQ(PushNumTeamsCallInst->getArgOperand(3), NumTeamsUpper);
+ EXPECT_EQ(PushNumTeamsCallInst->getArgOperand(4), Builder.getInt32(0));
+
+ // Verifying that the next instruction to execute is kmpc_fork_teams
+ BranchInst *BrInst =
+ dyn_cast<BranchInst>(PushNumTeamsCallInst->getNextNonDebugInstruction());
+ ASSERT_NE(BrInst, nullptr);
+ ASSERT_EQ(BrInst->getNumSuccessors(), 1U);
+ Instruction *NextInstruction =
+ BrInst->getSuccessor(0)->getFirstNonPHIOrDbgOrLifetime();
+ CallInst *ForkTeamsCI = dyn_cast_if_present<CallInst>(NextInstruction);
+ ASSERT_NE(ForkTeamsCI, nullptr);
+ EXPECT_EQ(ForkTeamsCI->getCalledFunction(),
+ OMPBuilder.getOrCreateRuntimeFunctionPtr(OMPRTL___kmpc_fork_teams));
+}
+
+TEST_F(OpenMPIRBuilderTest, CreateTeamsWithNumTeamsAndThreadLimit) {
+ using InsertPointTy = OpenMPIRBuilder::InsertPointTy;
+ OpenMPIRBuilder OMPBuilder(*M);
+ OMPBuilder.initialize();
+ F->setName("func");
+ IRBuilder<> &Builder = OMPBuilder.Builder;
+ Builder.SetInsertPoint(BB);
+
+ BasicBlock *CodegenBB = splitBB(Builder, true);
+ Builder.SetInsertPoint(CodegenBB);
+
+ // Generate values for `num_teams` and `thread_limit` using the first argument
+ // of the testing function.
+ Value *NumTeamsLower =
+ Builder.CreateAdd(F->arg_begin(), Builder.getInt32(5), "numTeamsLower");
+ Value *NumTeamsUpper =
+ Builder.CreateAdd(F->arg_begin(), Builder.getInt32(10), "numTeamsUpper");
+ Value *ThreadLimit =
+ Builder.CreateAdd(F->arg_begin(), Builder.getInt32(20), "threadLimit");
+
+ Function *FakeFunction =
+ Function::Create(FunctionType::get(Builder.getVoidTy(), false),
+ GlobalValue::ExternalLinkage, "fakeFunction", M.get());
+
+ auto BodyGenCB = [&](InsertPointTy AllocaIP, InsertPointTy CodeGenIP) {
+ Builder.restoreIP(CodeGenIP);
+ Builder.CreateCall(FakeFunction, {});
+ };
+
+ OpenMPIRBuilder::LocationDescription Loc({Builder.saveIP(), DL});
+ Builder.restoreIP(OMPBuilder.createTeams(Builder, BodyGenCB, NumTeamsLower,
+ NumTeamsUpper, ThreadLimit));
+
+ Builder.CreateRetVoid();
+ OMPBuilder.finalize();
+
+ ASSERT_FALSE(verifyModule(*M));
+
+ CallInst *PushNumTeamsCallInst =
+ findSingleCall(F, OMPRTL___kmpc_push_num_teams_51, OMPBuilder);
+ ASSERT_NE(PushNumTeamsCallInst, nullptr);
+
+ EXPECT_EQ(PushNumTeamsCallInst->getArgOperand(2), NumTeamsLower);
+ EXPECT_EQ(PushNumTeamsCallInst->getArgOperand(3), NumTeamsUpper);
+ EXPECT_EQ(PushNumTeamsCallInst->getArgOperand(4), ThreadLimit);
+
+ // Verifying that the next instruction to execute is kmpc_fork_teams
+ BranchInst *BrInst =
+ dyn_cast<BranchInst>(PushNumTeamsCallInst->getNextNonDebugInstruction());
+ ASSERT_NE(BrInst, nullptr);
+ ASSERT_EQ(BrInst->getNumSuccessors(), 1U);
+ Instruction *NextInstruction =
+ BrInst->getSuccessor(0)->getFirstNonPHIOrDbgOrLifetime();
+ CallInst *ForkTeamsCI = dyn_cast_if_present<CallInst>(NextInstruction);
+ ASSERT_NE(ForkTeamsCI, nullptr);
+ EXPECT_EQ(ForkTeamsCI->getCalledFunction(),
+ OMPBuilder.getOrCreateRuntimeFunctionPtr(OMPRTL___kmpc_fork_teams));
+}
+
/// Returns the single instruction of InstTy type in BB that uses the value V.
/// If there is more than one such instruction, returns null.
template <typename InstTy>
More information about the llvm-commits
mailing list