Skip to content

Commit ba9e523

Browse files
author
tnowicki
committed
[Coroutines] Refactor CoroShape::buildFrom to support ABI
* Refactor buildFrom to separate the analysis, abi related operations, tidying and bailout. * In a follow-up PR the code in initABI will be moved to an ABI object init method. And the Shape constructor will no longer perform any lowering, instead it will just call analysis. This will make the Shape object a bit more useful because it can be constructed and used anywhere. It may even be useful to make it an analysis pass. * 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. * Note, that it was necessary to introduce two new SmallVectors, one to track CoroFrames and the other for UnusedCoroSaves. The tidyCoroutine method requires both, while invalidateCoroutine (bailout) method just requires the former. See RFC for more info: https://discourse.llvm.org/t/rfc-abi-objects-for-coroutines/81057
1 parent 4c040c0 commit ba9e523

File tree

2 files changed

+117
-75
lines changed

2 files changed

+117
-75
lines changed

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

Lines changed: 47 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -50,15 +50,49 @@ 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 invalidateCoroutine() and tidyCoroutine()
62+
SmallVector<CoroFrameInst *, 8> CoroFrames;
63+
SmallVector<CoroSaveInst *, 2> UnusedCoroSaves;
64+
65+
// Values invalidated by replaceSwiftErrorOps()
66+
SmallVector<CallInst *, 2> SwiftErrorOps;
67+
68+
void clear() {
69+
CoroBegin = nullptr;
70+
CoroEnds.clear();
71+
CoroSizes.clear();
72+
CoroAligns.clear();
73+
CoroSuspends.clear();
74+
CoroAwaitSuspends.clear();
75+
SymmetricTransfers.clear();
76+
77+
CoroFrames.clear();
78+
UnusedCoroSaves.clear();
79+
80+
SwiftErrorOps.clear();
81+
82+
FrameTy = nullptr;
83+
FramePtr = nullptr;
84+
AllocaSpillBlock = nullptr;
85+
}
86+
87+
// Scan the function and collect the above intrinsics for later processing
88+
void analyze(Function &F);
89+
// If for some reason, we were not able to find coro.begin, bailout.
90+
void invalidateCoroutine(Function &F);
91+
// Perform ABI related initial transformation
92+
void initABI();
93+
// Remove orphaned and unnecessary intrinsics
94+
void tidyCoroutine();
95+
6296
// Field indexes for special fields in the switch lowering.
6397
struct SwitchFieldIndex {
6498
enum {
@@ -76,11 +110,11 @@ struct LLVM_LIBRARY_VISIBILITY Shape {
76110

77111
coro::ABI ABI;
78112

79-
StructType *FrameTy;
113+
StructType *FrameTy = nullptr;
80114
Align FrameAlign;
81-
uint64_t FrameSize;
82-
Value *FramePtr;
83-
BasicBlock *AllocaSpillBlock;
115+
uint64_t FrameSize = 0;
116+
Value *FramePtr = nullptr;
117+
BasicBlock *AllocaSpillBlock = nullptr;
84118

85119
/// This would only be true if optimization are enabled.
86120
bool OptimizeFrame;
@@ -237,9 +271,14 @@ struct LLVM_LIBRARY_VISIBILITY Shape {
237271
Shape() = default;
238272
explicit Shape(Function &F, bool OptimizeFrame = false)
239273
: OptimizeFrame(OptimizeFrame) {
240-
buildFrom(F);
274+
analyze(F);
275+
if (!CoroBegin) {
276+
invalidateCoroutine(F);
277+
return;
278+
}
279+
initABI();
280+
tidyCoroutine();
241281
}
242-
void buildFrom(Function &F);
243282
};
244283

245284
} // end namespace coro

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

Lines changed: 70 additions & 67 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,12 @@ 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+
clear();
194+
204195
bool HasFinalSuspend = false;
205196
bool HasUnwindCoroEnd = false;
206197
size_t FinalSuspendIndex = 0;
207-
clear(*this);
208-
SmallVector<CoroFrameInst *, 8> CoroFrames;
209-
SmallVector<CoroSaveInst *, 2> UnusedCoroSaves;
210198

211199
for (Instruction &I : instructions(F)) {
212200
// FIXME: coro_await_suspend_* are not proper `IntrinisicInst`s
@@ -298,8 +286,58 @@ void coro::Shape::buildFrom(Function &F) {
298286
}
299287
}
300288

301-
// If for some reason, we were not able to find coro.begin, bailout.
302-
if (!CoroBegin) {
289+
// If there is no CoroBegin then this is not a coroutine.
290+
if (!CoroBegin)
291+
return;
292+
293+
// Determination of ABI and initializing lowering info
294+
auto Id = CoroBegin->getId();
295+
auto IntrID = Id->getIntrinsicID();
296+
if (IntrID == Intrinsic::coro_id) {
297+
ABI = coro::ABI::Switch;
298+
SwitchLowering.HasFinalSuspend = HasFinalSuspend;
299+
SwitchLowering.HasUnwindCoroEnd = HasUnwindCoroEnd;
300+
301+
auto SwitchId = getSwitchCoroId();
302+
SwitchLowering.ResumeSwitch = nullptr;
303+
SwitchLowering.PromiseAlloca = SwitchId->getPromise();
304+
SwitchLowering.ResumeEntryBlock = nullptr;
305+
306+
// Move final suspend to the last element in the CoroSuspends vector.
307+
if (SwitchLowering.HasFinalSuspend &&
308+
FinalSuspendIndex != CoroSuspends.size() - 1)
309+
std::swap(CoroSuspends[FinalSuspendIndex], CoroSuspends.back());
310+
} else if (IntrID == Intrinsic::coro_id_async) {
311+
ABI = coro::ABI::Async;
312+
auto *AsyncId = getAsyncCoroId();
313+
AsyncId->checkWellFormed();
314+
AsyncLowering.Context = AsyncId->getStorage();
315+
AsyncLowering.ContextArgNo = AsyncId->getStorageArgumentIndex();
316+
AsyncLowering.ContextHeaderSize = AsyncId->getStorageSize();
317+
AsyncLowering.ContextAlignment = AsyncId->getStorageAlignment().value();
318+
AsyncLowering.AsyncFuncPointer = AsyncId->getAsyncFunctionPointer();
319+
AsyncLowering.AsyncCC = F.getCallingConv();
320+
} else if (IntrID == Intrinsic::coro_id_retcon ||
321+
IntrID == Intrinsic::coro_id_retcon_once) {
322+
ABI = IntrID == Intrinsic::coro_id_retcon ? coro::ABI::Retcon
323+
: coro::ABI::RetconOnce;
324+
auto ContinuationId = getRetconCoroId();
325+
ContinuationId->checkWellFormed();
326+
auto Prototype = ContinuationId->getPrototype();
327+
RetconLowering.ResumePrototype = Prototype;
328+
RetconLowering.Alloc = ContinuationId->getAllocFunction();
329+
RetconLowering.Dealloc = ContinuationId->getDeallocFunction();
330+
RetconLowering.ReturnBlock = nullptr;
331+
RetconLowering.IsFrameInlineInStorage = false;
332+
} else {
333+
llvm_unreachable("coro.begin is not dependent on a coro.id call");
334+
}
335+
}
336+
337+
// If for some reason, we were not able to find coro.begin, bailout.
338+
void coro::Shape::invalidateCoroutine(Function &F) {
339+
assert(!CoroBegin);
340+
{
303341
// Replace coro.frame which are supposed to be lowered to the result of
304342
// coro.begin with undef.
305343
auto *Undef = UndefValue::get(PointerType::get(F.getContext(), 0));
@@ -320,21 +358,13 @@ void coro::Shape::buildFrom(Function &F) {
320358
// Replace all coro.ends with unreachable instruction.
321359
for (AnyCoroEndInst *CE : CoroEnds)
322360
changeToUnreachable(CE);
323-
324-
return;
325361
}
362+
}
326363

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-
364+
// Perform semantic checking and initialization of the ABI
365+
void coro::Shape::initABI() {
366+
switch (ABI) {
367+
case coro::ABI::Switch: {
338368
for (auto *AnySuspend : CoroSuspends) {
339369
auto Suspend = dyn_cast<CoroSuspendInst>(AnySuspend);
340370
if (!Suspend) {
@@ -349,33 +379,11 @@ void coro::Shape::buildFrom(Function &F) {
349379
}
350380
break;
351381
}
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();
382+
case coro::ABI::Async: {
363383
break;
364384
};
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-
385+
case coro::ABI::Retcon:
386+
case coro::ABI::RetconOnce: {
379387
// Determine the result value types, and make sure they match up with
380388
// the values passed to the suspends.
381389
auto ResultTys = getRetconResultTypes();
@@ -408,7 +416,7 @@ void coro::Shape::buildFrom(Function &F) {
408416

409417
#ifndef NDEBUG
410418
Suspend->dump();
411-
Prototype->getFunctionType()->dump();
419+
RetconLowering.ResumePrototype->getFunctionType()->dump();
412420
#endif
413421
report_fatal_error("argument to coro.suspend.retcon does not "
414422
"match corresponding prototype function result");
@@ -417,14 +425,14 @@ void coro::Shape::buildFrom(Function &F) {
417425
if (SI != SE || RI != RE) {
418426
#ifndef NDEBUG
419427
Suspend->dump();
420-
Prototype->getFunctionType()->dump();
428+
RetconLowering.ResumePrototype->getFunctionType()->dump();
421429
#endif
422430
report_fatal_error("wrong number of arguments to coro.suspend.retcon");
423431
}
424432

425433
// Check that the result type of the suspend matches the resume types.
426434
Type *SResultTy = Suspend->getType();
427-
ArrayRef<Type*> SuspendResultTys;
435+
ArrayRef<Type *> SuspendResultTys;
428436
if (SResultTy->isVoidTy()) {
429437
// leave as empty array
430438
} else if (auto SResultStructTy = dyn_cast<StructType>(SResultTy)) {
@@ -436,15 +444,15 @@ void coro::Shape::buildFrom(Function &F) {
436444
if (SuspendResultTys.size() != ResumeTys.size()) {
437445
#ifndef NDEBUG
438446
Suspend->dump();
439-
Prototype->getFunctionType()->dump();
447+
RetconLowering.ResumePrototype->getFunctionType()->dump();
440448
#endif
441449
report_fatal_error("wrong number of results from coro.suspend.retcon");
442450
}
443451
for (size_t I = 0, E = ResumeTys.size(); I != E; ++I) {
444452
if (SuspendResultTys[I] != ResumeTys[I]) {
445453
#ifndef NDEBUG
446454
Suspend->dump();
447-
Prototype->getFunctionType()->dump();
455+
RetconLowering.ResumePrototype->getFunctionType()->dump();
448456
#endif
449457
report_fatal_error("result from coro.suspend.retcon does not "
450458
"match corresponding prototype function param");
@@ -453,23 +461,18 @@ void coro::Shape::buildFrom(Function &F) {
453461
}
454462
break;
455463
}
456-
457464
default:
458465
llvm_unreachable("coro.begin is not dependent on a coro.id call");
459466
}
467+
}
460468

469+
void coro::Shape::tidyCoroutine() {
461470
// The coro.free intrinsic is always lowered to the result of coro.begin.
462471
for (CoroFrameInst *CF : CoroFrames) {
463472
CF->replaceAllUsesWith(CoroBegin);
464473
CF->eraseFromParent();
465474
}
466475

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());
472-
473476
// Remove orphaned coro.saves.
474477
for (CoroSaveInst *CoroSave : UnusedCoroSaves)
475478
CoroSave->eraseFromParent();

0 commit comments

Comments
 (0)