| //===- CoroElide.cpp - Coroutine Frame Allocation Elision Pass ------------===// |
| // |
| // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. |
| // See https://llvm.org/LICENSE.txt for license information. |
| // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
| // |
| //===----------------------------------------------------------------------===// |
| |
| #include "llvm/Transforms/Coroutines/CoroElide.h" |
| #include "CoroInternal.h" |
| #include "llvm/ADT/DenseMap.h" |
| #include "llvm/Analysis/AliasAnalysis.h" |
| #include "llvm/Analysis/InstructionSimplify.h" |
| #include "llvm/IR/Dominators.h" |
| #include "llvm/IR/InstIterator.h" |
| #include "llvm/InitializePasses.h" |
| #include "llvm/Pass.h" |
| #include "llvm/Support/ErrorHandling.h" |
| |
| using namespace llvm; |
| |
| #define DEBUG_TYPE "coro-elide" |
| |
| namespace { |
| // Created on demand if the coro-elide pass has work to do. |
| struct Lowerer : coro::LowererBase { |
| SmallVector<CoroIdInst *, 4> CoroIds; |
| SmallVector<CoroBeginInst *, 1> CoroBegins; |
| SmallVector<CoroAllocInst *, 1> CoroAllocs; |
| SmallVector<CoroSubFnInst *, 4> ResumeAddr; |
| DenseMap<CoroBeginInst *, SmallVector<CoroSubFnInst *, 4>> DestroyAddr; |
| SmallVector<CoroFreeInst *, 1> CoroFrees; |
| SmallPtrSet<const SwitchInst *, 4> CoroSuspendSwitches; |
| |
| Lowerer(Module &M) : LowererBase(M) {} |
| |
| void elideHeapAllocations(Function *F, uint64_t FrameSize, Align FrameAlign, |
| AAResults &AA); |
| bool shouldElide(Function *F, DominatorTree &DT) const; |
| void collectPostSplitCoroIds(Function *F); |
| bool processCoroId(CoroIdInst *, AAResults &AA, DominatorTree &DT); |
| bool hasEscapePath(const CoroBeginInst *, |
| const SmallPtrSetImpl<BasicBlock *> &) const; |
| }; |
| } // end anonymous namespace |
| |
| // Go through the list of coro.subfn.addr intrinsics and replace them with the |
| // provided constant. |
| static void replaceWithConstant(Constant *Value, |
| SmallVectorImpl<CoroSubFnInst *> &Users) { |
| if (Users.empty()) |
| return; |
| |
| // See if we need to bitcast the constant to match the type of the intrinsic |
| // being replaced. Note: All coro.subfn.addr intrinsics return the same type, |
| // so we only need to examine the type of the first one in the list. |
| Type *IntrTy = Users.front()->getType(); |
| Type *ValueTy = Value->getType(); |
| if (ValueTy != IntrTy) { |
| // May need to tweak the function type to match the type expected at the |
| // use site. |
| assert(ValueTy->isPointerTy() && IntrTy->isPointerTy()); |
| Value = ConstantExpr::getBitCast(Value, IntrTy); |
| } |
| |
| // Now the value type matches the type of the intrinsic. Replace them all! |
| for (CoroSubFnInst *I : Users) |
| replaceAndRecursivelySimplify(I, Value); |
| } |
| |
| // See if any operand of the call instruction references the coroutine frame. |
| static bool operandReferences(CallInst *CI, AllocaInst *Frame, AAResults &AA) { |
| for (Value *Op : CI->operand_values()) |
| if (AA.alias(Op, Frame) != NoAlias) |
| return true; |
| return false; |
| } |
| |
| // Look for any tail calls referencing the coroutine frame and remove tail |
| // attribute from them, since now coroutine frame resides on the stack and tail |
| // call implies that the function does not references anything on the stack. |
| // However if it's a musttail call, we cannot remove the tailcall attribute. |
| // It's safe to keep it there as the musttail call is for symmetric transfer, |
| // and by that point the frame should have been destroyed and hence not |
| // interfering with operands. |
| static void removeTailCallAttribute(AllocaInst *Frame, AAResults &AA) { |
| Function &F = *Frame->getFunction(); |
| for (Instruction &I : instructions(F)) |
| if (auto *Call = dyn_cast<CallInst>(&I)) |
| if (Call->isTailCall() && operandReferences(Call, Frame, AA) && |
| !Call->isMustTailCall()) |
| Call->setTailCall(false); |
| } |
| |
| // Given a resume function @f.resume(%f.frame* %frame), returns the size |
| // and expected alignment of %f.frame type. |
| static std::pair<uint64_t, Align> getFrameLayout(Function *Resume) { |
| // Prefer to pull information from the function attributes. |
| auto Size = Resume->getParamDereferenceableBytes(0); |
| auto Align = Resume->getParamAlign(0); |
| |
| // If those aren't given, extract them from the type. |
| if (Size == 0 || !Align) { |
| auto *FrameTy = Resume->arg_begin()->getType()->getPointerElementType(); |
| |
| const DataLayout &DL = Resume->getParent()->getDataLayout(); |
| if (!Size) Size = DL.getTypeAllocSize(FrameTy); |
| if (!Align) Align = DL.getABITypeAlign(FrameTy); |
| } |
| |
| return std::make_pair(Size, *Align); |
| } |
| |
| // Finds first non alloca instruction in the entry block of a function. |
| static Instruction *getFirstNonAllocaInTheEntryBlock(Function *F) { |
| for (Instruction &I : F->getEntryBlock()) |
| if (!isa<AllocaInst>(&I)) |
| return &I; |
| llvm_unreachable("no terminator in the entry block"); |
| } |
| |
| // To elide heap allocations we need to suppress code blocks guarded by |
| // llvm.coro.alloc and llvm.coro.free instructions. |
| void Lowerer::elideHeapAllocations(Function *F, uint64_t FrameSize, |
| Align FrameAlign, AAResults &AA) { |
| LLVMContext &C = F->getContext(); |
| auto *InsertPt = |
| getFirstNonAllocaInTheEntryBlock(CoroIds.front()->getFunction()); |
| |
| // Replacing llvm.coro.alloc with false will suppress dynamic |
| // allocation as it is expected for the frontend to generate the code that |
| // looks like: |
| // id = coro.id(...) |
| // mem = coro.alloc(id) ? malloc(coro.size()) : 0; |
| // coro.begin(id, mem) |
| auto *False = ConstantInt::getFalse(C); |
| for (auto *CA : CoroAllocs) { |
| CA->replaceAllUsesWith(False); |
| CA->eraseFromParent(); |
| } |
| |
| // FIXME: Design how to transmit alignment information for every alloca that |
| // is spilled into the coroutine frame and recreate the alignment information |
| // here. Possibly we will need to do a mini SROA here and break the coroutine |
| // frame into individual AllocaInst recreating the original alignment. |
| const DataLayout &DL = F->getParent()->getDataLayout(); |
| auto FrameTy = ArrayType::get(Type::getInt8Ty(C), FrameSize); |
| auto *Frame = new AllocaInst(FrameTy, DL.getAllocaAddrSpace(), "", InsertPt); |
| Frame->setAlignment(FrameAlign); |
| auto *FrameVoidPtr = |
| new BitCastInst(Frame, Type::getInt8PtrTy(C), "vFrame", InsertPt); |
| |
| for (auto *CB : CoroBegins) { |
| CB->replaceAllUsesWith(FrameVoidPtr); |
| CB->eraseFromParent(); |
| } |
| |
| // Since now coroutine frame lives on the stack we need to make sure that |
| // any tail call referencing it, must be made non-tail call. |
| removeTailCallAttribute(Frame, AA); |
| } |
| |
| bool Lowerer::hasEscapePath(const CoroBeginInst *CB, |
| const SmallPtrSetImpl<BasicBlock *> &TIs) const { |
| const auto &It = DestroyAddr.find(CB); |
| assert(It != DestroyAddr.end()); |
| |
| // Limit the number of blocks we visit. |
| unsigned Limit = 32 * (1 + It->second.size()); |
| |
| SmallVector<const BasicBlock *, 32> Worklist; |
| Worklist.push_back(CB->getParent()); |
| |
| SmallPtrSet<const BasicBlock *, 32> Visited; |
| // Consider basicblock of coro.destroy as visited one, so that we |
| // skip the path pass through coro.destroy. |
| for (auto *DA : It->second) |
| Visited.insert(DA->getParent()); |
| |
| do { |
| const auto *BB = Worklist.pop_back_val(); |
| if (!Visited.insert(BB).second) |
| continue; |
| if (TIs.count(BB)) |
| return true; |
| |
| // Conservatively say that there is potentially a path. |
| if (!--Limit) |
| return true; |
| |
| auto TI = BB->getTerminator(); |
| // Although the default dest of coro.suspend switches is suspend pointer |
| // which means a escape path to normal terminator, it is reasonable to skip |
| // it since coroutine frame doesn't change outside the coroutine body. |
| if (isa<SwitchInst>(TI) && |
| CoroSuspendSwitches.count(cast<SwitchInst>(TI))) { |
| Worklist.push_back(cast<SwitchInst>(TI)->getSuccessor(1)); |
| Worklist.push_back(cast<SwitchInst>(TI)->getSuccessor(2)); |
| } else |
| Worklist.append(succ_begin(BB), succ_end(BB)); |
| |
| } while (!Worklist.empty()); |
| |
| // We have exhausted all possible paths and are certain that coro.begin can |
| // not reach to any of terminators. |
| return false; |
| } |
| |
| bool Lowerer::shouldElide(Function *F, DominatorTree &DT) const { |
| // If no CoroAllocs, we cannot suppress allocation, so elision is not |
| // possible. |
| if (CoroAllocs.empty()) |
| return false; |
| |
| // Check that for every coro.begin there is at least one coro.destroy directly |
| // referencing the SSA value of that coro.begin along each |
| // non-exceptional path. |
| // If the value escaped, then coro.destroy would have been referencing a |
| // memory location storing that value and not the virtual register. |
| |
| SmallPtrSet<BasicBlock *, 8> Terminators; |
| // First gather all of the non-exceptional terminators for the function. |
| // Consider the final coro.suspend as the real terminator when the current |
| // function is a coroutine. |
| for (BasicBlock &B : *F) { |
| auto *TI = B.getTerminator(); |
| if (TI->getNumSuccessors() == 0 && !TI->isExceptionalTerminator() && |
| !isa<UnreachableInst>(TI)) |
| Terminators.insert(&B); |
| } |
| |
| // Filter out the coro.destroy that lie along exceptional paths. |
| SmallPtrSet<CoroBeginInst *, 8> ReferencedCoroBegins; |
| for (auto &It : DestroyAddr) { |
| for (Instruction *DA : It.second) { |
| for (BasicBlock *TI : Terminators) { |
| if (DT.dominates(DA, TI->getTerminator())) { |
| ReferencedCoroBegins.insert(It.first); |
| break; |
| } |
| } |
| } |
| |
| // Whether there is any paths from coro.begin to Terminators which not pass |
| // through any of the coro.destroys. |
| if (!ReferencedCoroBegins.count(It.first) && |
| !hasEscapePath(It.first, Terminators)) |
| ReferencedCoroBegins.insert(It.first); |
| } |
| |
| // If size of the set is the same as total number of coro.begin, that means we |
| // found a coro.free or coro.destroy referencing each coro.begin, so we can |
| // perform heap elision. |
| return ReferencedCoroBegins.size() == CoroBegins.size(); |
| } |
| |
| void Lowerer::collectPostSplitCoroIds(Function *F) { |
| CoroIds.clear(); |
| CoroSuspendSwitches.clear(); |
| for (auto &I : instructions(F)) { |
| if (auto *CII = dyn_cast<CoroIdInst>(&I)) |
| if (CII->getInfo().isPostSplit()) |
| // If it is the coroutine itself, don't touch it. |
| if (CII->getCoroutine() != CII->getFunction()) |
| CoroIds.push_back(CII); |
| |
| // Consider case like: |
| // %0 = call i8 @llvm.coro.suspend(...) |
| // switch i8 %0, label %suspend [i8 0, label %resume |
| // i8 1, label %cleanup] |
| // and collect the SwitchInsts which are used by escape analysis later. |
| if (auto *CSI = dyn_cast<CoroSuspendInst>(&I)) |
| if (CSI->hasOneUse() && isa<SwitchInst>(CSI->use_begin()->getUser())) { |
| SwitchInst *SWI = cast<SwitchInst>(CSI->use_begin()->getUser()); |
| if (SWI->getNumCases() == 2) |
| CoroSuspendSwitches.insert(SWI); |
| } |
| } |
| } |
| |
| bool Lowerer::processCoroId(CoroIdInst *CoroId, AAResults &AA, |
| DominatorTree &DT) { |
| CoroBegins.clear(); |
| CoroAllocs.clear(); |
| CoroFrees.clear(); |
| ResumeAddr.clear(); |
| DestroyAddr.clear(); |
| |
| // Collect all coro.begin and coro.allocs associated with this coro.id. |
| for (User *U : CoroId->users()) { |
| if (auto *CB = dyn_cast<CoroBeginInst>(U)) |
| CoroBegins.push_back(CB); |
| else if (auto *CA = dyn_cast<CoroAllocInst>(U)) |
| CoroAllocs.push_back(CA); |
| else if (auto *CF = dyn_cast<CoroFreeInst>(U)) |
| CoroFrees.push_back(CF); |
| } |
| |
| // Collect all coro.subfn.addrs associated with coro.begin. |
| // Note, we only devirtualize the calls if their coro.subfn.addr refers to |
| // coro.begin directly. If we run into cases where this check is too |
| // conservative, we can consider relaxing the check. |
| for (CoroBeginInst *CB : CoroBegins) { |
| for (User *U : CB->users()) |
| if (auto *II = dyn_cast<CoroSubFnInst>(U)) |
| switch (II->getIndex()) { |
| case CoroSubFnInst::ResumeIndex: |
| ResumeAddr.push_back(II); |
| break; |
| case CoroSubFnInst::DestroyIndex: |
| DestroyAddr[CB].push_back(II); |
| break; |
| default: |
| llvm_unreachable("unexpected coro.subfn.addr constant"); |
| } |
| } |
| |
| // PostSplit coro.id refers to an array of subfunctions in its Info |
| // argument. |
| ConstantArray *Resumers = CoroId->getInfo().Resumers; |
| assert(Resumers && "PostSplit coro.id Info argument must refer to an array" |
| "of coroutine subfunctions"); |
| auto *ResumeAddrConstant = |
| ConstantExpr::getExtractValue(Resumers, CoroSubFnInst::ResumeIndex); |
| |
| replaceWithConstant(ResumeAddrConstant, ResumeAddr); |
| |
| bool ShouldElide = shouldElide(CoroId->getFunction(), DT); |
| |
| auto *DestroyAddrConstant = ConstantExpr::getExtractValue( |
| Resumers, |
| ShouldElide ? CoroSubFnInst::CleanupIndex : CoroSubFnInst::DestroyIndex); |
| |
| for (auto &It : DestroyAddr) |
| replaceWithConstant(DestroyAddrConstant, It.second); |
| |
| if (ShouldElide) { |
| auto FrameSizeAndAlign = getFrameLayout(cast<Function>(ResumeAddrConstant)); |
| elideHeapAllocations(CoroId->getFunction(), FrameSizeAndAlign.first, |
| FrameSizeAndAlign.second, AA); |
| coro::replaceCoroFree(CoroId, /*Elide=*/true); |
| } |
| |
| return true; |
| } |
| |
| // See if there are any coro.subfn.addr instructions referring to coro.devirt |
| // trigger, if so, replace them with a direct call to devirt trigger function. |
| static bool replaceDevirtTrigger(Function &F) { |
| SmallVector<CoroSubFnInst *, 1> DevirtAddr; |
| for (auto &I : instructions(F)) |
| if (auto *SubFn = dyn_cast<CoroSubFnInst>(&I)) |
| if (SubFn->getIndex() == CoroSubFnInst::RestartTrigger) |
| DevirtAddr.push_back(SubFn); |
| |
| if (DevirtAddr.empty()) |
| return false; |
| |
| Module &M = *F.getParent(); |
| Function *DevirtFn = M.getFunction(CORO_DEVIRT_TRIGGER_FN); |
| assert(DevirtFn && "coro.devirt.fn not found"); |
| replaceWithConstant(DevirtFn, DevirtAddr); |
| |
| return true; |
| } |
| |
| static bool declaresCoroElideIntrinsics(Module &M) { |
| return coro::declaresIntrinsics(M, {"llvm.coro.id", "llvm.coro.id.async"}); |
| } |
| |
| PreservedAnalyses CoroElidePass::run(Function &F, FunctionAnalysisManager &AM) { |
| auto &M = *F.getParent(); |
| if (!declaresCoroElideIntrinsics(M)) |
| return PreservedAnalyses::all(); |
| |
| Lowerer L(M); |
| L.CoroIds.clear(); |
| L.collectPostSplitCoroIds(&F); |
| // If we did not find any coro.id, there is nothing to do. |
| if (L.CoroIds.empty()) |
| return PreservedAnalyses::all(); |
| |
| AAResults &AA = AM.getResult<AAManager>(F); |
| DominatorTree &DT = AM.getResult<DominatorTreeAnalysis>(F); |
| |
| bool Changed = false; |
| for (auto *CII : L.CoroIds) |
| Changed |= L.processCoroId(CII, AA, DT); |
| |
| return Changed ? PreservedAnalyses::none() : PreservedAnalyses::all(); |
| } |
| |
| namespace { |
| struct CoroElideLegacy : FunctionPass { |
| static char ID; |
| CoroElideLegacy() : FunctionPass(ID) { |
| initializeCoroElideLegacyPass(*PassRegistry::getPassRegistry()); |
| } |
| |
| std::unique_ptr<Lowerer> L; |
| |
| bool doInitialization(Module &M) override { |
| if (declaresCoroElideIntrinsics(M)) |
| L = std::make_unique<Lowerer>(M); |
| return false; |
| } |
| |
| bool runOnFunction(Function &F) override { |
| if (!L) |
| return false; |
| |
| bool Changed = false; |
| |
| if (F.hasFnAttribute(CORO_PRESPLIT_ATTR)) |
| Changed = replaceDevirtTrigger(F); |
| |
| L->CoroIds.clear(); |
| L->collectPostSplitCoroIds(&F); |
| // If we did not find any coro.id, there is nothing to do. |
| if (L->CoroIds.empty()) |
| return Changed; |
| |
| AAResults &AA = getAnalysis<AAResultsWrapperPass>().getAAResults(); |
| DominatorTree &DT = getAnalysis<DominatorTreeWrapperPass>().getDomTree(); |
| |
| for (auto *CII : L->CoroIds) |
| Changed |= L->processCoroId(CII, AA, DT); |
| |
| return Changed; |
| } |
| void getAnalysisUsage(AnalysisUsage &AU) const override { |
| AU.addRequired<AAResultsWrapperPass>(); |
| AU.addRequired<DominatorTreeWrapperPass>(); |
| } |
| StringRef getPassName() const override { return "Coroutine Elision"; } |
| }; |
| } |
| |
| char CoroElideLegacy::ID = 0; |
| INITIALIZE_PASS_BEGIN( |
| CoroElideLegacy, "coro-elide", |
| "Coroutine frame allocation elision and indirect calls replacement", false, |
| false) |
| INITIALIZE_PASS_DEPENDENCY(AAResultsWrapperPass) |
| INITIALIZE_PASS_END( |
| CoroElideLegacy, "coro-elide", |
| "Coroutine frame allocation elision and indirect calls replacement", false, |
| false) |
| |
| Pass *llvm::createCoroElideLegacyPass() { return new CoroElideLegacy(); } |