Skip to content

Commit b0e1dc9

Browse files
nate-chandlernate-chandler
committed
[Coro] Add variant of retcon.once ABI.
Like async coroutines, it's fixed-per-function-size frame is caller allocated--the size is stored in a global "coro function pointer". Like retcon coroutines, dynamic allocations are performed via intrinsic-provided allocation and deallocation functions. Unlike both, it takes an allocator struct as an argument which is forwarded to the allocation/deallocation functions.
1 parent 5dc8ea8 commit b0e1dc9

File tree

8 files changed

+335
-116
lines changed

8 files changed

+335
-116
lines changed

‎llvm/include/llvm/IR/Intrinsics.td‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1680,6 +1680,10 @@ def int_coro_id_retcon_once : Intrinsic<[llvm_token_ty],
16801680
[llvm_i32_ty, llvm_i32_ty, llvm_ptr_ty,
16811681
llvm_ptr_ty, llvm_ptr_ty, llvm_ptr_ty, llvm_vararg_ty],
16821682
[]>;
1683+
def int_coro_id_retcon_once_dynamic : Intrinsic<[llvm_token_ty],
1684+
[llvm_i32_ty, llvm_i32_ty, llvm_ptr_ty, llvm_ptr_ty, llvm_ptr_ty,
1685+
llvm_ptr_ty, llvm_ptr_ty, llvm_ptr_ty],
1686+
[]>;
16831687
def int_coro_alloc : Intrinsic<[llvm_i1_ty], [llvm_token_ty], []>;
16841688
def int_coro_id_async : Intrinsic<[llvm_token_ty],
16851689
[llvm_i32_ty, llvm_i32_ty, llvm_i32_ty, llvm_ptr_ty],

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,7 @@ bool Lowerer::lower(Function &F) {
6868
case Intrinsic::coro_id:
6969
case Intrinsic::coro_id_retcon:
7070
case Intrinsic::coro_id_retcon_once:
71+
case Intrinsic::coro_id_retcon_once_dynamic:
7172
case Intrinsic::coro_id_async:
7273
II->replaceAllUsesWith(ConstantTokenNone::get(Context));
7374
break;

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -213,6 +213,7 @@ void Lowerer::lowerEarlyIntrinsics(Function &F) {
213213
break;
214214
case Intrinsic::coro_id_retcon:
215215
case Intrinsic::coro_id_retcon_once:
216+
case Intrinsic::coro_id_retcon_once_dynamic:
216217
case Intrinsic::coro_id_async:
217218
F.setPresplitCoroutine();
218219
break;

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

Lines changed: 21 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1394,6 +1394,19 @@ static StructType *buildFrameType(Function &F, coro::Shape &Shape,
13941394
B.getStructAlign() <= Id->getStorageAlignment());
13951395
break;
13961396
}
1397+
case coro::ABI::RetconOnceDynamic: {
1398+
// In the dynamic retcon.once ABI, the frame is always inline in the
1399+
// storage.
1400+
Shape.RetconLowering.IsFrameInlineInStorage = true;
1401+
Shape.RetconLowering.ContextSize =
1402+
alignTo(Shape.FrameSize, Shape.RetconLowering.StorageAlignment);
1403+
if (Shape.RetconLowering.StorageAlignment < Shape.FrameAlign) {
1404+
report_fatal_error(
1405+
"The alignment requirment of frame variables cannot be higher than "
1406+
"the alignment of the coro function context");
1407+
}
1408+
break;
1409+
}
13971410
case coro::ABI::Async: {
13981411
Shape.AsyncLowering.FrameOffset =
13991412
alignTo(Shape.AsyncLowering.ContextHeaderSize, Shape.FrameAlign);
@@ -1969,7 +1982,8 @@ static void insertSpills(const FrameDataInfo &FrameData, coro::Shape &Shape) {
19691982

19701983
// retcon and retcon.once lowering assumes all uses have been sunk.
19711984
if (Shape.ABI == coro::ABI::Retcon || Shape.ABI == coro::ABI::RetconOnce ||
1972-
Shape.ABI == coro::ABI::Async) {
1985+
Shape.ABI == coro::ABI::Async ||
1986+
Shape.ABI == coro::ABI::RetconOnceDynamic) {
19731987
// If we found any allocas, replace all of their remaining uses with Geps.
19741988
Builder.SetInsertPoint(SpillBlock, SpillBlock->begin());
19751989
for (const auto &P : FrameData.Allocas) {
@@ -2857,7 +2871,8 @@ static void collectFrameAlloca(AllocaInst *AI, coro::Shape &Shape,
28572871
// code.
28582872
bool ShouldUseLifetimeStartInfo =
28592873
(Shape.ABI != coro::ABI::Async && Shape.ABI != coro::ABI::Retcon &&
2860-
Shape.ABI != coro::ABI::RetconOnce);
2874+
Shape.ABI != coro::ABI::RetconOnce &&
2875+
Shape.ABI != coro::ABI::RetconOnceDynamic);
28612876
AllocaUseVisitor Visitor{AI->getDataLayout(), DT, Shape, Checker,
28622877
ShouldUseLifetimeStartInfo};
28632878
Visitor.visitPtr(*AI);
@@ -3160,7 +3175,8 @@ void coro::buildCoroutineFrame(
31603175
SmallVector<CoroAllocaAllocInst*, 4> LocalAllocas;
31613176
SmallVector<Instruction*, 4> DeadInstructions;
31623177
if (Shape.ABI != coro::ABI::Async && Shape.ABI != coro::ABI::Retcon &&
3163-
Shape.ABI != coro::ABI::RetconOnce)
3178+
Shape.ABI != coro::ABI::RetconOnce &&
3179+
Shape.ABI != coro::ABI::RetconOnceDynamic)
31643180
sinkLifetimeStartMarkers(F, Shape, Checker, DT);
31653181

31663182
// Collect the spills for arguments and other not-materializable values.
@@ -3238,7 +3254,8 @@ void coro::buildCoroutineFrame(
32383254

32393255
LLVM_DEBUG(dumpSpills("Spills", FrameData.Spills));
32403256
if (Shape.ABI == coro::ABI::Retcon || Shape.ABI == coro::ABI::RetconOnce ||
3241-
Shape.ABI == coro::ABI::Async)
3257+
Shape.ABI == coro::ABI::Async ||
3258+
Shape.ABI == coro::ABI::RetconOnceDynamic)
32423259
sinkSpillUsesAfterCoroBegin(F, FrameData, Shape.CoroBegin);
32433260
Shape.FrameTy = buildFrameType(F, Shape, FrameData);
32443261
Shape.FramePtr = Shape.CoroBegin;

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

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -134,6 +134,7 @@ class LLVM_LIBRARY_VISIBILITY AnyCoroIdInst : public IntrinsicInst {
134134
auto ID = I->getIntrinsicID();
135135
return ID == Intrinsic::coro_id || ID == Intrinsic::coro_id_retcon ||
136136
ID == Intrinsic::coro_id_retcon_once ||
137+
ID == Intrinsic::coro_id_retcon_once_dynamic ||
137138
ID == Intrinsic::coro_id_async;
138139
}
139140

@@ -317,6 +318,72 @@ class LLVM_LIBRARY_VISIBILITY CoroIdRetconOnceInst
317318
}
318319
};
319320

321+
/// This represents the llvm.coro.id.retcon.once.dynamic instruction.
322+
class LLVM_LIBRARY_VISIBILITY CoroIdRetconOnceDynamicInst
323+
: public AnyCoroIdInst {
324+
enum {
325+
SizeArg,
326+
AlignArg,
327+
CoroFuncPtrArg,
328+
AllocatorArg,
329+
StorageArg,
330+
PrototypeArg,
331+
AllocArg,
332+
DeallocArg
333+
};
334+
335+
public:
336+
void checkWellFormed() const;
337+
338+
uint64_t getStorageSize() const {
339+
return cast<ConstantInt>(getArgOperand(SizeArg))->getZExtValue();
340+
}
341+
342+
Align getStorageAlignment() const {
343+
return cast<ConstantInt>(getArgOperand(AlignArg))->getAlignValue();
344+
}
345+
346+
Value *getStorage() const { return getArgOperand(StorageArg); }
347+
348+
/// Return the coro function pointer address. This should be the address of
349+
/// a coro function pointer struct for the current coro function.
350+
/// struct coro_function_pointer {
351+
/// uint32_t frame size;
352+
/// uint32_t relative_pointer(coro_function);
353+
/// };
354+
GlobalVariable *getCoroFunctionPointer() const {
355+
return cast<GlobalVariable>(
356+
getArgOperand(CoroFuncPtrArg)->stripPointerCasts());
357+
}
358+
359+
/// Return the prototype for the continuation function. The type,
360+
/// attributes, and calling convention of the continuation function(s)
361+
/// are taken from this declaration.
362+
Function *getPrototype() const {
363+
return cast<Function>(getArgOperand(PrototypeArg)->stripPointerCasts());
364+
}
365+
366+
/// Return the function to use for allocating memory.
367+
Function *getAllocFunction() const {
368+
return cast<Function>(getArgOperand(AllocArg)->stripPointerCasts());
369+
}
370+
371+
/// Return the function to use for deallocating memory.
372+
Function *getDeallocFunction() const {
373+
return cast<Function>(getArgOperand(DeallocArg)->stripPointerCasts());
374+
}
375+
376+
Value *getAllocator() const { return getArgOperand(AllocatorArg); }
377+
378+
// Methods to support type inquiry through isa, cast, and dyn_cast:
379+
static bool classof(const IntrinsicInst *I) {
380+
return I->getIntrinsicID() == Intrinsic::coro_id_retcon_once_dynamic;
381+
}
382+
static bool classof(const Value *V) {
383+
return isa<IntrinsicInst>(V) && classof(cast<IntrinsicInst>(V));
384+
}
385+
};
386+
320387
/// This represents the llvm.coro.id.async instruction.
321388
class LLVM_LIBRARY_VISIBILITY CoroIdAsyncInst : public AnyCoroIdInst {
322389
enum { SizeArg, AlignArg, StorageArg, AsyncFuncPtrArg };

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

Lines changed: 23 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,10 @@ enum class ABI {
7373
/// single continuation function. The continuation function is available as an
7474
/// intrinsic.
7575
Async,
76+
77+
/// The variant of RetconOnce which features a dynamically-sized caller
78+
/// allocation.
79+
RetconOnceDynamic,
7680
};
7781

7882
// Holds structural Coroutine Intrinsics for a particular function and other
@@ -128,9 +132,18 @@ struct LLVM_LIBRARY_VISIBILITY Shape {
128132
Function *ResumePrototype;
129133
Function *Alloc;
130134
Function *Dealloc;
135+
Value *Allocator;
131136
BasicBlock *ReturnBlock;
132137
bool IsFrameInlineInStorage;
133138
ConstantInt* TypeId;
139+
GlobalVariable *CoroFuncPointer;
140+
Value *Storage;
141+
uint64_t StorageSize;
142+
Align StorageAlignment;
143+
// computed during splitting:
144+
uint64_t ContextSize;
145+
146+
Align getStorageAlignment() const { return Align(StorageAlignment); }
134147
};
135148

136149
struct AsyncLoweringStorage {
@@ -196,6 +209,7 @@ struct LLVM_LIBRARY_VISIBILITY Shape {
196209
/*IsVarArg=*/false);
197210
case coro::ABI::Retcon:
198211
case coro::ABI::RetconOnce:
212+
case coro::ABI::RetconOnceDynamic:
199213
return RetconLowering.ResumePrototype->getFunctionType();
200214
case coro::ABI::Async:
201215
// Not used. The function type depends on the active suspend.
@@ -206,8 +220,8 @@ struct LLVM_LIBRARY_VISIBILITY Shape {
206220
}
207221

208222
ArrayRef<Type*> getRetconResultTypes() const {
209-
assert(ABI == coro::ABI::Retcon ||
210-
ABI == coro::ABI::RetconOnce);
223+
assert(ABI == coro::ABI::Retcon || ABI == coro::ABI::RetconOnce ||
224+
ABI == coro::ABI::RetconOnceDynamic);
211225
auto FTy = CoroBegin->getFunction()->getFunctionType();
212226

213227
// The safety of all this is checked by checkWFRetconPrototype.
@@ -219,8 +233,8 @@ struct LLVM_LIBRARY_VISIBILITY Shape {
219233
}
220234

221235
ArrayRef<Type*> getRetconResumeTypes() const {
222-
assert(ABI == coro::ABI::Retcon ||
223-
ABI == coro::ABI::RetconOnce);
236+
assert(ABI == coro::ABI::Retcon || ABI == coro::ABI::RetconOnce ||
237+
ABI == coro::ABI::RetconOnceDynamic);
224238

225239
// The safety of all this is checked by checkWFRetconPrototype.
226240
auto FTy = RetconLowering.ResumePrototype->getFunctionType();
@@ -234,6 +248,7 @@ struct LLVM_LIBRARY_VISIBILITY Shape {
234248

235249
case coro::ABI::Retcon:
236250
case coro::ABI::RetconOnce:
251+
case coro::ABI::RetconOnceDynamic:
237252
return RetconLowering.ResumePrototype->getCallingConv();
238253
case coro::ABI::Async:
239254
return AsyncLowering.AsyncCC;
@@ -266,12 +281,15 @@ struct LLVM_LIBRARY_VISIBILITY Shape {
266281
/// \param CG - if non-null, will be updated for the new call
267282
void emitDealloc(IRBuilder<> &Builder, Value *Ptr, CallGraph *CG) const;
268283

269-
Shape() = default;
284+
Shape() = delete;
270285
explicit Shape(Function &F, bool OptimizeFrame = false)
271286
: OptimizeFrame(OptimizeFrame) {
272287
buildFrom(F);
273288
}
274289
void buildFrom(Function &F);
290+
291+
private:
292+
void buildShapeFromRetconInst(AnyCoroIdInst *Id, Function *Prototype);
275293
};
276294

277295
bool defaultMaterializable(Instruction &V);

0 commit comments

Comments
 (0)