309 lines · cpp
1//===- CoroEarly.cpp - Coroutine Early Function 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/CoroEarly.h"10#include "CoroInternal.h"11#include "llvm/IR/DIBuilder.h"12#include "llvm/IR/Function.h"13#include "llvm/IR/IRBuilder.h"14#include "llvm/IR/InstIterator.h"15#include "llvm/IR/Module.h"16#include "llvm/Transforms/Coroutines/CoroShape.h"17 18using namespace llvm;19 20#define DEBUG_TYPE "coro-early"21 22namespace {23// Created on demand if the coro-early pass has work to do.24class Lowerer : public coro::LowererBase {25 IRBuilder<> Builder;26 PointerType *const AnyResumeFnPtrTy;27 Constant *NoopCoro = nullptr;28 29 void lowerResumeOrDestroy(CallBase &CB, CoroSubFnInst::ResumeKind);30 void lowerCoroPromise(CoroPromiseInst *Intrin);31 void lowerCoroDone(IntrinsicInst *II);32 void lowerCoroNoop(IntrinsicInst *II);33 void hidePromiseAlloca(CoroIdInst *CoroId, CoroBeginInst *CoroBegin);34 35public:36 Lowerer(Module &M)37 : LowererBase(M), Builder(Context),38 AnyResumeFnPtrTy(PointerType::getUnqual(Context)) {}39 void lowerEarlyIntrinsics(Function &F);40};41} // namespace42 43// Replace a direct call to coro.resume or coro.destroy with an indirect call to44// an address returned by coro.subfn.addr intrinsic. This is done so that45// CGPassManager recognizes devirtualization when CoroElide pass replaces a call46// to coro.subfn.addr with an appropriate function address.47void Lowerer::lowerResumeOrDestroy(CallBase &CB,48 CoroSubFnInst::ResumeKind Index) {49 Value *ResumeAddr = makeSubFnCall(CB.getArgOperand(0), Index, &CB);50 CB.setCalledOperand(ResumeAddr);51 CB.setCallingConv(CallingConv::Fast);52}53 54// Coroutine promise field is always at the fixed offset from the beginning of55// the coroutine frame. i8* coro.promise(i8*, i1 from) intrinsic adds an offset56// to a passed pointer to move from coroutine frame to coroutine promise and57// vice versa. Since we don't know exactly which coroutine frame it is, we build58// a coroutine frame mock up starting with two function pointers, followed by a59// properly aligned coroutine promise field.60// TODO: Handle the case when coroutine promise alloca has align override.61void Lowerer::lowerCoroPromise(CoroPromiseInst *Intrin) {62 Value *Operand = Intrin->getArgOperand(0);63 Align Alignment = Intrin->getAlignment();64 Type *Int8Ty = Builder.getInt8Ty();65 66 auto *SampleStruct =67 StructType::get(Context, {AnyResumeFnPtrTy, AnyResumeFnPtrTy, Int8Ty});68 const DataLayout &DL = TheModule.getDataLayout();69 int64_t Offset = alignTo(70 DL.getStructLayout(SampleStruct)->getElementOffset(2), Alignment);71 if (Intrin->isFromPromise())72 Offset = -Offset;73 74 Builder.SetInsertPoint(Intrin);75 Value *Replacement =76 Builder.CreateConstInBoundsGEP1_32(Int8Ty, Operand, Offset);77 78 Intrin->replaceAllUsesWith(Replacement);79 Intrin->eraseFromParent();80}81 82// When a coroutine reaches final suspend point, it zeros out ResumeFnAddr in83// the coroutine frame (it is UB to resume from a final suspend point).84// The llvm.coro.done intrinsic is used to check whether a coroutine is85// suspended at the final suspend point or not.86void Lowerer::lowerCoroDone(IntrinsicInst *II) {87 Value *Operand = II->getArgOperand(0);88 89 // ResumeFnAddr is the first pointer sized element of the coroutine frame.90 static_assert(coro::Shape::SwitchFieldIndex::Resume == 0,91 "resume function not at offset zero");92 auto *FrameTy = Int8Ptr;93 94 Builder.SetInsertPoint(II);95 auto *Load = Builder.CreateLoad(FrameTy, Operand);96 auto *Cond = Builder.CreateICmpEQ(Load, NullPtr);97 98 II->replaceAllUsesWith(Cond);99 II->eraseFromParent();100}101 102static void buildDebugInfoForNoopResumeDestroyFunc(Function *NoopFn) {103 Module &M = *NoopFn->getParent();104 if (M.debug_compile_units().empty())105 return;106 107 DICompileUnit *CU = *M.debug_compile_units_begin();108 DIBuilder DB(M, /*AllowUnresolved*/ false, CU);109 std::array<Metadata *, 2> Params{nullptr, nullptr};110 auto *SubroutineType =111 DB.createSubroutineType(DB.getOrCreateTypeArray(Params));112 StringRef Name = NoopFn->getName();113 auto *SP = DB.createFunction(114 CU, /*Name=*/Name, /*LinkageName=*/Name, /*File=*/ CU->getFile(),115 /*LineNo=*/0, SubroutineType, /*ScopeLine=*/0, DINode::FlagArtificial,116 DISubprogram::SPFlagDefinition);117 NoopFn->setSubprogram(SP);118 DB.finalize();119}120 121void Lowerer::lowerCoroNoop(IntrinsicInst *II) {122 if (!NoopCoro) {123 LLVMContext &C = Builder.getContext();124 Module &M = *II->getModule();125 126 // Create a noop.frame struct type.127 auto *FnTy = FunctionType::get(Type::getVoidTy(C), Builder.getPtrTy(0),128 /*isVarArg=*/false);129 auto *FnPtrTy = Builder.getPtrTy(0);130 StructType *FrameTy =131 StructType::create({FnPtrTy, FnPtrTy}, "NoopCoro.Frame");132 133 // Create a Noop function that does nothing.134 Function *NoopFn = Function::createWithDefaultAttr(135 FnTy, GlobalValue::LinkageTypes::InternalLinkage,136 M.getDataLayout().getProgramAddressSpace(), "__NoopCoro_ResumeDestroy",137 &M);138 NoopFn->setCallingConv(CallingConv::Fast);139 buildDebugInfoForNoopResumeDestroyFunc(NoopFn);140 auto *Entry = BasicBlock::Create(C, "entry", NoopFn);141 ReturnInst::Create(C, Entry);142 143 // Create a constant struct for the frame.144 Constant* Values[] = {NoopFn, NoopFn};145 Constant* NoopCoroConst = ConstantStruct::get(FrameTy, Values);146 NoopCoro = new GlobalVariable(M, NoopCoroConst->getType(), /*isConstant=*/true,147 GlobalVariable::PrivateLinkage, NoopCoroConst,148 "NoopCoro.Frame.Const");149 cast<GlobalVariable>(NoopCoro)->setNoSanitizeMetadata();150 }151 152 Builder.SetInsertPoint(II);153 auto *NoopCoroVoidPtr = Builder.CreateBitCast(NoopCoro, Int8Ptr);154 II->replaceAllUsesWith(NoopCoroVoidPtr);155 II->eraseFromParent();156}157 158// Later middle-end passes will assume promise alloca dead after coroutine159// suspend, leading to misoptimizations. We hide promise alloca using160// coro.promise and will lower it back to alloca at CoroSplit.161void Lowerer::hidePromiseAlloca(CoroIdInst *CoroId, CoroBeginInst *CoroBegin) {162 auto *PA = CoroId->getPromise();163 if (!PA || !CoroBegin)164 return;165 Builder.SetInsertPoint(*CoroBegin->getInsertionPointAfterDef());166 167 auto *Alignment = Builder.getInt32(PA->getAlign().value());168 auto *FromPromise = Builder.getInt1(false);169 SmallVector<Value *, 3> Arg{CoroBegin, Alignment, FromPromise};170 auto *PI = Builder.CreateIntrinsic(171 Builder.getPtrTy(), Intrinsic::coro_promise, Arg, {}, "promise.addr");172 PI->setCannotDuplicate();173 // Remove lifetime markers, as these are only allowed on allocas.174 for (User *U : make_early_inc_range(PA->users())) {175 auto *I = cast<Instruction>(U);176 if (I->isLifetimeStartOrEnd())177 I->eraseFromParent();178 }179 PA->replaceUsesWithIf(PI, [CoroId](Use &U) {180 bool IsBitcast = U == U.getUser()->stripPointerCasts();181 bool IsCoroId = U.getUser() == CoroId;182 return !IsBitcast && !IsCoroId;183 });184}185 186// Prior to CoroSplit, calls to coro.begin needs to be marked as NoDuplicate,187// as CoroSplit assumes there is exactly one coro.begin. After CoroSplit,188// NoDuplicate attribute will be removed from coro.begin otherwise, it will189// interfere with inlining.190static void setCannotDuplicate(CoroIdInst *CoroId) {191 for (User *U : CoroId->users())192 if (auto *CB = dyn_cast<CoroBeginInst>(U))193 CB->setCannotDuplicate();194}195 196void Lowerer::lowerEarlyIntrinsics(Function &F) {197 CoroIdInst *CoroId = nullptr;198 CoroBeginInst *CoroBegin = nullptr;199 SmallVector<CoroFreeInst *, 4> CoroFrees;200 bool HasCoroSuspend = false;201 for (Instruction &I : llvm::make_early_inc_range(instructions(F))) {202 auto *CB = dyn_cast<CallBase>(&I);203 if (!CB)204 continue;205 206 switch (CB->getIntrinsicID()) {207 default:208 continue;209 case Intrinsic::coro_begin:210 case Intrinsic::coro_begin_custom_abi:211 if (CoroBegin)212 report_fatal_error(213 "coroutine should have exactly one defining @llvm.coro.begin");214 CoroBegin = cast<CoroBeginInst>(&I);215 break;216 case Intrinsic::coro_free:217 CoroFrees.push_back(cast<CoroFreeInst>(&I));218 break;219 case Intrinsic::coro_suspend:220 // Make sure that final suspend point is not duplicated as CoroSplit221 // pass expects that there is at most one final suspend point.222 if (cast<CoroSuspendInst>(&I)->isFinal())223 CB->setCannotDuplicate();224 HasCoroSuspend = true;225 break;226 case Intrinsic::coro_end_async:227 case Intrinsic::coro_end:228 // Make sure that fallthrough coro.end is not duplicated as CoroSplit229 // pass expects that there is at most one fallthrough coro.end.230 if (cast<AnyCoroEndInst>(&I)->isFallthrough())231 CB->setCannotDuplicate();232 break;233 case Intrinsic::coro_noop:234 lowerCoroNoop(cast<IntrinsicInst>(&I));235 break;236 case Intrinsic::coro_id:237 if (auto *CII = cast<CoroIdInst>(&I)) {238 if (CII->getInfo().isPreSplit()) {239 assert(F.isPresplitCoroutine() &&240 "The frontend uses Switch-Resumed ABI should emit "241 "\"presplitcoroutine\" attribute for the coroutine.");242 setCannotDuplicate(CII);243 CII->setCoroutineSelf();244 CoroId = cast<CoroIdInst>(&I);245 }246 }247 break;248 case Intrinsic::coro_id_retcon:249 case Intrinsic::coro_id_retcon_once:250 case Intrinsic::coro_id_async:251 F.setPresplitCoroutine();252 break;253 case Intrinsic::coro_resume:254 lowerResumeOrDestroy(*CB, CoroSubFnInst::ResumeIndex);255 break;256 case Intrinsic::coro_destroy:257 lowerResumeOrDestroy(*CB, CoroSubFnInst::DestroyIndex);258 break;259 case Intrinsic::coro_promise:260 lowerCoroPromise(cast<CoroPromiseInst>(&I));261 break;262 case Intrinsic::coro_done:263 lowerCoroDone(cast<IntrinsicInst>(&I));264 break;265 }266 }267 268 if (CoroId) {269 // Make sure that all CoroFree reference the coro.id intrinsic.270 // Token type is not exposed through coroutine C/C++ builtins to plain C, so271 // we allow specifying none and fixing it up here.272 for (CoroFreeInst *CF : CoroFrees)273 CF->setArgOperand(0, CoroId);274 275 hidePromiseAlloca(CoroId, CoroBegin);276 }277 278 // Coroutine suspention could potentially lead to any argument modified279 // outside of the function, hence arguments should not have noalias280 // attributes.281 if (HasCoroSuspend)282 for (Argument &A : F.args())283 if (A.hasNoAliasAttr())284 A.removeAttr(Attribute::NoAlias);285}286 287static bool declaresCoroEarlyIntrinsics(const Module &M) {288 // coro_suspend omitted as it is overloaded.289 return coro::declaresIntrinsics(290 M, {Intrinsic::coro_id, Intrinsic::coro_id_retcon,291 Intrinsic::coro_id_retcon_once, Intrinsic::coro_id_async,292 Intrinsic::coro_destroy, Intrinsic::coro_done, Intrinsic::coro_end,293 Intrinsic::coro_end_async, Intrinsic::coro_noop, Intrinsic::coro_free,294 Intrinsic::coro_promise, Intrinsic::coro_resume});295}296 297PreservedAnalyses CoroEarlyPass::run(Module &M, ModuleAnalysisManager &) {298 if (!declaresCoroEarlyIntrinsics(M))299 return PreservedAnalyses::all();300 301 Lowerer L(M);302 for (auto &F : M)303 L.lowerEarlyIntrinsics(F);304 305 PreservedAnalyses PA;306 PA.preserveSet<CFGAnalyses>();307 return PA;308}309