diff --git a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h b/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h --- a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h +++ b/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h @@ -103,6 +103,11 @@ using BodyGenCallbackTy = function_ref; + // This is created primarily for sections construct as llvm::function_ref + // are not storable + using StorableBodyGenCallbackTy = + std::function; /// Callback type for loop body code generation. /// @@ -643,6 +648,34 @@ 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 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. @@ -762,16 +795,17 @@ /// 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 diff --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp --- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp +++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp @@ -865,6 +865,138 @@ emitTaskyieldImpl(Loc); } +OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::createSections( + const LocationDescription &Loc, InsertPointTy AllocaIP, + ArrayRef SectionCBs, PrivatizeCallbackTy PrivCB, + FinalizeCallbackTy FiniCB, bool IsCancellable, bool IsNowait) { + if (!updateToLocation(Loc)) + return Loc.IP; + + auto FiniCBWrapper = [&](InsertPointTy 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. + if (IP.getBlock()->end() == IP.getPoint()) { + // 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 in 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: + // ; + // break; + // ... + // case - 1: + // - 1]>; + // break; + // } + // ... + // section_loop.after: + // ; + 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(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) { + // 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. + if (IP.getBlock()->end() == IP.getPoint()) { + // 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, @@ -1643,10 +1775,10 @@ 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 @@ -1711,9 +1843,8 @@ 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(); @@ -1763,6 +1894,9 @@ Builder.SetInsertPoint(FiniBBTI); } + if (!ExitCall) + return Builder.saveIP(); + // place the Exitcall as last instruction before Finalization block terminator ExitCall->removeFromParent(); Builder.Insert(ExitCall); diff --git a/llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp b/llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp --- a/llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp +++ b/llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp @@ -1984,4 +1984,126 @@ 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 SectionCBVector; + llvm::SmallVector 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(Inst)) { + if (cast(&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(Inst)) { + if (cast(&Inst)->getCalledFunction()->getName() == + "__kmpc_for_static_fini") { + FoundForExit = true; + } + if (cast(&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(SwitchBB->getTerminator()), true); + Switch = cast(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