Handle dist_schedule in unique workshare loop.

It should be noted that at this stage, the wsloop is inserted at
the wrong location, and the insertion point needs to be handled to
ensure the `.ll` output is the same as Clang.

Author: Stylie777

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

@@ -1098,15 +1098,12 @@ class OpenMPIRBuilder {
   /// \param LoopType Type of workshare loop.
   /// \param HasDistSchedule Defines if the clause being lowered is
   /// dist_schedule as this is handled slightly differently
-  /// \param DistScheduleSchedType Defines the Schedule Type for the Distribute
-  /// loop. Defaults to None if no Distribute loop is present.
   ///
   /// \returns Point where to insert code after the workshare construct.
   InsertPointOrErrorTy applyStaticWorkshareLoop(
       DebugLoc DL, CanonicalLoopInfo *CLI, InsertPointTy AllocaIP,
       omp::WorksharingLoopType LoopType, bool NeedsBarrier,
-      bool HasDistSchedule = false,
-      omp::OMPScheduleType DistScheduleSchedType = omp::OMPScheduleType::None);
+      bool HasDistSchedule = false);
 
   /// Modifies the canonical loop a statically-scheduled workshare loop with a
   /// user-specified chunk size.
@@ -1121,20 +1118,13 @@ class OpenMPIRBuilder {
   /// \param ChunkSize    The user-specified chunk size.
   /// \param SchedType    Optional type of scheduling to be passed to the init
   /// function.
-  /// \param DistScheduleChunkSize    The size of dist_shcedule chunk considered
-  /// as a unit when
-  ///                 scheduling. If \p nullptr, defaults to 1.
-  /// \param DistScheduleSchedType Defines the Schedule Type for the Distribute
-  /// loop. Defaults to None if no Distribute loop is present.
   ///
   /// \returns Point where to insert code after the workshare construct.
   InsertPointOrErrorTy applyStaticChunkedWorkshareLoop(
       DebugLoc DL, CanonicalLoopInfo *CLI, InsertPointTy AllocaIP,
       bool NeedsBarrier, Value *ChunkSize,
       omp::OMPScheduleType SchedType =
-          omp::OMPScheduleType::UnorderedStaticChunked,
-      Value *DistScheduleChunkSize = nullptr,
-      omp::OMPScheduleType DistScheduleSchedType = omp::OMPScheduleType::None);
+          omp::OMPScheduleType::UnorderedStaticChunked);
 
   /// Modifies the canonical loop to be a dynamically-scheduled workshare loop.
   ///
