Skip to content

Commit d575252

Browse files
TylerNowickiTylerNowicki
authored
[Coroutines] Refactor CoroShape::buildFrom for future use by ABI objects (#108623)
* Refactor buildFrom to separate the analysis, abi related operations, cleaning and invalidating. * In a follow-up PR the code in initABI will be moved to an ABI object init method. * In a follow-up PR the OptimizeFrame flag will also be removed from the Shape and instead will be passed directly to buildCoroutineFrame (although it would be nice to find another way to trigger this optimization). This is the only thing that Shape cannot determine from the Function/Coroutine, but it is only needed within buildCoroutineFrame. See RFC for more info: https://discourse.llvm.org/t/rfc-abi-objects-for-coroutines/81057
1 parent 14c7632 commit d575252

File tree

2 files changed

+132
-76
lines changed

2 files changed

+132
-76
lines changed

‎llvm/lib/Transforms/Coroutines/CoroShape.h

Lines changed: 46 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -50,15 +50,45 @@ enum class ABI {
5050
// Holds structural Coroutine Intrinsics for a particular function and other
5151
// values used during CoroSplit pass.
5252
struct LLVM_LIBRARY_VISIBILITY Shape {
53-
CoroBeginInst *CoroBegin;
53+
CoroBeginInst *CoroBegin = nullptr;
5454
SmallVector<AnyCoroEndInst *, 4> CoroEnds;
5555
SmallVector<CoroSizeInst *, 2> CoroSizes;
5656
SmallVector<CoroAlignInst *, 2> CoroAligns;
5757
SmallVector<AnyCoroSuspendInst *, 4> CoroSuspends;
58-
SmallVector<CallInst *, 2> SwiftErrorOps;
5958
SmallVector<CoroAwaitSuspendInst *, 4> CoroAwaitSuspends;
6059
SmallVector<CallInst *, 2> SymmetricTransfers;
6160

61+
// Values invalidated by replaceSwiftErrorOps()
62+
SmallVector<CallInst *, 2> SwiftErrorOps;
63+
64+
void clear() {
65+
CoroBegin = nullptr;
66+
CoroEnds.clear();
67+
CoroSizes.clear();
68+
CoroAligns.clear();
69+
CoroSuspends.clear();
70+
CoroAwaitSuspends.clear();
71+
SymmetricTransfers.clear();
72+
73+
SwiftErrorOps.clear();
74+
75+
FrameTy = nullptr;
76+
FramePtr = nullptr;
77+
AllocaSpillBlock = nullptr;
78+
}
79+
80+
// Scan the function and collect the above intrinsics for later processing
81+
void analyze(Function &F, SmallVectorImpl<CoroFrameInst *> &CoroFrames,
82+
SmallVectorImpl<CoroSaveInst *> &UnusedCoroSaves);
83+
// If for some reason, we were not able to find coro.begin, bailout.
84+
void invalidateCoroutine(Function &F,
85+
SmallVectorImpl<CoroFrameInst *> &CoroFrames);
86+
// Perform ABI related initial transformation
87+
void initABI();
88+
// Remove orphaned and unnecessary intrinsics
89+
void cleanCoroutine(SmallVectorImpl<CoroFrameInst *> &CoroFrames,
90+
SmallVectorImpl<CoroSaveInst *> &UnusedCoroSaves);
91+
6292
// Field indexes for special fields in the switch lowering.
6393
struct SwitchFieldIndex {
6494
enum {
@@ -76,11 +106,11 @@ struct LLVM_LIBRARY_VISIBILITY Shape {
76106

77107
coro::ABI ABI;
78108

79-
StructType *FrameTy;
109+
StructType *FrameTy = nullptr;
80110
Align FrameAlign;
81-
uint64_t FrameSize;
82-
Value *FramePtr;
83-
BasicBlock *AllocaSpillBlock;
111+
uint64_t FrameSize = 0;
112+
Value *FramePtr = nullptr;
113+
BasicBlock *AllocaSpillBlock = nullptr;
84114

85115
/// This would only be true if optimization are enabled.
86116
bool OptimizeFrame;
@@ -237,9 +267,17 @@ struct LLVM_LIBRARY_VISIBILITY Shape {
237267
Shape() = default;
238268
explicit Shape(Function &F, bool OptimizeFrame = false)
239269
: OptimizeFrame(OptimizeFrame) {
240-
buildFrom(F);
270+
SmallVector<CoroFrameInst *, 8> CoroFrames;
271+
SmallVector<CoroSaveInst *, 2> UnusedCoroSaves;
272+
273+
analyze(F, CoroFrames, UnusedCoroSaves);
274+
if (!CoroBegin) {
275+
invalidateCoroutine(F, CoroFrames);
276+
return;
277+
}
278+
initABI();
279+
cleanCoroutine(CoroFrames, UnusedCoroSaves);
241280
}
242-
void buildFrom(Function &F);
243281
};
244282

245283
} // end namespace coro

‎llvm/lib/Transforms/Coroutines/Coroutines.cpp

Lines changed: 86 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -177,17 +177,6 @@ void coro::suppressCoroAllocs(LLVMContext &Context,
177177
}
178178
}
179179

180-
static void clear(coro::Shape &Shape) {
181-
Shape.CoroBegin = nullptr;
182-
Shape.CoroEnds.clear();
183-
Shape.CoroSizes.clear();
184-
Shape.CoroSuspends.clear();
185-
186-
Shape.FrameTy = nullptr;
187-
Shape.FramePtr = nullptr;
188-
Shape.AllocaSpillBlock = nullptr;
189-
}
190-
191180
static CoroSaveInst *createCoroSave(CoroBeginInst *CoroBegin,
192181
CoroSuspendInst *SuspendInst) {
193182
Module *M = SuspendInst->getModule();
@@ -200,13 +189,14 @@ static CoroSaveInst *createCoroSave(CoroBeginInst *CoroBegin,
200189
}
201190

202191
// Collect "interesting" coroutine intrinsics.
203-
void coro::Shape::buildFrom(Function &F) {
192+
void coro::Shape::analyze(Function &F,
193+
SmallVectorImpl<CoroFrameInst *> &CoroFrames,
194+
SmallVectorImpl<CoroSaveInst *> &UnusedCoroSaves) {
195+
clear();
196+
204197
bool HasFinalSuspend = false;
205198
bool HasUnwindCoroEnd = false;
206199
size_t FinalSuspendIndex = 0;
207-
clear(*this);
208-
SmallVector<CoroFrameInst *, 8> CoroFrames;
209-
SmallVector<CoroSaveInst *, 2> UnusedCoroSaves;
210200

211201
for (Instruction &I : instructions(F)) {
212202
// FIXME: coro_await_suspend_* are not proper `IntrinisicInst`s
@@ -298,15 +288,73 @@ void coro::Shape::buildFrom(Function &F) {
298288
}
299289
}
300290

301-
// If for some reason, we were not able to find coro.begin, bailout.
302-
if (!CoroBegin) {
291+
// If there is no CoroBegin then this is not a coroutine.
292+
if (!CoroBegin)
293+
return;
294+
295+
// Determination of ABI and initializing lowering info
296+
auto Id = CoroBegin->getId();
297+
switch (auto IntrID = Id->getIntrinsicID()) {
298+
case Intrinsic::coro_id: {
299+
ABI = coro::ABI::Switch;
300+
SwitchLowering.HasFinalSuspend = HasFinalSuspend;
301+
SwitchLowering.HasUnwindCoroEnd = HasUnwindCoroEnd;
302+
303+
auto SwitchId = getSwitchCoroId();
304+
SwitchLowering.ResumeSwitch = nullptr;
305+
SwitchLowering.PromiseAlloca = SwitchId->getPromise();
306+
SwitchLowering.ResumeEntryBlock = nullptr;
307+
308+
// Move final suspend to the last element in the CoroSuspends vector.
309+
if (SwitchLowering.HasFinalSuspend &&
310+
FinalSuspendIndex != CoroSuspends.size() - 1)
311+
std::swap(CoroSuspends[FinalSuspendIndex], CoroSuspends.back());
312+
break;
313+
}
314+
case Intrinsic::coro_id_async: {
315+
ABI = coro::ABI::Async;
316+
auto *AsyncId = getAsyncCoroId();
317+
AsyncId->checkWellFormed();
318+
AsyncLowering.Context = AsyncId->getStorage();
319+
AsyncLowering.ContextArgNo = AsyncId->getStorageArgumentIndex();
320+
AsyncLowering.ContextHeaderSize = AsyncId->getStorageSize();
321+
AsyncLowering.ContextAlignment = AsyncId->getStorageAlignment().value();
322+
AsyncLowering.AsyncFuncPointer = AsyncId->getAsyncFunctionPointer();
323+
AsyncLowering.AsyncCC = F.getCallingConv();
324+
break;
325+
}
326+
case Intrinsic::coro_id_retcon:
327+
case Intrinsic::coro_id_retcon_once: {
328+
ABI = IntrID == Intrinsic::coro_id_retcon ? coro::ABI::Retcon
329+
: coro::ABI::RetconOnce;
330+
auto ContinuationId = getRetconCoroId();
331+
ContinuationId->checkWellFormed();
332+
auto Prototype = ContinuationId->getPrototype();
333+
RetconLowering.ResumePrototype = Prototype;
334+
RetconLowering.Alloc = ContinuationId->getAllocFunction();
335+
RetconLowering.Dealloc = ContinuationId->getDeallocFunction();
336+
RetconLowering.ReturnBlock = nullptr;
337+
RetconLowering.IsFrameInlineInStorage = false;
338+
break;
339+
}
340+
default:
341+
llvm_unreachable("coro.begin is not dependent on a coro.id call");
342+
}
343+
}
344+
345+
// If for some reason, we were not able to find coro.begin, bailout.
346+
void coro::Shape::invalidateCoroutine(
347+
Function &F, SmallVectorImpl<CoroFrameInst *> &CoroFrames) {
348+
assert(!CoroBegin);
349+
{
303350
// Replace coro.frame which are supposed to be lowered to the result of
304351
// coro.begin with undef.
305352
auto *Undef = UndefValue::get(PointerType::get(F.getContext(), 0));
306353
for (CoroFrameInst *CF : CoroFrames) {
307354
CF->replaceAllUsesWith(Undef);
308355
CF->eraseFromParent();
309356
}
357+
CoroFrames.clear();
310358

311359
// Replace all coro.suspend with undef and remove related coro.saves if
312360
// present.
@@ -316,25 +364,18 @@ void coro::Shape::buildFrom(Function &F) {
316364
if (auto *CoroSave = CS->getCoroSave())
317365
CoroSave->eraseFromParent();
318366
}
367+
CoroSuspends.clear();
319368

320369
// Replace all coro.ends with unreachable instruction.
321370
for (AnyCoroEndInst *CE : CoroEnds)
322371
changeToUnreachable(CE);
323-
324-
return;
325372
}
373+
}
326374

327-
auto Id = CoroBegin->getId();
328-
switch (auto IdIntrinsic = Id->getIntrinsicID()) {
329-
case Intrinsic::coro_id: {
330-
auto SwitchId = cast<CoroIdInst>(Id);
331-
this->ABI = coro::ABI::Switch;
332-
this->SwitchLowering.HasFinalSuspend = HasFinalSuspend;
333-
this->SwitchLowering.HasUnwindCoroEnd = HasUnwindCoroEnd;
334-
this->SwitchLowering.ResumeSwitch = nullptr;
335-
this->SwitchLowering.PromiseAlloca = SwitchId->getPromise();
336-
this->SwitchLowering.ResumeEntryBlock = nullptr;
337-
375+
// Perform semantic checking and initialization of the ABI
376+
void coro::Shape::initABI() {
377+
switch (ABI) {
378+
case coro::ABI::Switch: {
338379
for (auto *AnySuspend : CoroSuspends) {
339380
auto Suspend = dyn_cast<CoroSuspendInst>(AnySuspend);
340381
if (!Suspend) {
@@ -349,33 +390,11 @@ void coro::Shape::buildFrom(Function &F) {
349390
}
350391
break;
351392
}
352-
case Intrinsic::coro_id_async: {
353-
auto *AsyncId = cast<CoroIdAsyncInst>(Id);
354-
AsyncId->checkWellFormed();
355-
this->ABI = coro::ABI::Async;
356-
this->AsyncLowering.Context = AsyncId->getStorage();
357-
this->AsyncLowering.ContextArgNo = AsyncId->getStorageArgumentIndex();
358-
this->AsyncLowering.ContextHeaderSize = AsyncId->getStorageSize();
359-
this->AsyncLowering.ContextAlignment =
360-
AsyncId->getStorageAlignment().value();
361-
this->AsyncLowering.AsyncFuncPointer = AsyncId->getAsyncFunctionPointer();
362-
this->AsyncLowering.AsyncCC = F.getCallingConv();
393+
case coro::ABI::Async: {
363394
break;
364395
};
365-
case Intrinsic::coro_id_retcon:
366-
case Intrinsic::coro_id_retcon_once: {
367-
auto ContinuationId = cast<AnyCoroIdRetconInst>(Id);
368-
ContinuationId->checkWellFormed();
369-
this->ABI = (IdIntrinsic == Intrinsic::coro_id_retcon
370-
? coro::ABI::Retcon
371-
: coro::ABI::RetconOnce);
372-
auto Prototype = ContinuationId->getPrototype();
373-
this->RetconLowering.ResumePrototype = Prototype;
374-
this->RetconLowering.Alloc = ContinuationId->getAllocFunction();
375-
this->RetconLowering.Dealloc = ContinuationId->getDeallocFunction();
376-
this->RetconLowering.ReturnBlock = nullptr;
377-
this->RetconLowering.IsFrameInlineInStorage = false;
378-
396+
case coro::ABI::Retcon:
397+
case coro::ABI::RetconOnce: {
379398
// Determine the result value types, and make sure they match up with
380399
// the values passed to the suspends.
381400
auto ResultTys = getRetconResultTypes();
@@ -408,7 +427,7 @@ void coro::Shape::buildFrom(Function &F) {
408427

409428
#ifndef NDEBUG
410429
Suspend->dump();
411-
Prototype->getFunctionType()->dump();
430+
RetconLowering.ResumePrototype->getFunctionType()->dump();
412431
#endif
413432
report_fatal_error("argument to coro.suspend.retcon does not "
414433
"match corresponding prototype function result");
@@ -417,14 +436,14 @@ void coro::Shape::buildFrom(Function &F) {
417436
if (SI != SE || RI != RE) {
418437
#ifndef NDEBUG
419438
Suspend->dump();
420-
Prototype->getFunctionType()->dump();
439+
RetconLowering.ResumePrototype->getFunctionType()->dump();
421440
#endif
422441
report_fatal_error("wrong number of arguments to coro.suspend.retcon");
423442
}
424443

425444
// Check that the result type of the suspend matches the resume types.
426445
Type *SResultTy = Suspend->getType();
427-
ArrayRef<Type*> SuspendResultTys;
446+
ArrayRef<Type *> SuspendResultTys;
428447
if (SResultTy->isVoidTy()) {
429448
// leave as empty array
430449
} else if (auto SResultStructTy = dyn_cast<StructType>(SResultTy)) {
@@ -436,15 +455,15 @@ void coro::Shape::buildFrom(Function &F) {
436455
if (SuspendResultTys.size() != ResumeTys.size()) {
437456
#ifndef NDEBUG
438457
Suspend->dump();
439-
Prototype->getFunctionType()->dump();
458+
RetconLowering.ResumePrototype->getFunctionType()->dump();
440459
#endif
441460
report_fatal_error("wrong number of results from coro.suspend.retcon");
442461
}
443462
for (size_t I = 0, E = ResumeTys.size(); I != E; ++I) {
444463
if (SuspendResultTys[I] != ResumeTys[I]) {
445464
#ifndef NDEBUG
446465
Suspend->dump();
447-
Prototype->getFunctionType()->dump();
466+
RetconLowering.ResumePrototype->getFunctionType()->dump();
448467
#endif
449468
report_fatal_error("result from coro.suspend.retcon does not "
450469
"match corresponding prototype function param");
@@ -453,26 +472,25 @@ void coro::Shape::buildFrom(Function &F) {
453472
}
454473
break;
455474
}
456-
457475
default:
458476
llvm_unreachable("coro.begin is not dependent on a coro.id call");
459477
}
478+
}
460479

461-
// The coro.free intrinsic is always lowered to the result of coro.begin.
480+
void coro::Shape::cleanCoroutine(
481+
SmallVectorImpl<CoroFrameInst *> &CoroFrames,
482+
SmallVectorImpl<CoroSaveInst *> &UnusedCoroSaves) {
483+
// The coro.frame intrinsic is always lowered to the result of coro.begin.
462484
for (CoroFrameInst *CF : CoroFrames) {
463485
CF->replaceAllUsesWith(CoroBegin);
464486
CF->eraseFromParent();
465487
}
466-
467-
// Move final suspend to be the last element in the CoroSuspends vector.
468-
if (ABI == coro::ABI::Switch &&
469-
SwitchLowering.HasFinalSuspend &&
470-
FinalSuspendIndex != CoroSuspends.size() - 1)
471-
std::swap(CoroSuspends[FinalSuspendIndex], CoroSuspends.back());
488+
CoroFrames.clear();
472489

473490
// Remove orphaned coro.saves.
474491
for (CoroSaveInst *CoroSave : UnusedCoroSaves)
475492
CoroSave->eraseFromParent();
493+
UnusedCoroSaves.clear();
476494
}
477495

478496
static void propagateCallAttrsFromCallee(CallInst *Call, Function *Callee) {

0 commit comments

Comments
 (0)