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

[Leporacanthicus]

Summary

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.

[ftynse]

The MLIR part LGTM after comments are addressed.

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

185–187

Haven't we covered all clauses?

266

llvm_unreachable

270–271

Could we have a test that at list hits this code path? Does not need to be detailed, just the check that the calls to runtime functions are actually emitted, similarly to what we have for static loops.

[jdoerfert]

2 nits, LGTM otherwise.

llvm/include/llvm/Frontend/OpenMP/OMPConstants.h

118

Can we move all of the entries here and use //< ... comments maybe.

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

1387–1388

[Leporacanthicus]

Will do some minor fixes tomorrow, thanks for the feedback!

llvm/include/llvm/Frontend/OpenMP/OMPConstants.h

118

I'll take a look at moving the others as a separate patch - will fix the comment style in this patch to remind myself that it needs fixing for the others! :)

llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp

1760

Should use SchedType rather than GetParam()

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

185–187

Good point. Probably good enough to cover that in the default of the switch at the bottom.

270–271

The test code that I modified does this - I added the test when adding support for Dynamic scheduling, and this only updated the test to take the scheduling type as a parameter - because that's the only difference from this side - the OpenMP runtime will do different things [I think!]

[Leporacanthicus]

Updates based on review comments:

• Fix style using PascalCase rather than camelCase.
• Add Doxygen comment for parameter.
• Change comment style for constants.
• Remove useless error output - there was no way to get to that error.
• Use llvm_unreachable instead of assert(0).

[Leporacanthicus]

Adding MLIR to LLVM-IR tests. The tests are very basic, just checking that
the correct runtime functions are called.

[Leporacanthicus]

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

270–271

I have added some MLIR to LLVM-IR tests that check that the correct runtime calls are made, as well as the above unit-test.

llvm/include/llvm/Frontend/OpenMP/OMPConstants.h

for (const AssumptionClauseMappingInfo &ACMI : AssumptionClauseMappings)
     S += (S.empty() ? "'" : "', '") + ACMI.Identifier.str();
   return S + "'";
 }
 
 /// \note This needs to be kept in sync with kmp.h enum sched_type.
 /// 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)
 };
 
 } // end namespace omp
 
 } // end namespace llvm
 
 #endif // LLVM_FRONTEND_OPENMP_OMPCONSTANTS_H

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

public:
   /// turn it into a workshare loop. In particular, it calls to an OpenMP
   /// runtime function in the preheader to obtain, and then in each iteration
   /// to update the loop counter.
   /// \param Loc      The source location description, the insertion location
   ///                 is not used.
   /// \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
   ///                 scheduling. If \p nullptr, defaults to 1.
   ///
   /// \returns Point where to insert code after the loop.
   InsertPointTy createDynamicWorkshareLoop(const LocationDescription &Loc,
                                            CanonicalLoopInfo *CLI,
                                            InsertPointTy AllocaIP,
+                                           omp::OMPScheduleType SchedType,
                                            bool NeedsBarrier,
                                            Value *Chunk = nullptr);
 
   /// Modifies the canonical loop to be a workshare loop.
   ///
   /// This takes a \p LoopInfo representing a canonical loop, such as the one
   /// created by \p createCanonicalLoop and emits additional instructions to
   /// turn it into a workshare loop. In particular, it calls to an OpenMP

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

getKmpcForDynamicNextForType(Type *Ty, Module &M, OpenMPIRBuilder &OMPBuilder) {
   if (Bitwidth == 64)
     return OMPBuilder.getOrCreateRuntimeFunction(
         M, omp::RuntimeFunction::OMPRTL___kmpc_dispatch_next_8u);
   llvm_unreachable("unknown OpenMP loop iterator bitwidth");
 }
 
 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);
 
   Constant *SrcLocStr = getOrCreateSrcLocStr(Loc);
   Value *SrcLoc = getOrCreateIdent(SrcLocStr);
 
   // Declare useful OpenMP runtime functions.
   Value *IV = CLI->getIndVar();
OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::createDynamicWorkshareLoop(
   // a valid canonical loop.
 
   if (!Chunk)
     Chunk = One;
 
   Value *ThreadNum = getOrCreateThreadID(SrcLoc);
 
   OMPScheduleType DynamicSchedType =
-      OMPScheduleType::DynamicChunked | OMPScheduleType::ModifierNonmonotonic;
+      SchedType | OMPScheduleType::ModifierNonmonotonic;
   Constant *SchedulingType =
       ConstantInt::get(I32Type, static_cast<int>(DynamicSchedType));
 
   // Call the "init" function.
   Builder.CreateCall(DynamicInit,
                      {SrcLoc, ThreadNum, SchedulingType, /* LowerBound */ One,
                       UpperBound, /* step */ One, Chunk});
 

llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp

protected:
 
   LLVMContext Ctx;
   std::unique_ptr<Module> M;
   Function *F;
   BasicBlock *BB;
   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.
 static Value *findStoredValue(Value *AllocaValue) {
   Instruction *Alloca = dyn_cast<AllocaInst>(AllocaValue);
   if (!Alloca)
     return nullptr;
   StoreInst *Store = nullptr;
TEST_F(OpenMPIRBuilderTest, StaticWorkShareLoop) {
   // The exit block should contain the "fini" call and the barrier call,
   // plus the call to obtain the thread ID.
   BasicBlock *ExitBlock = CLI->getExit();
   size_t NumCallsInExitBlock =
       count_if(*ExitBlock, [](Instruction &I) { return isa<CallInst>(I); });
   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(
       Loc, LoopBodyGen, StartVal, StopVal, StepVal,
       /*IsSigned=*/false, /*InclusiveStop=*/false);
 
   Builder.SetInsertPoint(BB, BB->getFirstInsertionPt());
   InsertPointTy AllocaIP = Builder.saveIP();
 
   // Collect all the info from CLI, as it isn't usable after the call to
   // createDynamicWorkshareLoop.
   InsertPointTy AfterIP = CLI->getAfterIP();
   BasicBlock *Preheader = CLI->getPreheader();
   BasicBlock *ExitBlock = CLI->getExit();
   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());
   ASSERT_EQ(EndIP.getPoint(), AfterIP.getPoint());
 
   auto AllocaIter = BB->begin();
   ASSERT_GE(std::distance(BB->begin(), BB->end()), 4);
   AllocaInst *PLastIter = dyn_cast<AllocaInst>(&*(AllocaIter++));
TEST_P(OpenMPIRBuilderTestWithParams, DynamicWorkShareLoop) {
 
   CallInst *InitCall = dyn_cast<CallInst>(&*PreheaderIter);
 
   ASSERT_NE(InitCall, nullptr);
   EXPECT_EQ(InitCall->getCalledFunction()->getName(),
             "__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());
   ConstantInt *OrigUpperBound =
       dyn_cast<ConstantInt>(UpperBoundStore->getValueOperand());
   ConstantInt *OrigStride =
       dyn_cast<ConstantInt>(StrideStore->getValueOperand());
   ASSERT_NE(OrigLowerBound, nullptr);
TEST_P(OpenMPIRBuilderTestWithParams, DynamicWorkShareLoop) {
 
   // Add a termination to our block and check that it is internally consistent.
   Builder.restoreIP(EndIP);
   Builder.CreateRetVoid();
   OMPBuilder.finalize();
   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);
   OMPBuilder.initialize();
   F->setName("func");
   IRBuilder<> Builder(BB);
 
   OpenMPIRBuilder::LocationDescription Loc({Builder.saveIP(), DL});

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

convertOmpWsLoop(Operation &opInst, llvm::IRBuilderBase &builder,
   auto loop = cast<omp::WsLoopOp>(opInst);
   // TODO: this should be in the op verifier instead.
   if (loop.lowerBound().empty())
     return failure();
 
   if (loop.getNumLoops() != 1)
     return opInst.emitOpError("collapsed loops not yet supported");
 
-  bool isStatic = true;
-
-  if (loop.schedule_val().hasValue()) {
-    auto schedule =
-        omp::symbolizeClauseScheduleKind(loop.schedule_val().getValue());
-    if (schedule != omp::ClauseScheduleKind::Static &&
-        schedule != omp::ClauseScheduleKind::Dynamic)
-      return opInst.emitOpError("only static (default) and dynamic loop "
-                                "schedule is currently supported");
-    isStatic = (schedule == omp::ClauseScheduleKind::Static);
-  }
+  // Static is the default.
+  omp::ClauseScheduleKind schedule = omp::ClauseScheduleKind::Static;
+  if (loop.schedule_val().hasValue())
+    schedule =
+        *omp::symbolizeClauseScheduleKind(loop.schedule_val().getValue());
 
   // Find the loop configuration.
   llvm::Value *lowerBound = moduleTranslation.lookupValue(loop.lowerBound()[0]);
   llvm::Value *upperBound = moduleTranslation.lookupValue(loop.upperBound()[0]);
   llvm::Value *step = moduleTranslation.lookupValue(loop.step()[0]);
   llvm::Type *ivType = step->getType();
   llvm::Value *chunk =
       loop.schedule_chunk_var()
           ? moduleTranslation.lookupValue(loop.schedule_chunk_var())
convertOmpWsLoop(Operation &opInst, llvm::IRBuilderBase &builder,
   // TODO: get the alloca insertion point from the parallel operation builder.
   // If we insert the at the top of the current function, they will be passed as
   // extra arguments into the function the parallel operation builder outlines.
   // Put them at the start of the current block for now.
   llvm::OpenMPIRBuilder::InsertPointTy allocaIP(
       insertBlock, insertBlock->getFirstInsertionPt());
   llvm::OpenMPIRBuilder::InsertPointTy afterIP;
   llvm::OpenMPIRBuilder *ompBuilder = moduleTranslation.getOpenMPBuilder();
-  if (isStatic) {
+  if (schedule == omp::ClauseScheduleKind::Static) {
     loopInfo = ompBuilder->createStaticWorkshareLoop(ompLoc, loopInfo, allocaIP,
                                                      !loop.nowait(), chunk);
     afterIP = loopInfo->getAfterIP();
   } else {
-    afterIP = ompBuilder->createDynamicWorkshareLoop(ompLoc, loopInfo, allocaIP,
-                                                     !loop.nowait(), chunk);
+    llvm::omp::OMPScheduleType schedType;
+    switch (schedule) {
+    case omp::ClauseScheduleKind::Dynamic:
+      schedType = llvm::omp::OMPScheduleType::DynamicChunked;
+      break;
+    case omp::ClauseScheduleKind::Guided:
+      schedType = llvm::omp::OMPScheduleType::GuidedChunked;
+      break;
+    case omp::ClauseScheduleKind::Auto:
+      schedType = llvm::omp::OMPScheduleType::Auto;
+      break;
+    case omp::ClauseScheduleKind::Runtime:
+      schedType = llvm::omp::OMPScheduleType::Runtime;
+      break;
+    default:
+      llvm_unreachable("Unknown schedule value");
+      break;
+    }
+
+    afterIP = ompBuilder->createDynamicWorkshareLoop(
+        ompLoc, loopInfo, allocaIP, schedType, !loop.nowait(), chunk);
   }
 
   // Continue building IR after the loop.
   builder.restoreIP(afterIP);
   return success();
 }
 
 namespace {

mlir/test/Target/LLVMIR/openmp-llvm.mlir

llvm.func @wsloop_inclusive_2(%arg0: !llvm.ptr<f32>) {
   ^bb0(%arg1: i64):
     %3 = llvm.mlir.constant(2.000000e+00 : f32) : f32
     %4 = llvm.getelementptr %arg0[%arg1] : (!llvm.ptr<f32>, i64) -> !llvm.ptr<f32>
     llvm.store %3, %4 : !llvm.ptr<f32>
     omp.yield
   }) {inclusive, operand_segment_sizes = dense<[1, 1, 1, 0, 0, 0, 0, 0, 0]> : vector<9xi32>} : (i64, i64, i64) -> ()
   llvm.return
 }
+
+llvm.func @body(i64)
+
+llvm.func @test_omp_wsloop_dynamic(%lb : i64, %ub : i64, %step : i64) -> () {
+ omp.wsloop (%iv) : i64 = (%lb) to (%ub) step (%step) schedule(dynamic) {
+  // CHECK: call void @__kmpc_dispatch_init_8u
+  // CHECK: %[[continue:.*]] = call i32 @__kmpc_dispatch_next_8u
+  // CHECK: %[[cond:.*]] = icmp ne i32 %[[continue]], 0
+  // CHECK  br i1 %[[cond]], label %omp_loop.header{{.*}}, label %omp_loop.exit{{.*}}
+   llvm.call @body(%iv) : (i64) -> ()
+   omp.yield
+ }
+ llvm.return
+}
+
+llvm.func @test_omp_wsloop_auto(%lb : i64, %ub : i64, %step : i64) -> () {
+ omp.wsloop (%iv) : i64 = (%lb) to (%ub) step (%step) schedule(auto) {
+  // CHECK: call void @__kmpc_dispatch_init_8u
+  // CHECK: %[[continue:.*]] = call i32 @__kmpc_dispatch_next_8u
+  // CHECK: %[[cond:.*]] = icmp ne i32 %[[continue]], 0
+  // CHECK  br i1 %[[cond]], label %omp_loop.header{{.*}}, label %omp_loop.exit{{.*}}
+   llvm.call @body(%iv) : (i64) -> ()
+   omp.yield
+ }
+ llvm.return
+}
+
+llvm.func @test_omp_wsloop_runtime(%lb : i64, %ub : i64, %step : i64) -> () {
+ omp.wsloop (%iv) : i64 = (%lb) to (%ub) step (%step) schedule(runtime) {
+  // CHECK: call void @__kmpc_dispatch_init_8u
+  // CHECK: %[[continue:.*]] = call i32 @__kmpc_dispatch_next_8u
+  // CHECK: %[[cond:.*]] = icmp ne i32 %[[continue]], 0
+  // CHECK  br i1 %[[cond]], label %omp_loop.header{{.*}}, label %omp_loop.exit{{.*}}
+   llvm.call @body(%iv) : (i64) -> ()
+   omp.yield
+ }
+ llvm.return
+}
+
+llvm.func @test_omp_wsloop_guided(%lb : i64, %ub : i64, %step : i64) -> () {
+ omp.wsloop (%iv) : i64 = (%lb) to (%ub) step (%step) schedule(guided) {
+  // CHECK: call void @__kmpc_dispatch_init_8u
+  // CHECK: %[[continue:.*]] = call i32 @__kmpc_dispatch_next_8u
+  // CHECK: %[[cond:.*]] = icmp ne i32 %[[continue]], 0
+  // CHECK  br i1 %[[cond]], label %omp_loop.header{{.*}}, label %omp_loop.exit{{.*}}
+   llvm.call @body(%iv) : (i64) -> ()
+   omp.yield
+ }
+ llvm.return
+}