@@ -1216,8 +1206,6 @@ class OpenMPIRBuilder {
   /// \param HasDistSchedule Defines if the clause being lowered is
   /// dist_schedule as this is handled slightly differently
   ///
-  /// \param ChunkSize The chunk size for dist_schedule loop
-  ///
   /// \returns Point where to insert code after the workshare construct.
   LLVM_ABI InsertPointOrErrorTy applyWorkshareLoop(
       DebugLoc DL, CanonicalLoopInfo *CLI, InsertPointTy AllocaIP,
@@ -1228,7 +1216,7 @@ class OpenMPIRBuilder {
       bool HasOrderedClause = false,
       omp::WorksharingLoopType LoopType =
           omp::WorksharingLoopType::ForStaticLoop,
-      bool HasDistSchedule = false, Value *DistScheduleChunkSize = nullptr);
+      bool HasDistSchedule = false);
 
   /// Tile a loop nest.
   ///

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

@@ -172,7 +172,7 @@ static const omp::GV &getGridValue(const Triple &T, Function *Kernel) {
 /// arguments.
 static OMPScheduleType
 getOpenMPBaseScheduleType(llvm::omp::ScheduleKind ClauseKind, bool HasChunks,
-                          bool HasSimdModifier, bool HasDistScheduleChunks) {
+                          bool HasSimdModifier) {
   // Currently, the default schedule it static.
   switch (ClauseKind) {
   case OMP_SCHEDULE_Default:
@@ -190,7 +190,7 @@ getOpenMPBaseScheduleType(llvm::omp::ScheduleKind ClauseKind, bool HasChunks,
     return HasSimdModifier ? OMPScheduleType::BaseRuntimeSimd
                            : OMPScheduleType::BaseRuntime;
   case OMP_SCHEDULE_Distribute:
-    return HasDistScheduleChunks ? OMPScheduleType::BaseDistributeChunked
+    return HasChunks ? OMPScheduleType::BaseDistributeChunked
                                  : OMPScheduleType::BaseDistribute;
   }
   llvm_unreachable("unhandled schedule clause argument");
@@ -260,10 +260,9 @@ getOpenMPMonotonicityScheduleType(OMPScheduleType ScheduleType,
 static OMPScheduleType
 computeOpenMPScheduleType(ScheduleKind ClauseKind, bool HasChunks,
                           bool HasSimdModifier, bool HasMonotonicModifier,
-                          bool HasNonmonotonicModifier, bool HasOrderedClause,
-                          bool HasDistScheduleChunks) {
+                          bool HasNonmonotonicModifier, bool HasOrderedClause) {
   OMPScheduleType BaseSchedule = getOpenMPBaseScheduleType(
-      ClauseKind, HasChunks, HasSimdModifier, HasDistScheduleChunks);
+      ClauseKind, HasChunks, HasSimdModifier);
   OMPScheduleType OrderedSchedule =
       getOpenMPOrderingScheduleType(BaseSchedule, HasOrderedClause);
   OMPScheduleType Result = getOpenMPMonotonicityScheduleType(
@@ -4643,8 +4642,7 @@ static FunctionCallee getKmpcForStaticInitForType(Type *Ty, Module &M,
 
 OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::applyStaticWorkshareLoop(
     DebugLoc DL, CanonicalLoopInfo *CLI, InsertPointTy AllocaIP,
-    WorksharingLoopType LoopType, bool NeedsBarrier, bool HasDistSchedule,
-    OMPScheduleType DistScheduleSchedType) {
+    WorksharingLoopType LoopType, bool NeedsBarrier, bool HasDistSchedule) {
   assert(CLI->isValid() && "Requires a valid canonical loop");
   assert(!isConflictIP(AllocaIP, CLI->getPreheaderIP()) &&
          "Require dedicated allocate IP");
@@ -4714,12 +4712,6 @@ OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::applyStaticWorkshareLoop(
     Builder.CreateCall(StaticInit, Args);
   };
   BuildInitCall(SchedulingType, Builder);
-  if (HasDistSchedule &&
-      LoopType != WorksharingLoopType::DistributeStaticLoop) {
-    Constant *DistScheduleSchedType = ConstantInt::get(
-        I32Type, static_cast<int>(omp::OMPScheduleType::OrderedDistribute));
-    BuildInitCall(DistScheduleSchedType, Builder);
-  }
   Value *LowerBound = Builder.CreateLoad(IVTy, PLowerBound);
   Value *InclusiveUpperBound = Builder.CreateLoad(IVTy, PUpperBound);
   Value *TripCountMinusOne = Builder.CreateSub(InclusiveUpperBound, LowerBound);
@@ -4792,10 +4784,9 @@ static void applyParallelAccessesMetadata(CanonicalLoopInfo *CLI,
 OpenMPIRBuilder::InsertPointOrErrorTy
 OpenMPIRBuilder::applyStaticChunkedWorkshareLoop(
     DebugLoc DL, CanonicalLoopInfo *CLI, InsertPointTy AllocaIP,
-    bool NeedsBarrier, Value *ChunkSize, OMPScheduleType SchedType,
-    Value *DistScheduleChunkSize, OMPScheduleType DistScheduleSchedType) {
+    bool NeedsBarrier, Value *ChunkSize, OMPScheduleType SchedType) {
   assert(CLI->isValid() && "Requires a valid canonical loop");
-  assert(ChunkSize || DistScheduleChunkSize && "Chunk size is required");
+  assert(ChunkSize && "Chunk size is required");
 
   LLVMContext &Ctx = CLI->getFunction()->getContext();
   Value *IV = CLI->getIndVar();
@@ -4817,7 +4808,7 @@ OpenMPIRBuilder::applyStaticChunkedWorkshareLoop(
   LoopInfo &&LI = LIA.run(*F, FAM);
   Loop *L = LI.getLoopFor(CLI->getHeader());
   SmallVector<Metadata *> LoopMDList;
-  if (ChunkSize || DistScheduleChunkSize)
+  if (ChunkSize)
     applyParallelAccessesMetadata(CLI, Ctx, L, LI, LoopMDList);
   addLoopMetadata(CLI, LoopMDList);
 
@@ -4839,22 +4830,17 @@ OpenMPIRBuilder::applyStaticChunkedWorkshareLoop(
   CLI->setLastIter(PLastIter);
 
   // Set up the source location value for the OpenMP runtime.
-  Builder.restoreIP(CLI->getPreheaderIP());
+  Builder.restoreIP(CLI->getPreheaderIP()); // -> sets insert point to omploop! Why?
   Builder.SetCurrentDebugLocation(DL);
 
   // TODO: Detect overflow in ubsan or max-out with current tripcount.
   Value *CastedChunkSize = Builder.CreateZExtOrTrunc(
       ChunkSize ? ChunkSize : Zero, InternalIVTy, "chunksize");
-  Value *CastestDistScheduleChunkSize = Builder.CreateZExtOrTrunc(
-      DistScheduleChunkSize ? DistScheduleChunkSize : Zero, InternalIVTy,
-      "distschedulechunksize");
   Value *CastedTripCount =
       Builder.CreateZExt(OrigTripCount, InternalIVTy, "tripcount");
 
   Constant *SchedulingType =
       ConstantInt::get(I32Type, static_cast<int>(SchedType));
-  Constant *DistSchedulingType =
-      ConstantInt::get(I32Type, static_cast<int>(DistScheduleSchedType));
   Builder.CreateStore(Zero, PLowerBound);
   Value *OrigUpperBound = Builder.CreateSub(CastedTripCount, One);
   Builder.CreateStore(OrigUpperBound, PUpperBound);
@@ -4877,14 +4863,6 @@ OpenMPIRBuilder::applyStaticChunkedWorkshareLoop(
                          /*chunk=*/ChunkSize});
       };
   BuildInitCall(SchedulingType, CastedChunkSize, Builder);
-  if (DistScheduleSchedType != OMPScheduleType::None &&
-      SchedType != OMPScheduleType::OrderedDistributeChunked &&
-      SchedType != OMPScheduleType::OrderedDistribute) {
-    // We want to emit a second init function call for the dist_schedule clause
-    // to the Distribute construct. This should only be done however if a
-    // Workshare Loop is nested within a Distribute Construct
-    BuildInitCall(DistSchedulingType, CastestDistScheduleChunkSize, Builder);
-  }
 
   // Load values written by the "init" function.
   Value *FirstChunkStart =
@@ -5208,35 +5186,27 @@ OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::applyWorkshareLoop(
     bool NeedsBarrier, omp::ScheduleKind SchedKind, Value *ChunkSize,
     bool HasSimdModifier, bool HasMonotonicModifier,
     bool HasNonmonotonicModifier, bool HasOrderedClause,
-    WorksharingLoopType LoopType, bool HasDistSchedule,
-    Value *DistScheduleChunkSize) {
+    WorksharingLoopType LoopType, bool HasDistSchedule) {
   if (Config.isTargetDevice())
     return applyWorkshareLoopTarget(DL, CLI, AllocaIP, LoopType);
   OMPScheduleType EffectiveScheduleType = computeOpenMPScheduleType(
       SchedKind, ChunkSize, HasSimdModifier, HasMonotonicModifier,
-      HasNonmonotonicModifier, HasOrderedClause, DistScheduleChunkSize);
+      HasNonmonotonicModifier, HasOrderedClause);
 
   bool IsOrdered = (EffectiveScheduleType & OMPScheduleType::ModifierOrdered) ==
                    OMPScheduleType::ModifierOrdered;
-  OMPScheduleType DistScheduleSchedType = OMPScheduleType::None;
-  if (HasDistSchedule) {
-    DistScheduleSchedType = DistScheduleChunkSize
-                                ? OMPScheduleType::OrderedDistributeChunked
-                                : OMPScheduleType::OrderedDistribute;
-  }
   switch (EffectiveScheduleType & ~OMPScheduleType::ModifierMask) {
   case OMPScheduleType::BaseStatic:
   case OMPScheduleType::BaseDistribute:
-    assert(!ChunkSize || !DistScheduleChunkSize &&
+    assert(!ChunkSize &&
                              "No chunk size with static-chunked schedule");
     if (IsOrdered && !HasDistSchedule)
       return applyDynamicWorkshareLoop(DL, CLI, AllocaIP, EffectiveScheduleType,
                                        NeedsBarrier, ChunkSize);
     // FIXME: Monotonicity ignored?
-    if (DistScheduleChunkSize)
+    if (ChunkSize)
       return applyStaticChunkedWorkshareLoop(
-          DL, CLI, AllocaIP, NeedsBarrier, ChunkSize, EffectiveScheduleType,
-          DistScheduleChunkSize, DistScheduleSchedType);
+          DL, CLI, AllocaIP, NeedsBarrier, ChunkSize, EffectiveScheduleType);
     return applyStaticWorkshareLoop(DL, CLI, AllocaIP, LoopType, NeedsBarrier,
                                     HasDistSchedule);
 
@@ -5247,8 +5217,7 @@ OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::applyWorkshareLoop(
                                        NeedsBarrier, ChunkSize);
     // FIXME: Monotonicity ignored?
     return applyStaticChunkedWorkshareLoop(
-        DL, CLI, AllocaIP, NeedsBarrier, ChunkSize, EffectiveScheduleType,
-        DistScheduleChunkSize, DistScheduleSchedType);
+        DL, CLI, AllocaIP, NeedsBarrier, ChunkSize, EffectiveScheduleType);
 
   case OMPScheduleType::BaseRuntime:
   case OMPScheduleType::BaseAuto:

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

@@ -2464,6 +2464,7 @@ convertOmpTaskwaitOp(omp::TaskwaitOp twOp, llvm::IRBuilderBase &builder,
 static LogicalResult
 convertOmpWsloop(Operation &opInst, llvm::IRBuilderBase &builder,
                  LLVM::ModuleTranslation &moduleTranslation) {
+  printf("CONVERTING WSLOOP\n");
   llvm::OpenMPIRBuilder *ompBuilder = moduleTranslation.getOpenMPBuilder();
   auto wsloopOp = cast<omp::WsloopOp>(opInst);
   if (failed(checkImplementationStatus(opInst)))
@@ -2487,19 +2488,6 @@ convertOmpWsloop(Operation &opInst, llvm::IRBuilderBase &builder,
     chunk = builder.CreateSExtOrTrunc(chunkVar, ivType);
   }
 
-  omp::DistributeOp distributeOp = nullptr;
-  llvm::Value *distScheduleChunk = nullptr;
-  bool hasDistSchedule = false;
-  if (llvm::isa_and_present<omp::DistributeOp>(opInst.getParentOp())) {
-    distributeOp = cast<omp::DistributeOp>(opInst.getParentOp());
-    hasDistSchedule = distributeOp.getDistScheduleStatic();
-    if (distributeOp.getDistScheduleChunkSize()) {
-      llvm::Value *chunkVar = moduleTranslation.lookupValue(
-          distributeOp.getDistScheduleChunkSize());
-      distScheduleChunk = builder.CreateSExtOrTrunc(chunkVar, ivType);
-    }
-  }
-
   PrivateVarsInfo privateVarsInfo(wsloopOp);
 
   SmallVector<omp::DeclareReductionOp> reductionDecls;
@@ -2600,13 +2588,15 @@ convertOmpWsloop(Operation &opInst, llvm::IRBuilderBase &builder,
   }
 
   builder.SetInsertPoint(*regionBlock, (*regionBlock)->begin());
+      printf("loopInfo Address: %p\n", loopInfo);
+      printf("Applying omp.wloop Workshare Loop\n");
   llvm::OpenMPIRBuilder::InsertPointOrErrorTy wsloopIP =
       ompBuilder->applyWorkshareLoop(
           ompLoc.DL, loopInfo, allocaIP, loopNeedsBarrier,
           convertToScheduleKind(schedule), chunk, isSimd,
           scheduleMod == omp::ScheduleModifier::monotonic,
           scheduleMod == omp::ScheduleModifier::nonmonotonic, isOrdered,
-          workshareLoopType, hasDistSchedule, distScheduleChunk);
+          workshareLoopType);
 
   if (failed(handleError(wsloopIP, opInst)))
     return failure();
@@ -3052,6 +3042,7 @@ convertOmpLoopNest(Operation &opInst, llvm::IRBuilderBase &builder,
 
   // Update the stack frame created for this loop to point to the resulting loop
   // after applying transformations.
+  printf("Applying loopInfo\n");
   moduleTranslation.stackWalk<OpenMPLoopInfoStackFrame>(
       [&](OpenMPLoopInfoStackFrame &frame) {
         frame.loopInfo = ompBuilder->collapseLoops(ompLoc.DL, loopInfos, {});
@@ -4767,6 +4758,7 @@ convertOmpTargetData(Operation *op, llvm::IRBuilderBase &builder,
 static LogicalResult
 convertOmpDistribute(Operation &opInst, llvm::IRBuilderBase &builder,
                      LLVM::ModuleTranslation &moduleTranslation) {
+  printf("CONVERTING DISTRIBUTE\n");
   llvm::OpenMPIRBuilder *ompBuilder = moduleTranslation.getOpenMPBuilder();
   auto distributeOp = cast<omp::DistributeOp>(opInst);
   if (failed(checkImplementationStatus(opInst)))
@@ -4835,21 +4827,20 @@ convertOmpDistribute(Operation &opInst, llvm::IRBuilderBase &builder,
     llvm::OpenMPIRBuilder::LocationDescription ompLoc(builder);
     llvm::Expected<llvm::BasicBlock *> regionBlock =
         convertOmpOpRegions(distributeOp.getRegion(), "omp.distribute.region",
-                            builder, moduleTranslation);
+                            builder, moduleTranslation); // -> this is causing Schedule to be emitted first.
     if (!regionBlock)
       return regionBlock.takeError();
     builder.SetInsertPoint(*regionBlock, (*regionBlock)->begin());
 
     // Skip applying a workshare loop below when translating 'distribute
     // parallel do' (it's been already handled by this point while translating
     // the nested omp.wsloop).
-    if (!isa_and_present<omp::WsloopOp>(distributeOp.getNestedWrapper())) {
+    if (!isa_and_present<omp::WsloopOp>(distributeOp.getNestedWrapper()) || distributeOp.getDistScheduleStatic()) {
       // TODO: Add support for clauses which are valid for DISTRIBUTE
       // constructs. Static schedule is the default.
       bool hasDistSchedule = distributeOp.getDistScheduleStatic();
       auto schedule = hasDistSchedule ? omp::ClauseScheduleKind::Distribute
                                       : omp::ClauseScheduleKind::Static;
-      // dist_schedule clauses are ordered - otherise this should be false
       bool isOrdered = hasDistSchedule;
       std::optional<omp::ScheduleModifier> scheduleMod;
       bool isSimd = false;
@@ -4859,14 +4850,17 @@ convertOmpDistribute(Operation &opInst, llvm::IRBuilderBase &builder,
       llvm::Value *chunk = moduleTranslation.lookupValue(
           distributeOp.getDistScheduleChunkSize());
       llvm::CanonicalLoopInfo *loopInfo =
-          findCurrentLoopInfo(moduleTranslation);
+          findCurrentLoopInfo(moduleTranslation); // Do we need a new loop info here?
+      printf("loopInfo Address: %p\n", loopInfo);
+      printf("InsertPoint Name : %s\n", builder.GetInsertBlock()->getName().str().c_str());
+      printf("Applying omp.ditribute Workshare Loop\n");
       llvm::OpenMPIRBuilder::InsertPointOrErrorTy wsloopIP =
           ompBuilder->applyWorkshareLoop(
               ompLoc.DL, loopInfo, allocaIP, loopNeedsBarrier,
               convertToScheduleKind(schedule), chunk, isSimd,
               scheduleMod == omp::ScheduleModifier::monotonic,
               scheduleMod == omp::ScheduleModifier::nonmonotonic, isOrdered,
-              workshareLoopType, hasDistSchedule, chunk);
+              workshareLoopType, hasDistSchedule);
 
       if (!wsloopIP)
         return wsloopIP.takeError();
@@ -5907,7 +5901,7 @@ convertHostOrTargetOperation(Operation *op, llvm::IRBuilderBase &builder,
       !dyn_cast_if_present<omp::LoopWrapperInterface>(op->getParentOp());
 
   if (isOutermostLoopWrapper)
-    moduleTranslation.stackPush<OpenMPLoopInfoStackFrame>();
+    moduleTranslation.stackPush<OpenMPLoopInfoStackFrame>(); // -> Need another one of these when Distribute AND WSLoop is present?
 
   auto result =
       llvm::TypeSwitch<Operation *, LogicalResult>(op)