[OpenMPIRBuilder] introduce createStaticWorkshareLoop

[ftynse]

Summary

Introduce a function that creates a statically-scheduled workshare loop
out of a canonical loop created earlier by the OpenMPIRBuilder. This
basically amounts to injecting runtime calls to the preheader and the
after block and updating the trip count. Static scheduling kind is
currently hardcoded and needs to be extracted from the runtime library
into common TableGen definitions.

[jdoerfert]

Only Nits from me. @Meinersbur wdyt?

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

270

The preheader (only) solution is for static schedules but that is OK for now.

667

nice!

751

@Meinersbur: Eventually we might want to keep pointers identifying these different parts as we will add instructions to the blocks and therefore might break these invariants. Not necessary now though.

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

1006

Nit: Ty or IVType

1019

I'm not a fan of LoopInfo as a name, CLI could be our "go to" acronym?

1089

Add a TODO above this block stating this should eventually move to the CLI or a "CLIUpdate" interface.

[Meinersbur]

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

270

The method;s name is "createStaticWorkshareLoop"

281

Return the modified CanonicalLoopInfo instead. The IP can be determined using getAfterIP().

Mostly for consistency with additional createXYZLoop methods, which create a new CanonicalLoopInfo instead of modifying an existing one.

746

Maybe make this member private (to the OpenMPIRBuilder implementation)? Other canonical loops such as tiling or OpenMPIRBuilder itself assuming a specific trip count during finialize() might depend on a specific trip count.

There's nothing like this now nor do I plan to add such things, but until we have a use case why front-end should directly modify the trip-cout, I'd hide this from the API.

750

add a call to assertOK() here?

751

Uses of it within the CanonicalLoop still needs to be identified, for changes like this. This makes explicitly storing it redundant. Same would also apply to most BasicBlocks which have a fixed position relative to the header block, so we could traverse the CFG to find them instead of storing them.

I don't have a good reason for either direction, I suggest to do what's easier for now.

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

1019

+1

LoopInfo clashes with llvm::LoopInfo and clang::LoopInfo types. The latter already makes it impossible to include some LLVM headers into CGLoopInfo.cpp

1092–1093

Could you emit the fini call into the getExit block instead? getAfterIP is meant for user code that executes, well, after the runtime boilerplate has finished.

1095

Call LoopInfo->assertOK() when the changes are done?

[kiranchandramohan]

Thanks @ftynse for this patch. This is great and will speed us up by a month.

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

282

No change required but just asking a question for my understanding,
Can createWorkshareLoop call createCanonicalLoop internally?

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

1094

A barrier insertion with a TODO to look at nowait attribute?

[Meinersbur]

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

1088

Did you consider using replaceAllUsesWith? setTripCount should already have removed the instances ub cond and latch.

[Meinersbur]

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

282

In the patch I am working on, I extracted most of createCanonicalLoop into a private createCanonicalLoopSkeleton to be used by other createXYZLoop methods (in my case: tiling).

However, this patch modifies an existing CanonicalLoopInfo, instead of creating a new one.

[ftynse]

Address reviews

[ftynse]

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

270

I only implemented the static schedule for now. My idea was to have createStatic, createDynamic, etc and eventually add a createWorkshareLoop that additionally takes a schedule kind flag and dispatches to the right function.

282

If I understood @Meinersbur's comment on another patch correctly, the overall design is to create canonical loops and transform them later (e.g., tile or workshare) with dedicated functions. createStaticWorkshareLoop follows this idea and therefore does _not_ create the loop, only the worksharing constructs.

746

Moved to implementation.

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

1088

This replaces the uses of the induction variable (phi of the header block) with induction-variable + lower-bound-from-runtime-call, setTripCount updated the trip count (another argument in icmp in the condition block).

Switched to replaceAllUsesWithIf though.

1094

Added a flag to the function. Let's have the caller decide if they need a barrier or not. For example, clang doesn't insert a barrier after the last loop in the "parallel" segment even if the loop doesn't have "nowait".

[SouraVX]

@ftynse would you mind having a look at D87247 ?

[kiranchandramohan]

LGTM. Please wait for approval from @Meinersbur.

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

1106–1108

Nit: There is a create barrier function in the OpenMPIRBuilder.
CreateBarrier(Builder.saveIP(), llvm::omp::OMPD_barrier);

[fghanim]

Thanks for the patch. Just a couple of nits

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

1006

