473 lines · cpp
1//===- CoroElide.cpp - Coroutine Frame Allocation Elision Pass ------------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8 9#include "llvm/Transforms/Coroutines/CoroElide.h"10#include "CoroInternal.h"11#include "llvm/ADT/DenseMap.h"12#include "llvm/ADT/Statistic.h"13#include "llvm/Analysis/AliasAnalysis.h"14#include "llvm/Analysis/InstructionSimplify.h"15#include "llvm/Analysis/OptimizationRemarkEmitter.h"16#include "llvm/IR/Dominators.h"17#include "llvm/IR/InstIterator.h"18#include "llvm/Support/ErrorHandling.h"19#include "llvm/Support/FileSystem.h"20#include <optional>21 22using namespace llvm;23 24#define DEBUG_TYPE "coro-elide"25 26STATISTIC(NumOfCoroElided, "The # of coroutine get elided.");27 28#ifndef NDEBUG29static cl::opt<std::string> CoroElideInfoOutputFilename(30 "coro-elide-info-output-file", cl::value_desc("filename"),31 cl::desc("File to record the coroutines got elided"), cl::Hidden);32#endif33 34namespace {35// Created on demand if the coro-elide pass has work to do.36class FunctionElideInfo {37public:38 FunctionElideInfo(Function *F) : ContainingFunction(F) {39 this->collectPostSplitCoroIds();40 }41 42 bool hasCoroIds() const { return !CoroIds.empty(); }43 44 const SmallVectorImpl<CoroIdInst *> &getCoroIds() const { return CoroIds; }45 46private:47 Function *ContainingFunction;48 SmallVector<CoroIdInst *, 4> CoroIds;49 // Used in canCoroBeginEscape to distinguish coro.suspend switchs.50 SmallPtrSet<const SwitchInst *, 4> CoroSuspendSwitches;51 52 void collectPostSplitCoroIds();53 friend class CoroIdElider;54};55 56class CoroIdElider {57public:58 CoroIdElider(CoroIdInst *CoroId, FunctionElideInfo &FEI, AAResults &AA,59 DominatorTree &DT, OptimizationRemarkEmitter &ORE);60 void elideHeapAllocations(uint64_t FrameSize, Align FrameAlign);61 bool lifetimeEligibleForElide() const;62 bool attemptElide();63 bool canCoroBeginEscape(const CoroBeginInst *,64 const SmallPtrSetImpl<BasicBlock *> &) const;65 66private:67 CoroIdInst *CoroId;68 FunctionElideInfo &FEI;69 AAResults &AA;70 DominatorTree &DT;71 OptimizationRemarkEmitter &ORE;72 73 SmallVector<CoroBeginInst *, 1> CoroBegins;74 SmallVector<CoroAllocInst *, 1> CoroAllocs;75 SmallVector<CoroSubFnInst *, 4> ResumeAddr;76 DenseMap<CoroBeginInst *, SmallVector<CoroSubFnInst *, 4>> DestroyAddr;77};78} // end anonymous namespace79 80// Go through the list of coro.subfn.addr intrinsics and replace them with the81// provided constant.82static void replaceWithConstant(Constant *Value,83 SmallVectorImpl<CoroSubFnInst *> &Users) {84 for (CoroSubFnInst *I : Users)85 replaceAndRecursivelySimplify(I, Value);86}87 88// See if any operand of the call instruction references the coroutine frame.89static bool operandReferences(CallInst *CI, AllocaInst *Frame, AAResults &AA) {90 for (Value *Op : CI->operand_values())91 if (Op->getType()->isPointerTy() && !AA.isNoAlias(Op, Frame))92 return true;93 return false;94}95 96// Look for any tail calls referencing the coroutine frame and remove tail97// attribute from them, since now coroutine frame resides on the stack and tail98// call implies that the function does not references anything on the stack.99// However if it's a musttail call, we cannot remove the tailcall attribute.100// It's safe to keep it there as the musttail call is for symmetric transfer,101// and by that point the frame should have been destroyed and hence not102// interfering with operands.103static void removeTailCallAttribute(AllocaInst *Frame, AAResults &AA) {104 Function &F = *Frame->getFunction();105 for (Instruction &I : instructions(F))106 if (auto *Call = dyn_cast<CallInst>(&I))107 if (Call->isTailCall() && operandReferences(Call, Frame, AA) &&108 !Call->isMustTailCall())109 Call->setTailCall(false);110}111 112// Given a resume function @f.resume(%f.frame* %frame), returns the size113// and expected alignment of %f.frame type.114static std::optional<std::pair<uint64_t, Align>>115getFrameLayout(Function *Resume) {116 // Pull information from the function attributes.117 auto Size = Resume->getParamDereferenceableBytes(0);118 if (!Size)119 return std::nullopt;120 return std::make_pair(Size, Resume->getParamAlign(0).valueOrOne());121}122 123// Finds first non alloca instruction in the entry block of a function.124static Instruction *getFirstNonAllocaInTheEntryBlock(Function *F) {125 for (Instruction &I : F->getEntryBlock())126 if (!isa<AllocaInst>(&I))127 return &I;128 llvm_unreachable("no terminator in the entry block");129}130 131#ifndef NDEBUG132static std::unique_ptr<raw_fd_ostream> getOrCreateLogFile() {133 assert(!CoroElideInfoOutputFilename.empty() &&134 "coro-elide-info-output-file shouldn't be empty");135 std::error_code EC;136 auto Result = std::make_unique<raw_fd_ostream>(CoroElideInfoOutputFilename,137 EC, sys::fs::OF_Append);138 if (!EC)139 return Result;140 llvm::errs() << "Error opening coro-elide-info-output-file '"141 << CoroElideInfoOutputFilename << " for appending!\n";142 return std::make_unique<raw_fd_ostream>(2, false); // stderr.143}144#endif145 146void FunctionElideInfo::collectPostSplitCoroIds() {147 for (auto &I : instructions(this->ContainingFunction)) {148 if (auto *CII = dyn_cast<CoroIdInst>(&I))149 if (CII->getInfo().isPostSplit())150 // If it is the coroutine itself, don't touch it.151 if (CII->getCoroutine() != CII->getFunction())152 CoroIds.push_back(CII);153 154 // Consider case like:155 // %0 = call i8 @llvm.coro.suspend(...)156 // switch i8 %0, label %suspend [i8 0, label %resume157 // i8 1, label %cleanup]158 // and collect the SwitchInsts which are used by escape analysis later.159 if (auto *CSI = dyn_cast<CoroSuspendInst>(&I))160 if (CSI->hasOneUse() && isa<SwitchInst>(CSI->use_begin()->getUser())) {161 SwitchInst *SWI = cast<SwitchInst>(CSI->use_begin()->getUser());162 if (SWI->getNumCases() == 2)163 CoroSuspendSwitches.insert(SWI);164 }165 }166}167 168CoroIdElider::CoroIdElider(CoroIdInst *CoroId, FunctionElideInfo &FEI,169 AAResults &AA, DominatorTree &DT,170 OptimizationRemarkEmitter &ORE)171 : CoroId(CoroId), FEI(FEI), AA(AA), DT(DT), ORE(ORE) {172 // Collect all coro.begin and coro.allocs associated with this coro.id.173 for (User *U : CoroId->users()) {174 if (auto *CB = dyn_cast<CoroBeginInst>(U))175 CoroBegins.push_back(CB);176 else if (auto *CA = dyn_cast<CoroAllocInst>(U))177 CoroAllocs.push_back(CA);178 }179 180 // Collect all coro.subfn.addrs associated with coro.begin.181 // Note, we only devirtualize the calls if their coro.subfn.addr refers to182 // coro.begin directly. If we run into cases where this check is too183 // conservative, we can consider relaxing the check.184 for (CoroBeginInst *CB : CoroBegins) {185 for (User *U : CB->users())186 if (auto *II = dyn_cast<CoroSubFnInst>(U))187 switch (II->getIndex()) {188 case CoroSubFnInst::ResumeIndex:189 ResumeAddr.push_back(II);190 break;191 case CoroSubFnInst::DestroyIndex:192 DestroyAddr[CB].push_back(II);193 break;194 default:195 llvm_unreachable("unexpected coro.subfn.addr constant");196 }197 }198}199 200// To elide heap allocations we need to suppress code blocks guarded by201// llvm.coro.alloc and llvm.coro.free instructions.202void CoroIdElider::elideHeapAllocations(uint64_t FrameSize, Align FrameAlign) {203 LLVMContext &C = FEI.ContainingFunction->getContext();204 BasicBlock::iterator InsertPt =205 getFirstNonAllocaInTheEntryBlock(FEI.ContainingFunction)->getIterator();206 207 // Replacing llvm.coro.alloc with false will suppress dynamic208 // allocation as it is expected for the frontend to generate the code that209 // looks like:210 // id = coro.id(...)211 // mem = coro.alloc(id) ? malloc(coro.size()) : 0;212 // coro.begin(id, mem)213 auto *False = ConstantInt::getFalse(C);214 for (auto *CA : CoroAllocs) {215 CA->replaceAllUsesWith(False);216 CA->eraseFromParent();217 }218 219 // FIXME: Design how to transmit alignment information for every alloca that220 // is spilled into the coroutine frame and recreate the alignment information221 // here. Possibly we will need to do a mini SROA here and break the coroutine222 // frame into individual AllocaInst recreating the original alignment.223 const DataLayout &DL = FEI.ContainingFunction->getDataLayout();224 auto FrameTy = ArrayType::get(Type::getInt8Ty(C), FrameSize);225 auto *Frame = new AllocaInst(FrameTy, DL.getAllocaAddrSpace(), "", InsertPt);226 Frame->setAlignment(FrameAlign);227 auto *FrameVoidPtr =228 new BitCastInst(Frame, PointerType::getUnqual(C), "vFrame", InsertPt);229 230 for (auto *CB : CoroBegins) {231 CB->replaceAllUsesWith(FrameVoidPtr);232 CB->eraseFromParent();233 }234 235 // Since now coroutine frame lives on the stack we need to make sure that236 // any tail call referencing it, must be made non-tail call.237 removeTailCallAttribute(Frame, AA);238}239 240bool CoroIdElider::canCoroBeginEscape(241 const CoroBeginInst *CB, const SmallPtrSetImpl<BasicBlock *> &TIs) const {242 const auto &It = DestroyAddr.find(CB);243 assert(It != DestroyAddr.end());244 245 // Limit the number of blocks we visit.246 unsigned Limit = 32 * (1 + It->second.size());247 248 SmallVector<const BasicBlock *, 32> Worklist;249 Worklist.push_back(CB->getParent());250 251 SmallPtrSet<const BasicBlock *, 32> Visited;252 // Consider basicblock of coro.destroy as visited one, so that we253 // skip the path pass through coro.destroy.254 for (auto *DA : It->second)255 Visited.insert(DA->getParent());256 257 SmallPtrSet<const BasicBlock *, 32> EscapingBBs;258 for (auto *U : CB->users()) {259 // The use from coroutine intrinsics are not a problem.260 if (isa<CoroFreeInst, CoroSubFnInst, CoroSaveInst>(U))261 continue;262 263 // Think all other usages may be an escaping candidate conservatively.264 //265 // Note that the major user of switch ABI coroutine (the C++) will store266 // resume.fn, destroy.fn and the index to the coroutine frame immediately.267 // So the parent of the coro.begin in C++ will be always escaping.268 // Then we can't get any performance benefits for C++ by improving the269 // precision of the method.270 //271 // The reason why we still judge it is we want to make LLVM Coroutine in272 // switch ABIs to be self contained as much as possible instead of a273 // by-product of C++20 Coroutines.274 EscapingBBs.insert(cast<Instruction>(U)->getParent());275 }276 277 bool PotentiallyEscaped = false;278 279 do {280 const auto *BB = Worklist.pop_back_val();281 if (!Visited.insert(BB).second)282 continue;283 284 // A Path insensitive marker to test whether the coro.begin escapes.285 // It is intentional to make it path insensitive while it may not be286 // precise since we don't want the process to be too slow.287 PotentiallyEscaped |= EscapingBBs.count(BB);288 289 if (TIs.count(BB)) {290 if (isa<ReturnInst>(BB->getTerminator()) || PotentiallyEscaped)291 return true;292 293 // If the function ends with the exceptional terminator, the memory used294 // by the coroutine frame can be released by stack unwinding295 // automatically. So we can think the coro.begin doesn't escape if it296 // exits the function by exceptional terminator.297 298 continue;299 }300 301 // Conservatively say that there is potentially a path.302 if (!--Limit)303 return true;304 305 auto TI = BB->getTerminator();306 // Although the default dest of coro.suspend switches is suspend pointer307 // which means a escape path to normal terminator, it is reasonable to skip308 // it since coroutine frame doesn't change outside the coroutine body.309 if (isa<SwitchInst>(TI) &&310 FEI.CoroSuspendSwitches.count(cast<SwitchInst>(TI))) {311 Worklist.push_back(cast<SwitchInst>(TI)->getSuccessor(1));312 Worklist.push_back(cast<SwitchInst>(TI)->getSuccessor(2));313 } else314 Worklist.append(succ_begin(BB), succ_end(BB));315 316 } while (!Worklist.empty());317 318 // We have exhausted all possible paths and are certain that coro.begin can319 // not reach to any of terminators.320 return false;321}322 323bool CoroIdElider::lifetimeEligibleForElide() const {324 // If no CoroAllocs, we cannot suppress allocation, so elision is not325 // possible.326 if (CoroAllocs.empty())327 return false;328 329 // Check that for every coro.begin there is at least one coro.destroy directly330 // referencing the SSA value of that coro.begin along each331 // non-exceptional path.332 //333 // If the value escaped, then coro.destroy would have been referencing a334 // memory location storing that value and not the virtual register.335 336 SmallPtrSet<BasicBlock *, 8> Terminators;337 // First gather all of the terminators for the function.338 // Consider the final coro.suspend as the real terminator when the current339 // function is a coroutine.340 for (BasicBlock &B : *FEI.ContainingFunction) {341 auto *TI = B.getTerminator();342 343 if (TI->getNumSuccessors() != 0 || isa<UnreachableInst>(TI))344 continue;345 346 Terminators.insert(&B);347 }348 349 // Filter out the coro.destroy that lie along exceptional paths.350 for (const auto *CB : CoroBegins) {351 auto It = DestroyAddr.find(CB);352 353 // FIXME: If we have not found any destroys for this coro.begin, we354 // disqualify this elide.355 if (It == DestroyAddr.end())356 return false;357 358 const auto &CorrespondingDestroyAddrs = It->second;359 360 // If every terminators is dominated by coro.destroy, we could know the361 // corresponding coro.begin wouldn't escape.362 auto DominatesTerminator = [&](auto *TI) {363 return llvm::any_of(CorrespondingDestroyAddrs, [&](auto *Destroy) {364 return DT.dominates(Destroy, TI->getTerminator());365 });366 };367 368 if (llvm::all_of(Terminators, DominatesTerminator))369 continue;370 371 // Otherwise canCoroBeginEscape would decide whether there is any paths from372 // coro.begin to Terminators which not pass through any of the373 // coro.destroys. This is a slower analysis.374 //375 // canCoroBeginEscape is relatively slow, so we avoid to run it as much as376 // possible.377 if (canCoroBeginEscape(CB, Terminators))378 return false;379 }380 381 // We have checked all CoroBegins and their paths to the terminators without382 // finding disqualifying code patterns, so we can perform heap allocations.383 return true;384}385 386bool CoroIdElider::attemptElide() {387 // PostSplit coro.id refers to an array of subfunctions in its Info388 // argument.389 ConstantArray *Resumers = CoroId->getInfo().Resumers;390 assert(Resumers && "PostSplit coro.id Info argument must refer to an array"391 "of coroutine subfunctions");392 auto *ResumeAddrConstant =393 Resumers->getAggregateElement(CoroSubFnInst::ResumeIndex);394 395 replaceWithConstant(ResumeAddrConstant, ResumeAddr);396 397 bool EligibleForElide = lifetimeEligibleForElide();398 399 auto *DestroyAddrConstant = Resumers->getAggregateElement(400 EligibleForElide ? CoroSubFnInst::CleanupIndex401 : CoroSubFnInst::DestroyIndex);402 403 for (auto &It : DestroyAddr)404 replaceWithConstant(DestroyAddrConstant, It.second);405 406 auto FrameSizeAndAlign = getFrameLayout(cast<Function>(ResumeAddrConstant));407 408 auto CallerFunctionName = FEI.ContainingFunction->getName();409 auto CalleeCoroutineName = CoroId->getCoroutine()->getName();410 411 if (EligibleForElide && FrameSizeAndAlign) {412 elideHeapAllocations(FrameSizeAndAlign->first, FrameSizeAndAlign->second);413 coro::replaceCoroFree(CoroId, /*Elide=*/true);414 NumOfCoroElided++;415 416#ifndef NDEBUG417 if (!CoroElideInfoOutputFilename.empty())418 *getOrCreateLogFile() << "Elide " << CalleeCoroutineName << " in "419 << FEI.ContainingFunction->getName() << "\n";420#endif421 422 ORE.emit([&]() {423 return OptimizationRemark(DEBUG_TYPE, "CoroElide", CoroId)424 << "'" << ore::NV("callee", CalleeCoroutineName)425 << "' elided in '" << ore::NV("caller", CallerFunctionName)426 << "' (frame_size="427 << ore::NV("frame_size", FrameSizeAndAlign->first) << ", align="428 << ore::NV("align", FrameSizeAndAlign->second.value()) << ")";429 });430 } else {431 ORE.emit([&]() {432 auto Remark = OptimizationRemarkMissed(DEBUG_TYPE, "CoroElide", CoroId)433 << "'" << ore::NV("callee", CalleeCoroutineName)434 << "' not elided in '"435 << ore::NV("caller", CallerFunctionName);436 437 if (FrameSizeAndAlign)438 return Remark << "' (frame_size="439 << ore::NV("frame_size", FrameSizeAndAlign->first)440 << ", align="441 << ore::NV("align", FrameSizeAndAlign->second.value())442 << ")";443 else444 return Remark << "' (frame_size=unknown, align=unknown)";445 });446 }447 448 return true;449}450 451PreservedAnalyses CoroElidePass::run(Function &F, FunctionAnalysisManager &AM) {452 auto &M = *F.getParent();453 if (!coro::declaresIntrinsics(M, Intrinsic::coro_id))454 return PreservedAnalyses::all();455 456 FunctionElideInfo FEI{&F};457 // Elide is not necessary if there's no coro.id within the function.458 if (!FEI.hasCoroIds())459 return PreservedAnalyses::all();460 461 AAResults &AA = AM.getResult<AAManager>(F);462 DominatorTree &DT = AM.getResult<DominatorTreeAnalysis>(F);463 auto &ORE = AM.getResult<OptimizationRemarkEmitterAnalysis>(F);464 465 bool Changed = false;466 for (auto *CII : FEI.getCoroIds()) {467 CoroIdElider CIE(CII, FEI, AA, DT, ORE);468 Changed |= CIE.attemptElide();469 }470 471 return Changed ? PreservedAnalyses::none() : PreservedAnalyses::all();472}473