2455 lines · cpp
1//===- CoroSplit.cpp - Converts a coroutine into a state machine ----------===//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// This pass builds the coroutine frame and outlines resume and destroy parts9// of the coroutine into separate functions.10//11// We present a coroutine to an LLVM as an ordinary function with suspension12// points marked up with intrinsics. We let the optimizer party on the coroutine13// as a single function for as long as possible. Shortly before the coroutine is14// eligible to be inlined into its callers, we split up the coroutine into parts15// corresponding to an initial, resume and destroy invocations of the coroutine,16// add them to the current SCC and restart the IPO pipeline to optimize the17// coroutine subfunctions we extracted before proceeding to the caller of the18// coroutine.19//===----------------------------------------------------------------------===//20 21#include "llvm/Transforms/Coroutines/CoroSplit.h"22#include "CoroCloner.h"23#include "CoroInternal.h"24#include "llvm/ADT/DenseMap.h"25#include "llvm/ADT/PriorityWorklist.h"26#include "llvm/ADT/STLExtras.h"27#include "llvm/ADT/SmallPtrSet.h"28#include "llvm/ADT/SmallVector.h"29#include "llvm/ADT/StringExtras.h"30#include "llvm/ADT/StringRef.h"31#include "llvm/ADT/Twine.h"32#include "llvm/Analysis/CFG.h"33#include "llvm/Analysis/CallGraph.h"34#include "llvm/Analysis/ConstantFolding.h"35#include "llvm/Analysis/LazyCallGraph.h"36#include "llvm/Analysis/OptimizationRemarkEmitter.h"37#include "llvm/Analysis/TargetTransformInfo.h"38#include "llvm/BinaryFormat/Dwarf.h"39#include "llvm/IR/Argument.h"40#include "llvm/IR/Attributes.h"41#include "llvm/IR/BasicBlock.h"42#include "llvm/IR/CFG.h"43#include "llvm/IR/CallingConv.h"44#include "llvm/IR/Constants.h"45#include "llvm/IR/DIBuilder.h"46#include "llvm/IR/DataLayout.h"47#include "llvm/IR/DebugInfo.h"48#include "llvm/IR/DerivedTypes.h"49#include "llvm/IR/Dominators.h"50#include "llvm/IR/GlobalValue.h"51#include "llvm/IR/GlobalVariable.h"52#include "llvm/IR/InstIterator.h"53#include "llvm/IR/InstrTypes.h"54#include "llvm/IR/Instruction.h"55#include "llvm/IR/Instructions.h"56#include "llvm/IR/IntrinsicInst.h"57#include "llvm/IR/LLVMContext.h"58#include "llvm/IR/Module.h"59#include "llvm/IR/Type.h"60#include "llvm/IR/Value.h"61#include "llvm/IR/Verifier.h"62#include "llvm/Support/Casting.h"63#include "llvm/Support/Debug.h"64#include "llvm/Support/PrettyStackTrace.h"65#include "llvm/Support/raw_ostream.h"66#include "llvm/Transforms/Coroutines/MaterializationUtils.h"67#include "llvm/Transforms/Scalar.h"68#include "llvm/Transforms/Utils/BasicBlockUtils.h"69#include "llvm/Transforms/Utils/CallGraphUpdater.h"70#include "llvm/Transforms/Utils/Cloning.h"71#include "llvm/Transforms/Utils/Local.h"72#include <cassert>73#include <cstddef>74#include <cstdint>75#include <initializer_list>76#include <iterator>77 78using namespace llvm;79 80#define DEBUG_TYPE "coro-split"81 82// FIXME:83// Lower the intrinisc in CoroEarly phase if coroutine frame doesn't escape84// and it is known that other transformations, for example, sanitizers85// won't lead to incorrect code.86static void lowerAwaitSuspend(IRBuilder<> &Builder, CoroAwaitSuspendInst *CB,87 coro::Shape &Shape) {88 auto Wrapper = CB->getWrapperFunction();89 auto Awaiter = CB->getAwaiter();90 auto FramePtr = CB->getFrame();91 92 Builder.SetInsertPoint(CB);93 94 CallBase *NewCall = nullptr;95 // await_suspend has only 2 parameters, awaiter and handle.96 // Copy parameter attributes from the intrinsic call, but remove the last,97 // because the last parameter now becomes the function that is being called.98 AttributeList NewAttributes =99 CB->getAttributes().removeParamAttributes(CB->getContext(), 2);100 101 if (auto Invoke = dyn_cast<InvokeInst>(CB)) {102 auto WrapperInvoke =103 Builder.CreateInvoke(Wrapper, Invoke->getNormalDest(),104 Invoke->getUnwindDest(), {Awaiter, FramePtr});105 106 WrapperInvoke->setCallingConv(Invoke->getCallingConv());107 std::copy(Invoke->bundle_op_info_begin(), Invoke->bundle_op_info_end(),108 WrapperInvoke->bundle_op_info_begin());109 WrapperInvoke->setAttributes(NewAttributes);110 WrapperInvoke->setDebugLoc(Invoke->getDebugLoc());111 NewCall = WrapperInvoke;112 } else if (auto Call = dyn_cast<CallInst>(CB)) {113 auto WrapperCall = Builder.CreateCall(Wrapper, {Awaiter, FramePtr});114 115 WrapperCall->setAttributes(NewAttributes);116 WrapperCall->setDebugLoc(Call->getDebugLoc());117 NewCall = WrapperCall;118 } else {119 llvm_unreachable("Unexpected coro_await_suspend invocation method");120 }121 122 if (CB->getCalledFunction()->getIntrinsicID() ==123 Intrinsic::coro_await_suspend_handle) {124 // Follow the lowered await_suspend call above with a lowered resume call125 // to the returned coroutine.126 if (auto *Invoke = dyn_cast<InvokeInst>(CB)) {127 // If the await_suspend call is an invoke, we continue in the next block.128 Builder.SetInsertPoint(Invoke->getNormalDest()->getFirstInsertionPt());129 }130 131 coro::LowererBase LB(*Wrapper->getParent());132 auto *ResumeAddr = LB.makeSubFnCall(NewCall, CoroSubFnInst::ResumeIndex,133 &*Builder.GetInsertPoint());134 135 LLVMContext &Ctx = Builder.getContext();136 FunctionType *ResumeTy = FunctionType::get(137 Type::getVoidTy(Ctx), PointerType::getUnqual(Ctx), false);138 auto *ResumeCall = Builder.CreateCall(ResumeTy, ResumeAddr, {NewCall});139 ResumeCall->setCallingConv(CallingConv::Fast);140 141 // We can't insert the 'ret' instruction and adjust the cc until the142 // function has been split, so remember this for later.143 Shape.SymmetricTransfers.push_back(ResumeCall);144 145 NewCall = ResumeCall;146 }147 148 CB->replaceAllUsesWith(NewCall);149 CB->eraseFromParent();150}151 152static void lowerAwaitSuspends(Function &F, coro::Shape &Shape) {153 IRBuilder<> Builder(F.getContext());154 for (auto *AWS : Shape.CoroAwaitSuspends)155 lowerAwaitSuspend(Builder, AWS, Shape);156}157 158static void maybeFreeRetconStorage(IRBuilder<> &Builder,159 const coro::Shape &Shape, Value *FramePtr,160 CallGraph *CG) {161 assert(Shape.ABI == coro::ABI::Retcon || Shape.ABI == coro::ABI::RetconOnce);162 if (Shape.RetconLowering.IsFrameInlineInStorage)163 return;164 165 Shape.emitDealloc(Builder, FramePtr, CG);166}167 168/// Replace an llvm.coro.end.async.169/// Will inline the must tail call function call if there is one.170/// \returns true if cleanup of the coro.end block is needed, false otherwise.171static bool replaceCoroEndAsync(AnyCoroEndInst *End) {172 IRBuilder<> Builder(End);173 174 auto *EndAsync = dyn_cast<CoroAsyncEndInst>(End);175 if (!EndAsync) {176 Builder.CreateRetVoid();177 return true /*needs cleanup of coro.end block*/;178 }179 180 auto *MustTailCallFunc = EndAsync->getMustTailCallFunction();181 if (!MustTailCallFunc) {182 Builder.CreateRetVoid();183 return true /*needs cleanup of coro.end block*/;184 }185 186 // Move the must tail call from the predecessor block into the end block.187 auto *CoroEndBlock = End->getParent();188 auto *MustTailCallFuncBlock = CoroEndBlock->getSinglePredecessor();189 assert(MustTailCallFuncBlock && "Must have a single predecessor block");190 auto It = MustTailCallFuncBlock->getTerminator()->getIterator();191 auto *MustTailCall = cast<CallInst>(&*std::prev(It));192 CoroEndBlock->splice(End->getIterator(), MustTailCallFuncBlock,193 MustTailCall->getIterator());194 195 // Insert the return instruction.196 Builder.SetInsertPoint(End);197 Builder.CreateRetVoid();198 InlineFunctionInfo FnInfo;199 200 // Remove the rest of the block, by splitting it into an unreachable block.201 auto *BB = End->getParent();202 BB->splitBasicBlock(End);203 BB->getTerminator()->eraseFromParent();204 205 auto InlineRes = InlineFunction(*MustTailCall, FnInfo);206 assert(InlineRes.isSuccess() && "Expected inlining to succeed");207 (void)InlineRes;208 209 // We have cleaned up the coro.end block above.210 return false;211}212 213/// Replace a non-unwind call to llvm.coro.end.214static void replaceFallthroughCoroEnd(AnyCoroEndInst *End,215 const coro::Shape &Shape, Value *FramePtr,216 bool InRamp, CallGraph *CG) {217 // Start inserting right before the coro.end.218 IRBuilder<> Builder(End);219 220 // Create the return instruction.221 switch (Shape.ABI) {222 // The cloned functions in switch-lowering always return void.223 case coro::ABI::Switch:224 assert(!cast<CoroEndInst>(End)->hasResults() &&225 "switch coroutine should not return any values");226 // coro.end doesn't immediately end the coroutine in the main function227 // in this lowering, because we need to deallocate the coroutine.228 if (InRamp)229 return;230 Builder.CreateRetVoid();231 break;232 233 // In async lowering this returns.234 case coro::ABI::Async: {235 bool CoroEndBlockNeedsCleanup = replaceCoroEndAsync(End);236 if (!CoroEndBlockNeedsCleanup)237 return;238 break;239 }240 241 // In unique continuation lowering, the continuations always return void.242 // But we may have implicitly allocated storage.243 case coro::ABI::RetconOnce: {244 maybeFreeRetconStorage(Builder, Shape, FramePtr, CG);245 auto *CoroEnd = cast<CoroEndInst>(End);246 auto *RetTy = Shape.getResumeFunctionType()->getReturnType();247 248 if (!CoroEnd->hasResults()) {249 assert(RetTy->isVoidTy());250 Builder.CreateRetVoid();251 break;252 }253 254 auto *CoroResults = CoroEnd->getResults();255 unsigned NumReturns = CoroResults->numReturns();256 257 if (auto *RetStructTy = dyn_cast<StructType>(RetTy)) {258 assert(RetStructTy->getNumElements() == NumReturns &&259 "numbers of returns should match resume function singature");260 Value *ReturnValue = PoisonValue::get(RetStructTy);261 unsigned Idx = 0;262 for (Value *RetValEl : CoroResults->return_values())263 ReturnValue = Builder.CreateInsertValue(ReturnValue, RetValEl, Idx++);264 Builder.CreateRet(ReturnValue);265 } else if (NumReturns == 0) {266 assert(RetTy->isVoidTy());267 Builder.CreateRetVoid();268 } else {269 assert(NumReturns == 1);270 Builder.CreateRet(*CoroResults->retval_begin());271 }272 CoroResults->replaceAllUsesWith(273 ConstantTokenNone::get(CoroResults->getContext()));274 CoroResults->eraseFromParent();275 break;276 }277 278 // In non-unique continuation lowering, we signal completion by returning279 // a null continuation.280 case coro::ABI::Retcon: {281 assert(!cast<CoroEndInst>(End)->hasResults() &&282 "retcon coroutine should not return any values");283 maybeFreeRetconStorage(Builder, Shape, FramePtr, CG);284 auto RetTy = Shape.getResumeFunctionType()->getReturnType();285 auto RetStructTy = dyn_cast<StructType>(RetTy);286 PointerType *ContinuationTy =287 cast<PointerType>(RetStructTy ? RetStructTy->getElementType(0) : RetTy);288 289 Value *ReturnValue = ConstantPointerNull::get(ContinuationTy);290 if (RetStructTy) {291 ReturnValue = Builder.CreateInsertValue(PoisonValue::get(RetStructTy),292 ReturnValue, 0);293 }294 Builder.CreateRet(ReturnValue);295 break;296 }297 }298 299 // Remove the rest of the block, by splitting it into an unreachable block.300 auto *BB = End->getParent();301 BB->splitBasicBlock(End);302 BB->getTerminator()->eraseFromParent();303}304 305// Mark a coroutine as done, which implies that the coroutine is finished and306// never gets resumed.307//308// In resume-switched ABI, the done state is represented by storing zero in309// ResumeFnAddr.310//311// NOTE: We couldn't omit the argument `FramePtr`. It is necessary because the312// pointer to the frame in splitted function is not stored in `Shape`.313static void markCoroutineAsDone(IRBuilder<> &Builder, const coro::Shape &Shape,314 Value *FramePtr) {315 assert(316 Shape.ABI == coro::ABI::Switch &&317 "markCoroutineAsDone is only supported for Switch-Resumed ABI for now.");318 auto *GepIndex = Builder.CreateStructGEP(319 Shape.FrameTy, FramePtr, coro::Shape::SwitchFieldIndex::Resume,320 "ResumeFn.addr");321 auto *NullPtr = ConstantPointerNull::get(cast<PointerType>(322 Shape.FrameTy->getTypeAtIndex(coro::Shape::SwitchFieldIndex::Resume)));323 Builder.CreateStore(NullPtr, GepIndex);324 325 // If the coroutine don't have unwind coro end, we could omit the store to326 // the final suspend point since we could infer the coroutine is suspended327 // at the final suspend point by the nullness of ResumeFnAddr.328 // However, we can't skip it if the coroutine have unwind coro end. Since329 // the coroutine reaches unwind coro end is considered suspended at the330 // final suspend point (the ResumeFnAddr is null) but in fact the coroutine331 // didn't complete yet. We need the IndexVal for the final suspend point332 // to make the states clear.333 if (Shape.SwitchLowering.HasUnwindCoroEnd &&334 Shape.SwitchLowering.HasFinalSuspend) {335 assert(cast<CoroSuspendInst>(Shape.CoroSuspends.back())->isFinal() &&336 "The final suspend should only live in the last position of "337 "CoroSuspends.");338 ConstantInt *IndexVal = Shape.getIndex(Shape.CoroSuspends.size() - 1);339 auto *FinalIndex = Builder.CreateStructGEP(340 Shape.FrameTy, FramePtr, Shape.getSwitchIndexField(), "index.addr");341 342 Builder.CreateStore(IndexVal, FinalIndex);343 }344}345 346/// Replace an unwind call to llvm.coro.end.347static void replaceUnwindCoroEnd(AnyCoroEndInst *End, const coro::Shape &Shape,348 Value *FramePtr, bool InRamp, CallGraph *CG) {349 IRBuilder<> Builder(End);350 351 switch (Shape.ABI) {352 // In switch-lowering, this does nothing in the main function.353 case coro::ABI::Switch: {354 // In C++'s specification, the coroutine should be marked as done355 // if promise.unhandled_exception() throws. The frontend will356 // call coro.end(true) along this path.357 //358 // FIXME: We should refactor this once there is other language359 // which uses Switch-Resumed style other than C++.360 markCoroutineAsDone(Builder, Shape, FramePtr);361 if (InRamp)362 return;363 break;364 }365 // In async lowering this does nothing.366 case coro::ABI::Async:367 break;368 // In continuation-lowering, this frees the continuation storage.369 case coro::ABI::Retcon:370 case coro::ABI::RetconOnce:371 maybeFreeRetconStorage(Builder, Shape, FramePtr, CG);372 break;373 }374 375 // If coro.end has an associated bundle, add cleanupret instruction.376 if (auto Bundle = End->getOperandBundle(LLVMContext::OB_funclet)) {377 auto *FromPad = cast<CleanupPadInst>(Bundle->Inputs[0]);378 auto *CleanupRet = Builder.CreateCleanupRet(FromPad, nullptr);379 End->getParent()->splitBasicBlock(End);380 CleanupRet->getParent()->getTerminator()->eraseFromParent();381 }382}383 384static void replaceCoroEnd(AnyCoroEndInst *End, const coro::Shape &Shape,385 Value *FramePtr, bool InRamp, CallGraph *CG) {386 if (End->isUnwind())387 replaceUnwindCoroEnd(End, Shape, FramePtr, InRamp, CG);388 else389 replaceFallthroughCoroEnd(End, Shape, FramePtr, InRamp, CG);390 End->eraseFromParent();391}392 393// In the resume function, we remove the last case (when coro::Shape is built,394// the final suspend point (if present) is always the last element of395// CoroSuspends array) since it is an undefined behavior to resume a coroutine396// suspended at the final suspend point.397// In the destroy function, if it isn't possible that the ResumeFnAddr is NULL398// and the coroutine doesn't suspend at the final suspend point actually (this399// is possible since the coroutine is considered suspended at the final suspend400// point if promise.unhandled_exception() exits via an exception), we can401// remove the last case.402void coro::BaseCloner::handleFinalSuspend() {403 assert(Shape.ABI == coro::ABI::Switch &&404 Shape.SwitchLowering.HasFinalSuspend);405 406 if (isSwitchDestroyFunction() && Shape.SwitchLowering.HasUnwindCoroEnd)407 return;408 409 auto *Switch = cast<SwitchInst>(VMap[Shape.SwitchLowering.ResumeSwitch]);410 auto FinalCaseIt = std::prev(Switch->case_end());411 BasicBlock *ResumeBB = FinalCaseIt->getCaseSuccessor();412 Switch->removeCase(FinalCaseIt);413 if (isSwitchDestroyFunction()) {414 BasicBlock *OldSwitchBB = Switch->getParent();415 auto *NewSwitchBB = OldSwitchBB->splitBasicBlock(Switch, "Switch");416 Builder.SetInsertPoint(OldSwitchBB->getTerminator());417 418 if (NewF->isCoroOnlyDestroyWhenComplete()) {419 // When the coroutine can only be destroyed when complete, we don't need420 // to generate code for other cases.421 Builder.CreateBr(ResumeBB);422 } else {423 auto *GepIndex = Builder.CreateStructGEP(424 Shape.FrameTy, NewFramePtr, coro::Shape::SwitchFieldIndex::Resume,425 "ResumeFn.addr");426 auto *Load =427 Builder.CreateLoad(Shape.getSwitchResumePointerType(), GepIndex);428 auto *Cond = Builder.CreateIsNull(Load);429 Builder.CreateCondBr(Cond, ResumeBB, NewSwitchBB);430 }431 OldSwitchBB->getTerminator()->eraseFromParent();432 }433}434 435static FunctionType *436getFunctionTypeFromAsyncSuspend(AnyCoroSuspendInst *Suspend) {437 auto *AsyncSuspend = cast<CoroSuspendAsyncInst>(Suspend);438 auto *StructTy = cast<StructType>(AsyncSuspend->getType());439 auto &Context = Suspend->getParent()->getParent()->getContext();440 auto *VoidTy = Type::getVoidTy(Context);441 return FunctionType::get(VoidTy, StructTy->elements(), false);442}443 444static Function *createCloneDeclaration(Function &OrigF, coro::Shape &Shape,445 const Twine &Suffix,446 Module::iterator InsertBefore,447 AnyCoroSuspendInst *ActiveSuspend) {448 Module *M = OrigF.getParent();449 auto *FnTy = (Shape.ABI != coro::ABI::Async)450 ? Shape.getResumeFunctionType()451 : getFunctionTypeFromAsyncSuspend(ActiveSuspend);452 453 Function *NewF =454 Function::Create(FnTy, GlobalValue::LinkageTypes::InternalLinkage,455 OrigF.getName() + Suffix);456 457 M->getFunctionList().insert(InsertBefore, NewF);458 459 return NewF;460}461 462/// Replace uses of the active llvm.coro.suspend.retcon/async call with the463/// arguments to the continuation function.464///465/// This assumes that the builder has a meaningful insertion point.466void coro::BaseCloner::replaceRetconOrAsyncSuspendUses() {467 assert(Shape.ABI == coro::ABI::Retcon || Shape.ABI == coro::ABI::RetconOnce ||468 Shape.ABI == coro::ABI::Async);469 470 auto NewS = VMap[ActiveSuspend];471 if (NewS->use_empty())472 return;473 474 // Copy out all the continuation arguments after the buffer pointer into475 // an easily-indexed data structure for convenience.476 SmallVector<Value *, 8> Args;477 // The async ABI includes all arguments -- including the first argument.478 bool IsAsyncABI = Shape.ABI == coro::ABI::Async;479 for (auto I = IsAsyncABI ? NewF->arg_begin() : std::next(NewF->arg_begin()),480 E = NewF->arg_end();481 I != E; ++I)482 Args.push_back(&*I);483 484 // If the suspend returns a single scalar value, we can just do a simple485 // replacement.486 if (!isa<StructType>(NewS->getType())) {487 assert(Args.size() == 1);488 NewS->replaceAllUsesWith(Args.front());489 return;490 }491 492 // Try to peephole extracts of an aggregate return.493 for (Use &U : llvm::make_early_inc_range(NewS->uses())) {494 auto *EVI = dyn_cast<ExtractValueInst>(U.getUser());495 if (!EVI || EVI->getNumIndices() != 1)496 continue;497 498 EVI->replaceAllUsesWith(Args[EVI->getIndices().front()]);499 EVI->eraseFromParent();500 }501 502 // If we have no remaining uses, we're done.503 if (NewS->use_empty())504 return;505 506 // Otherwise, we need to create an aggregate.507 Value *Aggr = PoisonValue::get(NewS->getType());508 for (auto [Idx, Arg] : llvm::enumerate(Args))509 Aggr = Builder.CreateInsertValue(Aggr, Arg, Idx);510 511 NewS->replaceAllUsesWith(Aggr);512}513 514void coro::BaseCloner::replaceCoroSuspends() {515 Value *SuspendResult;516 517 switch (Shape.ABI) {518 // In switch lowering, replace coro.suspend with the appropriate value519 // for the type of function we're extracting.520 // Replacing coro.suspend with (0) will result in control flow proceeding to521 // a resume label associated with a suspend point, replacing it with (1) will522 // result in control flow proceeding to a cleanup label associated with this523 // suspend point.524 case coro::ABI::Switch:525 SuspendResult = Builder.getInt8(isSwitchDestroyFunction() ? 1 : 0);526 break;527 528 // In async lowering there are no uses of the result.529 case coro::ABI::Async:530 return;531 532 // In returned-continuation lowering, the arguments from earlier533 // continuations are theoretically arbitrary, and they should have been534 // spilled.535 case coro::ABI::RetconOnce:536 case coro::ABI::Retcon:537 return;538 }539 540 for (AnyCoroSuspendInst *CS : Shape.CoroSuspends) {541 // The active suspend was handled earlier.542 if (CS == ActiveSuspend)543 continue;544 545 auto *MappedCS = cast<AnyCoroSuspendInst>(VMap[CS]);546 MappedCS->replaceAllUsesWith(SuspendResult);547 MappedCS->eraseFromParent();548 }549}550 551void coro::BaseCloner::replaceCoroEnds() {552 for (AnyCoroEndInst *CE : Shape.CoroEnds) {553 // We use a null call graph because there's no call graph node for554 // the cloned function yet. We'll just be rebuilding that later.555 auto *NewCE = cast<AnyCoroEndInst>(VMap[CE]);556 replaceCoroEnd(NewCE, Shape, NewFramePtr, /*in ramp*/ false, nullptr);557 }558}559 560void coro::BaseCloner::replaceCoroIsInRamp() {561 auto &Ctx = OrigF.getContext();562 for (auto *II : Shape.CoroIsInRampInsts) {563 auto *NewII = cast<CoroIsInRampInst>(VMap[II]);564 NewII->replaceAllUsesWith(ConstantInt::getFalse(Ctx));565 NewII->eraseFromParent();566 }567}568 569static void replaceSwiftErrorOps(Function &F, coro::Shape &Shape,570 ValueToValueMapTy *VMap) {571 if (Shape.ABI == coro::ABI::Async && Shape.CoroSuspends.empty())572 return;573 Value *CachedSlot = nullptr;574 auto getSwiftErrorSlot = [&](Type *ValueTy) -> Value * {575 if (CachedSlot)576 return CachedSlot;577 578 // Check if the function has a swifterror argument.579 for (auto &Arg : F.args()) {580 if (Arg.isSwiftError()) {581 CachedSlot = &Arg;582 return &Arg;583 }584 }585 586 // Create a swifterror alloca.587 IRBuilder<> Builder(&F.getEntryBlock(),588 F.getEntryBlock().getFirstNonPHIOrDbg());589 auto Alloca = Builder.CreateAlloca(ValueTy);590 Alloca->setSwiftError(true);591 592 CachedSlot = Alloca;593 return Alloca;594 };595 596 for (CallInst *Op : Shape.SwiftErrorOps) {597 auto MappedOp = VMap ? cast<CallInst>((*VMap)[Op]) : Op;598 IRBuilder<> Builder(MappedOp);599 600 // If there are no arguments, this is a 'get' operation.601 Value *MappedResult;602 if (Op->arg_empty()) {603 auto ValueTy = Op->getType();604 auto Slot = getSwiftErrorSlot(ValueTy);605 MappedResult = Builder.CreateLoad(ValueTy, Slot);606 } else {607 assert(Op->arg_size() == 1);608 auto Value = MappedOp->getArgOperand(0);609 auto ValueTy = Value->getType();610 auto Slot = getSwiftErrorSlot(ValueTy);611 Builder.CreateStore(Value, Slot);612 MappedResult = Slot;613 }614 615 MappedOp->replaceAllUsesWith(MappedResult);616 MappedOp->eraseFromParent();617 }618 619 // If we're updating the original function, we've invalidated SwiftErrorOps.620 if (VMap == nullptr) {621 Shape.SwiftErrorOps.clear();622 }623}624 625/// Returns all debug records in F.626static SmallVector<DbgVariableRecord *>627collectDbgVariableRecords(Function &F) {628 SmallVector<DbgVariableRecord *> DbgVariableRecords;629 for (auto &I : instructions(F)) {630 for (DbgVariableRecord &DVR : filterDbgVars(I.getDbgRecordRange()))631 DbgVariableRecords.push_back(&DVR);632 }633 return DbgVariableRecords;634}635 636void coro::BaseCloner::replaceSwiftErrorOps() {637 ::replaceSwiftErrorOps(*NewF, Shape, &VMap);638}639 640void coro::BaseCloner::salvageDebugInfo() {641 auto DbgVariableRecords = collectDbgVariableRecords(*NewF);642 SmallDenseMap<Argument *, AllocaInst *, 4> ArgToAllocaMap;643 644 // Only 64-bit ABIs have a register we can refer to with the entry value.645 bool UseEntryValue = OrigF.getParent()->getTargetTriple().isArch64Bit();646 for (DbgVariableRecord *DVR : DbgVariableRecords)647 coro::salvageDebugInfo(ArgToAllocaMap, *DVR, UseEntryValue);648 649 // Remove all salvaged dbg.declare intrinsics that became650 // either unreachable or stale due to the CoroSplit transformation.651 DominatorTree DomTree(*NewF);652 auto IsUnreachableBlock = [&](BasicBlock *BB) {653 return !isPotentiallyReachable(&NewF->getEntryBlock(), BB, nullptr,654 &DomTree);655 };656 auto RemoveOne = [&](DbgVariableRecord *DVI) {657 if (IsUnreachableBlock(DVI->getParent()))658 DVI->eraseFromParent();659 else if (isa_and_nonnull<AllocaInst>(DVI->getVariableLocationOp(0))) {660 // Count all non-debuginfo uses in reachable blocks.661 unsigned Uses = 0;662 for (auto *User : DVI->getVariableLocationOp(0)->users())663 if (auto *I = dyn_cast<Instruction>(User))664 if (!isa<AllocaInst>(I) && !IsUnreachableBlock(I->getParent()))665 ++Uses;666 if (!Uses)667 DVI->eraseFromParent();668 }669 };670 for_each(DbgVariableRecords, RemoveOne);671}672 673void coro::BaseCloner::replaceEntryBlock() {674 // In the original function, the AllocaSpillBlock is a block immediately675 // following the allocation of the frame object which defines GEPs for676 // all the allocas that have been moved into the frame, and it ends by677 // branching to the original beginning of the coroutine. Make this678 // the entry block of the cloned function.679 auto *Entry = cast<BasicBlock>(VMap[Shape.AllocaSpillBlock]);680 auto *OldEntry = &NewF->getEntryBlock();681 Entry->setName("entry" + Suffix);682 Entry->moveBefore(OldEntry);683 Entry->getTerminator()->eraseFromParent();684 685 // Clear all predecessors of the new entry block. There should be686 // exactly one predecessor, which we created when splitting out687 // AllocaSpillBlock to begin with.688 assert(Entry->hasOneUse());689 auto BranchToEntry = cast<BranchInst>(Entry->user_back());690 assert(BranchToEntry->isUnconditional());691 Builder.SetInsertPoint(BranchToEntry);692 Builder.CreateUnreachable();693 BranchToEntry->eraseFromParent();694 695 // Branch from the entry to the appropriate place.696 Builder.SetInsertPoint(Entry);697 switch (Shape.ABI) {698 case coro::ABI::Switch: {699 // In switch-lowering, we built a resume-entry block in the original700 // function. Make the entry block branch to this.701 auto *SwitchBB =702 cast<BasicBlock>(VMap[Shape.SwitchLowering.ResumeEntryBlock]);703 Builder.CreateBr(SwitchBB);704 SwitchBB->moveAfter(Entry);705 break;706 }707 case coro::ABI::Async:708 case coro::ABI::Retcon:709 case coro::ABI::RetconOnce: {710 // In continuation ABIs, we want to branch to immediately after the711 // active suspend point. Earlier phases will have put the suspend in its712 // own basic block, so just thread our jump directly to its successor.713 assert((Shape.ABI == coro::ABI::Async &&714 isa<CoroSuspendAsyncInst>(ActiveSuspend)) ||715 ((Shape.ABI == coro::ABI::Retcon ||716 Shape.ABI == coro::ABI::RetconOnce) &&717 isa<CoroSuspendRetconInst>(ActiveSuspend)));718 auto *MappedCS = cast<AnyCoroSuspendInst>(VMap[ActiveSuspend]);719 auto Branch = cast<BranchInst>(MappedCS->getNextNode());720 assert(Branch->isUnconditional());721 Builder.CreateBr(Branch->getSuccessor(0));722 break;723 }724 }725 726 // Any static alloca that's still being used but not reachable from the new727 // entry needs to be moved to the new entry.728 Function *F = OldEntry->getParent();729 DominatorTree DT{*F};730 for (Instruction &I : llvm::make_early_inc_range(instructions(F))) {731 auto *Alloca = dyn_cast<AllocaInst>(&I);732 if (!Alloca || I.use_empty())733 continue;734 if (DT.isReachableFromEntry(I.getParent()) ||735 !isa<ConstantInt>(Alloca->getArraySize()))736 continue;737 I.moveBefore(*Entry, Entry->getFirstInsertionPt());738 }739}740 741/// Derive the value of the new frame pointer.742Value *coro::BaseCloner::deriveNewFramePointer() {743 // Builder should be inserting to the front of the new entry block.744 745 switch (Shape.ABI) {746 // In switch-lowering, the argument is the frame pointer.747 case coro::ABI::Switch:748 return &*NewF->arg_begin();749 // In async-lowering, one of the arguments is an async context as determined750 // by the `llvm.coro.id.async` intrinsic. We can retrieve the async context of751 // the resume function from the async context projection function associated752 // with the active suspend. The frame is located as a tail to the async753 // context header.754 case coro::ABI::Async: {755 auto *ActiveAsyncSuspend = cast<CoroSuspendAsyncInst>(ActiveSuspend);756 auto ContextIdx = ActiveAsyncSuspend->getStorageArgumentIndex() & 0xff;757 auto *CalleeContext = NewF->getArg(ContextIdx);758 auto *ProjectionFunc =759 ActiveAsyncSuspend->getAsyncContextProjectionFunction();760 auto DbgLoc =761 cast<CoroSuspendAsyncInst>(VMap[ActiveSuspend])->getDebugLoc();762 // Calling i8* (i8*)763 auto *CallerContext = Builder.CreateCall(ProjectionFunc->getFunctionType(),764 ProjectionFunc, CalleeContext);765 CallerContext->setCallingConv(ProjectionFunc->getCallingConv());766 CallerContext->setDebugLoc(DbgLoc);767 // The frame is located after the async_context header.768 auto &Context = Builder.getContext();769 auto *FramePtrAddr = Builder.CreateConstInBoundsGEP1_32(770 Type::getInt8Ty(Context), CallerContext,771 Shape.AsyncLowering.FrameOffset, "async.ctx.frameptr");772 // Inline the projection function.773 InlineFunctionInfo InlineInfo;774 auto InlineRes = InlineFunction(*CallerContext, InlineInfo);775 assert(InlineRes.isSuccess());776 (void)InlineRes;777 return FramePtrAddr;778 }779 // In continuation-lowering, the argument is the opaque storage.780 case coro::ABI::Retcon:781 case coro::ABI::RetconOnce: {782 Argument *NewStorage = &*NewF->arg_begin();783 auto FramePtrTy = PointerType::getUnqual(Shape.FrameTy->getContext());784 785 // If the storage is inline, just bitcast to the storage to the frame type.786 if (Shape.RetconLowering.IsFrameInlineInStorage)787 return NewStorage;788 789 // Otherwise, load the real frame from the opaque storage.790 return Builder.CreateLoad(FramePtrTy, NewStorage);791 }792 }793 llvm_unreachable("bad ABI");794}795 796/// Adjust the scope line of the funclet to the first line number after the797/// suspend point. This avoids a jump in the line table from the function798/// declaration (where prologue instructions are attributed to) to the suspend799/// point.800/// Only adjust the scope line when the files are the same.801/// If no candidate line number is found, fallback to the line of ActiveSuspend.802static void updateScopeLine(Instruction *ActiveSuspend,803 DISubprogram &SPToUpdate) {804 if (!ActiveSuspend)805 return;806 807 // No subsequent instruction -> fallback to the location of ActiveSuspend.808 if (!ActiveSuspend->getNextNode()) {809 if (auto DL = ActiveSuspend->getDebugLoc())810 if (SPToUpdate.getFile() == DL->getFile())811 SPToUpdate.setScopeLine(DL->getLine());812 return;813 }814 815 BasicBlock::iterator Successor = ActiveSuspend->getNextNode()->getIterator();816 // Corosplit splits the BB around ActiveSuspend, so the meaningful817 // instructions are not in the same BB.818 if (auto *Branch = dyn_cast_or_null<BranchInst>(Successor);819 Branch && Branch->isUnconditional())820 Successor = Branch->getSuccessor(0)->getFirstNonPHIOrDbg();821 822 // Find the first successor of ActiveSuspend with a non-zero line location.823 // If that matches the file of ActiveSuspend, use it.824 BasicBlock *PBB = Successor->getParent();825 for (; Successor != PBB->end(); Successor = std::next(Successor)) {826 Successor = skipDebugIntrinsics(Successor);827 auto DL = Successor->getDebugLoc();828 if (!DL || DL.getLine() == 0)829 continue;830 831 if (SPToUpdate.getFile() == DL->getFile()) {832 SPToUpdate.setScopeLine(DL.getLine());833 return;834 }835 836 break;837 }838 839 // If the search above failed, fallback to the location of ActiveSuspend.840 if (auto DL = ActiveSuspend->getDebugLoc())841 if (SPToUpdate.getFile() == DL->getFile())842 SPToUpdate.setScopeLine(DL->getLine());843}844 845static void addFramePointerAttrs(AttributeList &Attrs, LLVMContext &Context,846 unsigned ParamIndex, uint64_t Size,847 Align Alignment, bool NoAlias) {848 AttrBuilder ParamAttrs(Context);849 ParamAttrs.addAttribute(Attribute::NonNull);850 ParamAttrs.addAttribute(Attribute::NoUndef);851 852 if (NoAlias)853 ParamAttrs.addAttribute(Attribute::NoAlias);854 855 ParamAttrs.addAlignmentAttr(Alignment);856 ParamAttrs.addDereferenceableAttr(Size);857 Attrs = Attrs.addParamAttributes(Context, ParamIndex, ParamAttrs);858}859 860static void addAsyncContextAttrs(AttributeList &Attrs, LLVMContext &Context,861 unsigned ParamIndex) {862 AttrBuilder ParamAttrs(Context);863 ParamAttrs.addAttribute(Attribute::SwiftAsync);864 Attrs = Attrs.addParamAttributes(Context, ParamIndex, ParamAttrs);865}866 867static void addSwiftSelfAttrs(AttributeList &Attrs, LLVMContext &Context,868 unsigned ParamIndex) {869 AttrBuilder ParamAttrs(Context);870 ParamAttrs.addAttribute(Attribute::SwiftSelf);871 Attrs = Attrs.addParamAttributes(Context, ParamIndex, ParamAttrs);872}873 874/// Rebuild an aggregate/expression constant from already-remapped operands.875static Constant *rebuildConstant(Constant *C, ArrayRef<Constant *> NewOps) {876 if (auto *CE = dyn_cast<ConstantExpr>(C))877 return CE->getWithOperands(NewOps);878 if (auto *CA = dyn_cast<ConstantArray>(C))879 return ConstantArray::get(CA->getType(), NewOps);880 if (auto *CS = dyn_cast<ConstantStruct>(C))881 return ConstantStruct::get(CS->getType(), NewOps);882 if (isa<ConstantVector>(C))883 return ConstantVector::get(NewOps);884 llvm_unreachable("unexpected constant kind in blockaddress-table remap");885}886 887/// Recursively rebuild \p C, replacing every `blockaddress(&OrigF, BB)` with888/// `blockaddress(&NewF, VMap[BB])` for blocks that were cloned into \p NewF.889/// Returns \p C unchanged if it contains no such blockaddress.890static Constant *891remapBlockAddressesInConstant(Constant *C, Function &OrigF, Function &NewF,892 ValueToValueMapTy &VMap,893 DenseMap<Constant *, Constant *> &Cache) {894 auto It = Cache.find(C);895 if (It != Cache.end())896 return It->second;897 898 Constant *Result = C;899 if (auto *BA = dyn_cast<BlockAddress>(C)) {900 if (BA->getFunction() == &OrigF) {901 auto *MappedBB =902 dyn_cast_or_null<BasicBlock>(VMap.lookup(BA->getBasicBlock()));903 if (MappedBB && MappedBB->getParent() == &NewF)904 Result = BlockAddress::get(&NewF, MappedBB);905 }906 } else if (isa<ConstantAggregate>(C) || isa<ConstantExpr>(C)) {907 SmallVector<Constant *, 16> NewOps;908 bool Changed = false;909 for (Use &Op : C->operands()) {910 auto *OpC = cast<Constant>(Op.get());911 Constant *NewOp =912 remapBlockAddressesInConstant(OpC, OrigF, NewF, VMap, Cache);913 Changed |= NewOp != OpC;914 NewOps.push_back(NewOp);915 }916 if (Changed)917 Result = rebuildConstant(C, NewOps);918 }919 920 Cache[C] = Result;921 return Result;922}923 924/// Recursively rebuild \p C, replacing any global in \p Map with its925/// replacement. Returns \p C unchanged if it references no mapped global.926static Constant *927replaceGlobalsInConstant(Constant *C,928 const DenseMap<GlobalVariable *, GlobalVariable *> &Map,929 DenseMap<Constant *, Constant *> &Cache) {930 if (auto *GV = dyn_cast<GlobalVariable>(C)) {931 auto It = Map.find(GV);932 return It != Map.end() ? It->second : C;933 }934 935 auto CIt = Cache.find(C);936 if (CIt != Cache.end())937 return CIt->second;938 939 Constant *Result = C;940 if (isa<ConstantAggregate>(C) || isa<ConstantExpr>(C)) {941 SmallVector<Constant *, 16> NewOps;942 bool Changed = false;943 for (Use &Op : C->operands()) {944 auto *OpC = cast<Constant>(Op.get());945 Constant *NewOp = replaceGlobalsInConstant(OpC, Map, Cache);946 Changed |= NewOp != OpC;947 NewOps.push_back(NewOp);948 }949 if (Changed)950 Result = rebuildConstant(C, NewOps);951 }952 953 Cache[C] = Result;954 return Result;955}956 957/// Relocate GNU computed-goto jump tables across a coroutine split.958///959/// When a coroutine takes the address of its own basic blocks -- e.g. glibc's960/// printf_positional drives format dispatch with a module-level constant array961/// of `blockaddress` values feeding an `indirectbr` -- CoroSplit clones the962/// address-taken blocks into the resume/destroy/cleanup clone, but the jump963/// table global still refers to the *ramp* function's blocks. When964/// postSplitCleanup (or later CFG simplification) deletes those now-unreachable965/// ramp blocks, BasicBlock::~BasicBlock RAUWs every `blockaddress` of a deleted966/// block to the `inttoptr(i32 1)` sentinel, poisoning the shared table. The967/// clone's `indirectbr` then reads the poisoned table, which the WebAssembly968/// backend lowers to `unreachable` -> a runtime trap on the resume path.969///970/// Fix: give each clone its own copy of any jump table that references this971/// coroutine's blocks, with the `blockaddress` constants remapped to the972/// clone's blocks, and rewrite the clone's references to use it. The ramp keeps973/// the original table (correct for any address-taken blocks it retains;974/// harmlessly dead if it retains none). This must run while OrigF's blocks975/// still exist (i.e. before postSplitCleanup), which holds for all callers.976static void relocateBlockAddressTablesForClone(Function &OrigF, Function &NewF,977 ValueToValueMapTy &VMap) {978 // Collect the GlobalVariables transitively referenced by NewF's body.979 SmallPtrSet<GlobalVariable *, 4> Referenced;980 SmallPtrSet<Constant *, 16> Visited;981 SmallVector<Constant *, 16> Worklist;982 for (Instruction &I : instructions(NewF))983 for (Use &Op : I.operands())984 if (auto *C = dyn_cast<Constant>(Op.get()))985 Worklist.push_back(C);986 while (!Worklist.empty()) {987 Constant *C = Worklist.pop_back_val();988 if (!Visited.insert(C).second)989 continue;990 if (auto *GV = dyn_cast<GlobalVariable>(C)) {991 Referenced.insert(GV);992 continue;993 }994 for (Use &Op : C->operands())995 if (auto *OpC = dyn_cast<Constant>(Op.get()))996 Worklist.push_back(OpC);997 }998 999 // Build a per-clone copy of each referenced table whose initializer holds a1000 // blockaddress of one of OrigF's (now cloned) blocks.1001 DenseMap<GlobalVariable *, GlobalVariable *> TableClones;1002 DenseMap<Constant *, Constant *> RemapCache;1003 for (GlobalVariable *GV : Referenced) {1004 if (!GV->hasInitializer())1005 continue;1006 Constant *Init = GV->getInitializer();1007 Constant *NewInit =1008 remapBlockAddressesInConstant(Init, OrigF, NewF, VMap, RemapCache);1009 if (NewInit == Init)1010 continue; // No blockaddress of this coroutine inside.1011 auto *NewGV = new GlobalVariable(1012 *GV->getParent(), GV->getValueType(), GV->isConstant(),1013 GlobalValue::PrivateLinkage, NewInit, GV->getName() + "." + NewF.getName(),1014 /*InsertBefore=*/nullptr, GV->getThreadLocalMode(),1015 GV->getType()->getAddressSpace());1016 NewGV->setAlignment(GV->getAlign());1017 NewGV->setUnnamedAddr(GV->getUnnamedAddr());1018 TableClones[GV] = NewGV;1019 }1020 if (TableClones.empty())1021 return;1022 1023 // Rewrite NewF's references to the original tables to use the clone's copies.1024 DenseMap<Constant *, Constant *> ReplaceCache;1025 for (Instruction &I : instructions(NewF))1026 for (Use &Op : I.operands())1027 if (auto *C = dyn_cast<Constant>(Op.get())) {1028 Constant *NewC = replaceGlobalsInConstant(C, TableClones, ReplaceCache);1029 if (NewC != C)1030 Op.set(NewC);1031 }1032}1033 1034/// Clone the body of the original function into a resume function of1035/// some sort.1036void coro::BaseCloner::create() {1037 assert(NewF);1038 1039 // Replace all args with dummy instructions. If an argument is the old frame1040 // pointer, the dummy will be replaced by the new frame pointer once it is1041 // computed below. Uses of all other arguments should have already been1042 // rewritten by buildCoroutineFrame() to use loads/stores on the coroutine1043 // frame.1044 SmallVector<Instruction *> DummyArgs;1045 for (Argument &A : OrigF.args()) {1046 DummyArgs.push_back(new FreezeInst(PoisonValue::get(A.getType())));1047 VMap[&A] = DummyArgs.back();1048 }1049 1050 SmallVector<ReturnInst *, 4> Returns;1051 1052 // Ignore attempts to change certain attributes of the function.1053 // TODO: maybe there should be a way to suppress this during cloning?1054 auto savedVisibility = NewF->getVisibility();1055 auto savedUnnamedAddr = NewF->getUnnamedAddr();1056 auto savedDLLStorageClass = NewF->getDLLStorageClass();1057 1058 // NewF's linkage (which CloneFunctionInto does *not* change) might not1059 // be compatible with the visibility of OrigF (which it *does* change),1060 // so protect against that.1061 auto savedLinkage = NewF->getLinkage();1062 NewF->setLinkage(llvm::GlobalValue::ExternalLinkage);1063 1064 CloneFunctionInto(NewF, &OrigF, VMap,1065 CloneFunctionChangeType::LocalChangesOnly, Returns);1066 1067 // If this coroutine takes the address of its own blocks (a GNU computed-goto1068 // jump table driving an indirectbr), relocate those tables into per-clone1069 // copies so deleting the ramp's address-taken blocks does not poison the1070 // clone's table with the inttoptr(i32 1) sentinel. Must run while OrigF's1071 // blocks still exist, which holds here (postSplitCleanup runs later).1072 relocateBlockAddressTablesForClone(OrigF, *NewF, VMap);1073 1074 auto &Context = NewF->getContext();1075 1076 if (DISubprogram *SP = NewF->getSubprogram()) {1077 assert(SP != OrigF.getSubprogram() && SP->isDistinct());1078 updateScopeLine(ActiveSuspend, *SP);1079 1080 // Update the linkage name and the function name to reflect the modified1081 // name.1082 MDString *NewLinkageName = MDString::get(Context, NewF->getName());1083 SP->replaceLinkageName(NewLinkageName);1084 if (DISubprogram *Decl = SP->getDeclaration()) {1085 TempDISubprogram NewDecl = Decl->clone();1086 NewDecl->replaceLinkageName(NewLinkageName);1087 SP->replaceDeclaration(MDNode::replaceWithUniqued(std::move(NewDecl)));1088 }1089 }1090 1091 NewF->setLinkage(savedLinkage);1092 NewF->setVisibility(savedVisibility);1093 NewF->setUnnamedAddr(savedUnnamedAddr);1094 NewF->setDLLStorageClass(savedDLLStorageClass);1095 // The function sanitizer metadata needs to match the signature of the1096 // function it is being attached to. However this does not hold for split1097 // functions here. Thus remove the metadata for split functions.1098 if (Shape.ABI == coro::ABI::Switch &&1099 NewF->hasMetadata(LLVMContext::MD_func_sanitize))1100 NewF->eraseMetadata(LLVMContext::MD_func_sanitize);1101 1102 // Replace the attributes of the new function:1103 auto OrigAttrs = NewF->getAttributes();1104 auto NewAttrs = AttributeList();1105 1106 switch (Shape.ABI) {1107 case coro::ABI::Switch:1108 // Bootstrap attributes by copying function attributes from the1109 // original function. This should include optimization settings and so on.1110 NewAttrs = NewAttrs.addFnAttributes(1111 Context, AttrBuilder(Context, OrigAttrs.getFnAttrs()));1112 1113 addFramePointerAttrs(NewAttrs, Context, 0, Shape.FrameSize,1114 Shape.FrameAlign, /*NoAlias=*/false);1115 break;1116 case coro::ABI::Async: {1117 auto *ActiveAsyncSuspend = cast<CoroSuspendAsyncInst>(ActiveSuspend);1118 if (OrigF.hasParamAttribute(Shape.AsyncLowering.ContextArgNo,1119 Attribute::SwiftAsync)) {1120 uint32_t ArgAttributeIndices =1121 ActiveAsyncSuspend->getStorageArgumentIndex();1122 auto ContextArgIndex = ArgAttributeIndices & 0xff;1123 addAsyncContextAttrs(NewAttrs, Context, ContextArgIndex);1124 1125 // `swiftasync` must preceed `swiftself` so 0 is not a valid index for1126 // `swiftself`.1127 auto SwiftSelfIndex = ArgAttributeIndices >> 8;1128 if (SwiftSelfIndex)1129 addSwiftSelfAttrs(NewAttrs, Context, SwiftSelfIndex);1130 }1131 1132 // Transfer the original function's attributes.1133 auto FnAttrs = OrigF.getAttributes().getFnAttrs();1134 NewAttrs = NewAttrs.addFnAttributes(Context, AttrBuilder(Context, FnAttrs));1135 break;1136 }1137 case coro::ABI::Retcon:1138 case coro::ABI::RetconOnce:1139 // If we have a continuation prototype, just use its attributes,1140 // full-stop.1141 NewAttrs = Shape.RetconLowering.ResumePrototype->getAttributes();1142 1143 /// FIXME: Is it really good to add the NoAlias attribute?1144 addFramePointerAttrs(NewAttrs, Context, 0,1145 Shape.getRetconCoroId()->getStorageSize(),1146 Shape.getRetconCoroId()->getStorageAlignment(),1147 /*NoAlias=*/true);1148 1149 break;1150 }1151 1152 switch (Shape.ABI) {1153 // In these ABIs, the cloned functions always return 'void', and the1154 // existing return sites are meaningless. Note that for unique1155 // continuations, this includes the returns associated with suspends;1156 // this is fine because we can't suspend twice.1157 case coro::ABI::Switch:1158 case coro::ABI::RetconOnce:1159 // Remove old returns.1160 for (ReturnInst *Return : Returns)1161 changeToUnreachable(Return);1162 break;1163 1164 // With multi-suspend continuations, we'll already have eliminated the1165 // original returns and inserted returns before all the suspend points,1166 // so we want to leave any returns in place.1167 case coro::ABI::Retcon:1168 break;1169 // Async lowering will insert musttail call functions at all suspend points1170 // followed by a return.1171 // Don't change returns to unreachable because that will trip up the verifier.1172 // These returns should be unreachable from the clone.1173 case coro::ABI::Async:1174 break;1175 }1176 1177 NewF->setAttributes(NewAttrs);1178 NewF->setCallingConv(Shape.getResumeFunctionCC());1179 1180 // Set up the new entry block.1181 replaceEntryBlock();1182 1183 // Turn symmetric transfers into musttail calls.1184 for (CallInst *ResumeCall : Shape.SymmetricTransfers) {1185 ResumeCall = cast<CallInst>(VMap[ResumeCall]);1186 if (TTI.supportsTailCallFor(ResumeCall)) {1187 // FIXME: Could we support symmetric transfer effectively without1188 // musttail?1189 ResumeCall->setTailCallKind(CallInst::TCK_MustTail);1190 }1191 1192 // Put a 'ret void' after the call, and split any remaining instructions to1193 // an unreachable block.1194 BasicBlock *BB = ResumeCall->getParent();1195 BB->splitBasicBlock(ResumeCall->getNextNode());1196 Builder.SetInsertPoint(BB->getTerminator());1197 Builder.CreateRetVoid();1198 BB->getTerminator()->eraseFromParent();1199 }1200 1201 Builder.SetInsertPoint(&NewF->getEntryBlock().front());1202 NewFramePtr = deriveNewFramePointer();1203 1204 // Remap frame pointer.1205 Value *OldFramePtr = VMap[Shape.FramePtr];1206 NewFramePtr->takeName(OldFramePtr);1207 OldFramePtr->replaceAllUsesWith(NewFramePtr);1208 1209 // Remap vFrame pointer.1210 auto *NewVFrame = Builder.CreateBitCast(1211 NewFramePtr, PointerType::getUnqual(Builder.getContext()), "vFrame");1212 Value *OldVFrame = cast<Value>(VMap[Shape.CoroBegin]);1213 if (OldVFrame != NewVFrame)1214 OldVFrame->replaceAllUsesWith(NewVFrame);1215 1216 // All uses of the arguments should have been resolved by this point,1217 // so we can safely remove the dummy values.1218 for (Instruction *DummyArg : DummyArgs) {1219 DummyArg->replaceAllUsesWith(PoisonValue::get(DummyArg->getType()));1220 DummyArg->deleteValue();1221 }1222 1223 switch (Shape.ABI) {1224 case coro::ABI::Switch:1225 // Rewrite final suspend handling as it is not done via switch (allows to1226 // remove final case from the switch, since it is undefined behavior to1227 // resume the coroutine suspended at the final suspend point.1228 if (Shape.SwitchLowering.HasFinalSuspend)1229 handleFinalSuspend();1230 break;1231 case coro::ABI::Async:1232 case coro::ABI::Retcon:1233 case coro::ABI::RetconOnce:1234 // Replace uses of the active suspend with the corresponding1235 // continuation-function arguments.1236 assert(ActiveSuspend != nullptr &&1237 "no active suspend when lowering a continuation-style coroutine");1238 replaceRetconOrAsyncSuspendUses();1239 break;1240 }1241 1242 // Handle suspends.1243 replaceCoroSuspends();1244 1245 // Handle swifterror.1246 replaceSwiftErrorOps();1247 1248 // Remove coro.end intrinsics.1249 replaceCoroEnds();1250 1251 replaceCoroIsInRamp();1252 1253 // Salvage debug info that points into the coroutine frame.1254 salvageDebugInfo();1255}1256 1257void coro::SwitchCloner::create() {1258 // Create a new function matching the original type1259 NewF = createCloneDeclaration(OrigF, Shape, Suffix, OrigF.getParent()->end(),1260 ActiveSuspend);1261 1262 // Clone the function1263 coro::BaseCloner::create();1264 1265 // Eliminate coro.free from the clones, replacing it with 'null' in cleanup,1266 // to suppress deallocation code.1267 coro::replaceCoroFree(cast<CoroIdInst>(VMap[Shape.CoroBegin->getId()]),1268 /*Elide=*/FKind == coro::CloneKind::SwitchCleanup);1269}1270 1271static void updateAsyncFuncPointerContextSize(coro::Shape &Shape) {1272 assert(Shape.ABI == coro::ABI::Async);1273 1274 auto *FuncPtrStruct = cast<ConstantStruct>(1275 Shape.AsyncLowering.AsyncFuncPointer->getInitializer());1276 auto *OrigRelativeFunOffset = FuncPtrStruct->getOperand(0);1277 auto *OrigContextSize = FuncPtrStruct->getOperand(1);1278 auto *NewContextSize = ConstantInt::get(OrigContextSize->getType(),1279 Shape.AsyncLowering.ContextSize);1280 auto *NewFuncPtrStruct = ConstantStruct::get(1281 FuncPtrStruct->getType(), OrigRelativeFunOffset, NewContextSize);1282 1283 Shape.AsyncLowering.AsyncFuncPointer->setInitializer(NewFuncPtrStruct);1284}1285 1286static TypeSize getFrameSizeForShape(coro::Shape &Shape) {1287 // In the same function all coro.sizes should have the same result type.1288 auto *SizeIntrin = Shape.CoroSizes.back();1289 Module *M = SizeIntrin->getModule();1290 const DataLayout &DL = M->getDataLayout();1291 return DL.getTypeAllocSize(Shape.FrameTy);1292}1293 1294static void replaceFrameSizeAndAlignment(coro::Shape &Shape) {1295 if (Shape.ABI == coro::ABI::Async)1296 updateAsyncFuncPointerContextSize(Shape);1297 1298 for (CoroAlignInst *CA : Shape.CoroAligns) {1299 CA->replaceAllUsesWith(1300 ConstantInt::get(CA->getType(), Shape.FrameAlign.value()));1301 CA->eraseFromParent();1302 }1303 1304 if (Shape.CoroSizes.empty())1305 return;1306 1307 // In the same function all coro.sizes should have the same result type.1308 auto *SizeIntrin = Shape.CoroSizes.back();1309 auto *SizeConstant =1310 ConstantInt::get(SizeIntrin->getType(), getFrameSizeForShape(Shape));1311 1312 for (CoroSizeInst *CS : Shape.CoroSizes) {1313 CS->replaceAllUsesWith(SizeConstant);1314 CS->eraseFromParent();1315 }1316}1317 1318static void postSplitCleanup(Function &F) {1319 removeUnreachableBlocks(F);1320 1321#ifndef NDEBUG1322 // For now, we do a mandatory verification step because we don't1323 // entirely trust this pass. Note that we don't want to add a verifier1324 // pass to FPM below because it will also verify all the global data.1325 if (verifyFunction(F, &errs()))1326 report_fatal_error("Broken function");1327#endif1328}1329 1330// Coroutine has no suspend points. Remove heap allocation for the coroutine1331// frame if possible.1332static void handleNoSuspendCoroutine(coro::Shape &Shape) {1333 auto *CoroBegin = Shape.CoroBegin;1334 switch (Shape.ABI) {1335 case coro::ABI::Switch: {1336 auto SwitchId = Shape.getSwitchCoroId();1337 auto *AllocInst = SwitchId->getCoroAlloc();1338 coro::replaceCoroFree(SwitchId, /*Elide=*/AllocInst != nullptr);1339 if (AllocInst) {1340 IRBuilder<> Builder(AllocInst);1341 auto *Frame = Builder.CreateAlloca(Shape.FrameTy);1342 Frame->setAlignment(Shape.FrameAlign);1343 AllocInst->replaceAllUsesWith(Builder.getFalse());1344 AllocInst->eraseFromParent();1345 CoroBegin->replaceAllUsesWith(Frame);1346 } else {1347 CoroBegin->replaceAllUsesWith(CoroBegin->getMem());1348 }1349 1350 break;1351 }1352 case coro::ABI::Async:1353 case coro::ABI::Retcon:1354 case coro::ABI::RetconOnce:1355 CoroBegin->replaceAllUsesWith(PoisonValue::get(CoroBegin->getType()));1356 break;1357 }1358 1359 CoroBegin->eraseFromParent();1360 Shape.CoroBegin = nullptr;1361}1362 1363// SimplifySuspendPoint needs to check that there is no calls between1364// coro_save and coro_suspend, since any of the calls may potentially resume1365// the coroutine and if that is the case we cannot eliminate the suspend point.1366static bool hasCallsInBlockBetween(iterator_range<BasicBlock::iterator> R) {1367 for (Instruction &I : R) {1368 // Assume that no intrinsic can resume the coroutine.1369 if (isa<IntrinsicInst>(I))1370 continue;1371 1372 if (isa<CallBase>(I))1373 return true;1374 }1375 return false;1376}1377 1378static bool hasCallsInBlocksBetween(BasicBlock *SaveBB, BasicBlock *ResDesBB) {1379 SmallPtrSet<BasicBlock *, 8> Set;1380 SmallVector<BasicBlock *, 8> Worklist;1381 1382 Set.insert(SaveBB);1383 Worklist.push_back(ResDesBB);1384 1385 // Accumulate all blocks between SaveBB and ResDesBB. Because CoroSaveIntr1386 // returns a token consumed by suspend instruction, all blocks in between1387 // will have to eventually hit SaveBB when going backwards from ResDesBB.1388 while (!Worklist.empty()) {1389 auto *BB = Worklist.pop_back_val();1390 Set.insert(BB);1391 for (auto *Pred : predecessors(BB))1392 if (!Set.contains(Pred))1393 Worklist.push_back(Pred);1394 }1395 1396 // SaveBB and ResDesBB are checked separately in hasCallsBetween.1397 Set.erase(SaveBB);1398 Set.erase(ResDesBB);1399 1400 for (auto *BB : Set)1401 if (hasCallsInBlockBetween({BB->getFirstNonPHIIt(), BB->end()}))1402 return true;1403 1404 return false;1405}1406 1407static bool hasCallsBetween(Instruction *Save, Instruction *ResumeOrDestroy) {1408 auto *SaveBB = Save->getParent();1409 auto *ResumeOrDestroyBB = ResumeOrDestroy->getParent();1410 BasicBlock::iterator SaveIt = Save->getIterator();1411 BasicBlock::iterator ResumeOrDestroyIt = ResumeOrDestroy->getIterator();1412 1413 if (SaveBB == ResumeOrDestroyBB)1414 return hasCallsInBlockBetween({std::next(SaveIt), ResumeOrDestroyIt});1415 1416 // Any calls from Save to the end of the block?1417 if (hasCallsInBlockBetween({std::next(SaveIt), SaveBB->end()}))1418 return true;1419 1420 // Any calls from begging of the block up to ResumeOrDestroy?1421 if (hasCallsInBlockBetween(1422 {ResumeOrDestroyBB->getFirstNonPHIIt(), ResumeOrDestroyIt}))1423 return true;1424 1425 // Any calls in all of the blocks between SaveBB and ResumeOrDestroyBB?1426 if (hasCallsInBlocksBetween(SaveBB, ResumeOrDestroyBB))1427 return true;1428 1429 return false;1430}1431 1432// If a SuspendIntrin is preceded by Resume or Destroy, we can eliminate the1433// suspend point and replace it with nornal control flow.1434static bool simplifySuspendPoint(CoroSuspendInst *Suspend,1435 CoroBeginInst *CoroBegin) {1436 Instruction *Prev = Suspend->getPrevNode();1437 if (!Prev) {1438 auto *Pred = Suspend->getParent()->getSinglePredecessor();1439 if (!Pred)1440 return false;1441 Prev = Pred->getTerminator();1442 }1443 1444 CallBase *CB = dyn_cast<CallBase>(Prev);1445 if (!CB)1446 return false;1447 1448 auto *Callee = CB->getCalledOperand()->stripPointerCasts();1449 1450 // See if the callsite is for resumption or destruction of the coroutine.1451 auto *SubFn = dyn_cast<CoroSubFnInst>(Callee);1452 if (!SubFn)1453 return false;1454 1455 // Does not refer to the current coroutine, we cannot do anything with it.1456 if (SubFn->getFrame() != CoroBegin)1457 return false;1458 1459 // See if the transformation is safe. Specifically, see if there are any1460 // calls in between Save and CallInstr. They can potenitally resume the1461 // coroutine rendering this optimization unsafe.1462 auto *Save = Suspend->getCoroSave();1463 if (hasCallsBetween(Save, CB))1464 return false;1465 1466 // Replace llvm.coro.suspend with the value that results in resumption over1467 // the resume or cleanup path.1468 Suspend->replaceAllUsesWith(SubFn->getRawIndex());1469 Suspend->eraseFromParent();1470 Save->eraseFromParent();1471 1472 // No longer need a call to coro.resume or coro.destroy.1473 if (auto *Invoke = dyn_cast<InvokeInst>(CB)) {1474 BranchInst::Create(Invoke->getNormalDest(), Invoke->getIterator());1475 }1476 1477 // Grab the CalledValue from CB before erasing the CallInstr.1478 auto *CalledValue = CB->getCalledOperand();1479 CB->eraseFromParent();1480 1481 // If no more users remove it. Usually it is a bitcast of SubFn.1482 if (CalledValue != SubFn && CalledValue->user_empty())1483 if (auto *I = dyn_cast<Instruction>(CalledValue))1484 I->eraseFromParent();1485 1486 // Now we are good to remove SubFn.1487 if (SubFn->user_empty())1488 SubFn->eraseFromParent();1489 1490 return true;1491}1492 1493// Remove suspend points that are simplified.1494static void simplifySuspendPoints(coro::Shape &Shape) {1495 // Currently, the only simplification we do is switch-lowering-specific.1496 if (Shape.ABI != coro::ABI::Switch)1497 return;1498 1499 auto &S = Shape.CoroSuspends;1500 size_t I = 0, N = S.size();1501 if (N == 0)1502 return;1503 1504 size_t ChangedFinalIndex = std::numeric_limits<size_t>::max();1505 while (true) {1506 auto SI = cast<CoroSuspendInst>(S[I]);1507 // Leave final.suspend to handleFinalSuspend since it is undefined behavior1508 // to resume a coroutine suspended at the final suspend point.1509 if (!SI->isFinal() && simplifySuspendPoint(SI, Shape.CoroBegin)) {1510 if (--N == I)1511 break;1512 1513 std::swap(S[I], S[N]);1514 1515 if (cast<CoroSuspendInst>(S[I])->isFinal()) {1516 assert(Shape.SwitchLowering.HasFinalSuspend);1517 ChangedFinalIndex = I;1518 }1519 1520 continue;1521 }1522 if (++I == N)1523 break;1524 }1525 S.resize(N);1526 1527 // Maintain final.suspend in case final suspend was swapped.1528 // Due to we requrie the final suspend to be the last element of CoroSuspends.1529 if (ChangedFinalIndex < N) {1530 assert(cast<CoroSuspendInst>(S[ChangedFinalIndex])->isFinal());1531 std::swap(S[ChangedFinalIndex], S.back());1532 }1533}1534 1535namespace {1536 1537struct SwitchCoroutineSplitter {1538 static void split(Function &F, coro::Shape &Shape,1539 SmallVectorImpl<Function *> &Clones,1540 TargetTransformInfo &TTI) {1541 assert(Shape.ABI == coro::ABI::Switch);1542 1543 // Create a resume clone by cloning the body of the original function,1544 // setting new entry block and replacing coro.suspend an appropriate value1545 // to force resume or cleanup pass for every suspend point.1546 createResumeEntryBlock(F, Shape);1547 auto *ResumeClone = coro::SwitchCloner::createClone(1548 F, ".resume", Shape, coro::CloneKind::SwitchResume, TTI);1549 auto *DestroyClone = coro::SwitchCloner::createClone(1550 F, ".destroy", Shape, coro::CloneKind::SwitchUnwind, TTI);1551 auto *CleanupClone = coro::SwitchCloner::createClone(1552 F, ".cleanup", Shape, coro::CloneKind::SwitchCleanup, TTI);1553 1554 postSplitCleanup(*ResumeClone);1555 postSplitCleanup(*DestroyClone);1556 postSplitCleanup(*CleanupClone);1557 1558 // Store addresses resume/destroy/cleanup functions in the coroutine frame.1559 updateCoroFrame(Shape, ResumeClone, DestroyClone, CleanupClone);1560 1561 assert(Clones.empty());1562 Clones.push_back(ResumeClone);1563 Clones.push_back(DestroyClone);1564 Clones.push_back(CleanupClone);1565 1566 // Create a constant array referring to resume/destroy/clone functions1567 // pointed by the last argument of @llvm.coro.info, so that CoroElide pass1568 // can determined correct function to call.1569 setCoroInfo(F, Shape, Clones);1570 }1571 1572 // Create a variant of ramp function that does not perform heap allocation1573 // for a switch ABI coroutine.1574 //1575 // The newly split `.noalloc` ramp function has the following differences:1576 // - Has one additional frame pointer parameter in lieu of dynamic1577 // allocation.1578 // - Suppressed allocations by replacing coro.alloc and coro.free.1579 static Function *createNoAllocVariant(Function &F, coro::Shape &Shape,1580 SmallVectorImpl<Function *> &Clones) {1581 assert(Shape.ABI == coro::ABI::Switch);1582 auto *OrigFnTy = F.getFunctionType();1583 auto OldParams = OrigFnTy->params();1584 1585 SmallVector<Type *> NewParams;1586 NewParams.reserve(OldParams.size() + 1);1587 NewParams.append(OldParams.begin(), OldParams.end());1588 NewParams.push_back(PointerType::getUnqual(Shape.FrameTy->getContext()));1589 1590 auto *NewFnTy = FunctionType::get(OrigFnTy->getReturnType(), NewParams,1591 OrigFnTy->isVarArg());1592 Function *NoAllocF =1593 Function::Create(NewFnTy, F.getLinkage(), F.getName() + ".noalloc");1594 1595 ValueToValueMapTy VMap;1596 unsigned int Idx = 0;1597 for (const auto &I : F.args()) {1598 VMap[&I] = NoAllocF->getArg(Idx++);1599 }1600 // We just appended the frame pointer as the last argument of the new1601 // function.1602 auto FrameIdx = NoAllocF->arg_size() - 1;1603 SmallVector<ReturnInst *, 4> Returns;1604 CloneFunctionInto(NoAllocF, &F, VMap,1605 CloneFunctionChangeType::LocalChangesOnly, Returns);1606 1607 if (Shape.CoroBegin) {1608 auto *NewCoroBegin =1609 cast_if_present<CoroBeginInst>(VMap[Shape.CoroBegin]);1610 auto *NewCoroId = cast<CoroIdInst>(NewCoroBegin->getId());1611 coro::replaceCoroFree(NewCoroId, /*Elide=*/true);1612 coro::suppressCoroAllocs(NewCoroId);1613 NewCoroBegin->replaceAllUsesWith(NoAllocF->getArg(FrameIdx));1614 NewCoroBegin->eraseFromParent();1615 }1616 1617 Module *M = F.getParent();1618 M->getFunctionList().insert(M->end(), NoAllocF);1619 1620 removeUnreachableBlocks(*NoAllocF);1621 auto NewAttrs = NoAllocF->getAttributes();1622 // When we elide allocation, we read these attributes to determine the1623 // frame size and alignment.1624 addFramePointerAttrs(NewAttrs, NoAllocF->getContext(), FrameIdx,1625 Shape.FrameSize, Shape.FrameAlign,1626 /*NoAlias=*/false);1627 1628 NoAllocF->setAttributes(NewAttrs);1629 1630 Clones.push_back(NoAllocF);1631 // Reset the original function's coro info, make the new noalloc variant1632 // connected to the original ramp function.1633 setCoroInfo(F, Shape, Clones);1634 // After copying, set the linkage to internal linkage. Original function1635 // may have different linkage, but optimization dependent on this function1636 // generally relies on LTO.1637 NoAllocF->setLinkage(llvm::GlobalValue::InternalLinkage);1638 return NoAllocF;1639 }1640 1641private:1642 // Create an entry block for a resume function with a switch that will jump to1643 // suspend points.1644 static void createResumeEntryBlock(Function &F, coro::Shape &Shape) {1645 LLVMContext &C = F.getContext();1646 1647 DIBuilder DBuilder(*F.getParent(), /*AllowUnresolved*/ false);1648 DISubprogram *DIS = F.getSubprogram();1649 // If there is no DISubprogram for F, it implies the function is compiled1650 // without debug info. So we also don't generate debug info for the1651 // suspension points.1652 bool AddDebugLabels = DIS && DIS->getUnit() &&1653 (DIS->getUnit()->getEmissionKind() ==1654 DICompileUnit::DebugEmissionKind::FullDebug);1655 1656 // resume.entry:1657 // %index.addr = getelementptr inbounds %f.Frame, %f.Frame* %FramePtr, i321658 // 0, i32 2 % index = load i32, i32* %index.addr switch i32 %index, label1659 // %unreachable [1660 // i32 0, label %resume.01661 // i32 1, label %resume.11662 // ...1663 // ]1664 1665 auto *NewEntry = BasicBlock::Create(C, "resume.entry", &F);1666 auto *UnreachBB = BasicBlock::Create(C, "unreachable", &F);1667 1668 IRBuilder<> Builder(NewEntry);1669 auto *FramePtr = Shape.FramePtr;1670 auto *FrameTy = Shape.FrameTy;1671 auto *GepIndex = Builder.CreateStructGEP(1672 FrameTy, FramePtr, Shape.getSwitchIndexField(), "index.addr");1673 auto *Index = Builder.CreateLoad(Shape.getIndexType(), GepIndex, "index");1674 auto *Switch =1675 Builder.CreateSwitch(Index, UnreachBB, Shape.CoroSuspends.size());1676 Shape.SwitchLowering.ResumeSwitch = Switch;1677 1678 // Split all coro.suspend calls1679 size_t SuspendIndex = 0;1680 for (auto *AnyS : Shape.CoroSuspends) {1681 auto *S = cast<CoroSuspendInst>(AnyS);1682 ConstantInt *IndexVal = Shape.getIndex(SuspendIndex);1683 1684 // Replace CoroSave with a store to Index:1685 // %index.addr = getelementptr %f.frame... (index field number)1686 // store i32 %IndexVal, i32* %index.addr11687 auto *Save = S->getCoroSave();1688 Builder.SetInsertPoint(Save);1689 if (S->isFinal()) {1690 // The coroutine should be marked done if it reaches the final suspend1691 // point.1692 markCoroutineAsDone(Builder, Shape, FramePtr);1693 } else {1694 auto *GepIndex = Builder.CreateStructGEP(1695 FrameTy, FramePtr, Shape.getSwitchIndexField(), "index.addr");1696 Builder.CreateStore(IndexVal, GepIndex);1697 }1698 1699 Save->replaceAllUsesWith(ConstantTokenNone::get(C));1700 Save->eraseFromParent();1701 1702 // Split block before and after coro.suspend and add a jump from an entry1703 // switch:1704 //1705 // whateverBB:1706 // whatever1707 // %0 = call i8 @llvm.coro.suspend(token none, i1 false)1708 // switch i8 %0, label %suspend[i8 0, label %resume1709 // i8 1, label %cleanup]1710 // becomes:1711 //1712 // whateverBB:1713 // whatever1714 // br label %resume.0.landing1715 //1716 // resume.0: ; <--- jump from the switch in the resume.entry1717 // #dbg_label(...) ; <--- artificial label for debuggers1718 // %0 = tail call i8 @llvm.coro.suspend(token none, i1 false)1719 // br label %resume.0.landing1720 //1721 // resume.0.landing:1722 // %1 = phi i8[-1, %whateverBB], [%0, %resume.0]1723 // switch i8 % 1, label %suspend [i8 0, label %resume1724 // i8 1, label %cleanup]1725 1726 auto *SuspendBB = S->getParent();1727 auto *ResumeBB =1728 SuspendBB->splitBasicBlock(S, "resume." + Twine(SuspendIndex));1729 auto *LandingBB = ResumeBB->splitBasicBlock(1730 S->getNextNode(), ResumeBB->getName() + Twine(".landing"));1731 Switch->addCase(IndexVal, ResumeBB);1732 1733 cast<BranchInst>(SuspendBB->getTerminator())->setSuccessor(0, LandingBB);1734 auto *PN = PHINode::Create(Builder.getInt8Ty(), 2, "");1735 PN->insertBefore(LandingBB->begin());1736 S->replaceAllUsesWith(PN);1737 PN->addIncoming(Builder.getInt8(-1), SuspendBB);1738 PN->addIncoming(S, ResumeBB);1739 1740 if (AddDebugLabels) {1741 if (DebugLoc SuspendLoc = S->getDebugLoc()) {1742 std::string LabelName =1743 ("__coro_resume_" + Twine(SuspendIndex)).str();1744 // Take the "inlined at" location recursively, if present. This is1745 // mandatory as the DILabel insertion checks that the scopes of label1746 // and the attached location match. This is not the case when the1747 // suspend location has been inlined due to pointing to the original1748 // scope.1749 DILocation *DILoc = SuspendLoc;1750 while (DILocation *InlinedAt = DILoc->getInlinedAt())1751 DILoc = InlinedAt;1752 1753 DILabel *ResumeLabel =1754 DBuilder.createLabel(DIS, LabelName, DILoc->getFile(),1755 SuspendLoc.getLine(), SuspendLoc.getCol(),1756 /*IsArtificial=*/true,1757 /*CoroSuspendIdx=*/SuspendIndex,1758 /*AlwaysPreserve=*/false);1759 DBuilder.insertLabel(ResumeLabel, DILoc, ResumeBB->begin());1760 }1761 }1762 1763 ++SuspendIndex;1764 }1765 1766 Builder.SetInsertPoint(UnreachBB);1767 Builder.CreateUnreachable();1768 DBuilder.finalize();1769 1770 Shape.SwitchLowering.ResumeEntryBlock = NewEntry;1771 }1772 1773 // Store addresses of Resume/Destroy/Cleanup functions in the coroutine frame.1774 static void updateCoroFrame(coro::Shape &Shape, Function *ResumeFn,1775 Function *DestroyFn, Function *CleanupFn) {1776 IRBuilder<> Builder(&*Shape.getInsertPtAfterFramePtr());1777 1778 auto *ResumeAddr = Builder.CreateStructGEP(1779 Shape.FrameTy, Shape.FramePtr, coro::Shape::SwitchFieldIndex::Resume,1780 "resume.addr");1781 Builder.CreateStore(ResumeFn, ResumeAddr);1782 1783 Value *DestroyOrCleanupFn = DestroyFn;1784 1785 CoroIdInst *CoroId = Shape.getSwitchCoroId();1786 if (CoroAllocInst *CA = CoroId->getCoroAlloc()) {1787 // If there is a CoroAlloc and it returns false (meaning we elide the1788 // allocation, use CleanupFn instead of DestroyFn).1789 DestroyOrCleanupFn = Builder.CreateSelect(CA, DestroyFn, CleanupFn);1790 }1791 1792 auto *DestroyAddr = Builder.CreateStructGEP(1793 Shape.FrameTy, Shape.FramePtr, coro::Shape::SwitchFieldIndex::Destroy,1794 "destroy.addr");1795 Builder.CreateStore(DestroyOrCleanupFn, DestroyAddr);1796 }1797 1798 // Create a global constant array containing pointers to functions provided1799 // and set Info parameter of CoroBegin to point at this constant. Example:1800 //1801 // @f.resumers = internal constant [2 x void(%f.frame*)*]1802 // [void(%f.frame*)* @f.resume, void(%f.frame*)*1803 // @f.destroy]1804 // define void @f() {1805 // ...1806 // call i8* @llvm.coro.begin(i8* null, i32 0, i8* null,1807 // i8* bitcast([2 x void(%f.frame*)*] * @f.resumers to1808 // i8*))1809 //1810 // Assumes that all the functions have the same signature.1811 static void setCoroInfo(Function &F, coro::Shape &Shape,1812 ArrayRef<Function *> Fns) {1813 // This only works under the switch-lowering ABI because coro elision1814 // only works on the switch-lowering ABI.1815 SmallVector<Constant *, 4> Args(Fns);1816 assert(!Args.empty());1817 Function *Part = *Fns.begin();1818 Module *M = Part->getParent();1819 auto *ArrTy = ArrayType::get(Part->getType(), Args.size());1820 1821 auto *ConstVal = ConstantArray::get(ArrTy, Args);1822 auto *GV = new GlobalVariable(*M, ConstVal->getType(), /*isConstant=*/true,1823 GlobalVariable::PrivateLinkage, ConstVal,1824 F.getName() + Twine(".resumers"));1825 1826 // Update coro.begin instruction to refer to this constant.1827 LLVMContext &C = F.getContext();1828 auto *BC = ConstantExpr::getPointerCast(GV, PointerType::getUnqual(C));1829 Shape.getSwitchCoroId()->setInfo(BC);1830 }1831};1832 1833} // namespace1834 1835static void replaceAsyncResumeFunction(CoroSuspendAsyncInst *Suspend,1836 Value *Continuation) {1837 auto *ResumeIntrinsic = Suspend->getResumeFunction();1838 auto &Context = Suspend->getParent()->getParent()->getContext();1839 auto *Int8PtrTy = PointerType::getUnqual(Context);1840 1841 IRBuilder<> Builder(ResumeIntrinsic);1842 auto *Val = Builder.CreateBitOrPointerCast(Continuation, Int8PtrTy);1843 ResumeIntrinsic->replaceAllUsesWith(Val);1844 ResumeIntrinsic->eraseFromParent();1845 Suspend->setOperand(CoroSuspendAsyncInst::ResumeFunctionArg,1846 PoisonValue::get(Int8PtrTy));1847}1848 1849/// Coerce the arguments in \p FnArgs according to \p FnTy in \p CallArgs.1850static void coerceArguments(IRBuilder<> &Builder, FunctionType *FnTy,1851 ArrayRef<Value *> FnArgs,1852 SmallVectorImpl<Value *> &CallArgs) {1853 size_t ArgIdx = 0;1854 for (auto *paramTy : FnTy->params()) {1855 assert(ArgIdx < FnArgs.size());1856 if (paramTy != FnArgs[ArgIdx]->getType())1857 CallArgs.push_back(1858 Builder.CreateBitOrPointerCast(FnArgs[ArgIdx], paramTy));1859 else1860 CallArgs.push_back(FnArgs[ArgIdx]);1861 ++ArgIdx;1862 }1863}1864 1865CallInst *coro::createMustTailCall(DebugLoc Loc, Function *MustTailCallFn,1866 TargetTransformInfo &TTI,1867 ArrayRef<Value *> Arguments,1868 IRBuilder<> &Builder) {1869 auto *FnTy = MustTailCallFn->getFunctionType();1870 // Coerce the arguments, llvm optimizations seem to ignore the types in1871 // vaarg functions and throws away casts in optimized mode.1872 SmallVector<Value *, 8> CallArgs;1873 coerceArguments(Builder, FnTy, Arguments, CallArgs);1874 1875 auto *TailCall = Builder.CreateCall(FnTy, MustTailCallFn, CallArgs);1876 // Skip targets which don't support tail call.1877 if (TTI.supportsTailCallFor(TailCall)) {1878 TailCall->setTailCallKind(CallInst::TCK_MustTail);1879 }1880 TailCall->setDebugLoc(Loc);1881 TailCall->setCallingConv(MustTailCallFn->getCallingConv());1882 return TailCall;1883}1884 1885void coro::AsyncABI::splitCoroutine(Function &F, coro::Shape &Shape,1886 SmallVectorImpl<Function *> &Clones,1887 TargetTransformInfo &TTI) {1888 assert(Shape.ABI == coro::ABI::Async);1889 assert(Clones.empty());1890 // Reset various things that the optimizer might have decided it1891 // "knows" about the coroutine function due to not seeing a return.1892 F.removeFnAttr(Attribute::NoReturn);1893 F.removeRetAttr(Attribute::NoAlias);1894 F.removeRetAttr(Attribute::NonNull);1895 1896 auto &Context = F.getContext();1897 auto *Int8PtrTy = PointerType::getUnqual(Context);1898 1899 auto *Id = Shape.getAsyncCoroId();1900 IRBuilder<> Builder(Id);1901 1902 auto *FramePtr = Id->getStorage();1903 FramePtr = Builder.CreateBitOrPointerCast(FramePtr, Int8PtrTy);1904 FramePtr = Builder.CreateConstInBoundsGEP1_32(1905 Type::getInt8Ty(Context), FramePtr, Shape.AsyncLowering.FrameOffset,1906 "async.ctx.frameptr");1907 1908 // Map all uses of llvm.coro.begin to the allocated frame pointer.1909 {1910 // Make sure we don't invalidate Shape.FramePtr.1911 TrackingVH<Value> Handle(Shape.FramePtr);1912 Shape.CoroBegin->replaceAllUsesWith(FramePtr);1913 Shape.FramePtr = Handle.getValPtr();1914 }1915 1916 // Create all the functions in order after the main function.1917 auto NextF = std::next(F.getIterator());1918 1919 // Create a continuation function for each of the suspend points.1920 Clones.reserve(Shape.CoroSuspends.size());1921 for (auto [Idx, CS] : llvm::enumerate(Shape.CoroSuspends)) {1922 auto *Suspend = cast<CoroSuspendAsyncInst>(CS);1923 1924 // Create the clone declaration.1925 auto ResumeNameSuffix = ".resume.";1926 auto ProjectionFunctionName =1927 Suspend->getAsyncContextProjectionFunction()->getName();1928 bool UseSwiftMangling = false;1929 if (ProjectionFunctionName == "__swift_async_resume_project_context") {1930 ResumeNameSuffix = "TQ";1931 UseSwiftMangling = true;1932 } else if (ProjectionFunctionName == "__swift_async_resume_get_context") {1933 ResumeNameSuffix = "TY";1934 UseSwiftMangling = true;1935 }1936 auto *Continuation = createCloneDeclaration(1937 F, Shape,1938 UseSwiftMangling ? ResumeNameSuffix + Twine(Idx) + "_"1939 : ResumeNameSuffix + Twine(Idx),1940 NextF, Suspend);1941 Clones.push_back(Continuation);1942 1943 // Insert a branch to a new return block immediately before the suspend1944 // point.1945 auto *SuspendBB = Suspend->getParent();1946 auto *NewSuspendBB = SuspendBB->splitBasicBlock(Suspend);1947 auto *Branch = cast<BranchInst>(SuspendBB->getTerminator());1948 1949 // Place it before the first suspend.1950 auto *ReturnBB =1951 BasicBlock::Create(F.getContext(), "coro.return", &F, NewSuspendBB);1952 Branch->setSuccessor(0, ReturnBB);1953 1954 IRBuilder<> Builder(ReturnBB);1955 1956 // Insert the call to the tail call function and inline it.1957 auto *Fn = Suspend->getMustTailCallFunction();1958 SmallVector<Value *, 8> Args(Suspend->args());1959 auto FnArgs = ArrayRef<Value *>(Args).drop_front(1960 CoroSuspendAsyncInst::MustTailCallFuncArg + 1);1961 auto *TailCall = coro::createMustTailCall(Suspend->getDebugLoc(), Fn, TTI,1962 FnArgs, Builder);1963 Builder.CreateRetVoid();1964 InlineFunctionInfo FnInfo;1965 (void)InlineFunction(*TailCall, FnInfo);1966 1967 // Replace the lvm.coro.async.resume intrisic call.1968 replaceAsyncResumeFunction(Suspend, Continuation);1969 }1970 1971 assert(Clones.size() == Shape.CoroSuspends.size());1972 1973 for (auto [Idx, CS] : llvm::enumerate(Shape.CoroSuspends)) {1974 auto *Suspend = CS;1975 auto *Clone = Clones[Idx];1976 1977 coro::BaseCloner::createClone(F, "resume." + Twine(Idx), Shape, Clone,1978 Suspend, TTI);1979 }1980}1981 1982void coro::AnyRetconABI::splitCoroutine(Function &F, coro::Shape &Shape,1983 SmallVectorImpl<Function *> &Clones,1984 TargetTransformInfo &TTI) {1985 assert(Shape.ABI == coro::ABI::Retcon || Shape.ABI == coro::ABI::RetconOnce);1986 assert(Clones.empty());1987 1988 // Reset various things that the optimizer might have decided it1989 // "knows" about the coroutine function due to not seeing a return.1990 F.removeFnAttr(Attribute::NoReturn);1991 F.removeRetAttr(Attribute::NoAlias);1992 F.removeRetAttr(Attribute::NonNull);1993 1994 // Allocate the frame.1995 auto *Id = Shape.getRetconCoroId();1996 Value *RawFramePtr;1997 if (Shape.RetconLowering.IsFrameInlineInStorage) {1998 RawFramePtr = Id->getStorage();1999 } else {2000 IRBuilder<> Builder(Id);2001 2002 // Determine the size of the frame.2003 const DataLayout &DL = F.getDataLayout();2004 auto Size = DL.getTypeAllocSize(Shape.FrameTy);2005 2006 // Allocate. We don't need to update the call graph node because we're2007 // going to recompute it from scratch after splitting.2008 // FIXME: pass the required alignment2009 RawFramePtr = Shape.emitAlloc(Builder, Builder.getInt64(Size), nullptr);2010 RawFramePtr =2011 Builder.CreateBitCast(RawFramePtr, Shape.CoroBegin->getType());2012 2013 // Stash the allocated frame pointer in the continuation storage.2014 Builder.CreateStore(RawFramePtr, Id->getStorage());2015 }2016 2017 // Map all uses of llvm.coro.begin to the allocated frame pointer.2018 {2019 // Make sure we don't invalidate Shape.FramePtr.2020 TrackingVH<Value> Handle(Shape.FramePtr);2021 Shape.CoroBegin->replaceAllUsesWith(RawFramePtr);2022 Shape.FramePtr = Handle.getValPtr();2023 }2024 2025 // Create a unique return block.2026 BasicBlock *ReturnBB = nullptr;2027 PHINode *ContinuationPhi = nullptr;2028 SmallVector<PHINode *, 4> ReturnPHIs;2029 2030 // Create all the functions in order after the main function.2031 auto NextF = std::next(F.getIterator());2032 2033 // Create a continuation function for each of the suspend points.2034 Clones.reserve(Shape.CoroSuspends.size());2035 for (auto [Idx, CS] : llvm::enumerate(Shape.CoroSuspends)) {2036 auto Suspend = cast<CoroSuspendRetconInst>(CS);2037 2038 // Create the clone declaration.2039 auto Continuation = createCloneDeclaration(2040 F, Shape, ".resume." + Twine(Idx), NextF, nullptr);2041 Clones.push_back(Continuation);2042 2043 // Insert a branch to the unified return block immediately before2044 // the suspend point.2045 auto SuspendBB = Suspend->getParent();2046 auto NewSuspendBB = SuspendBB->splitBasicBlock(Suspend);2047 auto Branch = cast<BranchInst>(SuspendBB->getTerminator());2048 2049 // Create the unified return block.2050 if (!ReturnBB) {2051 // Place it before the first suspend.2052 ReturnBB =2053 BasicBlock::Create(F.getContext(), "coro.return", &F, NewSuspendBB);2054 Shape.RetconLowering.ReturnBlock = ReturnBB;2055 2056 IRBuilder<> Builder(ReturnBB);2057 2058 // First, the continuation.2059 ContinuationPhi =2060 Builder.CreatePHI(Continuation->getType(), Shape.CoroSuspends.size());2061 2062 // Create PHIs for all other return values.2063 assert(ReturnPHIs.empty());2064 2065 // Next, all the directly-yielded values.2066 for (auto *ResultTy : Shape.getRetconResultTypes())2067 ReturnPHIs.push_back(2068 Builder.CreatePHI(ResultTy, Shape.CoroSuspends.size()));2069 2070 // Build the return value.2071 auto RetTy = F.getReturnType();2072 2073 // Cast the continuation value if necessary.2074 // We can't rely on the types matching up because that type would2075 // have to be infinite.2076 auto CastedContinuationTy =2077 (ReturnPHIs.empty() ? RetTy : RetTy->getStructElementType(0));2078 auto *CastedContinuation =2079 Builder.CreateBitCast(ContinuationPhi, CastedContinuationTy);2080 2081 Value *RetV = CastedContinuation;2082 if (!ReturnPHIs.empty()) {2083 auto ValueIdx = 0;2084 RetV = PoisonValue::get(RetTy);2085 RetV = Builder.CreateInsertValue(RetV, CastedContinuation, ValueIdx++);2086 2087 for (auto Phi : ReturnPHIs)2088 RetV = Builder.CreateInsertValue(RetV, Phi, ValueIdx++);2089 }2090 2091 Builder.CreateRet(RetV);2092 }2093 2094 // Branch to the return block.2095 Branch->setSuccessor(0, ReturnBB);2096 assert(ContinuationPhi);2097 ContinuationPhi->addIncoming(Continuation, SuspendBB);2098 for (auto [Phi, VUse] :2099 llvm::zip_equal(ReturnPHIs, Suspend->value_operands()))2100 Phi->addIncoming(VUse, SuspendBB);2101 }2102 2103 assert(Clones.size() == Shape.CoroSuspends.size());2104 2105 for (auto [Idx, CS] : llvm::enumerate(Shape.CoroSuspends)) {2106 auto Suspend = CS;2107 auto Clone = Clones[Idx];2108 2109 coro::BaseCloner::createClone(F, "resume." + Twine(Idx), Shape, Clone,2110 Suspend, TTI);2111 }2112}2113 2114namespace {2115class PrettyStackTraceFunction : public PrettyStackTraceEntry {2116 Function &F;2117 2118public:2119 PrettyStackTraceFunction(Function &F) : F(F) {}2120 void print(raw_ostream &OS) const override {2121 OS << "While splitting coroutine ";2122 F.printAsOperand(OS, /*print type*/ false, F.getParent());2123 OS << "\n";2124 }2125};2126} // namespace2127 2128/// Remove calls to llvm.coro.end in the original function.2129static void removeCoroEndsFromRampFunction(const coro::Shape &Shape) {2130 if (Shape.ABI != coro::ABI::Switch) {2131 for (auto *End : Shape.CoroEnds) {2132 replaceCoroEnd(End, Shape, Shape.FramePtr, /*in ramp*/ true, nullptr);2133 }2134 } else {2135 for (llvm::AnyCoroEndInst *End : Shape.CoroEnds)2136 End->eraseFromParent();2137 }2138}2139 2140static void removeCoroIsInRampFromRampFunction(const coro::Shape &Shape) {2141 for (auto *II : Shape.CoroIsInRampInsts) {2142 auto &Ctx = II->getContext();2143 II->replaceAllUsesWith(ConstantInt::getTrue(Ctx));2144 II->eraseFromParent();2145 }2146}2147 2148static bool hasSafeElideCaller(Function &F) {2149 for (auto *U : F.users()) {2150 if (auto *CB = dyn_cast<CallBase>(U)) {2151 auto *Caller = CB->getFunction();2152 if (Caller && Caller->isPresplitCoroutine() &&2153 CB->hasFnAttr(llvm::Attribute::CoroElideSafe))2154 return true;2155 }2156 }2157 return false;2158}2159 2160void coro::SwitchABI::splitCoroutine(Function &F, coro::Shape &Shape,2161 SmallVectorImpl<Function *> &Clones,2162 TargetTransformInfo &TTI) {2163 SwitchCoroutineSplitter::split(F, Shape, Clones, TTI);2164}2165 2166static void doSplitCoroutine(Function &F, SmallVectorImpl<Function *> &Clones,2167 coro::BaseABI &ABI, TargetTransformInfo &TTI,2168 bool OptimizeFrame) {2169 PrettyStackTraceFunction prettyStackTrace(F);2170 2171 auto &Shape = ABI.Shape;2172 assert(Shape.CoroBegin);2173 2174 lowerAwaitSuspends(F, Shape);2175 2176 simplifySuspendPoints(Shape);2177 2178 normalizeCoroutine(F, Shape, TTI);2179 ABI.buildCoroutineFrame(OptimizeFrame);2180 replaceFrameSizeAndAlignment(Shape);2181 2182 bool isNoSuspendCoroutine = Shape.CoroSuspends.empty();2183 2184 bool shouldCreateNoAllocVariant =2185 !isNoSuspendCoroutine && Shape.ABI == coro::ABI::Switch &&2186 hasSafeElideCaller(F) && !F.hasFnAttribute(llvm::Attribute::NoInline);2187 2188 // If there are no suspend points, no split required, just remove2189 // the allocation and deallocation blocks, they are not needed.2190 if (isNoSuspendCoroutine) {2191 handleNoSuspendCoroutine(Shape);2192 } else {2193 ABI.splitCoroutine(F, Shape, Clones, TTI);2194 }2195 2196 // Replace all the swifterror operations in the original function.2197 // This invalidates SwiftErrorOps in the Shape.2198 replaceSwiftErrorOps(F, Shape, nullptr);2199 2200 // Salvage debug intrinsics that point into the coroutine frame in the2201 // original function. The Cloner has already salvaged debug info in the new2202 // coroutine funclets.2203 SmallDenseMap<Argument *, AllocaInst *, 4> ArgToAllocaMap;2204 auto DbgVariableRecords = collectDbgVariableRecords(F);2205 for (DbgVariableRecord *DVR : DbgVariableRecords)2206 coro::salvageDebugInfo(ArgToAllocaMap, *DVR, false /*UseEntryValue*/);2207 2208 removeCoroEndsFromRampFunction(Shape);2209 removeCoroIsInRampFromRampFunction(Shape);2210 2211 if (shouldCreateNoAllocVariant)2212 SwitchCoroutineSplitter::createNoAllocVariant(F, Shape, Clones);2213}2214 2215static LazyCallGraph::SCC &updateCallGraphAfterCoroutineSplit(2216 LazyCallGraph::Node &N, const coro::Shape &Shape,2217 const SmallVectorImpl<Function *> &Clones, LazyCallGraph::SCC &C,2218 LazyCallGraph &CG, CGSCCAnalysisManager &AM, CGSCCUpdateResult &UR,2219 FunctionAnalysisManager &FAM) {2220 2221 auto *CurrentSCC = &C;2222 if (!Clones.empty()) {2223 switch (Shape.ABI) {2224 case coro::ABI::Switch:2225 // Each clone in the Switch lowering is independent of the other clones.2226 // Let the LazyCallGraph know about each one separately.2227 for (Function *Clone : Clones)2228 CG.addSplitFunction(N.getFunction(), *Clone);2229 break;2230 case coro::ABI::Async:2231 case coro::ABI::Retcon:2232 case coro::ABI::RetconOnce:2233 // Each clone in the Async/Retcon lowering references of the other clones.2234 // Let the LazyCallGraph know about all of them at once.2235 if (!Clones.empty())2236 CG.addSplitRefRecursiveFunctions(N.getFunction(), Clones);2237 break;2238 }2239 2240 // Let the CGSCC infra handle the changes to the original function.2241 CurrentSCC = &updateCGAndAnalysisManagerForCGSCCPass(CG, *CurrentSCC, N, AM,2242 UR, FAM);2243 }2244 2245 // Do some cleanup and let the CGSCC infra see if we've cleaned up any edges2246 // to the split functions.2247 postSplitCleanup(N.getFunction());2248 CurrentSCC = &updateCGAndAnalysisManagerForFunctionPass(CG, *CurrentSCC, N,2249 AM, UR, FAM);2250 return *CurrentSCC;2251}2252 2253/// Replace a call to llvm.coro.prepare.retcon.2254static void replacePrepare(CallInst *Prepare, LazyCallGraph &CG,2255 LazyCallGraph::SCC &C) {2256 auto CastFn = Prepare->getArgOperand(0); // as an i8*2257 auto Fn = CastFn->stripPointerCasts(); // as its original type2258 2259 // Attempt to peephole this pattern:2260 // %0 = bitcast [[TYPE]] @some_function to i8*2261 // %1 = call @llvm.coro.prepare.retcon(i8* %0)2262 // %2 = bitcast %1 to [[TYPE]]2263 // ==>2264 // %2 = @some_function2265 for (Use &U : llvm::make_early_inc_range(Prepare->uses())) {2266 // Look for bitcasts back to the original function type.2267 auto *Cast = dyn_cast<BitCastInst>(U.getUser());2268 if (!Cast || Cast->getType() != Fn->getType())2269 continue;2270 2271 // Replace and remove the cast.2272 Cast->replaceAllUsesWith(Fn);2273 Cast->eraseFromParent();2274 }2275 2276 // Replace any remaining uses with the function as an i8*.2277 // This can never directly be a callee, so we don't need to update CG.2278 Prepare->replaceAllUsesWith(CastFn);2279 Prepare->eraseFromParent();2280 2281 // Kill dead bitcasts.2282 while (auto *Cast = dyn_cast<BitCastInst>(CastFn)) {2283 if (!Cast->use_empty())2284 break;2285 CastFn = Cast->getOperand(0);2286 Cast->eraseFromParent();2287 }2288}2289 2290static bool replaceAllPrepares(Function *PrepareFn, LazyCallGraph &CG,2291 LazyCallGraph::SCC &C) {2292 bool Changed = false;2293 for (Use &P : llvm::make_early_inc_range(PrepareFn->uses())) {2294 // Intrinsics can only be used in calls.2295 auto *Prepare = cast<CallInst>(P.getUser());2296 replacePrepare(Prepare, CG, C);2297 Changed = true;2298 }2299 2300 return Changed;2301}2302 2303static void addPrepareFunction(const Module &M,2304 SmallVectorImpl<Function *> &Fns,2305 StringRef Name) {2306 auto *PrepareFn = M.getFunction(Name);2307 if (PrepareFn && !PrepareFn->use_empty())2308 Fns.push_back(PrepareFn);2309}2310 2311static std::unique_ptr<coro::BaseABI>2312CreateNewABI(Function &F, coro::Shape &S,2313 std::function<bool(Instruction &)> IsMatCallback,2314 const SmallVector<CoroSplitPass::BaseABITy> GenCustomABIs) {2315 if (S.CoroBegin->hasCustomABI()) {2316 unsigned CustomABI = S.CoroBegin->getCustomABI();2317 if (CustomABI >= GenCustomABIs.size())2318 llvm_unreachable("Custom ABI not found amoung those specified");2319 return GenCustomABIs[CustomABI](F, S);2320 }2321 2322 switch (S.ABI) {2323 case coro::ABI::Switch:2324 return std::make_unique<coro::SwitchABI>(F, S, IsMatCallback);2325 case coro::ABI::Async:2326 return std::make_unique<coro::AsyncABI>(F, S, IsMatCallback);2327 case coro::ABI::Retcon:2328 return std::make_unique<coro::AnyRetconABI>(F, S, IsMatCallback);2329 case coro::ABI::RetconOnce:2330 return std::make_unique<coro::AnyRetconABI>(F, S, IsMatCallback);2331 }2332 llvm_unreachable("Unknown ABI");2333}2334 2335CoroSplitPass::CoroSplitPass(bool OptimizeFrame)2336 : CreateAndInitABI([](Function &F, coro::Shape &S) {2337 std::unique_ptr<coro::BaseABI> ABI =2338 CreateNewABI(F, S, coro::isTriviallyMaterializable, {});2339 ABI->init();2340 return ABI;2341 }),2342 OptimizeFrame(OptimizeFrame) {}2343 2344CoroSplitPass::CoroSplitPass(2345 SmallVector<CoroSplitPass::BaseABITy> GenCustomABIs, bool OptimizeFrame)2346 : CreateAndInitABI([=](Function &F, coro::Shape &S) {2347 std::unique_ptr<coro::BaseABI> ABI =2348 CreateNewABI(F, S, coro::isTriviallyMaterializable, GenCustomABIs);2349 ABI->init();2350 return ABI;2351 }),2352 OptimizeFrame(OptimizeFrame) {}2353 2354// For back compatibility, constructor takes a materializable callback and2355// creates a generator for an ABI with a modified materializable callback.2356CoroSplitPass::CoroSplitPass(std::function<bool(Instruction &)> IsMatCallback,2357 bool OptimizeFrame)2358 : CreateAndInitABI([=](Function &F, coro::Shape &S) {2359 std::unique_ptr<coro::BaseABI> ABI =2360 CreateNewABI(F, S, IsMatCallback, {});2361 ABI->init();2362 return ABI;2363 }),2364 OptimizeFrame(OptimizeFrame) {}2365 2366// For back compatibility, constructor takes a materializable callback and2367// creates a generator for an ABI with a modified materializable callback.2368CoroSplitPass::CoroSplitPass(2369 std::function<bool(Instruction &)> IsMatCallback,2370 SmallVector<CoroSplitPass::BaseABITy> GenCustomABIs, bool OptimizeFrame)2371 : CreateAndInitABI([=](Function &F, coro::Shape &S) {2372 std::unique_ptr<coro::BaseABI> ABI =2373 CreateNewABI(F, S, IsMatCallback, GenCustomABIs);2374 ABI->init();2375 return ABI;2376 }),2377 OptimizeFrame(OptimizeFrame) {}2378 2379PreservedAnalyses CoroSplitPass::run(LazyCallGraph::SCC &C,2380 CGSCCAnalysisManager &AM,2381 LazyCallGraph &CG, CGSCCUpdateResult &UR) {2382 // NB: One invariant of a valid LazyCallGraph::SCC is that it must contain a2383 // non-zero number of nodes, so we assume that here and grab the first2384 // node's function's module.2385 Module &M = *C.begin()->getFunction().getParent();2386 auto &FAM =2387 AM.getResult<FunctionAnalysisManagerCGSCCProxy>(C, CG).getManager();2388 2389 // Check for uses of llvm.coro.prepare.retcon/async.2390 SmallVector<Function *, 2> PrepareFns;2391 addPrepareFunction(M, PrepareFns, "llvm.coro.prepare.retcon");2392 addPrepareFunction(M, PrepareFns, "llvm.coro.prepare.async");2393 2394 // Find coroutines for processing.2395 SmallVector<LazyCallGraph::Node *> Coroutines;2396 for (LazyCallGraph::Node &N : C)2397 if (N.getFunction().isPresplitCoroutine())2398 Coroutines.push_back(&N);2399 2400 if (Coroutines.empty() && PrepareFns.empty())2401 return PreservedAnalyses::all();2402 2403 auto *CurrentSCC = &C;2404 // Split all the coroutines.2405 for (LazyCallGraph::Node *N : Coroutines) {2406 Function &F = N->getFunction();2407 LLVM_DEBUG(dbgs() << "CoroSplit: Processing coroutine '" << F.getName()2408 << "\n");2409 2410 // The suspend-crossing algorithm in buildCoroutineFrame gets tripped up2411 // by unreachable blocks, so remove them as a first pass. Remove the2412 // unreachable blocks before collecting intrinsics into Shape.2413 removeUnreachableBlocks(F);2414 2415 coro::Shape Shape(F);2416 if (!Shape.CoroBegin)2417 continue;2418 2419 F.setSplittedCoroutine();2420 2421 std::unique_ptr<coro::BaseABI> ABI = CreateAndInitABI(F, Shape);2422 2423 SmallVector<Function *, 4> Clones;2424 auto &TTI = FAM.getResult<TargetIRAnalysis>(F);2425 doSplitCoroutine(F, Clones, *ABI, TTI, OptimizeFrame);2426 CurrentSCC = &updateCallGraphAfterCoroutineSplit(2427 *N, Shape, Clones, *CurrentSCC, CG, AM, UR, FAM);2428 2429 auto &ORE = FAM.getResult<OptimizationRemarkEmitterAnalysis>(F);2430 ORE.emit([&]() {2431 return OptimizationRemark(DEBUG_TYPE, "CoroSplit", &F)2432 << "Split '" << ore::NV("function", F.getName())2433 << "' (frame_size=" << ore::NV("frame_size", Shape.FrameSize)2434 << ", align=" << ore::NV("align", Shape.FrameAlign.value()) << ")";2435 });2436 2437 if (!Shape.CoroSuspends.empty()) {2438 // Run the CGSCC pipeline on the original and newly split functions.2439 UR.CWorklist.insert(CurrentSCC);2440 for (Function *Clone : Clones)2441 UR.CWorklist.insert(CG.lookupSCC(CG.get(*Clone)));2442 } else if (Shape.ABI == coro::ABI::Async) {2443 // Reprocess the function to inline the tail called return function of2444 // coro.async.end.2445 UR.CWorklist.insert(&C);2446 }2447 }2448 2449 for (auto *PrepareFn : PrepareFns) {2450 replaceAllPrepares(PrepareFn, CG, *CurrentSCC);2451 }2452 2453 return PreservedAnalyses::none();2454}2455