I strongly suspect that @jdoerfert was referring to argument Type* type

1032

Nit: To keep a consistent behavior throughout the builder, please change to:

if (!updateToLocation(Loc))
    return Loc.IP;

1047–1049

NIT: please name all alloca's above - easier to track in IR.

[ftynse]

Address more reviews.

[ftynse]

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

1032

It's already inconsistent, createCanonicalLoop doesn't follow this pattern, and we cannot return Loc.IP here because the return type is CanonicalLoopInfo *. Returning nullptr instead.

1106–1108

Yes, but it would emit another call to get_thread_num and OMPD_barrier could misled somebody to think the emitted code comes from a barrier directive, which it is not.

[jdoerfert]

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

1106–1108

but it would emit another call to get_thread_num

That is not a problem we address here. They are deduplicated later.

OMPD_barrier could misled somebody to think the emitted code comes from a barrier directive, which it is not.

The OpenMP runtime is integrated with existing tooling. Barriers "know" what caused them and that is available to profilers, debuggers, etc. We should not break with this tooling capability. The worksharing type for barriers is available via:
createBarrier(Loc, OMPD_for, /* ForceSimpleCall */ false, /* CheckCancelFlag */ false)
The last argument is false only for now. The cancel flag capabilities will be needed later as well, at least from the OpenMP frontends.

[ftynse]

Address even more review

[Meinersbur]

LGTM

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

public:
   ///
   /// \returns An object representing the created control flow structure which
   ///          can be used for loop-associated directives.
   CanonicalLoopInfo *createCanonicalLoop(const LocationDescription &Loc,
                                          LoopBodyGenCallbackTy BodyGenCB,
                                          Value *Start, Value *Stop, Value *Step,
                                          bool IsSigned, bool InclusiveStop);
 
+  /// Modifies the canonical loop to be a statically-scheduled 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
+  /// runtime function in the preheader to obtain the loop bounds to be used in
+  /// the current thread, updates the relevant instructions in the canonical
+  /// loop and calls to an OpenMP runtime finalization function after the loop.
+  ///
+  /// \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 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 Updated CanonicalLoopInfo.
+  CanonicalLoopInfo *createStaticWorkshareLoop(const LocationDescription &Loc,
+                                               CanonicalLoopInfo *CLI,
+                                               InsertPointTy AllocaIP,
+                                               bool NeedsBarrier,
+                                               Value *Chunk = nullptr);
+
   /// Generator for '#omp flush'
   ///
   /// \param Loc The location where the flush directive was encountered
   void createFlush(const LocationDescription &Loc);
 
   /// Generator for '#omp taskwait'
   ///
   /// \param Loc The location where the taskwait directive was encountered.

 ///
 ///     Preheader
 ///        |
 ///  /-> Header
 ///  |     |
 ///  |    Cond---\
 ///  |     |     |
 ///  |    Body   |
-///  |     |     |
+///  |    | |    |
+///  |   <...>   |
+///  |    | |    |
 ///   \--Latch   |
 ///              |
 ///             Exit
 ///              |
 ///            After
 ///
 /// Code in the header, condition block, latch and exit block must not have any
-/// side-effect.
+/// side-effect. The body block is the single entry point into the loop body,
+/// which may contain arbitrary control flow as long as all control paths
+/// eventually branch to the latch block.
 ///
 /// Defined outside OpenMPIRBuilder because one cannot forward-declare nested
 /// classes.
 class CanonicalLoopInfo {
   friend class OpenMPIRBuilder;
 
 private:
   /// Whether this object currently represents a loop.
public:
 
   /// The after block is intended for clean-up code such as lifetime end
   /// markers. It is separate from the exit block to ensure, analogous to the
   /// preheader, it having just a single entry edge and being free from PHI
   /// nodes should there be multiple loop exits (such as from break
   /// statements/cancellations).
   BasicBlock *getAfter() const { return After; }
 
-  /// Returns the llvm::Value containing the number of loop iterations. I must
+  /// Returns the llvm::Value containing the number of loop iterations. It must
   /// be valid in the preheader and always interpreted as an unsigned integer of
   /// any bit-width.
   Value *getTripCount() const {
     Instruction *CmpI = &Cond->front();
     assert(isa<CmpInst>(CmpI) && "First inst must compare IV with TripCount");
     return CmpI->getOperand(1);
   }
 
   /// Returns the instruction representing the current logical induction
   /// variable. Always unsigned, always starting at 0 with an increment of one.
   Instruction *getIndVar() const {
     Instruction *IndVarPHI = &Header->front();
     assert(isa<PHINode>(IndVarPHI) && "First inst must be the IV PHI");
     return IndVarPHI;
   }
 
   /// Return the insertion point for user code after the loop.
   OpenMPIRBuilder::InsertPointTy getAfterIP() const {
     return {After, After->begin()};
   };
 
   /// Consistency self-check.
   void assertOK() const;
 };
 
 } // end namespace llvm
 
 #endif // LLVM_IR_IRBUILDER_H

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

