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