[OpenMP][MLIR]Add support for guided, auto and runtime scheduling

When using parallel loop construct, the OpenMP specification allows for
guided, auto and runtime as scheduling variants (as well as static and
dynamic which are already supported).

This adds the translation from MLIR to LLVM-IR for these scheduling
variants.

Reviewed By: jdoerfert

Differential Revision: https://reviews.llvm.org/D101435

include/llvm/Frontend/OpenMP/OMPConstants.h

@@ -111,10 +111,15 @@ inline std::string getAllAssumeClauseOptions() {
 /// Todo: Update kmp.h to include this file, and remove the enums in kmp.h
 ///       To complete this, more enum values will need to be moved here.
 enum class OMPScheduleType {
-  Static = 34, /**< static unspecialized */
+  Static = 34, //< static unspecialized
   DynamicChunked = 35,
+  GuidedChunked = 36, //< guided unspecialized
+  Runtime = 37,
+  Auto = 38, //< auto
+
   ModifierNonmonotonic =
       (1 << 30), /**< Set if the nonmonotonic schedule modifier was present */
+
   LLVM_MARK_AS_BITMASK_ENUM(/* LargestValue */ ModifierNonmonotonic)
 };
 

include/llvm/Frontend/OpenMP/OMPIRBuilder.h

@@ -387,6 +387,7 @@ public:
   /// \param CLI      A descriptor of the canonical loop to workshare.
   /// \param AllocaIP An insertion point for Alloca instructions usable in the
   ///                 preheader of the loop.
+  /// \param SchedType Type of scheduling to be passed to the init function.
   /// \param NeedsBarrier Indicates whether a barrier must be insterted after
   ///                     the loop.
   /// \param Chunk    The size of loop chunk considered as a unit when
@@ -396,6 +397,7 @@ public:
   InsertPointTy createDynamicWorkshareLoop(const LocationDescription &Loc,
                                            CanonicalLoopInfo *CLI,
                                            InsertPointTy AllocaIP,
+                                           omp::OMPScheduleType SchedType,
                                            bool NeedsBarrier,
                                            Value *Chunk = nullptr);
 

lib/Frontend/OpenMP/OMPIRBuilder.cpp

@@ -1384,7 +1384,8 @@ getKmpcForDynamicNextForType(Type *Ty, Module &M, OpenMPIRBuilder &OMPBuilder) {
 
 OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::createDynamicWorkshareLoop(
     const LocationDescription &Loc, CanonicalLoopInfo *CLI,
-    InsertPointTy AllocaIP, bool NeedsBarrier, Value *Chunk) {
+    InsertPointTy AllocaIP, OMPScheduleType SchedType, bool NeedsBarrier,
+    Value *Chunk) {
   // Set up the source location value for OpenMP runtime.
   Builder.SetCurrentDebugLocation(Loc.DL);
 
@@ -1431,7 +1432,7 @@ OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::createDynamicWorkshareLoop(
   Value *ThreadNum = getOrCreateThreadID(SrcLoc);
 
   OMPScheduleType DynamicSchedType =
-      OMPScheduleType::DynamicChunked | OMPScheduleType::ModifierNonmonotonic;
+      SchedType | OMPScheduleType::ModifierNonmonotonic;
   Constant *SchedulingType =
       ConstantInt::get(I32Type, static_cast<int>(DynamicSchedType));
 

unittests/Frontend/OpenMPIRBuilderTest.cpp

@@ -149,6 +149,10 @@ protected:
   DebugLoc DL;
 };
 
+class OpenMPIRBuilderTestWithParams
+    : public OpenMPIRBuilderTest,
+      public ::testing::WithParamInterface<omp::OMPScheduleType> {};
+
 // Returns the value stored in the given allocation. Returns null if the given
 // value is not a result of an allocation, if no value is stored or if there is
 // more than one store.
@@ -1708,18 +1712,34 @@ TEST_F(OpenMPIRBuilderTest, StaticWorkShareLoop) {
   EXPECT_EQ(NumCallsInExitBlock, 3u);
 }
 
-TEST_F(OpenMPIRBuilderTest, DynamicWorkShareLoop) {
+TEST_P(OpenMPIRBuilderTestWithParams, DynamicWorkShareLoop) {
   using InsertPointTy = OpenMPIRBuilder::InsertPointTy;
   OpenMPIRBuilder OMPBuilder(*M);
   OMPBuilder.initialize();
   IRBuilder<> Builder(BB);
   OpenMPIRBuilder::LocationDescription Loc({Builder.saveIP(), DL});
 
+  omp::OMPScheduleType SchedType = GetParam();
+  uint32_t ChunkSize = 1;
+  switch (SchedType) {
+  case omp::OMPScheduleType::DynamicChunked:
+  case omp::OMPScheduleType::GuidedChunked:
+    ChunkSize = 7;
+    break;
+  case omp::OMPScheduleType::Auto:
+  case omp::OMPScheduleType::Runtime:
+    ChunkSize = 1;
+    break;
+  default:
+    assert(0 && "unknown type for this test");
+    break;
+  }
+
   Type *LCTy = Type::getInt32Ty(Ctx);
   Value *StartVal = ConstantInt::get(LCTy, 10);
   Value *StopVal = ConstantInt::get(LCTy, 52);
   Value *StepVal = ConstantInt::get(LCTy, 2);
-  Value *ChunkVal = ConstantInt::get(LCTy, 7);
+  Value *ChunkVal = ConstantInt::get(LCTy, ChunkSize);
   auto LoopBodyGen = [&](InsertPointTy, llvm::Value *) {};
 
   CanonicalLoopInfo *CLI = OMPBuilder.createCanonicalLoop(
@@ -1737,7 +1757,7 @@ TEST_F(OpenMPIRBuilderTest, DynamicWorkShareLoop) {
   Value *IV = CLI->getIndVar();
 
   InsertPointTy EndIP =
-      OMPBuilder.createDynamicWorkshareLoop(Loc, CLI, AllocaIP,
+      OMPBuilder.createDynamicWorkshareLoop(Loc, CLI, AllocaIP, SchedType,
                                             /*NeedsBarrier=*/true, ChunkVal);
   // The returned value should be the "after" point.
   ASSERT_EQ(EndIP.getBlock(), AfterIP.getBlock());
@@ -1775,7 +1795,7 @@ TEST_F(OpenMPIRBuilderTest, DynamicWorkShareLoop) {
             "__kmpc_dispatch_init_4u");
   EXPECT_EQ(InitCall->getNumArgOperands(), 7U);
   EXPECT_EQ(InitCall->getArgOperand(6),
-            ConstantInt::get(Type::getInt32Ty(Ctx), 7));
+            ConstantInt::get(Type::getInt32Ty(Ctx), ChunkSize));
 
   ConstantInt *OrigLowerBound =
       dyn_cast<ConstantInt>(LowerBoundStore->getValueOperand());
@@ -1807,6 +1827,13 @@ TEST_F(OpenMPIRBuilderTest, DynamicWorkShareLoop) {
   EXPECT_FALSE(verifyModule(*M, &errs()));
 }
 
+INSTANTIATE_TEST_CASE_P(OpenMPWSLoopSchedulingTypes,
+                        OpenMPIRBuilderTestWithParams,
+                        ::testing::Values(omp::OMPScheduleType::DynamicChunked,
+                                          omp::OMPScheduleType::GuidedChunked,
+                                          omp::OMPScheduleType::Auto,
+                                          omp::OMPScheduleType::Runtime));
+
 TEST_F(OpenMPIRBuilderTest, MasterDirective) {
   using InsertPointTy = OpenMPIRBuilder::InsertPointTy;
   OpenMPIRBuilder OMPBuilder(*M);