auto BodyGen = [=](InsertPointTy CodeGenIP, Value *IV) {
     Builder.restoreIP(CodeGenIP);
     Value *Span = Builder.CreateMul(IV, Step);
     Value *IndVar = Builder.CreateAdd(Span, Start);
     BodyGenCB(Builder.saveIP(), IndVar);
   };
   return createCanonicalLoop(Builder.saveIP(), BodyGen, TripCount);
 }
 
+// Returns an LLVM function to call for initializing loop bounds using OpenMP
+// static scheduling depending on `type`. Only i32 and i64 are supported by the
+// runtime. Always interpret integers as unsigned similarly to
+// CanonicalLoopInfo.
+static FunctionCallee getKmpcForStaticInitForType(Type *Ty, Module &M,
+                                                  OpenMPIRBuilder &OMPBuilder) {
+  unsigned Bitwidth = Ty->getIntegerBitWidth();
+  if (Bitwidth == 32)
+    return OMPBuilder.getOrCreateRuntimeFunction(
+        M, omp::RuntimeFunction::OMPRTL___kmpc_for_static_init_4u);
+  if (Bitwidth == 64)
+    return OMPBuilder.getOrCreateRuntimeFunction(
+        M, omp::RuntimeFunction::OMPRTL___kmpc_for_static_init_8u);
+  llvm_unreachable("unknown OpenMP loop iterator bitwidth");
+}
+
+// Sets the number of loop iterations to the given value. This value must be
+// valid in the condition block (i.e., defined in the preheader) and is
+// interpreted as an unsigned integer.
+void setCanonicalLoopTripCount(CanonicalLoopInfo *CLI, Value *TripCount) {
+  Instruction *CmpI = &CLI->getCond()->front();
+  assert(isa<CmpInst>(CmpI) && "First inst must compare IV with TripCount");
+  CmpI->setOperand(1, TripCount);
+  CLI->assertOK();
+}
+
+CanonicalLoopInfo *OpenMPIRBuilder::createStaticWorkshareLoop(
+    const LocationDescription &Loc, CanonicalLoopInfo *CLI,
+    InsertPointTy AllocaIP, bool NeedsBarrier, Value *Chunk) {
+  // Set up the source location value for OpenMP runtime.
+  if (!updateToLocation(Loc))
+    return nullptr;
+
+  Constant *SrcLocStr = getOrCreateSrcLocStr(Loc);
+  Value *SrcLoc = getOrCreateIdent(SrcLocStr);
+
+  // Declare useful OpenMP runtime functions.
+  Value *IV = CLI->getIndVar();
+  Type *IVTy = IV->getType();
+  FunctionCallee StaticInit = getKmpcForStaticInitForType(IVTy, M, *this);
+  FunctionCallee StaticFini =
+      getOrCreateRuntimeFunction(M, omp::OMPRTL___kmpc_for_static_fini);
+
+  // Allocate space for computed loop bounds as expected by the "init" function.
+  Builder.restoreIP(AllocaIP);
+  Type *I32Type = Type::getInt32Ty(M.getContext());
+  Value *PLastIter = Builder.CreateAlloca(I32Type, nullptr, "p.lastiter");
+  Value *PLowerBound = Builder.CreateAlloca(IVTy, nullptr, "p.lowerbound");
+  Value *PUpperBound = Builder.CreateAlloca(IVTy, nullptr, "p.upperbound");
+  Value *PStride = Builder.CreateAlloca(IVTy, nullptr, "p.stride");
+
+  // At the end of the preheader, prepare for calling the "init" function by
+  // storing the current loop bounds into the allocated space. A canonical loop
+  // always iterates from 0 to trip-count with step 1. Note that "init" expects
+  // and produces an inclusive upper bound.
+  Builder.SetInsertPoint(CLI->getPreheader()->getTerminator());
+  Constant *Zero = ConstantInt::get(IVTy, 0);
+  Constant *One = ConstantInt::get(IVTy, 1);
+  Builder.CreateStore(Zero, PLowerBound);
+  Value *UpperBound = Builder.CreateSub(CLI->getTripCount(), One);
+  Builder.CreateStore(UpperBound, PUpperBound);
+  Builder.CreateStore(One, PStride);
+
+  if (!Chunk)
+    Chunk = One;
+
+  Value *ThreadNum = getOrCreateThreadID(SrcLoc);
+
+  // TODO: extract scheduling type and map it to OMP constant. This is curently
+  // happening in kmp.h and its ilk and needs to be moved to OpenMP.td first.
+  constexpr int StaticSchedType = 34;
+  Constant *SchedulingType = ConstantInt::get(I32Type, StaticSchedType);
+
+  // Call the "init" function and update the trip count of the loop with the
+  // value it produced.
+  Builder.CreateCall(StaticInit,
+                     {SrcLoc, ThreadNum, SchedulingType, PLastIter, PLowerBound,
+                      PUpperBound, PStride, One, Chunk});
+  Value *LowerBound = Builder.CreateLoad(PLowerBound);
+  Value *InclusiveUpperBound = Builder.CreateLoad(PUpperBound);
+  Value *TripCountMinusOne = Builder.CreateSub(InclusiveUpperBound, LowerBound);
+  Value *TripCount = Builder.CreateAdd(TripCountMinusOne, One);
+  setCanonicalLoopTripCount(CLI, TripCount);
+
+  // Update all uses of the induction variable except the one in the condition
+  // block that compares it with the actual upper bound, and the increment in
+  // the latch block.
+  // TODO: this can eventually move to CanonicalLoopInfo or to a new
+  // CanonicalLoopInfoUpdater interface.
+  Builder.SetInsertPoint(CLI->getBody(), CLI->getBody()->getFirstInsertionPt());
+  Value *UpdatedIV = Builder.CreateAdd(IV, LowerBound);
+  IV->replaceUsesWithIf(UpdatedIV, [&](Use &U) {
+    auto *Instr = dyn_cast<Instruction>(U.getUser());
+    return !Instr ||
+           (Instr->getParent() != CLI->getCond() &&
+            Instr->getParent() != CLI->getLatch() && Instr != UpdatedIV);
+  });
+
+  // In the "exit" block, call the "fini" function.
+  Builder.SetInsertPoint(CLI->getExit(),
+                         CLI->getExit()->getTerminator()->getIterator());
+  Builder.CreateCall(StaticFini, {SrcLoc, ThreadNum});
+
+  // Add the barrier if requested.
+  if (NeedsBarrier)
+    createBarrier(Loc, omp::Directive::OMPD_for, /* ForceSimpleCall */ false,
+                  /* CheckCancelFlag */ false);
+
+  CLI->assertOK();
+  return CLI;
+}
+
 void CanonicalLoopInfo::eraseFromParent() {
   assert(IsValid && "can only erase previously valid loop cfg");
   IsValid = false;
 
   SmallVector<BasicBlock *, 5> BBsToRemove{Header, Cond, Latch, Exit};
   SmallVector<Instruction *, 16> InstsToRemove;
 
   // Only remove preheader if not re-purposed somewhere else.

llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp

TEST_F(OpenMPIRBuilderTest, CanonicalLoopBounds) {
   ASSERT_EQ(EvalTripCount(0x7FFF, -1, INT16_MIN, true, true), 2);
 
   // Finalize the function and verify it.
   Builder.CreateRetVoid();
   OMPBuilder.finalize();
   EXPECT_FALSE(verifyModule(*M, &errs()));
 }
 
+TEST_F(OpenMPIRBuilderTest, StaticWorkShareLoop) {
+  using InsertPointTy = OpenMPIRBuilder::InsertPointTy;
+  OpenMPIRBuilder OMPBuilder(*M);
+  OMPBuilder.initialize();
+  IRBuilder<> Builder(BB);
+  OpenMPIRBuilder::LocationDescription Loc({Builder.saveIP(), DL});
+
+  Type *LCTy = Type::getInt32Ty(Ctx);
+  Value *StartVal = ConstantInt::get(LCTy, 10);
+  Value *StopVal = ConstantInt::get(LCTy, 52);
+  Value *StepVal = ConstantInt::get(LCTy, 2);
+  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();
+
+  CLI = OMPBuilder.createStaticWorkshareLoop(Loc, CLI, AllocaIP,
+                                             /*NeedsBarrier=*/true);
+  auto AllocaIter = BB->begin();
+  ASSERT_GE(std::distance(BB->begin(), BB->end()), 4);
+  AllocaInst *PLastIter = dyn_cast<AllocaInst>(&*(AllocaIter++));
+  AllocaInst *PLowerBound = dyn_cast<AllocaInst>(&*(AllocaIter++));
+  AllocaInst *PUpperBound = dyn_cast<AllocaInst>(&*(AllocaIter++));
+  AllocaInst *PStride = dyn_cast<AllocaInst>(&*(AllocaIter++));
+  EXPECT_NE(PLastIter, nullptr);
+  EXPECT_NE(PLowerBound, nullptr);
+  EXPECT_NE(PUpperBound, nullptr);
+  EXPECT_NE(PStride, nullptr);
+
+  auto PreheaderIter = CLI->getPreheader()->begin();
+  ASSERT_GE(
+      std::distance(CLI->getPreheader()->begin(), CLI->getPreheader()->end()),
+      7);
+  StoreInst *LowerBoundStore = dyn_cast<StoreInst>(&*(PreheaderIter++));
+  StoreInst *UpperBoundStore = dyn_cast<StoreInst>(&*(PreheaderIter++));
+  StoreInst *StrideStore = dyn_cast<StoreInst>(&*(PreheaderIter++));
+  ASSERT_NE(LowerBoundStore, nullptr);
+  ASSERT_NE(UpperBoundStore, nullptr);
+  ASSERT_NE(StrideStore, nullptr);
+
+  auto *OrigLowerBound =
+      dyn_cast<ConstantInt>(LowerBoundStore->getValueOperand());
+  auto *OrigUpperBound =
+      dyn_cast<ConstantInt>(UpperBoundStore->getValueOperand());
+  auto *OrigStride = dyn_cast<ConstantInt>(StrideStore->getValueOperand());
+  ASSERT_NE(OrigLowerBound, nullptr);
+  ASSERT_NE(OrigUpperBound, nullptr);
+  ASSERT_NE(OrigStride, nullptr);
+  EXPECT_EQ(OrigLowerBound->getValue(), 0);
+  EXPECT_EQ(OrigUpperBound->getValue(), 20);
+  EXPECT_EQ(OrigStride->getValue(), 1);
+
+  // Check that the loop IV is updated to account for the lower bound returned
+  // by the OpenMP runtime call.
+  BinaryOperator *Add = dyn_cast<BinaryOperator>(&CLI->getBody()->front());
+  EXPECT_EQ(Add->getOperand(0), CLI->getIndVar());
+  auto *LoadedLowerBound = dyn_cast<LoadInst>(Add->getOperand(1));
+  ASSERT_NE(LoadedLowerBound, nullptr);
+  EXPECT_EQ(LoadedLowerBound->getPointerOperand(), PLowerBound);
+
+  // Check that the trip count is updated to account for the lower and upper
+  // bounds return by the OpenMP runtime call.
+  auto *AddOne = dyn_cast<Instruction>(CLI->getTripCount());
+  ASSERT_NE(AddOne, nullptr);
+  ASSERT_TRUE(AddOne->isBinaryOp());
+  auto *One = dyn_cast<ConstantInt>(AddOne->getOperand(1));
+  ASSERT_NE(One, nullptr);
+  EXPECT_EQ(One->getValue(), 1);
+  auto *Difference = dyn_cast<Instruction>(AddOne->getOperand(0));
+  ASSERT_NE(Difference, nullptr);
+  ASSERT_TRUE(Difference->isBinaryOp());
+  EXPECT_EQ(Difference->getOperand(1), LoadedLowerBound);
+  auto *LoadedUpperBound = dyn_cast<LoadInst>(Difference->getOperand(0));
+  ASSERT_NE(LoadedUpperBound, nullptr);
+  EXPECT_EQ(LoadedUpperBound->getPointerOperand(), PUpperBound);
+
+  // The original loop iterator should only be used in the condition, in the
+  // increment and in the statement that adds the lower bound to it.
+  Value *IV = CLI->getIndVar();
+  EXPECT_EQ(std::distance(IV->use_begin(), IV->use_end()), 3);
+}
+
 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});