1 //===- CoroFrame.cpp - Builds and manipulates coroutine frame -------------===//
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 file contains classes used to discover if for a particular value
9 // there from sue to definition that crosses a suspend block.
10 //
11 // Using the information discovered we form a Coroutine Frame structure to
12 // contain those values. All uses of those values are replaced with appropriate
13 // GEP + load from the coroutine frame. At the point of the definition we spill
14 // the value into the coroutine frame.
15 //
16 // TODO: pack values tightly using liveness info.
17 //===----------------------------------------------------------------------===//
18
19 #include "CoroInternal.h"
20 #include "llvm/ADT/BitVector.h"
21 #include "llvm/Transforms/Utils/Local.h"
22 #include "llvm/Config/llvm-config.h"
23 #include "llvm/IR/CFG.h"
24 #include "llvm/IR/Dominators.h"
25 #include "llvm/IR/IRBuilder.h"
26 #include "llvm/IR/InstIterator.h"
27 #include "llvm/Support/Debug.h"
28 #include "llvm/Support/MathExtras.h"
29 #include "llvm/Support/circular_raw_ostream.h"
30 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
31
32 using namespace llvm;
33
34 // The "coro-suspend-crossing" flag is very noisy. There is another debug type,
35 // "coro-frame", which results in leaner debug spew.
36 #define DEBUG_TYPE "coro-suspend-crossing"
37
38 enum { SmallVectorThreshold = 32 };
39
40 // Provides two way mapping between the blocks and numbers.
41 namespace {
42 class BlockToIndexMapping {
43 SmallVector<BasicBlock *, SmallVectorThreshold> V;
44
45 public:
size() const46 size_t size() const { return V.size(); }
47
BlockToIndexMapping(Function & F)48 BlockToIndexMapping(Function &F) {
49 for (BasicBlock &BB : F)
50 V.push_back(&BB);
51 llvm::sort(V);
52 }
53
blockToIndex(BasicBlock * BB) const54 size_t blockToIndex(BasicBlock *BB) const {
55 auto *I = llvm::lower_bound(V, BB);
56 assert(I != V.end() && *I == BB && "BasicBlockNumberng: Unknown block");
57 return I - V.begin();
58 }
59
indexToBlock(unsigned Index) const60 BasicBlock *indexToBlock(unsigned Index) const { return V[Index]; }
61 };
62 } // end anonymous namespace
63
64 // The SuspendCrossingInfo maintains data that allows to answer a question
65 // whether given two BasicBlocks A and B there is a path from A to B that
66 // passes through a suspend point.
67 //
68 // For every basic block 'i' it maintains a BlockData that consists of:
69 // Consumes: a bit vector which contains a set of indices of blocks that can
70 // reach block 'i'
71 // Kills: a bit vector which contains a set of indices of blocks that can
72 // reach block 'i', but one of the path will cross a suspend point
73 // Suspend: a boolean indicating whether block 'i' contains a suspend point.
74 // End: a boolean indicating whether block 'i' contains a coro.end intrinsic.
75 //
76 namespace {
77 struct SuspendCrossingInfo {
78 BlockToIndexMapping Mapping;
79
80 struct BlockData {
81 BitVector Consumes;
82 BitVector Kills;
83 bool Suspend = false;
84 bool End = false;
85 };
86 SmallVector<BlockData, SmallVectorThreshold> Block;
87
successors__anon1ea1874c0311::SuspendCrossingInfo88 iterator_range<succ_iterator> successors(BlockData const &BD) const {
89 BasicBlock *BB = Mapping.indexToBlock(&BD - &Block[0]);
90 return llvm::successors(BB);
91 }
92
getBlockData__anon1ea1874c0311::SuspendCrossingInfo93 BlockData &getBlockData(BasicBlock *BB) {
94 return Block[Mapping.blockToIndex(BB)];
95 }
96
97 void dump() const;
98 void dump(StringRef Label, BitVector const &BV) const;
99
100 SuspendCrossingInfo(Function &F, coro::Shape &Shape);
101
hasPathCrossingSuspendPoint__anon1ea1874c0311::SuspendCrossingInfo102 bool hasPathCrossingSuspendPoint(BasicBlock *DefBB, BasicBlock *UseBB) const {
103 size_t const DefIndex = Mapping.blockToIndex(DefBB);
104 size_t const UseIndex = Mapping.blockToIndex(UseBB);
105
106 assert(Block[UseIndex].Consumes[DefIndex] && "use must consume def");
107 bool const Result = Block[UseIndex].Kills[DefIndex];
108 LLVM_DEBUG(dbgs() << UseBB->getName() << " => " << DefBB->getName()
109 << " answer is " << Result << "\n");
110 return Result;
111 }
112
isDefinitionAcrossSuspend__anon1ea1874c0311::SuspendCrossingInfo113 bool isDefinitionAcrossSuspend(BasicBlock *DefBB, User *U) const {
114 auto *I = cast<Instruction>(U);
115
116 // We rewrote PHINodes, so that only the ones with exactly one incoming
117 // value need to be analyzed.
118 if (auto *PN = dyn_cast<PHINode>(I))
119 if (PN->getNumIncomingValues() > 1)
120 return false;
121
122 BasicBlock *UseBB = I->getParent();
123 return hasPathCrossingSuspendPoint(DefBB, UseBB);
124 }
125
isDefinitionAcrossSuspend__anon1ea1874c0311::SuspendCrossingInfo126 bool isDefinitionAcrossSuspend(Argument &A, User *U) const {
127 return isDefinitionAcrossSuspend(&A.getParent()->getEntryBlock(), U);
128 }
129
isDefinitionAcrossSuspend__anon1ea1874c0311::SuspendCrossingInfo130 bool isDefinitionAcrossSuspend(Instruction &I, User *U) const {
131 return isDefinitionAcrossSuspend(I.getParent(), U);
132 }
133 };
134 } // end anonymous namespace
135
136 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
dump(StringRef Label,BitVector const & BV) const137 LLVM_DUMP_METHOD void SuspendCrossingInfo::dump(StringRef Label,
138 BitVector const &BV) const {
139 dbgs() << Label << ":";
140 for (size_t I = 0, N = BV.size(); I < N; ++I)
141 if (BV[I])
142 dbgs() << " " << Mapping.indexToBlock(I)->getName();
143 dbgs() << "\n";
144 }
145
dump() const146 LLVM_DUMP_METHOD void SuspendCrossingInfo::dump() const {
147 for (size_t I = 0, N = Block.size(); I < N; ++I) {
148 BasicBlock *const B = Mapping.indexToBlock(I);
149 dbgs() << B->getName() << ":\n";
150 dump(" Consumes", Block[I].Consumes);
151 dump(" Kills", Block[I].Kills);
152 }
153 dbgs() << "\n";
154 }
155 #endif
156
SuspendCrossingInfo(Function & F,coro::Shape & Shape)157 SuspendCrossingInfo::SuspendCrossingInfo(Function &F, coro::Shape &Shape)
158 : Mapping(F) {
159 const size_t N = Mapping.size();
160 Block.resize(N);
161
162 // Initialize every block so that it consumes itself
163 for (size_t I = 0; I < N; ++I) {
164 auto &B = Block[I];
165 B.Consumes.resize(N);
166 B.Kills.resize(N);
167 B.Consumes.set(I);
168 }
169
170 // Mark all CoroEnd Blocks. We do not propagate Kills beyond coro.ends as
171 // the code beyond coro.end is reachable during initial invocation of the
172 // coroutine.
173 for (auto *CE : Shape.CoroEnds)
174 getBlockData(CE->getParent()).End = true;
175
176 // Mark all suspend blocks and indicate that they kill everything they
177 // consume. Note, that crossing coro.save also requires a spill, as any code
178 // between coro.save and coro.suspend may resume the coroutine and all of the
179 // state needs to be saved by that time.
180 auto markSuspendBlock = [&](IntrinsicInst *BarrierInst) {
181 BasicBlock *SuspendBlock = BarrierInst->getParent();
182 auto &B = getBlockData(SuspendBlock);
183 B.Suspend = true;
184 B.Kills |= B.Consumes;
185 };
186 for (CoroSuspendInst *CSI : Shape.CoroSuspends) {
187 markSuspendBlock(CSI);
188 markSuspendBlock(CSI->getCoroSave());
189 }
190
191 // Iterate propagating consumes and kills until they stop changing.
192 int Iteration = 0;
193 (void)Iteration;
194
195 bool Changed;
196 do {
197 LLVM_DEBUG(dbgs() << "iteration " << ++Iteration);
198 LLVM_DEBUG(dbgs() << "==============\n");
199
200 Changed = false;
201 for (size_t I = 0; I < N; ++I) {
202 auto &B = Block[I];
203 for (BasicBlock *SI : successors(B)) {
204
205 auto SuccNo = Mapping.blockToIndex(SI);
206
207 // Saved Consumes and Kills bitsets so that it is easy to see
208 // if anything changed after propagation.
209 auto &S = Block[SuccNo];
210 auto SavedConsumes = S.Consumes;
211 auto SavedKills = S.Kills;
212
213 // Propagate Kills and Consumes from block B into its successor S.
214 S.Consumes |= B.Consumes;
215 S.Kills |= B.Kills;
216
217 // If block B is a suspend block, it should propagate kills into the
218 // its successor for every block B consumes.
219 if (B.Suspend) {
220 S.Kills |= B.Consumes;
221 }
222 if (S.Suspend) {
223 // If block S is a suspend block, it should kill all of the blocks it
224 // consumes.
225 S.Kills |= S.Consumes;
226 } else if (S.End) {
227 // If block S is an end block, it should not propagate kills as the
228 // blocks following coro.end() are reached during initial invocation
229 // of the coroutine while all the data are still available on the
230 // stack or in the registers.
231 S.Kills.reset();
232 } else {
233 // This is reached when S block it not Suspend nor coro.end and it
234 // need to make sure that it is not in the kill set.
235 S.Kills.reset(SuccNo);
236 }
237
238 // See if anything changed.
239 Changed |= (S.Kills != SavedKills) || (S.Consumes != SavedConsumes);
240
241 if (S.Kills != SavedKills) {
242 LLVM_DEBUG(dbgs() << "\nblock " << I << " follower " << SI->getName()
243 << "\n");
244 LLVM_DEBUG(dump("S.Kills", S.Kills));
245 LLVM_DEBUG(dump("SavedKills", SavedKills));
246 }
247 if (S.Consumes != SavedConsumes) {
248 LLVM_DEBUG(dbgs() << "\nblock " << I << " follower " << SI << "\n");
249 LLVM_DEBUG(dump("S.Consume", S.Consumes));
250 LLVM_DEBUG(dump("SavedCons", SavedConsumes));
251 }
252 }
253 }
254 } while (Changed);
255 LLVM_DEBUG(dump());
256 }
257
258 #undef DEBUG_TYPE // "coro-suspend-crossing"
259 #define DEBUG_TYPE "coro-frame"
260
261 // We build up the list of spills for every case where a use is separated
262 // from the definition by a suspend point.
263
264 namespace {
265 class Spill {
266 Value *Def = nullptr;
267 Instruction *User = nullptr;
268 unsigned FieldNo = 0;
269
270 public:
Spill(Value * Def,llvm::User * U)271 Spill(Value *Def, llvm::User *U) : Def(Def), User(cast<Instruction>(U)) {}
272
def() const273 Value *def() const { return Def; }
user() const274 Instruction *user() const { return User; }
userBlock() const275 BasicBlock *userBlock() const { return User->getParent(); }
276
277 // Note that field index is stored in the first SpillEntry for a particular
278 // definition. Subsequent mentions of a defintion do not have fieldNo
279 // assigned. This works out fine as the users of Spills capture the info about
280 // the definition the first time they encounter it. Consider refactoring
281 // SpillInfo into two arrays to normalize the spill representation.
fieldIndex() const282 unsigned fieldIndex() const {
283 assert(FieldNo && "Accessing unassigned field");
284 return FieldNo;
285 }
setFieldIndex(unsigned FieldNumber)286 void setFieldIndex(unsigned FieldNumber) {
287 assert(!FieldNo && "Reassigning field number");
288 FieldNo = FieldNumber;
289 }
290 };
291 } // namespace
292
293 // Note that there may be more than one record with the same value of Def in
294 // the SpillInfo vector.
295 using SpillInfo = SmallVector<Spill, 8>;
296
297 #ifndef NDEBUG
dump(StringRef Title,SpillInfo const & Spills)298 static void dump(StringRef Title, SpillInfo const &Spills) {
299 dbgs() << "------------- " << Title << "--------------\n";
300 Value *CurrentValue = nullptr;
301 for (auto const &E : Spills) {
302 if (CurrentValue != E.def()) {
303 CurrentValue = E.def();
304 CurrentValue->dump();
305 }
306 dbgs() << " user: ";
307 E.user()->dump();
308 }
309 }
310 #endif
311
312 namespace {
313 // We cannot rely solely on natural alignment of a type when building a
314 // coroutine frame and if the alignment specified on the Alloca instruction
315 // differs from the natural alignment of the alloca type we will need to insert
316 // padding.
317 struct PaddingCalculator {
318 const DataLayout &DL;
319 LLVMContext &Context;
320 unsigned StructSize = 0;
321
PaddingCalculator__anon1ea1874c0611::PaddingCalculator322 PaddingCalculator(LLVMContext &Context, DataLayout const &DL)
323 : DL(DL), Context(Context) {}
324
325 // Replicate the logic from IR/DataLayout.cpp to match field offset
326 // computation for LLVM structs.
addType__anon1ea1874c0611::PaddingCalculator327 void addType(Type *Ty) {
328 unsigned TyAlign = DL.getABITypeAlignment(Ty);
329 if ((StructSize & (TyAlign - 1)) != 0)
330 StructSize = alignTo(StructSize, TyAlign);
331
332 StructSize += DL.getTypeAllocSize(Ty); // Consume space for this data item.
333 }
334
addTypes__anon1ea1874c0611::PaddingCalculator335 void addTypes(SmallVectorImpl<Type *> const &Types) {
336 for (auto *Ty : Types)
337 addType(Ty);
338 }
339
computePadding__anon1ea1874c0611::PaddingCalculator340 unsigned computePadding(Type *Ty, unsigned ForcedAlignment) {
341 unsigned TyAlign = DL.getABITypeAlignment(Ty);
342 auto Natural = alignTo(StructSize, TyAlign);
343 auto Forced = alignTo(StructSize, ForcedAlignment);
344
345 // Return how many bytes of padding we need to insert.
346 if (Natural != Forced)
347 return std::max(Natural, Forced) - StructSize;
348
349 // Rely on natural alignment.
350 return 0;
351 }
352
353 // If padding required, return the padding field type to insert.
getPaddingType__anon1ea1874c0611::PaddingCalculator354 ArrayType *getPaddingType(Type *Ty, unsigned ForcedAlignment) {
355 if (auto Padding = computePadding(Ty, ForcedAlignment))
356 return ArrayType::get(Type::getInt8Ty(Context), Padding);
357
358 return nullptr;
359 }
360 };
361 } // namespace
362
363 // Build a struct that will keep state for an active coroutine.
364 // struct f.frame {
365 // ResumeFnTy ResumeFnAddr;
366 // ResumeFnTy DestroyFnAddr;
367 // int ResumeIndex;
368 // ... promise (if present) ...
369 // ... spills ...
370 // };
buildFrameType(Function & F,coro::Shape & Shape,SpillInfo & Spills)371 static StructType *buildFrameType(Function &F, coro::Shape &Shape,
372 SpillInfo &Spills) {
373 LLVMContext &C = F.getContext();
374 const DataLayout &DL = F.getParent()->getDataLayout();
375 PaddingCalculator Padder(C, DL);
376 SmallString<32> Name(F.getName());
377 Name.append(".Frame");
378 StructType *FrameTy = StructType::create(C, Name);
379 auto *FramePtrTy = FrameTy->getPointerTo();
380 auto *FnTy = FunctionType::get(Type::getVoidTy(C), FramePtrTy,
381 /*isVarArg=*/false);
382 auto *FnPtrTy = FnTy->getPointerTo();
383
384 // Figure out how wide should be an integer type storing the suspend index.
385 unsigned IndexBits = std::max(1U, Log2_64_Ceil(Shape.CoroSuspends.size()));
386 Type *PromiseType = Shape.PromiseAlloca
387 ? Shape.PromiseAlloca->getType()->getElementType()
388 : Type::getInt1Ty(C);
389 SmallVector<Type *, 8> Types{FnPtrTy, FnPtrTy, PromiseType,
390 Type::getIntNTy(C, IndexBits)};
391 Value *CurrentDef = nullptr;
392
393 Padder.addTypes(Types);
394
395 // Create an entry for every spilled value.
396 for (auto &S : Spills) {
397 if (CurrentDef == S.def())
398 continue;
399
400 CurrentDef = S.def();
401 // PromiseAlloca was already added to Types array earlier.
402 if (CurrentDef == Shape.PromiseAlloca)
403 continue;
404
405 uint64_t Count = 1;
406 Type *Ty = nullptr;
407 if (auto *AI = dyn_cast<AllocaInst>(CurrentDef)) {
408 Ty = AI->getAllocatedType();
409 if (unsigned AllocaAlignment = AI->getAlignment()) {
410 // If alignment is specified in alloca, see if we need to insert extra
411 // padding.
412 if (auto PaddingTy = Padder.getPaddingType(Ty, AllocaAlignment)) {
413 Types.push_back(PaddingTy);
414 Padder.addType(PaddingTy);
415 }
416 }
417 if (auto *CI = dyn_cast<ConstantInt>(AI->getArraySize()))
418 Count = CI->getValue().getZExtValue();
419 else
420 report_fatal_error("Coroutines cannot handle non static allocas yet");
421 } else {
422 Ty = CurrentDef->getType();
423 }
424 S.setFieldIndex(Types.size());
425 if (Count == 1)
426 Types.push_back(Ty);
427 else
428 Types.push_back(ArrayType::get(Ty, Count));
429 Padder.addType(Ty);
430 }
431 FrameTy->setBody(Types);
432
433 return FrameTy;
434 }
435
436 // We need to make room to insert a spill after initial PHIs, but before
437 // catchswitch instruction. Placing it before violates the requirement that
438 // catchswitch, like all other EHPads must be the first nonPHI in a block.
439 //
440 // Split away catchswitch into a separate block and insert in its place:
441 //
442 // cleanuppad <InsertPt> cleanupret.
443 //
444 // cleanupret instruction will act as an insert point for the spill.
splitBeforeCatchSwitch(CatchSwitchInst * CatchSwitch)445 static Instruction *splitBeforeCatchSwitch(CatchSwitchInst *CatchSwitch) {
446 BasicBlock *CurrentBlock = CatchSwitch->getParent();
447 BasicBlock *NewBlock = CurrentBlock->splitBasicBlock(CatchSwitch);
448 CurrentBlock->getTerminator()->eraseFromParent();
449
450 auto *CleanupPad =
451 CleanupPadInst::Create(CatchSwitch->getParentPad(), {}, "", CurrentBlock);
452 auto *CleanupRet =
453 CleanupReturnInst::Create(CleanupPad, NewBlock, CurrentBlock);
454 return CleanupRet;
455 }
456
457 // Replace all alloca and SSA values that are accessed across suspend points
458 // with GetElementPointer from coroutine frame + loads and stores. Create an
459 // AllocaSpillBB that will become the new entry block for the resume parts of
460 // the coroutine:
461 //
462 // %hdl = coro.begin(...)
463 // whatever
464 //
465 // becomes:
466 //
467 // %hdl = coro.begin(...)
468 // %FramePtr = bitcast i8* hdl to %f.frame*
469 // br label %AllocaSpillBB
470 //
471 // AllocaSpillBB:
472 // ; geps corresponding to allocas that were moved to coroutine frame
473 // br label PostSpill
474 //
475 // PostSpill:
476 // whatever
477 //
478 //
insertSpills(SpillInfo & Spills,coro::Shape & Shape)479 static Instruction *insertSpills(SpillInfo &Spills, coro::Shape &Shape) {
480 auto *CB = Shape.CoroBegin;
481 LLVMContext &C = CB->getContext();
482 IRBuilder<> Builder(CB->getNextNode());
483 StructType *FrameTy = Shape.FrameTy;
484 PointerType *FramePtrTy = FrameTy->getPointerTo();
485 auto *FramePtr =
486 cast<Instruction>(Builder.CreateBitCast(CB, FramePtrTy, "FramePtr"));
487
488 Value *CurrentValue = nullptr;
489 BasicBlock *CurrentBlock = nullptr;
490 Value *CurrentReload = nullptr;
491 unsigned Index = 0; // Proper field number will be read from field definition.
492
493 // We need to keep track of any allocas that need "spilling"
494 // since they will live in the coroutine frame now, all access to them
495 // need to be changed, not just the access across suspend points
496 // we remember allocas and their indices to be handled once we processed
497 // all the spills.
498 SmallVector<std::pair<AllocaInst *, unsigned>, 4> Allocas;
499 // Promise alloca (if present) has a fixed field number (Shape::PromiseField)
500 if (Shape.PromiseAlloca)
501 Allocas.emplace_back(Shape.PromiseAlloca, coro::Shape::PromiseField);
502
503 // Create a GEP with the given index into the coroutine frame for the original
504 // value Orig. Appends an extra 0 index for array-allocas, preserving the
505 // original type.
506 auto GetFramePointer = [&](uint32_t Index, Value *Orig) -> Value * {
507 SmallVector<Value *, 3> Indices = {
508 ConstantInt::get(Type::getInt32Ty(C), 0),
509 ConstantInt::get(Type::getInt32Ty(C), Index),
510 };
511
512 if (auto *AI = dyn_cast<AllocaInst>(Orig)) {
513 if (auto *CI = dyn_cast<ConstantInt>(AI->getArraySize())) {
514 auto Count = CI->getValue().getZExtValue();
515 if (Count > 1) {
516 Indices.push_back(ConstantInt::get(Type::getInt32Ty(C), 0));
517 }
518 } else {
519 report_fatal_error("Coroutines cannot handle non static allocas yet");
520 }
521 }
522
523 return Builder.CreateInBoundsGEP(FrameTy, FramePtr, Indices);
524 };
525
526 // Create a load instruction to reload the spilled value from the coroutine
527 // frame.
528 auto CreateReload = [&](Instruction *InsertBefore) {
529 assert(Index && "accessing unassigned field number");
530 Builder.SetInsertPoint(InsertBefore);
531
532 auto *G = GetFramePointer(Index, CurrentValue);
533 G->setName(CurrentValue->getName() + Twine(".reload.addr"));
534
535 return isa<AllocaInst>(CurrentValue)
536 ? G
537 : Builder.CreateLoad(FrameTy->getElementType(Index), G,
538 CurrentValue->getName() + Twine(".reload"));
539 };
540
541 for (auto const &E : Spills) {
542 // If we have not seen the value, generate a spill.
543 if (CurrentValue != E.def()) {
544 CurrentValue = E.def();
545 CurrentBlock = nullptr;
546 CurrentReload = nullptr;
547
548 Index = E.fieldIndex();
549
550 if (auto *AI = dyn_cast<AllocaInst>(CurrentValue)) {
551 // Spilled AllocaInst will be replaced with GEP from the coroutine frame
552 // there is no spill required.
553 Allocas.emplace_back(AI, Index);
554 if (!AI->isStaticAlloca())
555 report_fatal_error("Coroutines cannot handle non static allocas yet");
556 } else {
557 // Otherwise, create a store instruction storing the value into the
558 // coroutine frame.
559
560 Instruction *InsertPt = nullptr;
561 if (isa<Argument>(CurrentValue)) {
562 // For arguments, we will place the store instruction right after
563 // the coroutine frame pointer instruction, i.e. bitcast of
564 // coro.begin from i8* to %f.frame*.
565 InsertPt = FramePtr->getNextNode();
566 } else if (auto *II = dyn_cast<InvokeInst>(CurrentValue)) {
567 // If we are spilling the result of the invoke instruction, split the
568 // normal edge and insert the spill in the new block.
569 auto NewBB = SplitEdge(II->getParent(), II->getNormalDest());
570 InsertPt = NewBB->getTerminator();
571 } else if (dyn_cast<PHINode>(CurrentValue)) {
572 // Skip the PHINodes and EH pads instructions.
573 BasicBlock *DefBlock = cast<Instruction>(E.def())->getParent();
574 if (auto *CSI = dyn_cast<CatchSwitchInst>(DefBlock->getTerminator()))
575 InsertPt = splitBeforeCatchSwitch(CSI);
576 else
577 InsertPt = &*DefBlock->getFirstInsertionPt();
578 } else {
579 // For all other values, the spill is placed immediately after
580 // the definition.
581 assert(!cast<Instruction>(E.def())->isTerminator() &&
582 "unexpected terminator");
583 InsertPt = cast<Instruction>(E.def())->getNextNode();
584 }
585
586 Builder.SetInsertPoint(InsertPt);
587 auto *G = Builder.CreateConstInBoundsGEP2_32(
588 FrameTy, FramePtr, 0, Index,
589 CurrentValue->getName() + Twine(".spill.addr"));
590 Builder.CreateStore(CurrentValue, G);
591 }
592 }
593
594 // If we have not seen the use block, generate a reload in it.
595 if (CurrentBlock != E.userBlock()) {
596 CurrentBlock = E.userBlock();
597 CurrentReload = CreateReload(&*CurrentBlock->getFirstInsertionPt());
598 }
599
600 // If we have a single edge PHINode, remove it and replace it with a reload
601 // from the coroutine frame. (We already took care of multi edge PHINodes
602 // by rewriting them in the rewritePHIs function).
603 if (auto *PN = dyn_cast<PHINode>(E.user())) {
604 assert(PN->getNumIncomingValues() == 1 && "unexpected number of incoming "
605 "values in the PHINode");
606 PN->replaceAllUsesWith(CurrentReload);
607 PN->eraseFromParent();
608 continue;
609 }
610
611 // Replace all uses of CurrentValue in the current instruction with reload.
612 E.user()->replaceUsesOfWith(CurrentValue, CurrentReload);
613 }
614
615 BasicBlock *FramePtrBB = FramePtr->getParent();
616 Shape.AllocaSpillBlock =
617 FramePtrBB->splitBasicBlock(FramePtr->getNextNode(), "AllocaSpillBB");
618 Shape.AllocaSpillBlock->splitBasicBlock(&Shape.AllocaSpillBlock->front(),
619 "PostSpill");
620
621 Builder.SetInsertPoint(&Shape.AllocaSpillBlock->front());
622 // If we found any allocas, replace all of their remaining uses with Geps.
623 for (auto &P : Allocas) {
624 auto *G = GetFramePointer(P.second, P.first);
625
626 // We are not using ReplaceInstWithInst(P.first, cast<Instruction>(G)) here,
627 // as we are changing location of the instruction.
628 G->takeName(P.first);
629 P.first->replaceAllUsesWith(G);
630 P.first->eraseFromParent();
631 }
632 return FramePtr;
633 }
634
635 // Sets the unwind edge of an instruction to a particular successor.
setUnwindEdgeTo(Instruction * TI,BasicBlock * Succ)636 static void setUnwindEdgeTo(Instruction *TI, BasicBlock *Succ) {
637 if (auto *II = dyn_cast<InvokeInst>(TI))
638 II->setUnwindDest(Succ);
639 else if (auto *CS = dyn_cast<CatchSwitchInst>(TI))
640 CS->setUnwindDest(Succ);
641 else if (auto *CR = dyn_cast<CleanupReturnInst>(TI))
642 CR->setUnwindDest(Succ);
643 else
644 llvm_unreachable("unexpected terminator instruction");
645 }
646
647 // Replaces all uses of OldPred with the NewPred block in all PHINodes in a
648 // block.
updatePhiNodes(BasicBlock * DestBB,BasicBlock * OldPred,BasicBlock * NewPred,PHINode * LandingPadReplacement)649 static void updatePhiNodes(BasicBlock *DestBB, BasicBlock *OldPred,
650 BasicBlock *NewPred,
651 PHINode *LandingPadReplacement) {
652 unsigned BBIdx = 0;
653 for (BasicBlock::iterator I = DestBB->begin(); isa<PHINode>(I); ++I) {
654 PHINode *PN = cast<PHINode>(I);
655
656 // We manually update the LandingPadReplacement PHINode and it is the last
657 // PHI Node. So, if we find it, we are done.
658 if (LandingPadReplacement == PN)
659 break;
660
661 // Reuse the previous value of BBIdx if it lines up. In cases where we
662 // have multiple phi nodes with *lots* of predecessors, this is a speed
663 // win because we don't have to scan the PHI looking for TIBB. This
664 // happens because the BB list of PHI nodes are usually in the same
665 // order.
666 if (PN->getIncomingBlock(BBIdx) != OldPred)
667 BBIdx = PN->getBasicBlockIndex(OldPred);
668
669 assert(BBIdx != (unsigned)-1 && "Invalid PHI Index!");
670 PN->setIncomingBlock(BBIdx, NewPred);
671 }
672 }
673
674 // Uses SplitEdge unless the successor block is an EHPad, in which case do EH
675 // specific handling.
ehAwareSplitEdge(BasicBlock * BB,BasicBlock * Succ,LandingPadInst * OriginalPad,PHINode * LandingPadReplacement)676 static BasicBlock *ehAwareSplitEdge(BasicBlock *BB, BasicBlock *Succ,
677 LandingPadInst *OriginalPad,
678 PHINode *LandingPadReplacement) {
679 auto *PadInst = Succ->getFirstNonPHI();
680 if (!LandingPadReplacement && !PadInst->isEHPad())
681 return SplitEdge(BB, Succ);
682
683 auto *NewBB = BasicBlock::Create(BB->getContext(), "", BB->getParent(), Succ);
684 setUnwindEdgeTo(BB->getTerminator(), NewBB);
685 updatePhiNodes(Succ, BB, NewBB, LandingPadReplacement);
686
687 if (LandingPadReplacement) {
688 auto *NewLP = OriginalPad->clone();
689 auto *Terminator = BranchInst::Create(Succ, NewBB);
690 NewLP->insertBefore(Terminator);
691 LandingPadReplacement->addIncoming(NewLP, NewBB);
692 return NewBB;
693 }
694 Value *ParentPad = nullptr;
695 if (auto *FuncletPad = dyn_cast<FuncletPadInst>(PadInst))
696 ParentPad = FuncletPad->getParentPad();
697 else if (auto *CatchSwitch = dyn_cast<CatchSwitchInst>(PadInst))
698 ParentPad = CatchSwitch->getParentPad();
699 else
700 llvm_unreachable("handling for other EHPads not implemented yet");
701
702 auto *NewCleanupPad = CleanupPadInst::Create(ParentPad, {}, "", NewBB);
703 CleanupReturnInst::Create(NewCleanupPad, Succ, NewBB);
704 return NewBB;
705 }
706
rewritePHIs(BasicBlock & BB)707 static void rewritePHIs(BasicBlock &BB) {
708 // For every incoming edge we will create a block holding all
709 // incoming values in a single PHI nodes.
710 //
711 // loop:
712 // %n.val = phi i32[%n, %entry], [%inc, %loop]
713 //
714 // It will create:
715 //
716 // loop.from.entry:
717 // %n.loop.pre = phi i32 [%n, %entry]
718 // br %label loop
719 // loop.from.loop:
720 // %inc.loop.pre = phi i32 [%inc, %loop]
721 // br %label loop
722 //
723 // After this rewrite, further analysis will ignore any phi nodes with more
724 // than one incoming edge.
725
726 // TODO: Simplify PHINodes in the basic block to remove duplicate
727 // predecessors.
728
729 LandingPadInst *LandingPad = nullptr;
730 PHINode *ReplPHI = nullptr;
731 if ((LandingPad = dyn_cast_or_null<LandingPadInst>(BB.getFirstNonPHI()))) {
732 // ehAwareSplitEdge will clone the LandingPad in all the edge blocks.
733 // We replace the original landing pad with a PHINode that will collect the
734 // results from all of them.
735 ReplPHI = PHINode::Create(LandingPad->getType(), 1, "", LandingPad);
736 ReplPHI->takeName(LandingPad);
737 LandingPad->replaceAllUsesWith(ReplPHI);
738 // We will erase the original landing pad at the end of this function after
739 // ehAwareSplitEdge cloned it in the transition blocks.
740 }
741
742 SmallVector<BasicBlock *, 8> Preds(pred_begin(&BB), pred_end(&BB));
743 for (BasicBlock *Pred : Preds) {
744 auto *IncomingBB = ehAwareSplitEdge(Pred, &BB, LandingPad, ReplPHI);
745 IncomingBB->setName(BB.getName() + Twine(".from.") + Pred->getName());
746 auto *PN = cast<PHINode>(&BB.front());
747 do {
748 int Index = PN->getBasicBlockIndex(IncomingBB);
749 Value *V = PN->getIncomingValue(Index);
750 PHINode *InputV = PHINode::Create(
751 V->getType(), 1, V->getName() + Twine(".") + BB.getName(),
752 &IncomingBB->front());
753 InputV->addIncoming(V, Pred);
754 PN->setIncomingValue(Index, InputV);
755 PN = dyn_cast<PHINode>(PN->getNextNode());
756 } while (PN != ReplPHI); // ReplPHI is either null or the PHI that replaced
757 // the landing pad.
758 }
759
760 if (LandingPad) {
761 // Calls to ehAwareSplitEdge function cloned the original lading pad.
762 // No longer need it.
763 LandingPad->eraseFromParent();
764 }
765 }
766
rewritePHIs(Function & F)767 static void rewritePHIs(Function &F) {
768 SmallVector<BasicBlock *, 8> WorkList;
769
770 for (BasicBlock &BB : F)
771 if (auto *PN = dyn_cast<PHINode>(&BB.front()))
772 if (PN->getNumIncomingValues() > 1)
773 WorkList.push_back(&BB);
774
775 for (BasicBlock *BB : WorkList)
776 rewritePHIs(*BB);
777 }
778
779 // Check for instructions that we can recreate on resume as opposed to spill
780 // the result into a coroutine frame.
materializable(Instruction & V)781 static bool materializable(Instruction &V) {
782 return isa<CastInst>(&V) || isa<GetElementPtrInst>(&V) ||
783 isa<BinaryOperator>(&V) || isa<CmpInst>(&V) || isa<SelectInst>(&V);
784 }
785
786 // Check for structural coroutine intrinsics that should not be spilled into
787 // the coroutine frame.
isCoroutineStructureIntrinsic(Instruction & I)788 static bool isCoroutineStructureIntrinsic(Instruction &I) {
789 return isa<CoroIdInst>(&I) || isa<CoroSaveInst>(&I) ||
790 isa<CoroSuspendInst>(&I);
791 }
792
793 // For every use of the value that is across suspend point, recreate that value
794 // after a suspend point.
rewriteMaterializableInstructions(IRBuilder<> & IRB,SpillInfo const & Spills)795 static void rewriteMaterializableInstructions(IRBuilder<> &IRB,
796 SpillInfo const &Spills) {
797 BasicBlock *CurrentBlock = nullptr;
798 Instruction *CurrentMaterialization = nullptr;
799 Instruction *CurrentDef = nullptr;
800
801 for (auto const &E : Spills) {
802 // If it is a new definition, update CurrentXXX variables.
803 if (CurrentDef != E.def()) {
804 CurrentDef = cast<Instruction>(E.def());
805 CurrentBlock = nullptr;
806 CurrentMaterialization = nullptr;
807 }
808
809 // If we have not seen this block, materialize the value.
810 if (CurrentBlock != E.userBlock()) {
811 CurrentBlock = E.userBlock();
812 CurrentMaterialization = cast<Instruction>(CurrentDef)->clone();
813 CurrentMaterialization->setName(CurrentDef->getName());
814 CurrentMaterialization->insertBefore(
815 &*CurrentBlock->getFirstInsertionPt());
816 }
817
818 if (auto *PN = dyn_cast<PHINode>(E.user())) {
819 assert(PN->getNumIncomingValues() == 1 && "unexpected number of incoming "
820 "values in the PHINode");
821 PN->replaceAllUsesWith(CurrentMaterialization);
822 PN->eraseFromParent();
823 continue;
824 }
825
826 // Replace all uses of CurrentDef in the current instruction with the
827 // CurrentMaterialization for the block.
828 E.user()->replaceUsesOfWith(CurrentDef, CurrentMaterialization);
829 }
830 }
831
832 // Move early uses of spilled variable after CoroBegin.
833 // For example, if a parameter had address taken, we may end up with the code
834 // like:
835 // define @f(i32 %n) {
836 // %n.addr = alloca i32
837 // store %n, %n.addr
838 // ...
839 // call @coro.begin
840 // we need to move the store after coro.begin
moveSpillUsesAfterCoroBegin(Function & F,SpillInfo const & Spills,CoroBeginInst * CoroBegin)841 static void moveSpillUsesAfterCoroBegin(Function &F, SpillInfo const &Spills,
842 CoroBeginInst *CoroBegin) {
843 DominatorTree DT(F);
844 SmallVector<Instruction *, 8> NeedsMoving;
845
846 Value *CurrentValue = nullptr;
847
848 for (auto const &E : Spills) {
849 if (CurrentValue == E.def())
850 continue;
851
852 CurrentValue = E.def();
853
854 for (User *U : CurrentValue->users()) {
855 Instruction *I = cast<Instruction>(U);
856 if (!DT.dominates(CoroBegin, I)) {
857 LLVM_DEBUG(dbgs() << "will move: " << *I << "\n");
858
859 // TODO: Make this more robust. Currently if we run into a situation
860 // where simple instruction move won't work we panic and
861 // report_fatal_error.
862 for (User *UI : I->users()) {
863 if (!DT.dominates(CoroBegin, cast<Instruction>(UI)))
864 report_fatal_error("cannot move instruction since its users are not"
865 " dominated by CoroBegin");
866 }
867
868 NeedsMoving.push_back(I);
869 }
870 }
871 }
872
873 Instruction *InsertPt = CoroBegin->getNextNode();
874 for (Instruction *I : NeedsMoving)
875 I->moveBefore(InsertPt);
876 }
877
878 // Splits the block at a particular instruction unless it is the first
879 // instruction in the block with a single predecessor.
splitBlockIfNotFirst(Instruction * I,const Twine & Name)880 static BasicBlock *splitBlockIfNotFirst(Instruction *I, const Twine &Name) {
881 auto *BB = I->getParent();
882 if (&BB->front() == I) {
883 if (BB->getSinglePredecessor()) {
884 BB->setName(Name);
885 return BB;
886 }
887 }
888 return BB->splitBasicBlock(I, Name);
889 }
890
891 // Split above and below a particular instruction so that it
892 // will be all alone by itself in a block.
splitAround(Instruction * I,const Twine & Name)893 static void splitAround(Instruction *I, const Twine &Name) {
894 splitBlockIfNotFirst(I, Name);
895 splitBlockIfNotFirst(I->getNextNode(), "After" + Name);
896 }
897
buildCoroutineFrame(Function & F,Shape & Shape)898 void coro::buildCoroutineFrame(Function &F, Shape &Shape) {
899 // Lower coro.dbg.declare to coro.dbg.value, since we are going to rewrite
900 // access to local variables.
901 LowerDbgDeclare(F);
902
903 Shape.PromiseAlloca = Shape.CoroBegin->getId()->getPromise();
904 if (Shape.PromiseAlloca) {
905 Shape.CoroBegin->getId()->clearPromise();
906 }
907
908 // Make sure that all coro.save, coro.suspend and the fallthrough coro.end
909 // intrinsics are in their own blocks to simplify the logic of building up
910 // SuspendCrossing data.
911 for (CoroSuspendInst *CSI : Shape.CoroSuspends) {
912 splitAround(CSI->getCoroSave(), "CoroSave");
913 splitAround(CSI, "CoroSuspend");
914 }
915
916 // Put CoroEnds into their own blocks.
917 for (CoroEndInst *CE : Shape.CoroEnds)
918 splitAround(CE, "CoroEnd");
919
920 // Transforms multi-edge PHI Nodes, so that any value feeding into a PHI will
921 // never has its definition separated from the PHI by the suspend point.
922 rewritePHIs(F);
923
924 // Build suspend crossing info.
925 SuspendCrossingInfo Checker(F, Shape);
926
927 IRBuilder<> Builder(F.getContext());
928 SpillInfo Spills;
929
930 for (int Repeat = 0; Repeat < 4; ++Repeat) {
931 // See if there are materializable instructions across suspend points.
932 for (Instruction &I : instructions(F))
933 if (materializable(I))
934 for (User *U : I.users())
935 if (Checker.isDefinitionAcrossSuspend(I, U))
936 Spills.emplace_back(&I, U);
937
938 if (Spills.empty())
939 break;
940
941 // Rewrite materializable instructions to be materialized at the use point.
942 LLVM_DEBUG(dump("Materializations", Spills));
943 rewriteMaterializableInstructions(Builder, Spills);
944 Spills.clear();
945 }
946
947 // Collect the spills for arguments and other not-materializable values.
948 for (Argument &A : F.args())
949 for (User *U : A.users())
950 if (Checker.isDefinitionAcrossSuspend(A, U))
951 Spills.emplace_back(&A, U);
952
953 for (Instruction &I : instructions(F)) {
954 // Values returned from coroutine structure intrinsics should not be part
955 // of the Coroutine Frame.
956 if (isCoroutineStructureIntrinsic(I) || &I == Shape.CoroBegin)
957 continue;
958 // The Coroutine Promise always included into coroutine frame, no need to
959 // check for suspend crossing.
960 if (Shape.PromiseAlloca == &I)
961 continue;
962
963 for (User *U : I.users())
964 if (Checker.isDefinitionAcrossSuspend(I, U)) {
965 // We cannot spill a token.
966 if (I.getType()->isTokenTy())
967 report_fatal_error(
968 "token definition is separated from the use by a suspend point");
969 Spills.emplace_back(&I, U);
970 }
971 }
972 LLVM_DEBUG(dump("Spills", Spills));
973 moveSpillUsesAfterCoroBegin(F, Spills, Shape.CoroBegin);
974 Shape.FrameTy = buildFrameType(F, Shape, Spills);
975 Shape.FramePtr = insertSpills(Spills, Shape);
976 }
977