[OpenMPIRBuilder] Add ThreadLimit and NumTeams clauses to teams construct #68364

shraiysh · 2023-10-05T23:04:05Z

This patch adds support for thread_limit and bounds on num_teams clause for the teams construct in OpenMP.

Added testcases for the same.

This patch adds `num_teams` (upperbound) and `thread_limit` clauses to `OpenMPIRBuilder`.

llvmbot · 2023-10-05T23:05:16Z

@llvm/pr-subscribers-flang-openmp

Changes

This patch adds support for thread_limit and upperbound on num_teams clause for the teams construct in OpenMP.

Added testcases for the same.

Full diff: https://github.com/llvm/llvm-project/pull/68364.diff

3 Files Affected:

(modified) llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h (+6-1)
(modified) llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp (+13-1)
(modified) llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp (+157)

diff --git a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h b/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
index 75da461cfd8d95e..95f30d4a0d5cd5d 100644
--- a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
+++ b/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
@@ -1917,8 +1917,13 @@ class OpenMPIRBuilder {
   ///
   /// \param Loc The location where the teams construct was encountered.
   /// \param BodyGenCB Callback that will generate the region code.
+  /// \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 *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/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
index 72e1af55fe63f60..b6d79e068f61e6e 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 *NumTeamsUpper,
+                             Value *ThreadLimit) {
   if (!updateToLocation(Loc))
     return InsertPointTy();
 
@@ -5771,6 +5772,17 @@ OpenMPIRBuilder::createTeams(const LocationDescription &Loc,
   BasicBlock *AllocaBB =
       splitBB(Builder, /*CreateBranch=*/true, "teams.alloca");
 
+  // Push num_teams
+  if (NumTeamsUpper || ThreadLimit) {
+    NumTeamsUpper =
+        NumTeamsUpper == nullptr ? Builder.getInt32(0) : NumTeamsUpper;
+    ThreadLimit = ThreadLimit == nullptr ? Builder.getInt32(0) : ThreadLimit;
+    Value *ThreadNum = getOrCreateThreadID(Ident);
+    Builder.CreateCall(
+        getOrCreateRuntimeFunctionPtr(OMPRTL___kmpc_push_num_teams),
+        {Ident, ThreadNum, NumTeamsUpper, ThreadLimit});
+  }
+
   OutlineInfo OI;
   OI.EntryBB = AllocaBB;
   OI.ExitBB = ExitBB;
diff --git a/llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp b/llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp
index fd524f6067ee0ea..fb87389023910c2 100644
--- a/llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp
+++ b/llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp
@@ -4082,6 +4082,163 @@ 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, {});
+  };
+
+  OpenMPIRBuilder::LocationDescription Loc({Builder.saveIP(), DL});
+  Builder.restoreIP(
+      OMPBuilder.createTeams(Builder, BodyGenCB, nullptr, F->arg_begin()));
+
+  Builder.CreateRetVoid();
+  OMPBuilder.finalize();
+
+  ASSERT_FALSE(verifyModule(*M));
+
+  Function *PushNumTeamsRTL =
+      OMPBuilder.getOrCreateRuntimeFunctionPtr(OMPRTL___kmpc_push_num_teams);
+  EXPECT_EQ(PushNumTeamsRTL->getNumUses(), 1U);
+
+  CallInst *PushNumTeamsCallInst =
+      findSingleCall(F, OMPRTL___kmpc_push_num_teams, OMPBuilder);
+  ASSERT_NE(PushNumTeamsCallInst, nullptr);
+
+  EXPECT_EQ(PushNumTeamsCallInst->getArgOperand(2), Builder.getInt32(0));
+  EXPECT_EQ(PushNumTeamsCallInst->getArgOperand(3), &*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, CreateTeamsWithNumTeams) {
+  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, {});
+  };
+
+  OpenMPIRBuilder::LocationDescription Loc({Builder.saveIP(), DL});
+  Builder.restoreIP(OMPBuilder.createTeams(Builder, BodyGenCB, F->arg_begin()));
+
+  Builder.CreateRetVoid();
+  OMPBuilder.finalize();
+
+  ASSERT_FALSE(verifyModule(*M));
+
+  Function *PushNumTeamsRTL =
+      OMPBuilder.getOrCreateRuntimeFunctionPtr(OMPRTL___kmpc_push_num_teams);
+  EXPECT_EQ(PushNumTeamsRTL->getNumUses(), 1U);
+
+  CallInst *PushNumTeamsCallInst =
+      findSingleCall(F, OMPRTL___kmpc_push_num_teams, OMPBuilder);
+  ASSERT_NE(PushNumTeamsCallInst, nullptr);
+
+  EXPECT_EQ(PushNumTeamsCallInst->getArgOperand(2), &*F->arg_begin());
+  EXPECT_EQ(PushNumTeamsCallInst->getArgOperand(3), 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);
+
+  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, NumTeamsUpper, ThreadLimit));
+
+  Builder.CreateRetVoid();
+  OMPBuilder.finalize();
+
+  ASSERT_FALSE(verifyModule(*M));
+
+  Function *PushNumTeamsRTL =
+      OMPBuilder.getOrCreateRuntimeFunctionPtr(OMPRTL___kmpc_push_num_teams);
+  EXPECT_EQ(PushNumTeamsRTL->getNumUses(), 1U);
+
+  CallInst *PushNumTeamsCallInst =
+      findSingleCall(F, OMPRTL___kmpc_push_num_teams, OMPBuilder);
+  ASSERT_NE(PushNumTeamsCallInst, nullptr);
+
+  EXPECT_EQ(PushNumTeamsCallInst->getArgOperand(2), NumTeamsUpper);
+  EXPECT_EQ(PushNumTeamsCallInst->getArgOperand(3), 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>

kiranchandramohan

Is this change for OMP 5.0 and before?

kiranchandramohan · 2023-10-09T10:25:24Z

llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h

  InsertPointTy createTeams(const LocationDescription &Loc,
-                            BodyGenCallbackTy BodyGenCB);
+                            BodyGenCallbackTy BodyGenCB,
+                            Value *NumTeamsUpper = nullptr,


What about the LowerBound?

Lowerbound isn't handled for clang either - it is not even recognized as syntactically correct. I don't know what the generated code should look like, as @__kmpc_push_num_teams does not accept a lowerbound on number of teams. So, I kept it unimplemented for future work.

So, there are two functions in the OpenMP runtime. It is OK to leave with a TODO, just write a status of the support in Clang and the OpenMP runtime.

KMP_EXPORT void __kmpc_push_num_teams(ident_t *loc, kmp_int32 global_tid, kmp_int32 num_teams, kmp_int32 num_threads); /* Function for OpenMP 5.1 num_teams clause */ KMP_EXPORT void __kmpc_push_num_teams_51(ident_t *loc, kmp_int32 global_tid, kmp_int32 num_teams_lb, kmp_int32 num_teams_ub, kmp_int32 num_threads);

Oh, I had not seen this. I think I will update the PR to handle this. Thank you for pointing this out! 👍

kiranchandramohan · 2023-10-09T10:41:23Z

llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp

+
+  OpenMPIRBuilder::LocationDescription Loc({Builder.saveIP(), DL});
+  Builder.restoreIP(
+      OMPBuilder.createTeams(Builder, BodyGenCB, nullptr, F->arg_begin()));


Please add a comment for what is nullptr and why using F->arg_begin() is fine. Same for the createTeams invocation in the other tests.

shraiysh · 2023-10-09T13:31:08Z

Is this change for OMP 5.0 and before?

No, it should be for all the standards since num_teams was introduced. Why do you think it is for OMP 5.0 and before, but not for 5.1 and 5.2? I will try to address it and clarify.

…o teams_clauses1

kiranchandramohan

This is OK with me.

The only question is whether we should use the __kmpc_push_num_teams in the case when the lowerbound is not specified.

Please wait for a day to see whether there is any opinion on this topic.

shraiysh and others added 4 commits October 4, 2023 15:57

[OpenMPIRBuilder] Add clauses to teams

2d3b344

This patch adds `num_teams` (upperbound) and `thread_limit` clauses to `OpenMPIRBuilder`.

Merge branch 'llvm:main' into teams_clauses1

0268bfa

Add testcases

8393f14

Formatting

9f36870

llvmbot added the flang:openmp label Oct 5, 2023

shraiysh requested review from kiranchandramohan, kiranktp, raghavendhra, dpalermo, jsjodin, NimishMishra, skatrak, DominikAdamski and jdoerfert October 6, 2023 02:35

kiranchandramohan reviewed Oct 9, 2023

View reviewed changes

shraiysh closed this Oct 9, 2023

shraiysh force-pushed the teams_clauses1 branch from 9f36870 to e01c867 Compare October 9, 2023 13:40

Merge branch 'teams_clauses1' of github.com:shraiysh/llvm-project int…

49fa1c7

…o teams_clauses1

shraiysh reopened this Oct 9, 2023

shraiysh added 2 commits October 9, 2023 09:10

Address comments

25870b3

Address comment about lower bound

c13e639

shraiysh changed the title ~~[OpenMPIRBuilder] Add ThreadLimit and NumTeamsUpper clauses to teams clause~~ [OpenMPIRBuilder] Add ThreadLimit and NumTeams clauses to teams construct Oct 9, 2023

kiranchandramohan approved these changes Oct 10, 2023

View reviewed changes

shraiysh merged commit e41eaf4 into llvm:main Oct 11, 2023

shraiysh deleted the teams_clauses1 branch October 11, 2023 15:36

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[OpenMPIRBuilder] Add ThreadLimit and NumTeams clauses to teams construct #68364

[OpenMPIRBuilder] Add ThreadLimit and NumTeams clauses to teams construct #68364

Uh oh!

shraiysh commented Oct 5, 2023 •

edited

Loading

Uh oh!

llvmbot commented Oct 5, 2023

Uh oh!

kiranchandramohan left a comment

Uh oh!

kiranchandramohan Oct 9, 2023

Uh oh!

shraiysh Oct 9, 2023

Uh oh!

kiranchandramohan Oct 9, 2023

Uh oh!

shraiysh Oct 9, 2023

Uh oh!

kiranchandramohan Oct 9, 2023

Uh oh!

shraiysh commented Oct 9, 2023

Uh oh!

kiranchandramohan left a comment

Uh oh!

Uh oh!

[OpenMPIRBuilder] Add ThreadLimit and NumTeams clauses to teams construct #68364

[OpenMPIRBuilder] Add ThreadLimit and NumTeams clauses to teams construct #68364

Uh oh!

Conversation

shraiysh commented Oct 5, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

llvmbot commented Oct 5, 2023

Uh oh!

kiranchandramohan left a comment

Choose a reason for hiding this comment

Uh oh!

kiranchandramohan Oct 9, 2023

Choose a reason for hiding this comment

Uh oh!

shraiysh Oct 9, 2023

Choose a reason for hiding this comment

Uh oh!

kiranchandramohan Oct 9, 2023

Choose a reason for hiding this comment

Uh oh!

shraiysh Oct 9, 2023

Choose a reason for hiding this comment

Uh oh!

kiranchandramohan Oct 9, 2023

Choose a reason for hiding this comment

Uh oh!

shraiysh commented Oct 9, 2023

Uh oh!

kiranchandramohan left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

shraiysh commented Oct 5, 2023 •

edited

Loading