[LLVM][OpenMP] Adding support for OpenMP sections construct in OpenMPIRBuilder

This patch adds section support in the OpenMP IRBuilder module, along with a test for the same.

Reviewed By: fghanim

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

include/llvm/Frontend/OpenMP/OMPIRBuilder.h

@@ -104,6 +104,14 @@ public:
       function_ref<void(InsertPointTy AllocaIP, InsertPointTy CodeGenIP,
                         BasicBlock &ContinuationBB)>;
 
+  // This is created primarily for sections construct as llvm::function_ref
+  // (BodyGenCallbackTy) is not storable (as described in the comments of
+  // function_ref class - function_ref contains non-ownable reference
+  // to the callable.
+  using StorableBodyGenCallbackTy =
+      std::function<void(InsertPointTy AllocaIP, InsertPointTy CodeGenIP,
+                         BasicBlock &ContinuationBB)>;
+
   /// Callback type for loop body code generation.
   ///
   /// \param CodeGenIP is the insertion point where the loop's body code must be
@@ -678,6 +686,34 @@ public:
                                FinalizeCallbackTy FiniCB,
                                StringRef CriticalName, Value *HintInst);
 
+  /// Generator for '#omp sections'
+  ///
+  /// \param Loc The insert and source location description.
+  /// \param AllocaIP The insertion points to be used for alloca instructions.
+  /// \param SectionCBs Callbacks that will generate body of each section.
+  /// \param PrivCB Callback to copy a given variable (think copy constructor).
+  /// \param FiniCB Callback to finalize variable copies.
+  /// \param IsCancellable Flag to indicate a cancellable parallel region.
+  /// \param IsNowait If true, barrier - to ensure all sections are executed
+  /// before moving forward will not be generated.
+  /// \returns The insertion position *after* the sections.
+  InsertPointTy createSections(const LocationDescription &Loc,
+                               InsertPointTy AllocaIP,
+                               ArrayRef<StorableBodyGenCallbackTy> SectionCBs,
+                               PrivatizeCallbackTy PrivCB,
+                               FinalizeCallbackTy FiniCB, bool IsCancellable,
+                               bool IsNowait);
+
+  /// Generator for '#omp section'
+  ///
+  /// \param Loc The insert and source location description.
+  /// \param BodyGenCB Callback that will generate the region body code.
+  /// \param FiniCB Callback to finalize variable copies.
+  /// \returns The insertion position *after* the section.
+  InsertPointTy createSection(const LocationDescription &Loc,
+                              BodyGenCallbackTy BodyGenCB,
+                              FinalizeCallbackTy FiniCB);
+
   /// Generate conditional branch and relevant BasicBlocks through which private
   /// threads copy the 'copyin' variables from Master copy to threadprivate
   /// copies.
@@ -797,16 +833,17 @@ private:
   ///        to evaluate a conditional of whether a thread will execute
   ///        body code or not.
   /// \param HasFinalize indicate if the directive will require finalization
-  ///         and has a finalization callback in the stack that
-  /// should        be called.
-  ///
+  ///        and has a finalization callback in the stack that
+  ///        should be called.
+  /// \param IsCancellable if HasFinalize is set to true, indicate if the
+  ///        the directive should be cancellable.
   /// \return The insertion point after the region
 
   InsertPointTy
   EmitOMPInlinedRegion(omp::Directive OMPD, Instruction *EntryCall,
                        Instruction *ExitCall, BodyGenCallbackTy BodyGenCB,
                        FinalizeCallbackTy FiniCB, bool Conditional = false,
-                       bool HasFinalize = true);
+                       bool HasFinalize = true, bool IsCancellable = false);
 
   /// Get the platform-specific name separator.
   /// \param Parts different parts of the final name that needs separation

lib/Frontend/OpenMP/OMPIRBuilder.cpp

@@ -877,6 +877,138 @@ void OpenMPIRBuilder::createTaskyield(const LocationDescription &Loc) {
   emitTaskyieldImpl(Loc);
 }
 
+OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::createSections(
+    const LocationDescription &Loc, InsertPointTy AllocaIP,
+    ArrayRef<StorableBodyGenCallbackTy> SectionCBs, PrivatizeCallbackTy PrivCB,
+    FinalizeCallbackTy FiniCB, bool IsCancellable, bool IsNowait) {
+  if (!updateToLocation(Loc))
+    return Loc.IP;
+
+  auto FiniCBWrapper = [&](InsertPointTy IP) {
+    if (IP.getBlock()->end() != IP.getPoint())
+      return FiniCB(IP);
+    // This must be done otherwise any nested constructs using FinalizeOMPRegion
+    // will fail because that function requires the Finalization Basic Block to
+    // have a terminator, which is already removed by EmitOMPRegionBody.
+    // IP is currently at cancelation block.
+    // We need to backtrack to the condition block to fetch
+    // the exit block and create a branch from cancelation
+    // to exit block.
+    IRBuilder<>::InsertPointGuard IPG(Builder);
+    Builder.restoreIP(IP);
+    auto *CaseBB = IP.getBlock()->getSinglePredecessor();
+    auto *CondBB = CaseBB->getSinglePredecessor()->getSinglePredecessor();
+    auto *ExitBB = CondBB->getTerminator()->getSuccessor(1);
+    Instruction *I = Builder.CreateBr(ExitBB);
+    IP = InsertPointTy(I->getParent(), I->getIterator());
+    return FiniCB(IP);
+  };
+
+  FinalizationStack.push_back({FiniCBWrapper, OMPD_sections, IsCancellable});
+
+  // Each section is emitted as a switch case
+  // Each finalization callback is handled from clang.EmitOMPSectionDirective()
+  // -> OMP.createSection() which generates the IR for each section
+  // Iterate through all sections and emit a switch construct:
+  // switch (IV) {
+  //   case 0:
+  //     <SectionStmt[0]>;
+  //     break;
+  // ...
+  //   case <NumSection> - 1:
+  //     <SectionStmt[<NumSection> - 1]>;
+  //     break;
+  // }
+  // ...
+  // section_loop.after:
+  // <FiniCB>;
+  auto LoopBodyGenCB = [&](InsertPointTy CodeGenIP, Value *IndVar) {
+    auto *CurFn = CodeGenIP.getBlock()->getParent();
+    auto *ForIncBB = CodeGenIP.getBlock()->getSingleSuccessor();
+    auto *ForExitBB = CodeGenIP.getBlock()
+                          ->getSinglePredecessor()
+                          ->getTerminator()
+                          ->getSuccessor(1);
+    SwitchInst *SwitchStmt = Builder.CreateSwitch(IndVar, ForIncBB);
+    Builder.restoreIP(CodeGenIP);
+    unsigned CaseNumber = 0;
+    for (auto SectionCB : SectionCBs) {
+      auto *CaseBB = BasicBlock::Create(M.getContext(),
+                                        "omp_section_loop.body.case", CurFn);
+      SwitchStmt->addCase(Builder.getInt32(CaseNumber), CaseBB);
+      Builder.SetInsertPoint(CaseBB);
+      SectionCB(InsertPointTy(), Builder.saveIP(), *ForExitBB);
+      CaseNumber++;
+    }
+    // remove the existing terminator from body BB since there can be no
+    // terminators after switch/case
+    CodeGenIP.getBlock()->getTerminator()->eraseFromParent();
+  };
+  // Loop body ends here
+  // LowerBound, UpperBound, and STride for createCanonicalLoop
+  Type *I32Ty = Type::getInt32Ty(M.getContext());
+  Value *LB = ConstantInt::get(I32Ty, 0);
+  Value *UB = ConstantInt::get(I32Ty, SectionCBs.size());
+  Value *ST = ConstantInt::get(I32Ty, 1);
+  llvm::CanonicalLoopInfo *LoopInfo = createCanonicalLoop(
+      Loc, LoopBodyGenCB, LB, UB, ST, true, false, AllocaIP, "section_loop");
+  LoopInfo = createStaticWorkshareLoop(Loc, LoopInfo, AllocaIP, true);
+  BasicBlock *LoopAfterBB = LoopInfo->getAfter();
+  Instruction *SplitPos = LoopAfterBB->getTerminator();
+  if (!isa_and_nonnull<BranchInst>(SplitPos))
+    SplitPos = new UnreachableInst(Builder.getContext(), LoopAfterBB);
+  // ExitBB after LoopAfterBB because LoopAfterBB is used for FinalizationCB,
+  // which requires a BB with branch
+  BasicBlock *ExitBB =
+      LoopAfterBB->splitBasicBlock(SplitPos, "omp_sections.end");
+  SplitPos->eraseFromParent();
+
+  // Apply the finalization callback in LoopAfterBB
+  auto FiniInfo = FinalizationStack.pop_back_val();
+  assert(FiniInfo.DK == OMPD_sections &&
+         "Unexpected finalization stack state!");
+  Builder.SetInsertPoint(LoopAfterBB->getTerminator());
+  FiniInfo.FiniCB(Builder.saveIP());
+  Builder.SetInsertPoint(ExitBB);
+
+  return Builder.saveIP();
+}
+
+OpenMPIRBuilder::InsertPointTy
+OpenMPIRBuilder::createSection(const LocationDescription &Loc,
+                               BodyGenCallbackTy BodyGenCB,
+                               FinalizeCallbackTy FiniCB) {
+  if (!updateToLocation(Loc))
+    return Loc.IP;
+
+  auto FiniCBWrapper = [&](InsertPointTy IP) {
+    if (IP.getBlock()->end() != IP.getPoint())
+      return FiniCB(IP);
+    // This must be done otherwise any nested constructs using FinalizeOMPRegion
+    // will fail because that function requires the Finalization Basic Block to
+    // have a terminator, which is already removed by EmitOMPRegionBody.
+    // IP is currently at cancelation block.
+    // We need to backtrack to the condition block to fetch
+    // the exit block and create a branch from cancelation
+    // to exit block.
+    IRBuilder<>::InsertPointGuard IPG(Builder);
+    Builder.restoreIP(IP);
+    auto *CaseBB = Loc.IP.getBlock();
+    auto *CondBB = CaseBB->getSinglePredecessor()->getSinglePredecessor();
+    auto *ExitBB = CondBB->getTerminator()->getSuccessor(1);
+    Instruction *I = Builder.CreateBr(ExitBB);
+    IP = InsertPointTy(I->getParent(), I->getIterator());
+    return FiniCB(IP);
+  };
+
+  Directive OMPD = Directive::OMPD_sections;
+  // Since we are using Finalization Callback here, HasFinalize
+  // and IsCancellable have to be true
+  return EmitOMPInlinedRegion(OMPD, nullptr, nullptr, BodyGenCB, FiniCBWrapper,
+                              /*Conditional*/ false, /*hasFinalize*/ true,
+                              /*IsCancellable*/ true);
+}
+
 OpenMPIRBuilder::InsertPointTy
 OpenMPIRBuilder::createMaster(const LocationDescription &Loc,
                               BodyGenCallbackTy BodyGenCB,
@@ -1819,10 +1951,10 @@ OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::createCritical(
 OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::EmitOMPInlinedRegion(
     Directive OMPD, Instruction *EntryCall, Instruction *ExitCall,
     BodyGenCallbackTy BodyGenCB, FinalizeCallbackTy FiniCB, bool Conditional,
-    bool HasFinalize) {
+    bool HasFinalize, bool IsCancellable) {
 
   if (HasFinalize)
-    FinalizationStack.push_back({FiniCB, OMPD, /*IsCancellable*/ false});
+    FinalizationStack.push_back({FiniCB, OMPD, IsCancellable});
 
   // Create inlined region's entry and body blocks, in preparation
   // for conditional creation
@@ -1887,9 +2019,8 @@ OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::EmitOMPInlinedRegion(
 
 OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::emitCommonDirectiveEntry(
     Directive OMPD, Value *EntryCall, BasicBlock *ExitBB, bool Conditional) {
-
   // if nothing to do, Return current insertion point.
-  if (!Conditional)
+  if (!Conditional || !EntryCall)
     return Builder.saveIP();
 
   BasicBlock *EntryBB = Builder.GetInsertBlock();
@@ -1939,6 +2070,9 @@ OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::emitCommonDirectiveExit(
     Builder.SetInsertPoint(FiniBBTI);
   }
 
+  if (!ExitCall)
+    return Builder.saveIP();
+
   // place the Exitcall as last instruction before Finalization block terminator
   ExitCall->removeFromParent();
   Builder.Insert(ExitCall);

unittests/Frontend/OpenMPIRBuilderTest.cpp

@@ -2165,4 +2165,126 @@ TEST_F(OpenMPIRBuilderTest, SingleDirective) {
   EXPECT_EQ(SingleEndCI->getArgOperand(1), SingleEntryCI->getArgOperand(1));
 }
 
+TEST_F(OpenMPIRBuilderTest, CreateSections) {
+  using InsertPointTy = OpenMPIRBuilder::InsertPointTy;
+  using BodyGenCallbackTy = llvm::OpenMPIRBuilder::StorableBodyGenCallbackTy;
+  OpenMPIRBuilder OMPBuilder(*M);
+  OMPBuilder.initialize();
+  F->setName("func");
+  IRBuilder<> Builder(BB);
+
+  OpenMPIRBuilder::LocationDescription Loc({Builder.saveIP(), DL});
+  llvm::SmallVector<BodyGenCallbackTy, 4> SectionCBVector;
+  llvm::SmallVector<BasicBlock *, 4> CaseBBs;
+
+  BasicBlock *SwitchBB = nullptr;
+  BasicBlock *ForExitBB = nullptr;
+  BasicBlock *ForIncBB = nullptr;
+  AllocaInst *PrivAI = nullptr;
+  SwitchInst *Switch = nullptr;
+
+  unsigned NumBodiesGenerated = 0;
+  unsigned NumFiniCBCalls = 0;
+  PrivAI = Builder.CreateAlloca(F->arg_begin()->getType());
+
+  auto FiniCB = [&](InsertPointTy IP) {
+    ++NumFiniCBCalls;
+    BasicBlock *IPBB = IP.getBlock();
+    EXPECT_NE(IPBB->end(), IP.getPoint());
+  };
+
+  auto SectionCB = [&](InsertPointTy AllocaIP, InsertPointTy CodeGenIP,
+                       BasicBlock &FiniBB) {
+    ++NumBodiesGenerated;
+    CaseBBs.push_back(CodeGenIP.getBlock());
+    SwitchBB = CodeGenIP.getBlock()->getSinglePredecessor();
+    Builder.restoreIP(CodeGenIP);
+    Builder.CreateStore(F->arg_begin(), PrivAI);
+    Value *PrivLoad =
+        Builder.CreateLoad(F->arg_begin()->getType(), PrivAI, "local.alloca");
+    Builder.CreateICmpNE(F->arg_begin(), PrivLoad);
+    Builder.CreateBr(&FiniBB);
+    ForIncBB =
+        CodeGenIP.getBlock()->getSinglePredecessor()->getSingleSuccessor();
+  };
+  auto PrivCB = [](InsertPointTy AllocaIP, InsertPointTy CodeGenIP,
+                   llvm::Value &, llvm::Value &Val, llvm::Value *&ReplVal) {
+    // TODO: Privatization not implemented yet
+    return CodeGenIP;
+  };
+
+  SectionCBVector.push_back(SectionCB);
+  SectionCBVector.push_back(SectionCB);
+
+  IRBuilder<>::InsertPoint AllocaIP(&F->getEntryBlock(),
+                                    F->getEntryBlock().getFirstInsertionPt());
+  Builder.restoreIP(OMPBuilder.createSections(Loc, AllocaIP, SectionCBVector,
+                                              PrivCB, FiniCB, false, false));
+  Builder.CreateRetVoid(); // Required at the end of the function
+
+  // Switch BB's predecessor is loop condition BB, whose successor at index 1 is
+  // loop's exit BB
+  ForExitBB =
+      SwitchBB->getSinglePredecessor()->getTerminator()->getSuccessor(1);
+  EXPECT_NE(ForExitBB, nullptr);
+
+  EXPECT_NE(PrivAI, nullptr);
+  Function *OutlinedFn = PrivAI->getFunction();
+  EXPECT_EQ(F, OutlinedFn);
+  EXPECT_FALSE(verifyModule(*M, &errs()));
+  EXPECT_EQ(OutlinedFn->arg_size(), 1U);
+  EXPECT_EQ(OutlinedFn->getBasicBlockList().size(), 11);
+
+  BasicBlock *LoopPreheaderBB =
+      OutlinedFn->getEntryBlock().getSingleSuccessor();
+  // loop variables are 5 - lower bound, upper bound, stride, islastiter, and
+  // iterator/counter
+  bool FoundForInit = false;
+  for (Instruction &Inst : *LoopPreheaderBB) {
+    if (isa<CallInst>(Inst)) {
+      if (cast<CallInst>(&Inst)->getCalledFunction()->getName() ==
+          "__kmpc_for_static_init_4u") {
+        FoundForInit = true;
+      }
+    }
+  }
+  EXPECT_EQ(FoundForInit, true);
+
+  bool FoundForExit = false;
+  bool FoundBarrier = false;
+  for (Instruction &Inst : *ForExitBB) {
+    if (isa<CallInst>(Inst)) {
+      if (cast<CallInst>(&Inst)->getCalledFunction()->getName() ==
+          "__kmpc_for_static_fini") {
+        FoundForExit = true;
+      }
+      if (cast<CallInst>(&Inst)->getCalledFunction()->getName() ==
+          "__kmpc_barrier") {
+        FoundBarrier = true;
+      }
+      if (FoundForExit && FoundBarrier)
+        break;
+    }
+  }
+  EXPECT_EQ(FoundForExit, true);
+  EXPECT_EQ(FoundBarrier, true);
+
+  EXPECT_NE(SwitchBB, nullptr);
+  EXPECT_NE(SwitchBB->getTerminator(), nullptr);
+  EXPECT_EQ(isa<SwitchInst>(SwitchBB->getTerminator()), true);
+  Switch = cast<SwitchInst>(SwitchBB->getTerminator());
+  EXPECT_EQ(Switch->getNumCases(), 2U);
+  EXPECT_NE(ForIncBB, nullptr);
+  EXPECT_EQ(Switch->getSuccessor(0), ForIncBB);
+
+  EXPECT_EQ(CaseBBs.size(), 2U);
+  for (auto *&CaseBB : CaseBBs) {
+    EXPECT_EQ(CaseBB->getParent(), OutlinedFn);
+    EXPECT_EQ(CaseBB->getSingleSuccessor(), ForExitBB);
+  }
+
+  ASSERT_EQ(NumBodiesGenerated, 2U);
+  ASSERT_EQ(NumFiniCBCalls, 1U);
+}
+
 } // namespace