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-exception
6//
7//===----------------------------------------------------------------------===//
8// This pass builds the coroutine frame and outlines resume and destroy parts
9// of the coroutine into separate functions.
10//
11// We present a coroutine to an LLVM as an ordinary function with suspension
12// points marked up with intrinsics. We let the optimizer party on the coroutine
13// as a single function for as long as possible. Shortly before the coroutine is
14// eligible to be inlined into its callers, we split up the coroutine into parts
15// 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 the
17// coroutine subfunctions we extracted before proceeding to the caller of the
18// 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 escape
84// and it is known that other transformations, for example, sanitizers
85// 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(C&: CB->getContext(), ArgNo: 2);
100
101 if (auto Invoke = dyn_cast<InvokeInst>(Val: CB)) {
102 auto WrapperInvoke =
103 Builder.CreateInvoke(Callee: Wrapper, NormalDest: Invoke->getNormalDest(),
104 UnwindDest: Invoke->getUnwindDest(), Args: {Awaiter, FramePtr});
105
106 WrapperInvoke->setCallingConv(Invoke->getCallingConv());
107 std::copy(first: Invoke->bundle_op_info_begin(), last: Invoke->bundle_op_info_end(),
108 result: 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>(Val: CB)) {
113 auto WrapperCall = Builder.CreateCall(Callee: Wrapper, Args: {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 call
125 // to the returned coroutine.
126 if (auto *Invoke = dyn_cast<InvokeInst>(Val: 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(Arg: NewCall, Index: CoroSubFnInst::ResumeIndex,
133 InsertPt: &*Builder.GetInsertPoint());
134
135 LLVMContext &Ctx = Builder.getContext();
136 FunctionType *ResumeTy = FunctionType::get(
137 Result: Type::getVoidTy(C&: Ctx), Params: PointerType::getUnqual(C&: Ctx), isVarArg: false);
138 auto *ResumeCall = Builder.CreateCall(FTy: ResumeTy, Callee: ResumeAddr, Args: {NewCall});
139
140 // We can't insert the 'ret' instruction and adjust the cc until the
141 // function has been split, so remember this for later.
142 Shape.SymmetricTransfers.push_back(Elt: ResumeCall);
143
144 NewCall = ResumeCall;
145 }
146
147 CB->replaceAllUsesWith(V: NewCall);
148 CB->eraseFromParent();
149}
150
151static void lowerAwaitSuspends(Function &F, coro::Shape &Shape) {
152 IRBuilder<> Builder(F.getContext());
153 for (auto *AWS : Shape.CoroAwaitSuspends)
154 lowerAwaitSuspend(Builder, CB: AWS, Shape);
155}
156
157static void maybeFreeRetconStorage(IRBuilder<> &Builder,
158 const coro::Shape &Shape, Value *FramePtr,
159 CallGraph *CG) {
160 assert(Shape.ABI == coro::ABI::Retcon || Shape.ABI == coro::ABI::RetconOnce);
161 if (Shape.RetconLowering.IsFrameInlineInStorage)
162 return;
163
164 Shape.emitDealloc(Builder, Ptr: FramePtr, CG);
165}
166
167/// Replace an llvm.coro.end.async.
168/// Will inline the must tail call function call if there is one.
169/// \returns true if cleanup of the coro.end block is needed, false otherwise.
170static bool replaceCoroEndAsync(AnyCoroEndInst *End) {
171 IRBuilder<> Builder(End);
172
173 auto *EndAsync = dyn_cast<CoroAsyncEndInst>(Val: End);
174 if (!EndAsync) {
175 Builder.CreateRetVoid();
176 return true /*needs cleanup of coro.end block*/;
177 }
178
179 auto *MustTailCallFunc = EndAsync->getMustTailCallFunction();
180 if (!MustTailCallFunc) {
181 Builder.CreateRetVoid();
182 return true /*needs cleanup of coro.end block*/;
183 }
184
185 // Move the must tail call from the predecessor block into the end block.
186 auto *CoroEndBlock = End->getParent();
187 auto *MustTailCallFuncBlock = CoroEndBlock->getSinglePredecessor();
188 assert(MustTailCallFuncBlock && "Must have a single predecessor block");
189 auto It = MustTailCallFuncBlock->getTerminator()->getIterator();
190 auto *MustTailCall = cast<CallInst>(Val: &*std::prev(x: It));
191 CoroEndBlock->splice(ToIt: End->getIterator(), FromBB: MustTailCallFuncBlock,
192 FromIt: MustTailCall->getIterator());
193
194 // Insert the return instruction.
195 Builder.SetInsertPoint(End);
196 Builder.CreateRetVoid();
197 InlineFunctionInfo FnInfo;
198
199 // Remove the rest of the block, by splitting it into an unreachable block.
200 auto *BB = End->getParent();
201 BB->splitBasicBlock(I: End);
202 BB->getTerminator()->eraseFromParent();
203
204 auto InlineRes = InlineFunction(CB&: *MustTailCall, IFI&: FnInfo);
205 assert(InlineRes.isSuccess() && "Expected inlining to succeed");
206 (void)InlineRes;
207
208 // We have cleaned up the coro.end block above.
209 return false;
210}
211
212/// Replace a non-unwind call to llvm.coro.end.
213static void replaceFallthroughCoroEnd(AnyCoroEndInst *End,
214 const coro::Shape &Shape, Value *FramePtr,
215 bool InRamp, CallGraph *CG) {
216 // Start inserting right before the coro.end.
217 IRBuilder<> Builder(End);
218
219 // Create the return instruction.
220 switch (Shape.ABI) {
221 // The cloned functions in switch-lowering always return void.
222 case coro::ABI::Switch:
223 assert(!cast<CoroEndInst>(End)->hasResults() &&
224 "switch coroutine should not return any values");
225 // coro.end doesn't immediately end the coroutine in the main function
226 // in this lowering, because we need to deallocate the coroutine.
227 if (InRamp)
228 return;
229 Builder.CreateRetVoid();
230 break;
231
232 // In async lowering this returns.
233 case coro::ABI::Async: {
234 bool CoroEndBlockNeedsCleanup = replaceCoroEndAsync(End);
235 if (!CoroEndBlockNeedsCleanup)
236 return;
237 break;
238 }
239
240 // In unique continuation lowering, the continuations always return void.
241 // But we may have implicitly allocated storage.
242 case coro::ABI::RetconOnce: {
243 maybeFreeRetconStorage(Builder, Shape, FramePtr, CG);
244 auto *CoroEnd = cast<CoroEndInst>(Val: End);
245 auto *RetTy = Shape.getResumeFunctionType()->getReturnType();
246
247 if (!CoroEnd->hasResults()) {
248 assert(RetTy->isVoidTy());
249 Builder.CreateRetVoid();
250 break;
251 }
252
253 auto *CoroResults = CoroEnd->getResults();
254 unsigned NumReturns = CoroResults->numReturns();
255
256 if (auto *RetStructTy = dyn_cast<StructType>(Val: RetTy)) {
257 assert(RetStructTy->getNumElements() == NumReturns &&
258 "numbers of returns should match resume function singature");
259 Value *ReturnValue = PoisonValue::get(T: RetStructTy);
260 unsigned Idx = 0;
261 for (Value *RetValEl : CoroResults->return_values())
262 ReturnValue = Builder.CreateInsertValue(Agg: ReturnValue, Val: RetValEl, Idxs: Idx++);
263 Builder.CreateRet(V: ReturnValue);
264 } else if (NumReturns == 0) {
265 assert(RetTy->isVoidTy());
266 Builder.CreateRetVoid();
267 } else {
268 assert(NumReturns == 1);
269 Builder.CreateRet(V: *CoroResults->retval_begin());
270 }
271 CoroResults->replaceAllUsesWith(
272 V: ConstantTokenNone::get(Context&: CoroResults->getContext()));
273 CoroResults->eraseFromParent();
274 break;
275 }
276
277 // In non-unique continuation lowering, we signal completion by returning
278 // a null continuation.
279 case coro::ABI::Retcon: {
280 assert(!cast<CoroEndInst>(End)->hasResults() &&
281 "retcon coroutine should not return any values");
282 maybeFreeRetconStorage(Builder, Shape, FramePtr, CG);
283 auto RetTy = Shape.getResumeFunctionType()->getReturnType();
284 auto RetStructTy = dyn_cast<StructType>(Val: RetTy);
285 PointerType *ContinuationTy =
286 cast<PointerType>(Val: RetStructTy ? RetStructTy->getElementType(N: 0) : RetTy);
287
288 Value *ReturnValue = ConstantPointerNull::get(T: ContinuationTy);
289 if (RetStructTy) {
290 ReturnValue = Builder.CreateInsertValue(Agg: PoisonValue::get(T: RetStructTy),
291 Val: ReturnValue, Idxs: 0);
292 }
293 Builder.CreateRet(V: ReturnValue);
294 break;
295 }
296 }
297
298 // Remove the rest of the block, by splitting it into an unreachable block.
299 auto *BB = End->getParent();
300 BB->splitBasicBlock(I: End);
301 BB->getTerminator()->eraseFromParent();
302}
303
304/// Create a pointer to the switch index field in the coroutine frame.
305static Value *createSwitchIndexPtr(const coro::Shape &Shape,
306 IRBuilder<> &Builder, Value *FramePtr) {
307 auto *Offset = ConstantInt::get(Ty: Type::getInt64Ty(C&: FramePtr->getContext()),
308 V: Shape.SwitchLowering.IndexOffset);
309 return Builder.CreateInBoundsPtrAdd(Ptr: FramePtr, Offset, Name: "index.addr");
310}
311
312// Mark a coroutine as done, which implies that the coroutine is finished and
313// never gets resumed.
314//
315// In resume-switched ABI, the done state is represented by storing zero in
316// ResumeFnAddr.
317//
318// NOTE: We couldn't omit the argument `FramePtr`. It is necessary because the
319// pointer to the frame in splitted function is not stored in `Shape`.
320static void markCoroutineAsDone(IRBuilder<> &Builder, const coro::Shape &Shape,
321 Value *FramePtr) {
322 assert(
323 Shape.ABI == coro::ABI::Switch &&
324 "markCoroutineAsDone is only supported for Switch-Resumed ABI for now.");
325 // Resume function pointer is always first
326 auto *NullPtr = ConstantPointerNull::get(T: Shape.getSwitchResumePointerType());
327 Builder.CreateStore(Val: NullPtr, Ptr: FramePtr);
328
329 // If the coroutine don't have unwind coro end, we could omit the store to
330 // the final suspend point since we could infer the coroutine is suspended
331 // at the final suspend point by the nullness of ResumeFnAddr.
332 // However, we can't skip it if the coroutine have unwind coro end. Since
333 // the coroutine reaches unwind coro end is considered suspended at the
334 // final suspend point (the ResumeFnAddr is null) but in fact the coroutine
335 // didn't complete yet. We need the IndexVal for the final suspend point
336 // to make the states clear.
337 if (Shape.SwitchLowering.HasUnwindCoroEnd &&
338 Shape.SwitchLowering.HasFinalSuspend) {
339 assert(cast<CoroSuspendInst>(Shape.CoroSuspends.back())->isFinal() &&
340 "The final suspend should only live in the last position of "
341 "CoroSuspends.");
342 ConstantInt *IndexVal = Shape.getIndex(Value: Shape.CoroSuspends.size() - 1);
343 Value *FinalIndex = createSwitchIndexPtr(Shape, Builder, FramePtr);
344 Builder.CreateStore(Val: IndexVal, Ptr: FinalIndex);
345 }
346}
347
348/// Replace an unwind call to llvm.coro.end.
349static void replaceUnwindCoroEnd(AnyCoroEndInst *End, const coro::Shape &Shape,
350 Value *FramePtr, bool InRamp, CallGraph *CG) {
351 IRBuilder<> Builder(End);
352
353 switch (Shape.ABI) {
354 // In switch-lowering, this does nothing in the main function.
355 case coro::ABI::Switch: {
356 // In C++'s specification, the coroutine should be marked as done
357 // if promise.unhandled_exception() throws. The frontend will
358 // call coro.end(true) along this path.
359 //
360 // FIXME: We should refactor this once there is other language
361 // which uses Switch-Resumed style other than C++.
362 markCoroutineAsDone(Builder, Shape, FramePtr);
363 if (InRamp)
364 return;
365 break;
366 }
367 // In async lowering this does nothing.
368 case coro::ABI::Async:
369 break;
370 // In continuation-lowering, this frees the continuation storage.
371 case coro::ABI::Retcon:
372 case coro::ABI::RetconOnce:
373 maybeFreeRetconStorage(Builder, Shape, FramePtr, CG);
374 break;
375 }
376
377 // If coro.end has an associated bundle, add cleanupret instruction.
378 if (auto Bundle = End->getOperandBundle(ID: LLVMContext::OB_funclet)) {
379 auto *FromPad = cast<CleanupPadInst>(Val: Bundle->Inputs[0]);
380 auto *CleanupRet = Builder.CreateCleanupRet(CleanupPad: FromPad, UnwindBB: nullptr);
381 End->getParent()->splitBasicBlock(I: End);
382 CleanupRet->getParent()->getTerminator()->eraseFromParent();
383 }
384}
385
386static void replaceCoroEnd(AnyCoroEndInst *End, const coro::Shape &Shape,
387 Value *FramePtr, bool InRamp, CallGraph *CG) {
388 if (End->isUnwind())
389 replaceUnwindCoroEnd(End, Shape, FramePtr, InRamp, CG);
390 else
391 replaceFallthroughCoroEnd(End, Shape, FramePtr, InRamp, CG);
392 End->eraseFromParent();
393}
394
395// In the resume function, we remove the last case (when coro::Shape is built,
396// the final suspend point (if present) is always the last element of
397// CoroSuspends array) since it is an undefined behavior to resume a coroutine
398// suspended at the final suspend point.
399// In the destroy function, if it isn't possible that the ResumeFnAddr is NULL
400// and the coroutine doesn't suspend at the final suspend point actually (this
401// is possible since the coroutine is considered suspended at the final suspend
402// point if promise.unhandled_exception() exits via an exception), we can
403// remove the last case.
404void coro::BaseCloner::handleFinalSuspend() {
405 assert(Shape.ABI == coro::ABI::Switch &&
406 Shape.SwitchLowering.HasFinalSuspend);
407
408 if (isSwitchDestroyFunction() && Shape.SwitchLowering.HasUnwindCoroEnd)
409 return;
410
411 auto *Switch = cast<SwitchInst>(Val&: VMap[Shape.SwitchLowering.ResumeSwitch]);
412 auto FinalCaseIt = std::prev(x: Switch->case_end());
413 BasicBlock *ResumeBB = FinalCaseIt->getCaseSuccessor();
414 Switch->removeCase(I: FinalCaseIt);
415 if (isSwitchDestroyFunction()) {
416 BasicBlock *OldSwitchBB = Switch->getParent();
417 auto *NewSwitchBB = OldSwitchBB->splitBasicBlock(I: Switch, BBName: "Switch");
418 Builder.SetInsertPoint(OldSwitchBB->getTerminator());
419
420 if (NewF->isCoroOnlyDestroyWhenComplete()) {
421 // When the coroutine can only be destroyed when complete, we don't need
422 // to generate code for other cases.
423 Builder.CreateBr(Dest: ResumeBB);
424 } else {
425 // Resume function pointer is always first
426 auto *Load =
427 Builder.CreateLoad(Ty: Shape.getSwitchResumePointerType(), Ptr: NewFramePtr);
428 auto *Cond = Builder.CreateIsNull(Arg: Load);
429 Builder.CreateCondBr(Cond, True: ResumeBB, False: NewSwitchBB);
430 }
431 OldSwitchBB->getTerminator()->eraseFromParent();
432 }
433}
434
435static FunctionType *
436getFunctionTypeFromAsyncSuspend(AnyCoroSuspendInst *Suspend) {
437 auto *AsyncSuspend = cast<CoroSuspendAsyncInst>(Val: Suspend);
438 auto *StructTy = cast<StructType>(Val: AsyncSuspend->getType());
439 auto &Context = Suspend->getParent()->getParent()->getContext();
440 auto *VoidTy = Type::getVoidTy(C&: Context);
441 return FunctionType::get(Result: VoidTy, Params: StructTy->elements(), isVarArg: 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(Suspend: ActiveSuspend);
452
453 Function *NewF =
454 Function::Create(Ty: FnTy, Linkage: GlobalValue::LinkageTypes::InternalLinkage,
455 AddrSpace: OrigF.getAddressSpace(), N: OrigF.getName() + Suffix);
456
457 M->getFunctionList().insert(where: InsertBefore, New: NewF);
458
459 return NewF;
460}
461
462/// Replace uses of the active llvm.coro.suspend.retcon/async call with the
463/// 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 into
475 // 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(x: NewF->arg_begin()),
480 E = NewF->arg_end();
481 I != E; ++I)
482 Args.push_back(Elt: &*I);
483
484 // If the suspend returns a single scalar value, we can just do a simple
485 // replacement.
486 if (!isa<StructType>(Val: NewS->getType())) {
487 assert(Args.size() == 1);
488 NewS->replaceAllUsesWith(V: Args.front());
489 return;
490 }
491
492 // Try to peephole extracts of an aggregate return.
493 for (Use &U : llvm::make_early_inc_range(Range: NewS->uses())) {
494 auto *EVI = dyn_cast<ExtractValueInst>(Val: U.getUser());
495 if (!EVI || EVI->getNumIndices() != 1)
496 continue;
497
498 EVI->replaceAllUsesWith(V: 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(T: NewS->getType());
508 for (auto [Idx, Arg] : llvm::enumerate(First&: Args))
509 Aggr = Builder.CreateInsertValue(Agg: Aggr, Val: Arg, Idxs: Idx);
510
511 NewS->replaceAllUsesWith(V: 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 value
519 // for the type of function we're extracting.
520 // Replacing coro.suspend with (0) will result in control flow proceeding to
521 // a resume label associated with a suspend point, replacing it with (1) will
522 // result in control flow proceeding to a cleanup label associated with this
523 // suspend point.
524 case coro::ABI::Switch:
525 SuspendResult = Builder.getInt8(C: 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 earlier
533 // continuations are theoretically arbitrary, and they should have been
534 // 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>(Val&: VMap[CS]);
546 MappedCS->replaceAllUsesWith(V: 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 for
554 // the cloned function yet. We'll just be rebuilding that later.
555 auto *NewCE = cast<AnyCoroEndInst>(Val&: VMap[CE]);
556 replaceCoroEnd(End: NewCE, Shape, FramePtr: NewFramePtr, /*in ramp*/ InRamp: false, CG: nullptr);
557 }
558}
559
560void coro::BaseCloner::replaceCoroIsInRamp() {
561 auto &Ctx = OrigF.getContext();
562 for (auto *II : Shape.CoroIsInRampInsts) {
563 auto *NewII = cast<CoroIsInRampInst>(Val&: VMap[II]);
564 NewII->replaceAllUsesWith(V: ConstantInt::getFalse(Context&: 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(Ty: 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>(Val&: (*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(Ty: ValueTy, Ptr: Slot);
606 } else {
607 assert(Op->arg_size() == 1);
608 auto Value = MappedOp->getArgOperand(i: 0);
609 auto ValueTy = Value->getType();
610 auto Slot = getSwiftErrorSlot(ValueTy);
611 Builder.CreateStore(Val: Value, Ptr: Slot);
612 MappedResult = Slot;
613 }
614
615 MappedOp->replaceAllUsesWith(V: 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(R: I.getDbgRecordRange()))
631 DbgVariableRecords.push_back(Elt: &DVR);
632 }
633 return DbgVariableRecords;
634}
635
636void coro::BaseCloner::replaceSwiftErrorOps() {
637 ::replaceSwiftErrorOps(F&: *NewF, Shape, VMap: &VMap);
638}
639
640void coro::BaseCloner::salvageDebugInfo() {
641 auto DbgVariableRecords = collectDbgVariableRecords(F&: *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&: *DVR, UseEntryValue);
648
649 // Remove all salvaged dbg.declare intrinsics that became
650 // either unreachable or stale due to the CoroSplit transformation.
651 DominatorTree DomTree(*NewF);
652 auto IsUnreachableBlock = [&](BasicBlock *BB) {
653 return !isPotentiallyReachable(From: &NewF->getEntryBlock(), To: BB, ExclusionSet: nullptr,
654 DT: &DomTree);
655 };
656 auto RemoveOne = [&](DbgVariableRecord *DVI) {
657 if (IsUnreachableBlock(DVI->getParent()))
658 DVI->eraseFromParent();
659 else if (isa_and_nonnull<AllocaInst>(Val: DVI->getVariableLocationOp(OpIdx: 0))) {
660 // Count all non-debuginfo uses in reachable blocks.
661 unsigned Uses = 0;
662 for (auto *User : DVI->getVariableLocationOp(OpIdx: 0)->users())
663 if (auto *I = dyn_cast<Instruction>(Val: User))
664 if (!isa<AllocaInst>(Val: I) && !IsUnreachableBlock(I->getParent()))
665 ++Uses;
666 if (!Uses)
667 DVI->eraseFromParent();
668 }
669 };
670 for_each(Range&: DbgVariableRecords, F: RemoveOne);
671}
672
673void coro::BaseCloner::replaceEntryBlock() {
674 // In the original function, the AllocaSpillBlock is a block immediately
675 // following the allocation of the frame object which defines GEPs for
676 // all the allocas that have been moved into the frame, and it ends by
677 // branching to the original beginning of the coroutine. Make this
678 // the entry block of the cloned function.
679 auto *Entry = cast<BasicBlock>(Val&: VMap[Shape.AllocaSpillBlock]);
680 auto *OldEntry = &NewF->getEntryBlock();
681 Entry->setName("entry" + Suffix);
682 Entry->moveBefore(MovePos: OldEntry);
683 Entry->getTerminator()->eraseFromParent();
684
685 // Clear all predecessors of the new entry block. There should be
686 // exactly one predecessor, which we created when splitting out
687 // AllocaSpillBlock to begin with.
688 assert(Entry->hasOneUse());
689 auto BranchToEntry = cast<UncondBrInst>(Val: Entry->user_back());
690 Builder.SetInsertPoint(BranchToEntry);
691 Builder.CreateUnreachable();
692 BranchToEntry->eraseFromParent();
693
694 // Branch from the entry to the appropriate place.
695 Builder.SetInsertPoint(Entry);
696 switch (Shape.ABI) {
697 case coro::ABI::Switch: {
698 // In switch-lowering, we built a resume-entry block in the original
699 // function. Make the entry block branch to this.
700 auto *SwitchBB =
701 cast<BasicBlock>(Val&: VMap[Shape.SwitchLowering.ResumeEntryBlock]);
702 Builder.CreateBr(Dest: SwitchBB);
703 SwitchBB->moveAfter(MovePos: Entry);
704 break;
705 }
706 case coro::ABI::Async:
707 case coro::ABI::Retcon:
708 case coro::ABI::RetconOnce: {
709 // In continuation ABIs, we want to branch to immediately after the
710 // active suspend point. Earlier phases will have put the suspend in its
711 // own basic block, so just thread our jump directly to its successor.
712 assert((Shape.ABI == coro::ABI::Async &&
713 isa<CoroSuspendAsyncInst>(ActiveSuspend)) ||
714 ((Shape.ABI == coro::ABI::Retcon ||
715 Shape.ABI == coro::ABI::RetconOnce) &&
716 isa<CoroSuspendRetconInst>(ActiveSuspend)));
717 auto *MappedCS = cast<AnyCoroSuspendInst>(Val&: VMap[ActiveSuspend]);
718 auto Branch = cast<UncondBrInst>(Val: MappedCS->getNextNode());
719 Builder.CreateBr(Dest: Branch->getSuccessor(i: 0));
720 break;
721 }
722 }
723
724 // Any static alloca that's still being used but not reachable from the new
725 // entry needs to be moved to the new entry.
726 Function *F = OldEntry->getParent();
727 DominatorTree DT{*F};
728 for (Instruction &I : llvm::make_early_inc_range(Range: instructions(F))) {
729 auto *Alloca = dyn_cast<AllocaInst>(Val: &I);
730 if (!Alloca || I.use_empty())
731 continue;
732 if (DT.isReachableFromEntry(A: I.getParent()) ||
733 !isa<ConstantInt>(Val: Alloca->getArraySize()))
734 continue;
735 I.moveBefore(BB&: *Entry, I: Entry->getFirstInsertionPt());
736 }
737}
738
739/// Derive the value of the new frame pointer.
740Value *coro::BaseCloner::deriveNewFramePointer() {
741 // Builder should be inserting to the front of the new entry block.
742
743 switch (Shape.ABI) {
744 // In switch-lowering, the argument is the frame pointer.
745 case coro::ABI::Switch:
746 return &*NewF->arg_begin();
747 // In async-lowering, one of the arguments is an async context as determined
748 // by the `llvm.coro.id.async` intrinsic. We can retrieve the async context of
749 // the resume function from the async context projection function associated
750 // with the active suspend. The frame is located as a tail to the async
751 // context header.
752 case coro::ABI::Async: {
753 auto *ActiveAsyncSuspend = cast<CoroSuspendAsyncInst>(Val: ActiveSuspend);
754 auto ContextIdx = ActiveAsyncSuspend->getStorageArgumentIndex() & 0xff;
755 auto *CalleeContext = NewF->getArg(i: ContextIdx);
756 auto *ProjectionFunc =
757 ActiveAsyncSuspend->getAsyncContextProjectionFunction();
758 auto DbgLoc =
759 cast<CoroSuspendAsyncInst>(Val&: VMap[ActiveSuspend])->getDebugLoc();
760 // Calling i8* (i8*)
761 auto *CallerContext = Builder.CreateCall(FTy: ProjectionFunc->getFunctionType(),
762 Callee: ProjectionFunc, Args: CalleeContext);
763 CallerContext->setCallingConv(ProjectionFunc->getCallingConv());
764 CallerContext->setDebugLoc(DbgLoc);
765 // The frame is located after the async_context header.
766 auto &Context = Builder.getContext();
767 auto *FramePtrAddr = Builder.CreateInBoundsPtrAdd(
768 Ptr: CallerContext,
769 Offset: ConstantInt::get(Ty: Type::getInt64Ty(C&: Context),
770 V: Shape.AsyncLowering.FrameOffset),
771 Name: "async.ctx.frameptr");
772 // Inline the projection function.
773 InlineFunctionInfo InlineInfo;
774 auto InlineRes = InlineFunction(CB&: *CallerContext, IFI&: 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(C&: Shape.FramePtr->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(Ty: FramePtrTy, Ptr: NewStorage);
791 }
792 }
793 llvm_unreachable("bad ABI");
794}
795
796/// Adjust the scope line of the funclet to the first line number after the
797/// suspend point. This avoids a jump in the line table from the function
798/// declaration (where prologue instructions are attributed to) to the suspend
799/// 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 meaningful
817 // instructions are not in the same BB.
818 // FIXME: remove this hardcoded number of tries.
819 for (unsigned Repeat = 0; Repeat < 2; Repeat++) {
820 auto *Branch = dyn_cast_or_null<UncondBrInst>(Val&: Successor);
821 if (!Branch)
822 break;
823 Successor = Branch->getSuccessor()->getFirstNonPHIOrDbg();
824 }
825
826 // Find the first successor of ActiveSuspend with a non-zero line location.
827 // If that matches the file of ActiveSuspend, use it.
828 BasicBlock *PBB = Successor->getParent();
829 for (; Successor != PBB->end(); Successor = std::next(x: Successor)) {
830 Successor = skipDebugIntrinsics(It: Successor);
831 auto DL = Successor->getDebugLoc();
832 if (!DL || DL.getLine() == 0)
833 continue;
834
835 if (SPToUpdate.getFile() == DL->getFile()) {
836 SPToUpdate.setScopeLine(DL.getLine());
837 return;
838 }
839
840 break;
841 }
842
843 // If the search above failed, fallback to the location of ActiveSuspend.
844 if (auto DL = ActiveSuspend->getDebugLoc())
845 if (SPToUpdate.getFile() == DL->getFile())
846 SPToUpdate.setScopeLine(DL->getLine());
847}
848
849static void addFramePointerAttrs(AttributeList &Attrs, LLVMContext &Context,
850 unsigned ParamIndex, uint64_t Size,
851 Align Alignment, bool NoAlias) {
852 AttrBuilder ParamAttrs(Context);
853 ParamAttrs.addAttribute(Val: Attribute::NonNull);
854 ParamAttrs.addAttribute(Val: Attribute::NoUndef);
855
856 if (NoAlias)
857 ParamAttrs.addAttribute(Val: Attribute::NoAlias);
858
859 ParamAttrs.addAlignmentAttr(Align: Alignment);
860 ParamAttrs.addDereferenceableAttr(Bytes: Size);
861 Attrs = Attrs.addParamAttributes(C&: Context, ArgNo: ParamIndex, B: ParamAttrs);
862}
863
864static void addAsyncContextAttrs(AttributeList &Attrs, LLVMContext &Context,
865 unsigned ParamIndex) {
866 AttrBuilder ParamAttrs(Context);
867 ParamAttrs.addAttribute(Val: Attribute::SwiftAsync);
868 Attrs = Attrs.addParamAttributes(C&: Context, ArgNo: ParamIndex, B: ParamAttrs);
869}
870
871static void addSwiftSelfAttrs(AttributeList &Attrs, LLVMContext &Context,
872 unsigned ParamIndex) {
873 AttrBuilder ParamAttrs(Context);
874 ParamAttrs.addAttribute(Val: Attribute::SwiftSelf);
875 Attrs = Attrs.addParamAttributes(C&: Context, ArgNo: ParamIndex, B: ParamAttrs);
876}
877
878/// Clone the body of the original function into a resume function of
879/// some sort.
880void coro::BaseCloner::create() {
881 assert(NewF);
882
883 // Replace all args with dummy instructions. If an argument is the old frame
884 // pointer, the dummy will be replaced by the new frame pointer once it is
885 // computed below. Uses of all other arguments should have already been
886 // rewritten by buildCoroutineFrame() to use loads/stores on the coroutine
887 // frame.
888 SmallVector<Instruction *> DummyArgs;
889 for (Argument &A : OrigF.args()) {
890 DummyArgs.push_back(Elt: new FreezeInst(PoisonValue::get(T: A.getType())));
891 VMap[&A] = DummyArgs.back();
892 }
893
894 SmallVector<ReturnInst *, 4> Returns;
895
896 // Ignore attempts to change certain attributes of the function.
897 // TODO: maybe there should be a way to suppress this during cloning?
898 auto savedVisibility = NewF->getVisibility();
899 auto savedUnnamedAddr = NewF->getUnnamedAddr();
900 auto savedDLLStorageClass = NewF->getDLLStorageClass();
901
902 // NewF's linkage (which CloneFunctionInto does *not* change) might not
903 // be compatible with the visibility of OrigF (which it *does* change),
904 // so protect against that.
905 auto savedLinkage = NewF->getLinkage();
906 NewF->setLinkage(llvm::GlobalValue::ExternalLinkage);
907
908 CloneFunctionInto(NewFunc: NewF, OldFunc: &OrigF, VMap,
909 Changes: CloneFunctionChangeType::LocalChangesOnly, Returns);
910
911 auto &Context = NewF->getContext();
912
913 if (DISubprogram *SP = NewF->getSubprogram()) {
914 assert(SP != OrigF.getSubprogram() && SP->isDistinct());
915 updateScopeLine(ActiveSuspend, SPToUpdate&: *SP);
916
917 // Update the linkage name and the function name to reflect the modified
918 // name.
919 MDString *NewLinkageName = MDString::get(Context, Str: NewF->getName());
920 SP->replaceLinkageName(LN: NewLinkageName);
921 if (DISubprogram *Decl = SP->getDeclaration()) {
922 TempDISubprogram NewDecl = Decl->clone();
923 NewDecl->replaceLinkageName(LN: NewLinkageName);
924 SP->replaceDeclaration(Decl: MDNode::replaceWithUniqued(N: std::move(NewDecl)));
925 }
926 }
927
928 NewF->setLinkage(savedLinkage);
929 NewF->setVisibility(savedVisibility);
930 NewF->setUnnamedAddr(savedUnnamedAddr);
931 NewF->setDLLStorageClass(savedDLLStorageClass);
932 // The function sanitizer metadata needs to match the signature of the
933 // function it is being attached to. However this does not hold for split
934 // functions here. Thus remove the metadata for split functions.
935 if (Shape.ABI == coro::ABI::Switch &&
936 NewF->hasMetadata(KindID: LLVMContext::MD_func_sanitize))
937 NewF->eraseMetadata(KindID: LLVMContext::MD_func_sanitize);
938
939 // Replace the attributes of the new function:
940 auto OrigAttrs = NewF->getAttributes();
941 auto NewAttrs = AttributeList();
942
943 switch (Shape.ABI) {
944 case coro::ABI::Switch:
945 // Bootstrap attributes by copying function attributes from the
946 // original function. This should include optimization settings and so on.
947 NewAttrs = NewAttrs.addFnAttributes(
948 C&: Context, B: AttrBuilder(Context, OrigAttrs.getFnAttrs()));
949
950 addFramePointerAttrs(Attrs&: NewAttrs, Context, ParamIndex: 0, Size: Shape.FrameSize,
951 Alignment: Shape.FrameAlign, /*NoAlias=*/false);
952 break;
953 case coro::ABI::Async: {
954 auto *ActiveAsyncSuspend = cast<CoroSuspendAsyncInst>(Val: ActiveSuspend);
955 if (OrigF.hasParamAttribute(ArgNo: Shape.AsyncLowering.ContextArgNo,
956 Kind: Attribute::SwiftAsync)) {
957 uint32_t ArgAttributeIndices =
958 ActiveAsyncSuspend->getStorageArgumentIndex();
959 auto ContextArgIndex = ArgAttributeIndices & 0xff;
960 addAsyncContextAttrs(Attrs&: NewAttrs, Context, ParamIndex: ContextArgIndex);
961
962 // `swiftasync` must preceed `swiftself` so 0 is not a valid index for
963 // `swiftself`.
964 auto SwiftSelfIndex = ArgAttributeIndices >> 8;
965 if (SwiftSelfIndex)
966 addSwiftSelfAttrs(Attrs&: NewAttrs, Context, ParamIndex: SwiftSelfIndex);
967 }
968
969 // Transfer the original function's attributes.
970 auto FnAttrs = OrigF.getAttributes().getFnAttrs();
971 NewAttrs = NewAttrs.addFnAttributes(C&: Context, B: AttrBuilder(Context, FnAttrs));
972 break;
973 }
974 case coro::ABI::Retcon:
975 case coro::ABI::RetconOnce:
976 // If we have a continuation prototype, just use its attributes,
977 // full-stop.
978 NewAttrs = Shape.RetconLowering.ResumePrototype->getAttributes();
979
980 /// FIXME: Is it really good to add the NoAlias attribute?
981 addFramePointerAttrs(Attrs&: NewAttrs, Context, ParamIndex: 0,
982 Size: Shape.getRetconCoroId()->getStorageSize(),
983 Alignment: Shape.getRetconCoroId()->getStorageAlignment(),
984 /*NoAlias=*/true);
985
986 break;
987 }
988
989 switch (Shape.ABI) {
990 // In these ABIs, the cloned functions always return 'void', and the
991 // existing return sites are meaningless. Note that for unique
992 // continuations, this includes the returns associated with suspends;
993 // this is fine because we can't suspend twice.
994 case coro::ABI::Switch:
995 case coro::ABI::RetconOnce:
996 // Remove old returns.
997 for (ReturnInst *Return : Returns)
998 changeToUnreachable(I: Return);
999 break;
1000
1001 // With multi-suspend continuations, we'll already have eliminated the
1002 // original returns and inserted returns before all the suspend points,
1003 // so we want to leave any returns in place.
1004 case coro::ABI::Retcon:
1005 break;
1006 // Async lowering will insert musttail call functions at all suspend points
1007 // followed by a return.
1008 // Don't change returns to unreachable because that will trip up the verifier.
1009 // These returns should be unreachable from the clone.
1010 case coro::ABI::Async:
1011 break;
1012 }
1013
1014 NewF->setAttributes(NewAttrs);
1015 NewF->setCallingConv(Shape.getResumeFunctionCC());
1016
1017 // Set up the new entry block.
1018 replaceEntryBlock();
1019
1020 // Turn symmetric transfers into musttail calls.
1021 for (CallInst *ResumeCall : Shape.SymmetricTransfers) {
1022 ResumeCall = cast<CallInst>(Val&: VMap[ResumeCall]);
1023 if (TTI.supportsTailCallFor(CB: ResumeCall)) {
1024 // FIXME: Could we support symmetric transfer effectively without
1025 // musttail?
1026 ResumeCall->setTailCallKind(CallInst::TCK_MustTail);
1027 }
1028
1029 // Put a 'ret void' after the call, and split any remaining instructions to
1030 // an unreachable block.
1031 BasicBlock *BB = ResumeCall->getParent();
1032 BB->splitBasicBlock(I: ResumeCall->getNextNode());
1033 Builder.SetInsertPoint(BB->getTerminator());
1034 Builder.CreateRetVoid();
1035 BB->getTerminator()->eraseFromParent();
1036 }
1037
1038 Builder.SetInsertPoint(&NewF->getEntryBlock().front());
1039 NewFramePtr = deriveNewFramePointer();
1040
1041 // Remap frame pointer.
1042 Value *OldFramePtr = VMap[Shape.FramePtr];
1043 NewFramePtr->takeName(V: OldFramePtr);
1044 OldFramePtr->replaceAllUsesWith(V: NewFramePtr);
1045
1046 // Remap vFrame pointer.
1047 auto *NewVFrame = Builder.CreateBitCast(
1048 V: NewFramePtr, DestTy: PointerType::getUnqual(C&: Builder.getContext()), Name: "vFrame");
1049 Value *OldVFrame = cast<Value>(Val&: VMap[Shape.CoroBegin]);
1050 if (OldVFrame != NewVFrame)
1051 OldVFrame->replaceAllUsesWith(V: NewVFrame);
1052
1053 // All uses of the arguments should have been resolved by this point,
1054 // so we can safely remove the dummy values.
1055 for (Instruction *DummyArg : DummyArgs) {
1056 DummyArg->replaceAllUsesWith(V: PoisonValue::get(T: DummyArg->getType()));
1057 DummyArg->deleteValue();
1058 }
1059
1060 switch (Shape.ABI) {
1061 case coro::ABI::Switch:
1062 // Rewrite final suspend handling as it is not done via switch (allows to
1063 // remove final case from the switch, since it is undefined behavior to
1064 // resume the coroutine suspended at the final suspend point.
1065 if (Shape.SwitchLowering.HasFinalSuspend)
1066 handleFinalSuspend();
1067 break;
1068 case coro::ABI::Async:
1069 case coro::ABI::Retcon:
1070 case coro::ABI::RetconOnce:
1071 // Replace uses of the active suspend with the corresponding
1072 // continuation-function arguments.
1073 assert(ActiveSuspend != nullptr &&
1074 "no active suspend when lowering a continuation-style coroutine");
1075 replaceRetconOrAsyncSuspendUses();
1076 break;
1077 }
1078
1079 // Handle suspends.
1080 replaceCoroSuspends();
1081
1082 // Handle swifterror.
1083 replaceSwiftErrorOps();
1084
1085 // Remove coro.end intrinsics.
1086 replaceCoroEnds();
1087
1088 replaceCoroIsInRamp();
1089
1090 // Salvage debug info that points into the coroutine frame.
1091 salvageDebugInfo();
1092}
1093
1094void coro::SwitchCloner::create() {
1095 // Create a new function matching the original type
1096 NewF = createCloneDeclaration(OrigF, Shape, Suffix, InsertBefore: OrigF.getParent()->end(),
1097 ActiveSuspend);
1098
1099 // Clone the function
1100 coro::BaseCloner::create();
1101
1102 // Replacing coro.free with 'null' in cleanup to suppress deallocation code.
1103 if (FKind == coro::CloneKind::SwitchCleanup)
1104 elideCoroFree(FramePtr: NewFramePtr);
1105}
1106
1107static void updateAsyncFuncPointerContextSize(coro::Shape &Shape) {
1108 assert(Shape.ABI == coro::ABI::Async);
1109
1110 auto *FuncPtrStruct = cast<ConstantStruct>(
1111 Val: Shape.AsyncLowering.AsyncFuncPointer->getInitializer());
1112 auto *OrigRelativeFunOffset = FuncPtrStruct->getOperand(i_nocapture: 0);
1113 auto *OrigContextSize = FuncPtrStruct->getOperand(i_nocapture: 1);
1114 auto *NewContextSize = ConstantInt::get(Ty: OrigContextSize->getType(),
1115 V: Shape.AsyncLowering.ContextSize);
1116 auto *NewFuncPtrStruct = ConstantStruct::get(
1117 T: FuncPtrStruct->getType(), Vs: OrigRelativeFunOffset, Vs: NewContextSize);
1118
1119 Shape.AsyncLowering.AsyncFuncPointer->setInitializer(NewFuncPtrStruct);
1120}
1121
1122static void replaceFrameSizeAndAlignment(coro::Shape &Shape) {
1123 if (Shape.ABI == coro::ABI::Async)
1124 updateAsyncFuncPointerContextSize(Shape);
1125
1126 for (CoroAlignInst *CA : Shape.CoroAligns) {
1127 CA->replaceAllUsesWith(
1128 V: ConstantInt::get(Ty: CA->getType(), V: Shape.FrameAlign.value()));
1129 CA->eraseFromParent();
1130 }
1131
1132 if (Shape.CoroSizes.empty())
1133 return;
1134
1135 // In the same function all coro.sizes should have the same result type.
1136 auto *SizeIntrin = Shape.CoroSizes.back();
1137 auto *SizeConstant = ConstantInt::get(Ty: SizeIntrin->getType(),
1138 V: TypeSize::getFixed(ExactSize: Shape.FrameSize));
1139
1140 for (CoroSizeInst *CS : Shape.CoroSizes) {
1141 CS->replaceAllUsesWith(V: SizeConstant);
1142 CS->eraseFromParent();
1143 }
1144}
1145
1146static void postSplitCleanup(Function &F) {
1147 removeUnreachableBlocks(F);
1148
1149#ifndef NDEBUG
1150 // For now, we do a mandatory verification step because we don't
1151 // entirely trust this pass. Note that we don't want to add a verifier
1152 // pass to FPM below because it will also verify all the global data.
1153 if (verifyFunction(F, &errs()))
1154 report_fatal_error("Broken function");
1155#endif
1156}
1157
1158// Coroutine has no suspend points. Remove heap allocation for the coroutine
1159// frame if possible.
1160static void handleNoSuspendCoroutine(coro::Shape &Shape) {
1161 auto *CoroBegin = Shape.CoroBegin;
1162 switch (Shape.ABI) {
1163 case coro::ABI::Switch: {
1164 if (auto *AllocInst = Shape.getSwitchCoroId()->getCoroAlloc()) {
1165 coro::elideCoroFree(FramePtr: CoroBegin);
1166
1167 IRBuilder<> Builder(AllocInst);
1168 // Create an alloca for a byte array of the frame size
1169 auto *FrameTy = ArrayType::get(ElementType: Type::getInt8Ty(C&: Builder.getContext()),
1170 NumElements: Shape.FrameSize);
1171 auto *Frame = Builder.CreateAlloca(
1172 Ty: FrameTy, ArraySize: nullptr, Name: AllocInst->getFunction()->getName() + ".Frame");
1173 Frame->setAlignment(Shape.FrameAlign);
1174 AllocInst->replaceAllUsesWith(V: Builder.getFalse());
1175 AllocInst->eraseFromParent();
1176 CoroBegin->replaceAllUsesWith(V: Frame);
1177 } else {
1178 CoroBegin->replaceAllUsesWith(V: CoroBegin->getMem());
1179 }
1180
1181 break;
1182 }
1183 case coro::ABI::Async:
1184 case coro::ABI::Retcon:
1185 case coro::ABI::RetconOnce:
1186 CoroBegin->replaceAllUsesWith(V: PoisonValue::get(T: CoroBegin->getType()));
1187 break;
1188 }
1189
1190 CoroBegin->eraseFromParent();
1191 Shape.CoroBegin = nullptr;
1192}
1193
1194// SimplifySuspendPoint needs to check that there is no calls between
1195// coro_save and coro_suspend, since any of the calls may potentially resume
1196// the coroutine and if that is the case we cannot eliminate the suspend point.
1197static bool hasCallsInBlockBetween(iterator_range<BasicBlock::iterator> R) {
1198 for (Instruction &I : R) {
1199 // Assume that no intrinsic can resume the coroutine.
1200 if (isa<IntrinsicInst>(Val: I))
1201 continue;
1202
1203 if (isa<CallBase>(Val: I))
1204 return true;
1205 }
1206 return false;
1207}
1208
1209static bool hasCallsInBlocksBetween(BasicBlock *SaveBB, BasicBlock *ResDesBB) {
1210 SmallPtrSet<BasicBlock *, 8> Set;
1211 SmallVector<BasicBlock *, 8> Worklist;
1212
1213 Set.insert(Ptr: SaveBB);
1214 Worklist.push_back(Elt: ResDesBB);
1215
1216 // Accumulate all blocks between SaveBB and ResDesBB. Because CoroSaveIntr
1217 // returns a token consumed by suspend instruction, all blocks in between
1218 // will have to eventually hit SaveBB when going backwards from ResDesBB.
1219 while (!Worklist.empty()) {
1220 auto *BB = Worklist.pop_back_val();
1221 Set.insert(Ptr: BB);
1222 for (auto *Pred : predecessors(BB))
1223 if (!Set.contains(Ptr: Pred))
1224 Worklist.push_back(Elt: Pred);
1225 }
1226
1227 // SaveBB and ResDesBB are checked separately in hasCallsBetween.
1228 Set.erase(Ptr: SaveBB);
1229 Set.erase(Ptr: ResDesBB);
1230
1231 for (auto *BB : Set)
1232 if (hasCallsInBlockBetween(R: {BB->getFirstNonPHIIt(), BB->end()}))
1233 return true;
1234
1235 return false;
1236}
1237
1238static bool hasCallsBetween(Instruction *Save, Instruction *ResumeOrDestroy) {
1239 auto *SaveBB = Save->getParent();
1240 auto *ResumeOrDestroyBB = ResumeOrDestroy->getParent();
1241 BasicBlock::iterator SaveIt = Save->getIterator();
1242 BasicBlock::iterator ResumeOrDestroyIt = ResumeOrDestroy->getIterator();
1243
1244 if (SaveBB == ResumeOrDestroyBB)
1245 return hasCallsInBlockBetween(R: {std::next(x: SaveIt), ResumeOrDestroyIt});
1246
1247 // Any calls from Save to the end of the block?
1248 if (hasCallsInBlockBetween(R: {std::next(x: SaveIt), SaveBB->end()}))
1249 return true;
1250
1251 // Any calls from begging of the block up to ResumeOrDestroy?
1252 if (hasCallsInBlockBetween(
1253 R: {ResumeOrDestroyBB->getFirstNonPHIIt(), ResumeOrDestroyIt}))
1254 return true;
1255
1256 // Any calls in all of the blocks between SaveBB and ResumeOrDestroyBB?
1257 if (hasCallsInBlocksBetween(SaveBB, ResDesBB: ResumeOrDestroyBB))
1258 return true;
1259
1260 return false;
1261}
1262
1263// If a SuspendIntrin is preceded by Resume or Destroy, we can eliminate the
1264// suspend point and replace it with nornal control flow.
1265static bool simplifySuspendPoint(CoroSuspendInst *Suspend,
1266 CoroBeginInst *CoroBegin) {
1267 Instruction *Prev = Suspend->getPrevNode();
1268 if (!Prev) {
1269 auto *Pred = Suspend->getParent()->getSinglePredecessor();
1270 if (!Pred)
1271 return false;
1272 Prev = Pred->getTerminator();
1273 }
1274
1275 CallBase *CB = dyn_cast<CallBase>(Val: Prev);
1276 if (!CB)
1277 return false;
1278
1279 auto *Callee = CB->getCalledOperand()->stripPointerCasts();
1280
1281 // See if the callsite is for resumption or destruction of the coroutine.
1282 auto *SubFn = dyn_cast<CoroSubFnInst>(Val: Callee);
1283 if (!SubFn)
1284 return false;
1285
1286 // Does not refer to the current coroutine, we cannot do anything with it.
1287 if (SubFn->getFrame() != CoroBegin)
1288 return false;
1289
1290 // See if the transformation is safe. Specifically, see if there are any
1291 // calls in between Save and CallInstr. They can potenitally resume the
1292 // coroutine rendering this optimization unsafe.
1293 auto *Save = Suspend->getCoroSave();
1294 if (hasCallsBetween(Save, ResumeOrDestroy: CB))
1295 return false;
1296
1297 // Replace llvm.coro.suspend with the value that results in resumption over
1298 // the resume or cleanup path.
1299 Suspend->replaceAllUsesWith(V: SubFn->getRawIndex());
1300 Suspend->eraseFromParent();
1301 Save->eraseFromParent();
1302
1303 // No longer need a call to coro.resume or coro.destroy.
1304 if (auto *Invoke = dyn_cast<InvokeInst>(Val: CB)) {
1305 UncondBrInst::Create(Target: Invoke->getNormalDest(), InsertBefore: Invoke->getIterator());
1306 }
1307
1308 // Grab the CalledValue from CB before erasing the CallInstr.
1309 auto *CalledValue = CB->getCalledOperand();
1310 CB->eraseFromParent();
1311
1312 // If no more users remove it. Usually it is a bitcast of SubFn.
1313 if (CalledValue != SubFn && CalledValue->user_empty())
1314 if (auto *I = dyn_cast<Instruction>(Val: CalledValue))
1315 I->eraseFromParent();
1316
1317 // Now we are good to remove SubFn.
1318 if (SubFn->user_empty())
1319 SubFn->eraseFromParent();
1320
1321 return true;
1322}
1323
1324// Remove suspend points that are simplified.
1325static void simplifySuspendPoints(coro::Shape &Shape) {
1326 // Currently, the only simplification we do is switch-lowering-specific.
1327 if (Shape.ABI != coro::ABI::Switch)
1328 return;
1329
1330 auto &S = Shape.CoroSuspends;
1331 size_t I = 0, N = S.size();
1332 if (N == 0)
1333 return;
1334
1335 size_t ChangedFinalIndex = std::numeric_limits<size_t>::max();
1336 while (true) {
1337 auto SI = cast<CoroSuspendInst>(Val: S[I]);
1338 // Leave final.suspend to handleFinalSuspend since it is undefined behavior
1339 // to resume a coroutine suspended at the final suspend point.
1340 if (!SI->isFinal() && simplifySuspendPoint(Suspend: SI, CoroBegin: Shape.CoroBegin)) {
1341 if (--N == I)
1342 break;
1343
1344 std::swap(a&: S[I], b&: S[N]);
1345
1346 if (cast<CoroSuspendInst>(Val: S[I])->isFinal()) {
1347 assert(Shape.SwitchLowering.HasFinalSuspend);
1348 ChangedFinalIndex = I;
1349 }
1350
1351 continue;
1352 }
1353 if (++I == N)
1354 break;
1355 }
1356 S.resize(N);
1357
1358 // Maintain final.suspend in case final suspend was swapped.
1359 // Due to we requrie the final suspend to be the last element of CoroSuspends.
1360 if (ChangedFinalIndex < N) {
1361 assert(cast<CoroSuspendInst>(S[ChangedFinalIndex])->isFinal());
1362 std::swap(a&: S[ChangedFinalIndex], b&: S.back());
1363 }
1364}
1365
1366namespace {
1367
1368struct SwitchCoroutineSplitter {
1369 static void split(Function &F, coro::Shape &Shape,
1370 SmallVectorImpl<Function *> &Clones,
1371 TargetTransformInfo &TTI) {
1372 assert(Shape.ABI == coro::ABI::Switch);
1373
1374 // Create a resume clone by cloning the body of the original function,
1375 // setting new entry block and replacing coro.suspend an appropriate value
1376 // to force resume or cleanup pass for every suspend point.
1377 createResumeEntryBlock(F, Shape);
1378 auto *ResumeClone = coro::SwitchCloner::createClone(
1379 OrigF&: F, Suffix: ".resume", Shape, FKind: coro::CloneKind::SwitchResume, TTI);
1380 auto *DestroyClone = coro::SwitchCloner::createClone(
1381 OrigF&: F, Suffix: ".destroy", Shape, FKind: coro::CloneKind::SwitchUnwind, TTI);
1382 auto *CleanupClone = coro::SwitchCloner::createClone(
1383 OrigF&: F, Suffix: ".cleanup", Shape, FKind: coro::CloneKind::SwitchCleanup, TTI);
1384
1385 postSplitCleanup(F&: *ResumeClone);
1386 postSplitCleanup(F&: *DestroyClone);
1387 postSplitCleanup(F&: *CleanupClone);
1388
1389 // Store addresses resume/destroy/cleanup functions in the coroutine frame.
1390 updateCoroFrame(Shape, ResumeFn: ResumeClone, DestroyFn: DestroyClone, CleanupFn: CleanupClone);
1391
1392 assert(Clones.empty());
1393 Clones.push_back(Elt: ResumeClone);
1394 Clones.push_back(Elt: DestroyClone);
1395 Clones.push_back(Elt: CleanupClone);
1396
1397 // Create a constant array referring to resume/destroy/clone functions
1398 // pointed by the last argument of @llvm.coro.info, so that CoroElide pass
1399 // can determined correct function to call.
1400 setCoroInfo(F, Shape, Fns: Clones);
1401 }
1402
1403 // Create a variant of ramp function that does not perform heap allocation
1404 // for a switch ABI coroutine.
1405 //
1406 // The newly split `.noalloc` ramp function has the following differences:
1407 // - Has one additional frame pointer parameter in lieu of dynamic
1408 // allocation.
1409 // - Suppressed allocations by replacing coro.alloc and coro.free.
1410 static Function *createNoAllocVariant(Function &F, coro::Shape &Shape,
1411 SmallVectorImpl<Function *> &Clones) {
1412 assert(Shape.ABI == coro::ABI::Switch);
1413 auto *OrigFnTy = F.getFunctionType();
1414 auto OldParams = OrigFnTy->params();
1415
1416 SmallVector<Type *> NewParams;
1417 NewParams.reserve(N: OldParams.size() + 1);
1418 NewParams.append(in_start: OldParams.begin(), in_end: OldParams.end());
1419 NewParams.push_back(Elt: PointerType::getUnqual(C&: Shape.FramePtr->getContext()));
1420
1421 auto *NewFnTy = FunctionType::get(Result: OrigFnTy->getReturnType(), Params: NewParams,
1422 isVarArg: OrigFnTy->isVarArg());
1423 Function *NoAllocF = Function::Create(
1424 Ty: NewFnTy, Linkage: F.getLinkage(), AddrSpace: F.getAddressSpace(), N: F.getName() + ".noalloc");
1425
1426 ValueToValueMapTy VMap;
1427 unsigned int Idx = 0;
1428 for (const auto &I : F.args()) {
1429 VMap[&I] = NoAllocF->getArg(i: Idx++);
1430 }
1431 // We just appended the frame pointer as the last argument of the new
1432 // function.
1433 auto FrameIdx = NoAllocF->arg_size() - 1;
1434 SmallVector<ReturnInst *, 4> Returns;
1435 CloneFunctionInto(NewFunc: NoAllocF, OldFunc: &F, VMap,
1436 Changes: CloneFunctionChangeType::LocalChangesOnly, Returns);
1437
1438 if (Shape.CoroBegin) {
1439 auto *NewCoroBegin =
1440 cast_if_present<CoroBeginInst>(Val&: VMap[Shape.CoroBegin]);
1441 coro::elideCoroFree(FramePtr: NewCoroBegin);
1442 coro::suppressCoroAllocs(CoroId: cast<CoroIdInst>(Val: NewCoroBegin->getId()));
1443 NewCoroBegin->replaceAllUsesWith(V: NoAllocF->getArg(i: FrameIdx));
1444 NewCoroBegin->eraseFromParent();
1445 }
1446
1447 Module *M = F.getParent();
1448 M->getFunctionList().insert(where: M->end(), New: NoAllocF);
1449
1450 removeUnreachableBlocks(F&: *NoAllocF);
1451 auto NewAttrs = NoAllocF->getAttributes();
1452 // When we elide allocation, we read these attributes to determine the
1453 // frame size and alignment.
1454 addFramePointerAttrs(Attrs&: NewAttrs, Context&: NoAllocF->getContext(), ParamIndex: FrameIdx,
1455 Size: Shape.FrameSize, Alignment: Shape.FrameAlign,
1456 /*NoAlias=*/false);
1457
1458 NoAllocF->setAttributes(NewAttrs);
1459
1460 Clones.push_back(Elt: NoAllocF);
1461 // Reset the original function's coro info, make the new noalloc variant
1462 // connected to the original ramp function.
1463 setCoroInfo(F, Shape, Fns: Clones);
1464 // After copying, set the linkage to internal linkage. Original function
1465 // may have different linkage, but optimization dependent on this function
1466 // generally relies on LTO.
1467 NoAllocF->setLinkage(llvm::GlobalValue::InternalLinkage);
1468 return NoAllocF;
1469 }
1470
1471private:
1472 // Create an entry block for a resume function with a switch that will jump to
1473 // suspend points.
1474 static void createResumeEntryBlock(Function &F, coro::Shape &Shape) {
1475 LLVMContext &C = F.getContext();
1476
1477 DIBuilder DBuilder(*F.getParent(), /*AllowUnresolved*/ false);
1478 DISubprogram *DIS = F.getSubprogram();
1479 // If there is no DISubprogram for F, it implies the function is compiled
1480 // without debug info. So we also don't generate debug info for the
1481 // suspension points.
1482 bool AddDebugLabels = DIS && DIS->getUnit() &&
1483 (DIS->getUnit()->getEmissionKind() ==
1484 DICompileUnit::DebugEmissionKind::FullDebug);
1485
1486 // resume.entry:
1487 // %index.addr = getelementptr inbounds %f.Frame, %f.Frame* %FramePtr, i32
1488 // 0, i32 2 % index = load i32, i32* %index.addr switch i32 %index, label
1489 // %unreachable [
1490 // i32 0, label %resume.0
1491 // i32 1, label %resume.1
1492 // ...
1493 // ]
1494
1495 auto *NewEntry = BasicBlock::Create(Context&: C, Name: "resume.entry", Parent: &F);
1496 auto *UnreachBB = BasicBlock::Create(Context&: C, Name: "unreachable", Parent: &F);
1497
1498 IRBuilder<> Builder(NewEntry);
1499 auto *FramePtr = Shape.FramePtr;
1500 Value *GepIndex = createSwitchIndexPtr(Shape, Builder, FramePtr);
1501 auto *Index = Builder.CreateLoad(Ty: Shape.getIndexType(), Ptr: GepIndex, Name: "index");
1502 auto *Switch =
1503 Builder.CreateSwitch(V: Index, Dest: UnreachBB, NumCases: Shape.CoroSuspends.size());
1504 Shape.SwitchLowering.ResumeSwitch = Switch;
1505
1506 // Split all coro.suspend calls
1507 size_t SuspendIndex = 0;
1508 for (auto *AnyS : Shape.CoroSuspends) {
1509 auto *S = cast<CoroSuspendInst>(Val: AnyS);
1510 ConstantInt *IndexVal = Shape.getIndex(Value: SuspendIndex);
1511
1512 // Replace CoroSave with a store to Index:
1513 // %index.addr = getelementptr %f.frame... (index field number)
1514 // store i32 %IndexVal, i32* %index.addr1
1515 auto *Save = S->getCoroSave();
1516 Builder.SetInsertPoint(Save);
1517 if (S->isFinal()) {
1518 // The coroutine should be marked done if it reaches the final suspend
1519 // point.
1520 markCoroutineAsDone(Builder, Shape, FramePtr);
1521 } else {
1522 Value *GepIndex = createSwitchIndexPtr(Shape, Builder, FramePtr);
1523 Builder.CreateStore(Val: IndexVal, Ptr: GepIndex);
1524 }
1525
1526 Save->replaceAllUsesWith(V: ConstantTokenNone::get(Context&: C));
1527 Save->eraseFromParent();
1528
1529 // Split block before and after coro.suspend and add a jump from an entry
1530 // switch:
1531 //
1532 // whateverBB:
1533 // whatever
1534 // %0 = call i8 @llvm.coro.suspend(token none, i1 false)
1535 // switch i8 %0, label %suspend[i8 0, label %resume
1536 // i8 1, label %cleanup]
1537 // becomes:
1538 //
1539 // whateverBB:
1540 // whatever
1541 // br label %resume.0.landing
1542 //
1543 // resume.0: ; <--- jump from the switch in the resume.entry
1544 // #dbg_label(...) ; <--- artificial label for debuggers
1545 // %0 = tail call i8 @llvm.coro.suspend(token none, i1 false)
1546 // br label %resume.0.landing
1547 //
1548 // resume.0.landing:
1549 // %1 = phi i8[-1, %whateverBB], [%0, %resume.0]
1550 // switch i8 % 1, label %suspend [i8 0, label %resume
1551 // i8 1, label %cleanup]
1552
1553 auto *SuspendBB = S->getParent();
1554 auto *ResumeBB =
1555 SuspendBB->splitBasicBlock(I: S, BBName: "resume." + Twine(SuspendIndex));
1556 auto *LandingBB = ResumeBB->splitBasicBlock(
1557 I: S->getNextNode(), BBName: ResumeBB->getName() + Twine(".landing"));
1558 Switch->addCase(OnVal: IndexVal, Dest: ResumeBB);
1559
1560 cast<UncondBrInst>(Val: SuspendBB->getTerminator())->setSuccessor(LandingBB);
1561 auto *PN = PHINode::Create(Ty: Builder.getInt8Ty(), NumReservedValues: 2, NameStr: "");
1562 PN->insertBefore(InsertPos: LandingBB->begin());
1563 S->replaceAllUsesWith(V: PN);
1564 PN->addIncoming(V: Builder.getInt8(C: -1), BB: SuspendBB);
1565 PN->addIncoming(V: S, BB: ResumeBB);
1566
1567 if (AddDebugLabels) {
1568 if (DebugLoc SuspendLoc = S->getDebugLoc()) {
1569 std::string LabelName =
1570 ("__coro_resume_" + Twine(SuspendIndex)).str();
1571 // Take the "inlined at" location recursively, if present. This is
1572 // mandatory as the DILabel insertion checks that the scopes of label
1573 // and the attached location match. This is not the case when the
1574 // suspend location has been inlined due to pointing to the original
1575 // scope.
1576 DILocation *DILoc = SuspendLoc;
1577 while (DILocation *InlinedAt = DILoc->getInlinedAt())
1578 DILoc = InlinedAt;
1579
1580 DILabel *ResumeLabel =
1581 DBuilder.createLabel(Scope: DIS, Name: LabelName, File: DILoc->getFile(),
1582 LineNo: SuspendLoc.getLine(), Column: SuspendLoc.getCol(),
1583 /*IsArtificial=*/true,
1584 /*CoroSuspendIdx=*/SuspendIndex,
1585 /*AlwaysPreserve=*/false);
1586 DBuilder.insertLabel(LabelInfo: ResumeLabel, DL: DILoc, InsertPt: ResumeBB->begin());
1587 }
1588 }
1589
1590 ++SuspendIndex;
1591 }
1592
1593 Builder.SetInsertPoint(UnreachBB);
1594 Builder.CreateUnreachable();
1595 DBuilder.finalize();
1596
1597 Shape.SwitchLowering.ResumeEntryBlock = NewEntry;
1598 }
1599
1600 // Store addresses of Resume/Destroy/Cleanup functions in the coroutine frame.
1601 static void updateCoroFrame(coro::Shape &Shape, Function *ResumeFn,
1602 Function *DestroyFn, Function *CleanupFn) {
1603 IRBuilder<> Builder(&*Shape.getInsertPtAfterFramePtr());
1604 LLVMContext &C = ResumeFn->getContext();
1605
1606 // Resume function pointer
1607 Value *ResumeAddr = Shape.FramePtr;
1608 Builder.CreateStore(Val: ResumeFn, Ptr: ResumeAddr);
1609
1610 Value *DestroyOrCleanupFn = DestroyFn;
1611
1612 CoroIdInst *CoroId = Shape.getSwitchCoroId();
1613 if (CoroAllocInst *CA = CoroId->getCoroAlloc()) {
1614 // If there is a CoroAlloc and it returns false (meaning we elide the
1615 // allocation, use CleanupFn instead of DestroyFn).
1616 DestroyOrCleanupFn = Builder.CreateSelect(C: CA, True: DestroyFn, False: CleanupFn);
1617 }
1618
1619 // Destroy function pointer
1620 Value *DestroyAddr = Builder.CreateInBoundsPtrAdd(
1621 Ptr: Shape.FramePtr,
1622 Offset: ConstantInt::get(Ty: Type::getInt64Ty(C),
1623 V: Shape.SwitchLowering.DestroyOffset),
1624 Name: "destroy.addr");
1625 Builder.CreateStore(Val: DestroyOrCleanupFn, Ptr: DestroyAddr);
1626 }
1627
1628 // Create a global constant array containing pointers to functions provided
1629 // and set Info parameter of CoroBegin to point at this constant. Example:
1630 //
1631 // @f.resumers = internal constant [2 x void(%f.frame*)*]
1632 // [void(%f.frame*)* @f.resume, void(%f.frame*)*
1633 // @f.destroy]
1634 // define void @f() {
1635 // ...
1636 // call i8* @llvm.coro.begin(i8* null, i32 0, i8* null,
1637 // i8* bitcast([2 x void(%f.frame*)*] * @f.resumers to
1638 // i8*))
1639 //
1640 // Assumes that all the functions have the same signature.
1641 static void setCoroInfo(Function &F, coro::Shape &Shape,
1642 ArrayRef<Function *> Fns) {
1643 // This only works under the switch-lowering ABI because coro elision
1644 // only works on the switch-lowering ABI.
1645 SmallVector<Constant *, 4> Args(Fns);
1646 assert(!Args.empty());
1647 Function *Part = *Fns.begin();
1648 Module *M = Part->getParent();
1649 auto *ArrTy = ArrayType::get(ElementType: Part->getType(), NumElements: Args.size());
1650
1651 auto *ConstVal = ConstantArray::get(T: ArrTy, V: Args);
1652 auto *GV = new GlobalVariable(*M, ConstVal->getType(), /*isConstant=*/true,
1653 GlobalVariable::PrivateLinkage, ConstVal,
1654 F.getName() + Twine(".resumers"));
1655
1656 // Update coro.begin instruction to refer to this constant.
1657 LLVMContext &C = F.getContext();
1658 auto *BC = ConstantExpr::getPointerCast(C: GV, Ty: PointerType::getUnqual(C));
1659 Shape.getSwitchCoroId()->setInfo(BC);
1660 }
1661};
1662
1663} // namespace
1664
1665static void replaceAsyncResumeFunction(CoroSuspendAsyncInst *Suspend,
1666 Value *Continuation) {
1667 auto *ResumeIntrinsic = Suspend->getResumeFunction();
1668 auto &Context = Suspend->getParent()->getParent()->getContext();
1669 auto *Int8PtrTy = PointerType::getUnqual(C&: Context);
1670
1671 IRBuilder<> Builder(ResumeIntrinsic);
1672 auto *Val = Builder.CreateBitOrPointerCast(V: Continuation, DestTy: Int8PtrTy);
1673 ResumeIntrinsic->replaceAllUsesWith(V: Val);
1674 ResumeIntrinsic->eraseFromParent();
1675 Suspend->setOperand(i_nocapture: CoroSuspendAsyncInst::ResumeFunctionArg,
1676 Val_nocapture: PoisonValue::get(T: Int8PtrTy));
1677}
1678
1679/// Coerce the arguments in \p FnArgs according to \p FnTy in \p CallArgs.
1680static void coerceArguments(IRBuilder<> &Builder, FunctionType *FnTy,
1681 ArrayRef<Value *> FnArgs,
1682 SmallVectorImpl<Value *> &CallArgs) {
1683 size_t ArgIdx = 0;
1684 for (auto *paramTy : FnTy->params()) {
1685 assert(ArgIdx < FnArgs.size());
1686 if (paramTy != FnArgs[ArgIdx]->getType())
1687 CallArgs.push_back(
1688 Elt: Builder.CreateBitOrPointerCast(V: FnArgs[ArgIdx], DestTy: paramTy));
1689 else
1690 CallArgs.push_back(Elt: FnArgs[ArgIdx]);
1691 ++ArgIdx;
1692 }
1693}
1694
1695CallInst *coro::createMustTailCall(DebugLoc Loc, Function *MustTailCallFn,
1696 TargetTransformInfo &TTI,
1697 ArrayRef<Value *> Arguments,
1698 IRBuilder<> &Builder) {
1699 auto *FnTy = MustTailCallFn->getFunctionType();
1700 // Coerce the arguments, llvm optimizations seem to ignore the types in
1701 // vaarg functions and throws away casts in optimized mode.
1702 SmallVector<Value *, 8> CallArgs;
1703 coerceArguments(Builder, FnTy, FnArgs: Arguments, CallArgs);
1704
1705 auto *TailCall = Builder.CreateCall(FTy: FnTy, Callee: MustTailCallFn, Args: CallArgs);
1706 // Skip targets which don't support tail call.
1707 if (TTI.supportsTailCallFor(CB: TailCall)) {
1708 TailCall->setTailCallKind(CallInst::TCK_MustTail);
1709 }
1710 TailCall->setDebugLoc(Loc);
1711 TailCall->setCallingConv(MustTailCallFn->getCallingConv());
1712 return TailCall;
1713}
1714
1715void coro::AsyncABI::splitCoroutine(Function &F, coro::Shape &Shape,
1716 SmallVectorImpl<Function *> &Clones,
1717 TargetTransformInfo &TTI) {
1718 assert(Shape.ABI == coro::ABI::Async);
1719 assert(Clones.empty());
1720 // Reset various things that the optimizer might have decided it
1721 // "knows" about the coroutine function due to not seeing a return.
1722 F.removeFnAttr(Kind: Attribute::NoReturn);
1723 F.removeRetAttr(Kind: Attribute::NoAlias);
1724 F.removeRetAttr(Kind: Attribute::NonNull);
1725
1726 auto &Context = F.getContext();
1727 auto *Int8PtrTy = PointerType::getUnqual(C&: Context);
1728
1729 auto *Id = Shape.getAsyncCoroId();
1730 IRBuilder<> Builder(Id);
1731
1732 auto *FramePtr = Id->getStorage();
1733 FramePtr = Builder.CreateBitOrPointerCast(V: FramePtr, DestTy: Int8PtrTy);
1734 FramePtr = Builder.CreateInBoundsPtrAdd(
1735 Ptr: FramePtr,
1736 Offset: ConstantInt::get(Ty: Type::getInt64Ty(C&: Context),
1737 V: Shape.AsyncLowering.FrameOffset),
1738 Name: "async.ctx.frameptr");
1739
1740 // Map all uses of llvm.coro.begin to the allocated frame pointer.
1741 {
1742 // Make sure we don't invalidate Shape.FramePtr.
1743 TrackingVH<Value> Handle(Shape.FramePtr);
1744 Shape.CoroBegin->replaceAllUsesWith(V: FramePtr);
1745 Shape.FramePtr = Handle.getValPtr();
1746 }
1747
1748 // Create all the functions in order after the main function.
1749 auto NextF = std::next(x: F.getIterator());
1750
1751 // Create a continuation function for each of the suspend points.
1752 Clones.reserve(N: Shape.CoroSuspends.size());
1753 for (auto [Idx, CS] : llvm::enumerate(First&: Shape.CoroSuspends)) {
1754 auto *Suspend = cast<CoroSuspendAsyncInst>(Val: CS);
1755
1756 // Create the clone declaration.
1757 auto ResumeNameSuffix = ".resume.";
1758 auto ProjectionFunctionName =
1759 Suspend->getAsyncContextProjectionFunction()->getName();
1760 bool UseSwiftMangling = false;
1761 if (ProjectionFunctionName == "__swift_async_resume_project_context") {
1762 ResumeNameSuffix = "TQ";
1763 UseSwiftMangling = true;
1764 } else if (ProjectionFunctionName == "__swift_async_resume_get_context") {
1765 ResumeNameSuffix = "TY";
1766 UseSwiftMangling = true;
1767 }
1768 auto *Continuation = createCloneDeclaration(
1769 OrigF&: F, Shape,
1770 Suffix: UseSwiftMangling ? ResumeNameSuffix + Twine(Idx) + "_"
1771 : ResumeNameSuffix + Twine(Idx),
1772 InsertBefore: NextF, ActiveSuspend: Suspend);
1773 Clones.push_back(Elt: Continuation);
1774
1775 // Insert a branch to a new return block immediately before the suspend
1776 // point.
1777 auto *SuspendBB = Suspend->getParent();
1778 auto *NewSuspendBB = SuspendBB->splitBasicBlock(I: Suspend);
1779 auto *Branch = cast<UncondBrInst>(Val: SuspendBB->getTerminator());
1780
1781 // Place it before the first suspend.
1782 auto *ReturnBB =
1783 BasicBlock::Create(Context&: F.getContext(), Name: "coro.return", Parent: &F, InsertBefore: NewSuspendBB);
1784 Branch->setSuccessor(idx: 0, NewSucc: ReturnBB);
1785
1786 IRBuilder<> Builder(ReturnBB);
1787
1788 // Insert the call to the tail call function and inline it.
1789 auto *Fn = Suspend->getMustTailCallFunction();
1790 SmallVector<Value *, 8> Args(Suspend->args());
1791 auto FnArgs = ArrayRef<Value *>(Args).drop_front(
1792 N: CoroSuspendAsyncInst::MustTailCallFuncArg + 1);
1793 auto *TailCall = coro::createMustTailCall(Loc: Suspend->getDebugLoc(), MustTailCallFn: Fn, TTI,
1794 Arguments: FnArgs, Builder);
1795 Builder.CreateRetVoid();
1796 InlineFunctionInfo FnInfo;
1797 (void)InlineFunction(CB&: *TailCall, IFI&: FnInfo);
1798
1799 // Replace the lvm.coro.async.resume intrisic call.
1800 replaceAsyncResumeFunction(Suspend, Continuation);
1801 }
1802
1803 assert(Clones.size() == Shape.CoroSuspends.size());
1804
1805 for (auto [Idx, CS] : llvm::enumerate(First&: Shape.CoroSuspends)) {
1806 auto *Suspend = CS;
1807 auto *Clone = Clones[Idx];
1808
1809 coro::BaseCloner::createClone(OrigF&: F, Suffix: "resume." + Twine(Idx), Shape, NewF: Clone,
1810 ActiveSuspend: Suspend, TTI);
1811 }
1812}
1813
1814void coro::AnyRetconABI::splitCoroutine(Function &F, coro::Shape &Shape,
1815 SmallVectorImpl<Function *> &Clones,
1816 TargetTransformInfo &TTI) {
1817 assert(Shape.ABI == coro::ABI::Retcon || Shape.ABI == coro::ABI::RetconOnce);
1818 assert(Clones.empty());
1819
1820 // Reset various things that the optimizer might have decided it
1821 // "knows" about the coroutine function due to not seeing a return.
1822 F.removeFnAttr(Kind: Attribute::NoReturn);
1823 F.removeRetAttr(Kind: Attribute::NoAlias);
1824 F.removeRetAttr(Kind: Attribute::NonNull);
1825
1826 // Allocate the frame.
1827 auto *Id = Shape.getRetconCoroId();
1828 Value *RawFramePtr;
1829 if (Shape.RetconLowering.IsFrameInlineInStorage) {
1830 RawFramePtr = Id->getStorage();
1831 } else {
1832 IRBuilder<> Builder(Id);
1833
1834 auto FrameSize = Builder.getInt64(C: Shape.FrameSize);
1835
1836 // Allocate. We don't need to update the call graph node because we're
1837 // going to recompute it from scratch after splitting.
1838 // FIXME: pass the required alignment
1839 RawFramePtr = Shape.emitAlloc(Builder, Size: FrameSize, CG: nullptr);
1840 RawFramePtr =
1841 Builder.CreateBitCast(V: RawFramePtr, DestTy: Shape.CoroBegin->getType());
1842
1843 // Stash the allocated frame pointer in the continuation storage.
1844 Builder.CreateStore(Val: RawFramePtr, Ptr: Id->getStorage());
1845 }
1846
1847 // Map all uses of llvm.coro.begin to the allocated frame pointer.
1848 {
1849 // Make sure we don't invalidate Shape.FramePtr.
1850 TrackingVH<Value> Handle(Shape.FramePtr);
1851 Shape.CoroBegin->replaceAllUsesWith(V: RawFramePtr);
1852 Shape.FramePtr = Handle.getValPtr();
1853 }
1854
1855 // Create a unique return block.
1856 BasicBlock *ReturnBB = nullptr;
1857 PHINode *ContinuationPhi = nullptr;
1858 SmallVector<PHINode *, 4> ReturnPHIs;
1859
1860 // Create all the functions in order after the main function.
1861 auto NextF = std::next(x: F.getIterator());
1862
1863 // Create a continuation function for each of the suspend points.
1864 Clones.reserve(N: Shape.CoroSuspends.size());
1865 for (auto [Idx, CS] : llvm::enumerate(First&: Shape.CoroSuspends)) {
1866 auto Suspend = cast<CoroSuspendRetconInst>(Val: CS);
1867
1868 // Create the clone declaration.
1869 auto Continuation = createCloneDeclaration(
1870 OrigF&: F, Shape, Suffix: ".resume." + Twine(Idx), InsertBefore: NextF, ActiveSuspend: nullptr);
1871 Clones.push_back(Elt: Continuation);
1872
1873 // Insert a branch to the unified return block immediately before
1874 // the suspend point.
1875 auto SuspendBB = Suspend->getParent();
1876 auto NewSuspendBB = SuspendBB->splitBasicBlock(I: Suspend);
1877 auto Branch = cast<UncondBrInst>(Val: SuspendBB->getTerminator());
1878
1879 // Create the unified return block.
1880 if (!ReturnBB) {
1881 // Place it before the first suspend.
1882 ReturnBB =
1883 BasicBlock::Create(Context&: F.getContext(), Name: "coro.return", Parent: &F, InsertBefore: NewSuspendBB);
1884 Shape.RetconLowering.ReturnBlock = ReturnBB;
1885
1886 IRBuilder<> Builder(ReturnBB);
1887
1888 // First, the continuation.
1889 ContinuationPhi =
1890 Builder.CreatePHI(Ty: Continuation->getType(), NumReservedValues: Shape.CoroSuspends.size());
1891
1892 // Create PHIs for all other return values.
1893 assert(ReturnPHIs.empty());
1894
1895 // Next, all the directly-yielded values.
1896 for (auto *ResultTy : Shape.getRetconResultTypes())
1897 ReturnPHIs.push_back(
1898 Elt: Builder.CreatePHI(Ty: ResultTy, NumReservedValues: Shape.CoroSuspends.size()));
1899
1900 // Build the return value.
1901 auto RetTy = F.getReturnType();
1902
1903 // Cast the continuation value if necessary.
1904 // We can't rely on the types matching up because that type would
1905 // have to be infinite.
1906 auto CastedContinuationTy =
1907 (ReturnPHIs.empty() ? RetTy : RetTy->getStructElementType(N: 0));
1908 auto *CastedContinuation =
1909 Builder.CreateBitCast(V: ContinuationPhi, DestTy: CastedContinuationTy);
1910
1911 Value *RetV = CastedContinuation;
1912 if (!ReturnPHIs.empty()) {
1913 auto ValueIdx = 0;
1914 RetV = PoisonValue::get(T: RetTy);
1915 RetV = Builder.CreateInsertValue(Agg: RetV, Val: CastedContinuation, Idxs: ValueIdx++);
1916
1917 for (auto Phi : ReturnPHIs)
1918 RetV = Builder.CreateInsertValue(Agg: RetV, Val: Phi, Idxs: ValueIdx++);
1919 }
1920
1921 Builder.CreateRet(V: RetV);
1922 }
1923
1924 // Branch to the return block.
1925 Branch->setSuccessor(idx: 0, NewSucc: ReturnBB);
1926 assert(ContinuationPhi);
1927 ContinuationPhi->addIncoming(V: Continuation, BB: SuspendBB);
1928 for (auto [Phi, VUse] :
1929 llvm::zip_equal(t&: ReturnPHIs, u: Suspend->value_operands()))
1930 Phi->addIncoming(V: VUse, BB: SuspendBB);
1931 }
1932
1933 assert(Clones.size() == Shape.CoroSuspends.size());
1934
1935 for (auto [Idx, CS] : llvm::enumerate(First&: Shape.CoroSuspends)) {
1936 auto Suspend = CS;
1937 auto Clone = Clones[Idx];
1938
1939 coro::BaseCloner::createClone(OrigF&: F, Suffix: "resume." + Twine(Idx), Shape, NewF: Clone,
1940 ActiveSuspend: Suspend, TTI);
1941 }
1942}
1943
1944namespace {
1945class PrettyStackTraceFunction : public PrettyStackTraceEntry {
1946 Function &F;
1947
1948public:
1949 PrettyStackTraceFunction(Function &F) : F(F) {}
1950 void print(raw_ostream &OS) const override {
1951 OS << "While splitting coroutine ";
1952 F.printAsOperand(O&: OS, /*print type*/ PrintType: false, M: F.getParent());
1953 OS << "\n";
1954 }
1955};
1956} // namespace
1957
1958/// Remove calls to llvm.coro.end in the original function.
1959static void removeCoroEndsFromRampFunction(const coro::Shape &Shape) {
1960 if (Shape.ABI != coro::ABI::Switch) {
1961 for (auto *End : Shape.CoroEnds) {
1962 replaceCoroEnd(End, Shape, FramePtr: Shape.FramePtr, /*in ramp*/ InRamp: true, CG: nullptr);
1963 }
1964 } else {
1965 for (llvm::AnyCoroEndInst *End : Shape.CoroEnds)
1966 End->eraseFromParent();
1967 }
1968}
1969
1970static void removeCoroIsInRampFromRampFunction(const coro::Shape &Shape) {
1971 for (auto *II : Shape.CoroIsInRampInsts) {
1972 auto &Ctx = II->getContext();
1973 II->replaceAllUsesWith(V: ConstantInt::getTrue(Context&: Ctx));
1974 II->eraseFromParent();
1975 }
1976}
1977
1978static bool hasSafeElideCaller(Function &F) {
1979 for (auto *U : F.users()) {
1980 if (auto *CB = dyn_cast<CallBase>(Val: U)) {
1981 auto *Caller = CB->getFunction();
1982 if (Caller && Caller->isPresplitCoroutine() &&
1983 CB->hasFnAttr(Kind: llvm::Attribute::CoroElideSafe))
1984 return true;
1985 }
1986 }
1987 return false;
1988}
1989
1990void coro::SwitchABI::splitCoroutine(Function &F, coro::Shape &Shape,
1991 SmallVectorImpl<Function *> &Clones,
1992 TargetTransformInfo &TTI) {
1993 SwitchCoroutineSplitter::split(F, Shape, Clones, TTI);
1994}
1995
1996static void doSplitCoroutine(Function &F, SmallVectorImpl<Function *> &Clones,
1997 coro::BaseABI &ABI, TargetTransformInfo &TTI,
1998 bool OptimizeFrame) {
1999 PrettyStackTraceFunction prettyStackTrace(F);
2000
2001 auto &Shape = ABI.Shape;
2002 assert(Shape.CoroBegin);
2003
2004 lowerAwaitSuspends(F, Shape);
2005
2006 simplifySuspendPoints(Shape);
2007
2008 normalizeCoroutine(F, Shape, TTI);
2009 ABI.buildCoroutineFrame(OptimizeFrame);
2010 replaceFrameSizeAndAlignment(Shape);
2011
2012 bool isNoSuspendCoroutine = Shape.CoroSuspends.empty();
2013
2014 bool shouldCreateNoAllocVariant =
2015 !isNoSuspendCoroutine && Shape.ABI == coro::ABI::Switch &&
2016 hasSafeElideCaller(F) && !F.hasFnAttribute(Kind: llvm::Attribute::NoInline);
2017
2018 // If there are no suspend points, no split required, just remove
2019 // the allocation and deallocation blocks, they are not needed.
2020 if (isNoSuspendCoroutine) {
2021 handleNoSuspendCoroutine(Shape);
2022 } else {
2023 ABI.splitCoroutine(F, Shape, Clones, TTI);
2024 }
2025
2026 // Replace all the swifterror operations in the original function.
2027 // This invalidates SwiftErrorOps in the Shape.
2028 replaceSwiftErrorOps(F, Shape, VMap: nullptr);
2029
2030 // Salvage debug intrinsics that point into the coroutine frame in the
2031 // original function. The Cloner has already salvaged debug info in the new
2032 // coroutine funclets.
2033 SmallDenseMap<Argument *, AllocaInst *, 4> ArgToAllocaMap;
2034 auto DbgVariableRecords = collectDbgVariableRecords(F);
2035 for (DbgVariableRecord *DVR : DbgVariableRecords)
2036 coro::salvageDebugInfo(ArgToAllocaMap, DVR&: *DVR, UseEntryValue: false /*UseEntryValue*/);
2037
2038 removeCoroEndsFromRampFunction(Shape);
2039 removeCoroIsInRampFromRampFunction(Shape);
2040
2041 if (shouldCreateNoAllocVariant)
2042 SwitchCoroutineSplitter::createNoAllocVariant(F, Shape, Clones);
2043}
2044
2045static LazyCallGraph::SCC &updateCallGraphAfterCoroutineSplit(
2046 LazyCallGraph::Node &N, const coro::Shape &Shape,
2047 const SmallVectorImpl<Function *> &Clones, LazyCallGraph::SCC &C,
2048 LazyCallGraph &CG, CGSCCAnalysisManager &AM, CGSCCUpdateResult &UR,
2049 FunctionAnalysisManager &FAM) {
2050
2051 auto *CurrentSCC = &C;
2052 if (!Clones.empty()) {
2053 switch (Shape.ABI) {
2054 case coro::ABI::Switch:
2055 // Each clone in the Switch lowering is independent of the other clones.
2056 // Let the LazyCallGraph know about each one separately.
2057 for (Function *Clone : Clones)
2058 CG.addSplitFunction(OriginalFunction&: N.getFunction(), NewFunction&: *Clone);
2059 break;
2060 case coro::ABI::Async:
2061 case coro::ABI::Retcon:
2062 case coro::ABI::RetconOnce:
2063 // Each clone in the Async/Retcon lowering references of the other clones.
2064 // Let the LazyCallGraph know about all of them at once.
2065 if (!Clones.empty())
2066 CG.addSplitRefRecursiveFunctions(OriginalFunction&: N.getFunction(), NewFunctions: Clones);
2067 break;
2068 }
2069
2070 // Let the CGSCC infra handle the changes to the original function.
2071 CurrentSCC = &updateCGAndAnalysisManagerForCGSCCPass(G&: CG, C&: *CurrentSCC, N, AM,
2072 UR, FAM);
2073 }
2074
2075 // Do some cleanup and let the CGSCC infra see if we've cleaned up any edges
2076 // to the split functions.
2077 postSplitCleanup(F&: N.getFunction());
2078 CurrentSCC = &updateCGAndAnalysisManagerForFunctionPass(G&: CG, C&: *CurrentSCC, N,
2079 AM, UR, FAM);
2080 return *CurrentSCC;
2081}
2082
2083/// Replace a call to llvm.coro.prepare.retcon.
2084static void replacePrepare(CallInst *Prepare, LazyCallGraph &CG,
2085 LazyCallGraph::SCC &C) {
2086 auto CastFn = Prepare->getArgOperand(i: 0); // as an i8*
2087 auto Fn = CastFn->stripPointerCasts(); // as its original type
2088
2089 // Attempt to peephole this pattern:
2090 // %0 = bitcast [[TYPE]] @some_function to i8*
2091 // %1 = call @llvm.coro.prepare.retcon(i8* %0)
2092 // %2 = bitcast %1 to [[TYPE]]
2093 // ==>
2094 // %2 = @some_function
2095 for (Use &U : llvm::make_early_inc_range(Range: Prepare->uses())) {
2096 // Look for bitcasts back to the original function type.
2097 auto *Cast = dyn_cast<BitCastInst>(Val: U.getUser());
2098 if (!Cast || Cast->getType() != Fn->getType())
2099 continue;
2100
2101 // Replace and remove the cast.
2102 Cast->replaceAllUsesWith(V: Fn);
2103 Cast->eraseFromParent();
2104 }
2105
2106 // Replace any remaining uses with the function as an i8*.
2107 // This can never directly be a callee, so we don't need to update CG.
2108 Prepare->replaceAllUsesWith(V: CastFn);
2109 Prepare->eraseFromParent();
2110
2111 // Kill dead bitcasts.
2112 while (auto *Cast = dyn_cast<BitCastInst>(Val: CastFn)) {
2113 if (!Cast->use_empty())
2114 break;
2115 CastFn = Cast->getOperand(i_nocapture: 0);
2116 Cast->eraseFromParent();
2117 }
2118}
2119
2120static bool replaceAllPrepares(Function *PrepareFn, LazyCallGraph &CG,
2121 LazyCallGraph::SCC &C) {
2122 bool Changed = false;
2123 for (Use &P : llvm::make_early_inc_range(Range: PrepareFn->uses())) {
2124 // Intrinsics can only be used in calls.
2125 auto *Prepare = cast<CallInst>(Val: P.getUser());
2126 replacePrepare(Prepare, CG, C);
2127 Changed = true;
2128 }
2129
2130 return Changed;
2131}
2132
2133static void addPrepareFunction(const Module &M,
2134 SmallVectorImpl<Function *> &Fns,
2135 StringRef Name) {
2136 auto *PrepareFn = M.getFunction(Name);
2137 if (PrepareFn && !PrepareFn->use_empty())
2138 Fns.push_back(Elt: PrepareFn);
2139}
2140
2141static std::unique_ptr<coro::BaseABI>
2142CreateNewABI(Function &F, coro::Shape &S,
2143 std::function<bool(Instruction &)> IsMatCallback,
2144 const SmallVector<CoroSplitPass::BaseABITy> GenCustomABIs) {
2145 if (S.CoroBegin->hasCustomABI()) {
2146 unsigned CustomABI = S.CoroBegin->getCustomABI();
2147 if (CustomABI >= GenCustomABIs.size())
2148 llvm_unreachable("Custom ABI not found amoung those specified");
2149 return GenCustomABIs[CustomABI](F, S);
2150 }
2151
2152 switch (S.ABI) {
2153 case coro::ABI::Switch:
2154 return std::make_unique<coro::SwitchABI>(args&: F, args&: S, args&: IsMatCallback);
2155 case coro::ABI::Async:
2156 return std::make_unique<coro::AsyncABI>(args&: F, args&: S, args&: IsMatCallback);
2157 case coro::ABI::Retcon:
2158 return std::make_unique<coro::AnyRetconABI>(args&: F, args&: S, args&: IsMatCallback);
2159 case coro::ABI::RetconOnce:
2160 return std::make_unique<coro::AnyRetconABI>(args&: F, args&: S, args&: IsMatCallback);
2161 }
2162 llvm_unreachable("Unknown ABI");
2163}
2164
2165CoroSplitPass::CoroSplitPass(bool OptimizeFrame)
2166 : CreateAndInitABI([](Function &F, coro::Shape &S) {
2167 std::unique_ptr<coro::BaseABI> ABI =
2168 CreateNewABI(F, S, IsMatCallback: coro::isTriviallyMaterializable, GenCustomABIs: {});
2169 ABI->init();
2170 return ABI;
2171 }),
2172 OptimizeFrame(OptimizeFrame) {}
2173
2174CoroSplitPass::CoroSplitPass(
2175 SmallVector<CoroSplitPass::BaseABITy> GenCustomABIs, bool OptimizeFrame)
2176 : CreateAndInitABI([=](Function &F, coro::Shape &S) {
2177 std::unique_ptr<coro::BaseABI> ABI =
2178 CreateNewABI(F, S, IsMatCallback: coro::isTriviallyMaterializable, GenCustomABIs);
2179 ABI->init();
2180 return ABI;
2181 }),
2182 OptimizeFrame(OptimizeFrame) {}
2183
2184// For back compatibility, constructor takes a materializable callback and
2185// creates a generator for an ABI with a modified materializable callback.
2186CoroSplitPass::CoroSplitPass(std::function<bool(Instruction &)> IsMatCallback,
2187 bool OptimizeFrame)
2188 : CreateAndInitABI([=](Function &F, coro::Shape &S) {
2189 std::unique_ptr<coro::BaseABI> ABI =
2190 CreateNewABI(F, S, IsMatCallback, GenCustomABIs: {});
2191 ABI->init();
2192 return ABI;
2193 }),
2194 OptimizeFrame(OptimizeFrame) {}
2195
2196// For back compatibility, constructor takes a materializable callback and
2197// creates a generator for an ABI with a modified materializable callback.
2198CoroSplitPass::CoroSplitPass(
2199 std::function<bool(Instruction &)> IsMatCallback,
2200 SmallVector<CoroSplitPass::BaseABITy> GenCustomABIs, bool OptimizeFrame)
2201 : CreateAndInitABI([=](Function &F, coro::Shape &S) {
2202 std::unique_ptr<coro::BaseABI> ABI =
2203 CreateNewABI(F, S, IsMatCallback, GenCustomABIs);
2204 ABI->init();
2205 return ABI;
2206 }),
2207 OptimizeFrame(OptimizeFrame) {}
2208
2209PreservedAnalyses CoroSplitPass::run(LazyCallGraph::SCC &C,
2210 CGSCCAnalysisManager &AM,
2211 LazyCallGraph &CG, CGSCCUpdateResult &UR) {
2212 // NB: One invariant of a valid LazyCallGraph::SCC is that it must contain a
2213 // non-zero number of nodes, so we assume that here and grab the first
2214 // node's function's module.
2215 Module &M = *C.begin()->getFunction().getParent();
2216 auto &FAM =
2217 AM.getResult<FunctionAnalysisManagerCGSCCProxy>(IR&: C, ExtraArgs&: CG).getManager();
2218
2219 // Check for uses of llvm.coro.prepare.retcon/async.
2220 SmallVector<Function *, 2> PrepareFns;
2221 addPrepareFunction(M, Fns&: PrepareFns, Name: "llvm.coro.prepare.retcon");
2222 addPrepareFunction(M, Fns&: PrepareFns, Name: "llvm.coro.prepare.async");
2223
2224 // Find coroutines for processing.
2225 SmallVector<LazyCallGraph::Node *> Coroutines;
2226 for (LazyCallGraph::Node &N : C)
2227 if (N.getFunction().isPresplitCoroutine())
2228 Coroutines.push_back(Elt: &N);
2229
2230 if (Coroutines.empty() && PrepareFns.empty())
2231 return PreservedAnalyses::all();
2232
2233 auto *CurrentSCC = &C;
2234 // Split all the coroutines.
2235 for (LazyCallGraph::Node *N : Coroutines) {
2236 Function &F = N->getFunction();
2237 LLVM_DEBUG(dbgs() << "CoroSplit: Processing coroutine '" << F.getName()
2238 << "\n");
2239
2240 // The suspend-crossing algorithm in buildCoroutineFrame gets tripped up
2241 // by unreachable blocks, so remove them as a first pass. Remove the
2242 // unreachable blocks before collecting intrinsics into Shape.
2243 removeUnreachableBlocks(F);
2244
2245 coro::Shape Shape(F);
2246 if (!Shape.CoroBegin)
2247 continue;
2248
2249 F.setSplittedCoroutine();
2250
2251 std::unique_ptr<coro::BaseABI> ABI = CreateAndInitABI(F, Shape);
2252
2253 SmallVector<Function *, 4> Clones;
2254 auto &TTI = FAM.getResult<TargetIRAnalysis>(IR&: F);
2255 doSplitCoroutine(F, Clones, ABI&: *ABI, TTI, OptimizeFrame);
2256 CurrentSCC = &updateCallGraphAfterCoroutineSplit(
2257 N&: *N, Shape, Clones, C&: *CurrentSCC, CG, AM, UR, FAM);
2258
2259 auto &ORE = FAM.getResult<OptimizationRemarkEmitterAnalysis>(IR&: F);
2260 ORE.emit(RemarkBuilder: [&]() {
2261 return OptimizationRemark(DEBUG_TYPE, "CoroSplit", &F)
2262 << "Split '" << ore::NV("function", F.getName())
2263 << "' (frame_size=" << ore::NV("frame_size", Shape.FrameSize)
2264 << ", align=" << ore::NV("align", Shape.FrameAlign.value()) << ")";
2265 });
2266
2267 if (!Shape.CoroSuspends.empty()) {
2268 // Run the CGSCC pipeline on the original and newly split functions.
2269 UR.CWorklist.insert(X: CurrentSCC);
2270 for (Function *Clone : Clones)
2271 UR.CWorklist.insert(X: CG.lookupSCC(N&: CG.get(F&: *Clone)));
2272 } else if (Shape.ABI == coro::ABI::Async) {
2273 // Reprocess the function to inline the tail called return function of
2274 // coro.async.end.
2275 UR.CWorklist.insert(X: &C);
2276 }
2277 }
2278
2279 for (auto *PrepareFn : PrepareFns) {
2280 replaceAllPrepares(PrepareFn, CG, C&: *CurrentSCC);
2281 }
2282
2283 return PreservedAnalyses::none();
2284}
2285