xref: /freebsd/contrib/llvm-project/llvm/lib/Transforms/Coroutines/SpillUtils.cpp (revision 700637cbb5e582861067a11aaca4d053546871d2)
1 //===- SpillUtils.cpp - Utilities for checking for spills ---------------===//
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 
9 #include "llvm/Transforms/Coroutines/SpillUtils.h"
10 #include "CoroInternal.h"
11 #include "llvm/Analysis/CFG.h"
12 #include "llvm/Analysis/PtrUseVisitor.h"
13 #include "llvm/IR/CFG.h"
14 #include "llvm/IR/DebugInfo.h"
15 #include "llvm/IR/Dominators.h"
16 #include "llvm/IR/InstIterator.h"
17 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
18 
19 namespace llvm {
20 
21 namespace coro {
22 
23 namespace {
24 
25 typedef SmallPtrSet<BasicBlock *, 8> VisitedBlocksSet;
26 
isNonSpilledIntrinsic(Instruction & I)27 static bool isNonSpilledIntrinsic(Instruction &I) {
28   // Structural coroutine intrinsics that should not be spilled into the
29   // coroutine frame.
30   return isa<CoroIdInst>(&I) || isa<CoroSaveInst>(&I);
31 }
32 
33 /// Does control flow starting at the given block ever reach a suspend
34 /// instruction before reaching a block in VisitedOrFreeBBs?
isSuspendReachableFrom(BasicBlock * From,VisitedBlocksSet & VisitedOrFreeBBs)35 static bool isSuspendReachableFrom(BasicBlock *From,
36                                    VisitedBlocksSet &VisitedOrFreeBBs) {
37   // Eagerly try to add this block to the visited set. If it's already
38   // there, stop recursing; this path doesn't reach a suspend before
39   // either looping or reaching a freeing block.
40   if (!VisitedOrFreeBBs.insert(From).second)
41     return false;
42 
43   // We assume that we'll already have split suspends into their own blocks.
44   if (coro::isSuspendBlock(From))
45     return true;
46 
47   // Recurse on the successors.
48   for (auto *Succ : successors(From)) {
49     if (isSuspendReachableFrom(Succ, VisitedOrFreeBBs))
50       return true;
51   }
52 
53   return false;
54 }
55 
56 /// Is the given alloca "local", i.e. bounded in lifetime to not cross a
57 /// suspend point?
isLocalAlloca(CoroAllocaAllocInst * AI)58 static bool isLocalAlloca(CoroAllocaAllocInst *AI) {
59   // Seed the visited set with all the basic blocks containing a free
60   // so that we won't pass them up.
61   VisitedBlocksSet VisitedOrFreeBBs;
62   for (auto *User : AI->users()) {
63     if (auto FI = dyn_cast<CoroAllocaFreeInst>(User))
64       VisitedOrFreeBBs.insert(FI->getParent());
65   }
66 
67   return !isSuspendReachableFrom(AI->getParent(), VisitedOrFreeBBs);
68 }
69 
70 /// Turn the given coro.alloca.alloc call into a dynamic allocation.
71 /// This happens during the all-instructions iteration, so it must not
72 /// delete the call.
73 static Instruction *
lowerNonLocalAlloca(CoroAllocaAllocInst * AI,const coro::Shape & Shape,SmallVectorImpl<Instruction * > & DeadInsts)74 lowerNonLocalAlloca(CoroAllocaAllocInst *AI, const coro::Shape &Shape,
75                     SmallVectorImpl<Instruction *> &DeadInsts) {
76   IRBuilder<> Builder(AI);
77   auto Alloc = Shape.emitAlloc(Builder, AI->getSize(), nullptr);
78 
79   for (User *U : AI->users()) {
80     if (isa<CoroAllocaGetInst>(U)) {
81       U->replaceAllUsesWith(Alloc);
82     } else {
83       auto FI = cast<CoroAllocaFreeInst>(U);
84       Builder.SetInsertPoint(FI);
85       Shape.emitDealloc(Builder, Alloc, nullptr);
86     }
87     DeadInsts.push_back(cast<Instruction>(U));
88   }
89 
90   // Push this on last so that it gets deleted after all the others.
91   DeadInsts.push_back(AI);
92 
93   // Return the new allocation value so that we can check for needed spills.
94   return cast<Instruction>(Alloc);
95 }
96 
97 // We need to make room to insert a spill after initial PHIs, but before
98 // catchswitch instruction. Placing it before violates the requirement that
99 // catchswitch, like all other EHPads must be the first nonPHI in a block.
100 //
101 // Split away catchswitch into a separate block and insert in its place:
102 //
103 //   cleanuppad <InsertPt> cleanupret.
104 //
105 // cleanupret instruction will act as an insert point for the spill.
splitBeforeCatchSwitch(CatchSwitchInst * CatchSwitch)106 static Instruction *splitBeforeCatchSwitch(CatchSwitchInst *CatchSwitch) {
107   BasicBlock *CurrentBlock = CatchSwitch->getParent();
108   BasicBlock *NewBlock = CurrentBlock->splitBasicBlock(CatchSwitch);
109   CurrentBlock->getTerminator()->eraseFromParent();
110 
111   auto *CleanupPad =
112       CleanupPadInst::Create(CatchSwitch->getParentPad(), {}, "", CurrentBlock);
113   auto *CleanupRet =
114       CleanupReturnInst::Create(CleanupPad, NewBlock, CurrentBlock);
115   return CleanupRet;
116 }
117 
118 // We use a pointer use visitor to track how an alloca is being used.
119 // The goal is to be able to answer the following three questions:
120 //   1. Should this alloca be allocated on the frame instead.
121 //   2. Could the content of the alloca be modified prior to CoroBegin, which
122 //      would require copying the data from the alloca to the frame after
123 //      CoroBegin.
124 //   3. Are there any aliases created for this alloca prior to CoroBegin, but
125 //      used after CoroBegin. In that case, we will need to recreate the alias
126 //      after CoroBegin based off the frame.
127 //
128 // To answer question 1, we track two things:
129 //   A. List of all BasicBlocks that use this alloca or any of the aliases of
130 //   the alloca. In the end, we check if there exists any two basic blocks that
131 //   cross suspension points. If so, this alloca must be put on the frame.
132 //   B. Whether the alloca or any alias of the alloca is escaped at some point,
133 //   either by storing the address somewhere, or the address is used in a
134 //   function call that might capture. If it's ever escaped, this alloca must be
135 //   put on the frame conservatively.
136 //
137 // To answer quetion 2, we track through the variable MayWriteBeforeCoroBegin.
138 // Whenever a potential write happens, either through a store instruction, a
139 // function call or any of the memory intrinsics, we check whether this
140 // instruction is prior to CoroBegin.
141 //
142 // To answer question 3, we track the offsets of all aliases created for the
143 // alloca prior to CoroBegin but used after CoroBegin. std::optional is used to
144 // be able to represent the case when the offset is unknown (e.g. when you have
145 // a PHINode that takes in different offset values). We cannot handle unknown
146 // offsets and will assert. This is the potential issue left out. An ideal
147 // solution would likely require a significant redesign.
148 
149 namespace {
150 struct AllocaUseVisitor : PtrUseVisitor<AllocaUseVisitor> {
151   using Base = PtrUseVisitor<AllocaUseVisitor>;
AllocaUseVisitorllvm::coro::__anondea4afbd0111::__anondea4afbd0211::AllocaUseVisitor152   AllocaUseVisitor(const DataLayout &DL, const DominatorTree &DT,
153                    const coro::Shape &CoroShape,
154                    const SuspendCrossingInfo &Checker,
155                    bool ShouldUseLifetimeStartInfo)
156       : PtrUseVisitor(DL), DT(DT), CoroShape(CoroShape), Checker(Checker),
157         ShouldUseLifetimeStartInfo(ShouldUseLifetimeStartInfo) {
158     for (AnyCoroSuspendInst *SuspendInst : CoroShape.CoroSuspends)
159       CoroSuspendBBs.insert(SuspendInst->getParent());
160   }
161 
visitllvm::coro::__anondea4afbd0111::__anondea4afbd0211::AllocaUseVisitor162   void visit(Instruction &I) {
163     Users.insert(&I);
164     Base::visit(I);
165     // If the pointer is escaped prior to CoroBegin, we have to assume it would
166     // be written into before CoroBegin as well.
167     if (PI.isEscaped() &&
168         !DT.dominates(CoroShape.CoroBegin, PI.getEscapingInst())) {
169       MayWriteBeforeCoroBegin = true;
170     }
171   }
172   // We need to provide this overload as PtrUseVisitor uses a pointer based
173   // visiting function.
visitllvm::coro::__anondea4afbd0111::__anondea4afbd0211::AllocaUseVisitor174   void visit(Instruction *I) { return visit(*I); }
175 
visitPHINodellvm::coro::__anondea4afbd0111::__anondea4afbd0211::AllocaUseVisitor176   void visitPHINode(PHINode &I) {
177     enqueueUsers(I);
178     handleAlias(I);
179   }
180 
visitSelectInstllvm::coro::__anondea4afbd0111::__anondea4afbd0211::AllocaUseVisitor181   void visitSelectInst(SelectInst &I) {
182     enqueueUsers(I);
183     handleAlias(I);
184   }
185 
visitStoreInstllvm::coro::__anondea4afbd0111::__anondea4afbd0211::AllocaUseVisitor186   void visitStoreInst(StoreInst &SI) {
187     // Regardless whether the alias of the alloca is the value operand or the
188     // pointer operand, we need to assume the alloca is been written.
189     handleMayWrite(SI);
190 
191     if (SI.getValueOperand() != U->get())
192       return;
193 
194     // We are storing the pointer into a memory location, potentially escaping.
195     // As an optimization, we try to detect simple cases where it doesn't
196     // actually escape, for example:
197     //   %ptr = alloca ..
198     //   %addr = alloca ..
199     //   store %ptr, %addr
200     //   %x = load %addr
201     //   ..
202     // If %addr is only used by loading from it, we could simply treat %x as
203     // another alias of %ptr, and not considering %ptr being escaped.
204     auto IsSimpleStoreThenLoad = [&]() {
205       auto *AI = dyn_cast<AllocaInst>(SI.getPointerOperand());
206       // If the memory location we are storing to is not an alloca, it
207       // could be an alias of some other memory locations, which is difficult
208       // to analyze.
209       if (!AI)
210         return false;
211       // StoreAliases contains aliases of the memory location stored into.
212       SmallVector<Instruction *, 4> StoreAliases = {AI};
213       while (!StoreAliases.empty()) {
214         Instruction *I = StoreAliases.pop_back_val();
215         for (User *U : I->users()) {
216           // If we are loading from the memory location, we are creating an
217           // alias of the original pointer.
218           if (auto *LI = dyn_cast<LoadInst>(U)) {
219             enqueueUsers(*LI);
220             handleAlias(*LI);
221             continue;
222           }
223           // If we are overriding the memory location, the pointer certainly
224           // won't escape.
225           if (auto *S = dyn_cast<StoreInst>(U))
226             if (S->getPointerOperand() == I)
227               continue;
228           if (isa<LifetimeIntrinsic>(U))
229             continue;
230           // BitCastInst creats aliases of the memory location being stored
231           // into.
232           if (auto *BI = dyn_cast<BitCastInst>(U)) {
233             StoreAliases.push_back(BI);
234             continue;
235           }
236           return false;
237         }
238       }
239 
240       return true;
241     };
242 
243     if (!IsSimpleStoreThenLoad())
244       PI.setEscaped(&SI);
245   }
246 
247   // All mem intrinsics modify the data.
visitMemIntrinsicllvm::coro::__anondea4afbd0111::__anondea4afbd0211::AllocaUseVisitor248   void visitMemIntrinsic(MemIntrinsic &MI) { handleMayWrite(MI); }
249 
visitBitCastInstllvm::coro::__anondea4afbd0111::__anondea4afbd0211::AllocaUseVisitor250   void visitBitCastInst(BitCastInst &BC) {
251     Base::visitBitCastInst(BC);
252     handleAlias(BC);
253   }
254 
visitAddrSpaceCastInstllvm::coro::__anondea4afbd0111::__anondea4afbd0211::AllocaUseVisitor255   void visitAddrSpaceCastInst(AddrSpaceCastInst &ASC) {
256     Base::visitAddrSpaceCastInst(ASC);
257     handleAlias(ASC);
258   }
259 
visitGetElementPtrInstllvm::coro::__anondea4afbd0111::__anondea4afbd0211::AllocaUseVisitor260   void visitGetElementPtrInst(GetElementPtrInst &GEPI) {
261     // The base visitor will adjust Offset accordingly.
262     Base::visitGetElementPtrInst(GEPI);
263     handleAlias(GEPI);
264   }
265 
visitIntrinsicInstllvm::coro::__anondea4afbd0111::__anondea4afbd0211::AllocaUseVisitor266   void visitIntrinsicInst(IntrinsicInst &II) {
267     // When we found the lifetime markers refers to a
268     // subrange of the original alloca, ignore the lifetime
269     // markers to avoid misleading the analysis.
270     if (!IsOffsetKnown || !Offset.isZero())
271       return Base::visitIntrinsicInst(II);
272     switch (II.getIntrinsicID()) {
273     default:
274       return Base::visitIntrinsicInst(II);
275     case Intrinsic::lifetime_start:
276       LifetimeStarts.insert(&II);
277       LifetimeStartBBs.push_back(II.getParent());
278       break;
279     case Intrinsic::lifetime_end:
280       LifetimeEndBBs.insert(II.getParent());
281       break;
282     }
283   }
284 
visitCallBasellvm::coro::__anondea4afbd0111::__anondea4afbd0211::AllocaUseVisitor285   void visitCallBase(CallBase &CB) {
286     for (unsigned Op = 0, OpCount = CB.arg_size(); Op < OpCount; ++Op)
287       if (U->get() == CB.getArgOperand(Op) && !CB.doesNotCapture(Op))
288         PI.setEscaped(&CB);
289     handleMayWrite(CB);
290   }
291 
getShouldLiveOnFramellvm::coro::__anondea4afbd0111::__anondea4afbd0211::AllocaUseVisitor292   bool getShouldLiveOnFrame() const {
293     if (!ShouldLiveOnFrame)
294       ShouldLiveOnFrame = computeShouldLiveOnFrame();
295     return *ShouldLiveOnFrame;
296   }
297 
getMayWriteBeforeCoroBeginllvm::coro::__anondea4afbd0111::__anondea4afbd0211::AllocaUseVisitor298   bool getMayWriteBeforeCoroBegin() const { return MayWriteBeforeCoroBegin; }
299 
getAliasesCopyllvm::coro::__anondea4afbd0111::__anondea4afbd0211::AllocaUseVisitor300   DenseMap<Instruction *, std::optional<APInt>> getAliasesCopy() const {
301     assert(getShouldLiveOnFrame() && "This method should only be called if the "
302                                      "alloca needs to live on the frame.");
303     for (const auto &P : AliasOffetMap)
304       if (!P.second)
305         report_fatal_error("Unable to handle an alias with unknown offset "
306                            "created before CoroBegin.");
307     return AliasOffetMap;
308   }
309 
310 private:
311   const DominatorTree &DT;
312   const coro::Shape &CoroShape;
313   const SuspendCrossingInfo &Checker;
314   // All alias to the original AllocaInst, created before CoroBegin and used
315   // after CoroBegin. Each entry contains the instruction and the offset in the
316   // original Alloca. They need to be recreated after CoroBegin off the frame.
317   DenseMap<Instruction *, std::optional<APInt>> AliasOffetMap{};
318   SmallPtrSet<Instruction *, 4> Users{};
319   SmallPtrSet<IntrinsicInst *, 2> LifetimeStarts{};
320   SmallVector<BasicBlock *> LifetimeStartBBs{};
321   SmallPtrSet<BasicBlock *, 2> LifetimeEndBBs{};
322   SmallPtrSet<const BasicBlock *, 2> CoroSuspendBBs{};
323   bool MayWriteBeforeCoroBegin{false};
324   bool ShouldUseLifetimeStartInfo{true};
325 
326   mutable std::optional<bool> ShouldLiveOnFrame{};
327 
computeShouldLiveOnFramellvm::coro::__anondea4afbd0111::__anondea4afbd0211::AllocaUseVisitor328   bool computeShouldLiveOnFrame() const {
329     // If lifetime information is available, we check it first since it's
330     // more precise. We look at every pair of lifetime.start intrinsic and
331     // every basic block that uses the pointer to see if they cross suspension
332     // points. The uses cover both direct uses as well as indirect uses.
333     if (ShouldUseLifetimeStartInfo && !LifetimeStarts.empty()) {
334       // If there is no explicit lifetime.end, then assume the address can
335       // cross suspension points.
336       if (LifetimeEndBBs.empty())
337         return true;
338 
339       // If there is a path from a lifetime.start to a suspend without a
340       // corresponding lifetime.end, then the alloca's lifetime persists
341       // beyond that suspension point and the alloca must go on the frame.
342       llvm::SmallVector<BasicBlock *> Worklist(LifetimeStartBBs);
343       if (isManyPotentiallyReachableFromMany(Worklist, CoroSuspendBBs,
344                                              &LifetimeEndBBs, &DT))
345         return true;
346 
347       // Addresses are guaranteed to be identical after every lifetime.start so
348       // we cannot use the local stack if the address escaped and there is a
349       // suspend point between lifetime markers. This should also cover the
350       // case of a single lifetime.start intrinsic in a loop with suspend point.
351       if (PI.isEscaped()) {
352         for (auto *A : LifetimeStarts) {
353           for (auto *B : LifetimeStarts) {
354             if (Checker.hasPathOrLoopCrossingSuspendPoint(A->getParent(),
355                                                           B->getParent()))
356               return true;
357           }
358         }
359       }
360       return false;
361     }
362     // FIXME: Ideally the isEscaped check should come at the beginning.
363     // However there are a few loose ends that need to be fixed first before
364     // we can do that. We need to make sure we are not over-conservative, so
365     // that the data accessed in-between await_suspend and symmetric transfer
366     // is always put on the stack, and also data accessed after coro.end is
367     // always put on the stack (esp the return object). To fix that, we need
368     // to:
369     //  1) Potentially treat sret as nocapture in calls
370     //  2) Special handle the return object and put it on the stack
371     //  3) Utilize lifetime.end intrinsic
372     if (PI.isEscaped())
373       return true;
374 
375     for (auto *U1 : Users)
376       for (auto *U2 : Users)
377         if (Checker.isDefinitionAcrossSuspend(*U1, U2))
378           return true;
379 
380     return false;
381   }
382 
handleMayWritellvm::coro::__anondea4afbd0111::__anondea4afbd0211::AllocaUseVisitor383   void handleMayWrite(const Instruction &I) {
384     if (!DT.dominates(CoroShape.CoroBegin, &I))
385       MayWriteBeforeCoroBegin = true;
386   }
387 
usedAfterCoroBeginllvm::coro::__anondea4afbd0111::__anondea4afbd0211::AllocaUseVisitor388   bool usedAfterCoroBegin(Instruction &I) {
389     for (auto &U : I.uses())
390       if (DT.dominates(CoroShape.CoroBegin, U))
391         return true;
392     return false;
393   }
394 
handleAliasllvm::coro::__anondea4afbd0111::__anondea4afbd0211::AllocaUseVisitor395   void handleAlias(Instruction &I) {
396     // We track all aliases created prior to CoroBegin but used after.
397     // These aliases may need to be recreated after CoroBegin if the alloca
398     // need to live on the frame.
399     if (DT.dominates(CoroShape.CoroBegin, &I) || !usedAfterCoroBegin(I))
400       return;
401 
402     if (!IsOffsetKnown) {
403       AliasOffetMap[&I].reset();
404     } else {
405       auto [Itr, Inserted] = AliasOffetMap.try_emplace(&I, Offset);
406       if (!Inserted && Itr->second && *Itr->second != Offset) {
407         // If we have seen two different possible values for this alias, we set
408         // it to empty.
409         Itr->second.reset();
410       }
411     }
412   }
413 };
414 } // namespace
415 
collectFrameAlloca(AllocaInst * AI,const coro::Shape & Shape,const SuspendCrossingInfo & Checker,SmallVectorImpl<AllocaInfo> & Allocas,const DominatorTree & DT)416 static void collectFrameAlloca(AllocaInst *AI, const coro::Shape &Shape,
417                                const SuspendCrossingInfo &Checker,
418                                SmallVectorImpl<AllocaInfo> &Allocas,
419                                const DominatorTree &DT) {
420   if (Shape.CoroSuspends.empty())
421     return;
422 
423   // The PromiseAlloca will be specially handled since it needs to be in a
424   // fixed position in the frame.
425   if (AI == Shape.SwitchLowering.PromiseAlloca)
426     return;
427 
428   // The __coro_gro alloca should outlive the promise, make sure we
429   // keep it outside the frame.
430   if (AI->hasMetadata(LLVMContext::MD_coro_outside_frame))
431     return;
432 
433   // The code that uses lifetime.start intrinsic does not work for functions
434   // with loops without exit. Disable it on ABIs we know to generate such
435   // code.
436   bool ShouldUseLifetimeStartInfo =
437       (Shape.ABI != coro::ABI::Async && Shape.ABI != coro::ABI::Retcon &&
438        Shape.ABI != coro::ABI::RetconOnce);
439   AllocaUseVisitor Visitor{AI->getDataLayout(), DT, Shape, Checker,
440                            ShouldUseLifetimeStartInfo};
441   Visitor.visitPtr(*AI);
442   if (!Visitor.getShouldLiveOnFrame())
443     return;
444   Allocas.emplace_back(AI, Visitor.getAliasesCopy(),
445                        Visitor.getMayWriteBeforeCoroBegin());
446 }
447 
448 } // namespace
449 
collectSpillsFromArgs(SpillInfo & Spills,Function & F,const SuspendCrossingInfo & Checker)450 void collectSpillsFromArgs(SpillInfo &Spills, Function &F,
451                            const SuspendCrossingInfo &Checker) {
452   // Collect the spills for arguments and other not-materializable values.
453   for (Argument &A : F.args())
454     for (User *U : A.users())
455       if (Checker.isDefinitionAcrossSuspend(A, U))
456         Spills[&A].push_back(cast<Instruction>(U));
457 }
458 
collectSpillsAndAllocasFromInsts(SpillInfo & Spills,SmallVector<AllocaInfo,8> & Allocas,SmallVector<Instruction *,4> & DeadInstructions,SmallVector<CoroAllocaAllocInst *,4> & LocalAllocas,Function & F,const SuspendCrossingInfo & Checker,const DominatorTree & DT,const coro::Shape & Shape)459 void collectSpillsAndAllocasFromInsts(
460     SpillInfo &Spills, SmallVector<AllocaInfo, 8> &Allocas,
461     SmallVector<Instruction *, 4> &DeadInstructions,
462     SmallVector<CoroAllocaAllocInst *, 4> &LocalAllocas, Function &F,
463     const SuspendCrossingInfo &Checker, const DominatorTree &DT,
464     const coro::Shape &Shape) {
465 
466   for (Instruction &I : instructions(F)) {
467     // Values returned from coroutine structure intrinsics should not be part
468     // of the Coroutine Frame.
469     if (isNonSpilledIntrinsic(I) || &I == Shape.CoroBegin)
470       continue;
471 
472     // Handle alloca.alloc specially here.
473     if (auto AI = dyn_cast<CoroAllocaAllocInst>(&I)) {
474       // Check whether the alloca's lifetime is bounded by suspend points.
475       if (isLocalAlloca(AI)) {
476         LocalAllocas.push_back(AI);
477         continue;
478       }
479 
480       // If not, do a quick rewrite of the alloca and then add spills of
481       // the rewritten value. The rewrite doesn't invalidate anything in
482       // Spills because the other alloca intrinsics have no other operands
483       // besides AI, and it doesn't invalidate the iteration because we delay
484       // erasing AI.
485       auto Alloc = lowerNonLocalAlloca(AI, Shape, DeadInstructions);
486 
487       for (User *U : Alloc->users()) {
488         if (Checker.isDefinitionAcrossSuspend(*Alloc, U))
489           Spills[Alloc].push_back(cast<Instruction>(U));
490       }
491       continue;
492     }
493 
494     // Ignore alloca.get; we process this as part of coro.alloca.alloc.
495     if (isa<CoroAllocaGetInst>(I))
496       continue;
497 
498     if (auto *AI = dyn_cast<AllocaInst>(&I)) {
499       collectFrameAlloca(AI, Shape, Checker, Allocas, DT);
500       continue;
501     }
502 
503     for (User *U : I.users())
504       if (Checker.isDefinitionAcrossSuspend(I, U)) {
505         // We cannot spill a token.
506         if (I.getType()->isTokenTy())
507           report_fatal_error(
508               "token definition is separated from the use by a suspend point");
509         Spills[&I].push_back(cast<Instruction>(U));
510       }
511   }
512 }
513 
collectSpillsFromDbgInfo(SpillInfo & Spills,Function & F,const SuspendCrossingInfo & Checker)514 void collectSpillsFromDbgInfo(SpillInfo &Spills, Function &F,
515                               const SuspendCrossingInfo &Checker) {
516   // We don't want the layout of coroutine frame to be affected
517   // by debug information. So we only choose to salvage DbgValueInst for
518   // whose value is already in the frame.
519   // We would handle the dbg.values for allocas specially
520   for (auto &Iter : Spills) {
521     auto *V = Iter.first;
522     SmallVector<DbgValueInst *, 16> DVIs;
523     SmallVector<DbgVariableRecord *, 16> DVRs;
524     findDbgValues(DVIs, V, &DVRs);
525     for (DbgValueInst *DVI : DVIs)
526       if (Checker.isDefinitionAcrossSuspend(*V, DVI))
527         Spills[V].push_back(DVI);
528     // Add the instructions which carry debug info that is in the frame.
529     for (DbgVariableRecord *DVR : DVRs)
530       if (Checker.isDefinitionAcrossSuspend(*V, DVR->Marker->MarkedInstr))
531         Spills[V].push_back(DVR->Marker->MarkedInstr);
532   }
533 }
534 
535 /// Async and Retcon{Once} conventions assume that all spill uses can be sunk
536 /// after the coro.begin intrinsic.
sinkSpillUsesAfterCoroBegin(const DominatorTree & Dom,CoroBeginInst * CoroBegin,coro::SpillInfo & Spills,SmallVectorImpl<coro::AllocaInfo> & Allocas)537 void sinkSpillUsesAfterCoroBegin(const DominatorTree &Dom,
538                                  CoroBeginInst *CoroBegin,
539                                  coro::SpillInfo &Spills,
540                                  SmallVectorImpl<coro::AllocaInfo> &Allocas) {
541   SmallSetVector<Instruction *, 32> ToMove;
542   SmallVector<Instruction *, 32> Worklist;
543 
544   // Collect all users that precede coro.begin.
545   auto collectUsers = [&](Value *Def) {
546     for (User *U : Def->users()) {
547       auto Inst = cast<Instruction>(U);
548       if (Inst->getParent() != CoroBegin->getParent() ||
549           Dom.dominates(CoroBegin, Inst))
550         continue;
551       if (ToMove.insert(Inst))
552         Worklist.push_back(Inst);
553     }
554   };
555   for (auto &I : Spills)
556     collectUsers(I.first);
557   for (auto &I : Allocas)
558     collectUsers(I.Alloca);
559 
560   // Recursively collect users before coro.begin.
561   while (!Worklist.empty()) {
562     auto *Def = Worklist.pop_back_val();
563     for (User *U : Def->users()) {
564       auto Inst = cast<Instruction>(U);
565       if (Dom.dominates(CoroBegin, Inst))
566         continue;
567       if (ToMove.insert(Inst))
568         Worklist.push_back(Inst);
569     }
570   }
571 
572   // Sort by dominance.
573   SmallVector<Instruction *, 64> InsertionList(ToMove.begin(), ToMove.end());
574   llvm::sort(InsertionList, [&Dom](Instruction *A, Instruction *B) -> bool {
575     // If a dominates b it should precede (<) b.
576     return Dom.dominates(A, B);
577   });
578 
579   Instruction *InsertPt = CoroBegin->getNextNode();
580   for (Instruction *Inst : InsertionList)
581     Inst->moveBefore(InsertPt->getIterator());
582 }
583 
getSpillInsertionPt(const coro::Shape & Shape,Value * Def,const DominatorTree & DT)584 BasicBlock::iterator getSpillInsertionPt(const coro::Shape &Shape, Value *Def,
585                                          const DominatorTree &DT) {
586   BasicBlock::iterator InsertPt;
587   if (auto *Arg = dyn_cast<Argument>(Def)) {
588     // For arguments, we will place the store instruction right after
589     // the coroutine frame pointer instruction, i.e. coro.begin.
590     InsertPt = Shape.getInsertPtAfterFramePtr();
591 
592     // If we're spilling an Argument, make sure we clear 'captures'
593     // from the coroutine function.
594     Arg->getParent()->removeParamAttr(Arg->getArgNo(), Attribute::Captures);
595   } else if (auto *CSI = dyn_cast<AnyCoroSuspendInst>(Def)) {
596     // Don't spill immediately after a suspend; splitting assumes
597     // that the suspend will be followed by a branch.
598     InsertPt = CSI->getParent()->getSingleSuccessor()->getFirstNonPHIIt();
599   } else {
600     auto *I = cast<Instruction>(Def);
601     if (!DT.dominates(Shape.CoroBegin, I)) {
602       // If it is not dominated by CoroBegin, then spill should be
603       // inserted immediately after CoroFrame is computed.
604       InsertPt = Shape.getInsertPtAfterFramePtr();
605     } else if (auto *II = dyn_cast<InvokeInst>(I)) {
606       // If we are spilling the result of the invoke instruction, split
607       // the normal edge and insert the spill in the new block.
608       auto *NewBB = SplitEdge(II->getParent(), II->getNormalDest());
609       InsertPt = NewBB->getTerminator()->getIterator();
610     } else if (isa<PHINode>(I)) {
611       // Skip the PHINodes and EH pads instructions.
612       BasicBlock *DefBlock = I->getParent();
613       if (auto *CSI = dyn_cast<CatchSwitchInst>(DefBlock->getTerminator()))
614         InsertPt = splitBeforeCatchSwitch(CSI)->getIterator();
615       else
616         InsertPt = DefBlock->getFirstInsertionPt();
617     } else {
618       assert(!I->isTerminator() && "unexpected terminator");
619       // For all other values, the spill is placed immediately after
620       // the definition.
621       InsertPt = I->getNextNode()->getIterator();
622     }
623   }
624 
625   return InsertPt;
626 }
627 
628 } // End namespace coro.
629 
630 } // End namespace llvm.
631