Index: llvm/docs/Coroutines.rst =================================================================== --- llvm/docs/Coroutines.rst +++ llvm/docs/Coroutines.rst @@ -752,6 +752,33 @@ The `coro.size` intrinsic is lowered to a constant representing the size of the coroutine frame. +.. _coro.sizeof: + +'llvm.coro.sizeof' Intrinsic +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +:: + + declare i32 @llvm.coro.sizeof.i32(i8* %func) + declare i64 @llvm.coro.sizeof.i64(i8* %func) + +Overview: +""""""""" + +The '``llvm.coro.sizeof``' is the same as `coro.size` except it returns the +coroutine frame size of the parameter. + +Arguments: +"""""""""" + +The func argument to '``llvm.coro.sizeof``' must be a constant bitcasted pointer +to a function defined in the current module. + +Semantics: +"""""""""" + +Same as `coro.size` except for the function passed as a parameter instead of +the current function. + .. _coro.begin: 'llvm.coro.begin' Intrinsic Index: llvm/include/llvm/IR/Intrinsics.td =================================================================== --- llvm/include/llvm/IR/Intrinsics.td +++ llvm/include/llvm/IR/Intrinsics.td @@ -769,6 +769,7 @@ def int_coro_frame : Intrinsic<[llvm_ptr_ty], [], [IntrNoMem]>; def int_coro_size : Intrinsic<[llvm_anyint_ty], [], [IntrNoMem]>; +def int_coro_sizeof : Intrinsic<[llvm_anyint_ty], [llvm_ptr_ty], [IntrNoMem]>; def int_coro_save : Intrinsic<[llvm_token_ty], [llvm_ptr_ty], []>; def int_coro_suspend : Intrinsic<[llvm_i8_ty], [llvm_token_ty, llvm_i1_ty], []>; Index: llvm/lib/Transforms/Coroutines/CoroEarly.cpp =================================================================== --- llvm/lib/Transforms/Coroutines/CoroEarly.cpp +++ llvm/lib/Transforms/Coroutines/CoroEarly.cpp @@ -123,6 +123,9 @@ switch (CS.getIntrinsicID()) { default: continue; + case Intrinsic::coro_sizeof: + F.addFnAttr(CORO_SIZEOF_ATTR, HAS_SIZEOF_CALL); + continue; case Intrinsic::coro_free: CoroFrees.push_back(cast(&I)); break; @@ -195,7 +198,8 @@ if (coro::declaresIntrinsics(M, {"llvm.coro.id", "llvm.coro.destroy", "llvm.coro.done", "llvm.coro.end", "llvm.coro.free", "llvm.coro.promise", - "llvm.coro.resume", "llvm.coro.suspend"})) + "llvm.coro.resume", "llvm.coro.suspend", + "llvm.coro.sizeof" })) L = llvm::make_unique(M); return false; } Index: llvm/lib/Transforms/Coroutines/CoroInstr.h =================================================================== --- llvm/lib/Transforms/Coroutines/CoroInstr.h +++ llvm/lib/Transforms/Coroutines/CoroInstr.h @@ -299,6 +299,18 @@ } }; +/// This represents the llvm.coro.sizeof instruction. +class LLVM_LIBRARY_VISIBILITY CoroSizeofInst : public IntrinsicInst { +public: + // Methods to support type inquiry through isa, cast, and dyn_cast: + static bool classof(const IntrinsicInst *I) { + return I->getIntrinsicID() == Intrinsic::coro_sizeof; + } + static bool classof(const Value *V) { + return isa(V) && classof(cast(V)); + } +}; + /// This represents the llvm.coro.end instruction. class LLVM_LIBRARY_VISIBILITY CoroEndInst : public IntrinsicInst { enum { FrameArg, UnwindArg }; Index: llvm/lib/Transforms/Coroutines/CoroInternal.h =================================================================== --- llvm/lib/Transforms/Coroutines/CoroInternal.h +++ llvm/lib/Transforms/Coroutines/CoroInternal.h @@ -39,6 +39,15 @@ #define UNPREPARED_FOR_SPLIT "0" #define PREPARED_FOR_SPLIT "1" +// CoroEarly pass marks functions that do coro.sizeof +// and then CoroSplit goes back and lowers this to the frame size. +#define CORO_SIZEOF_ATTR "coroutine.sizeof" +#define HAS_SIZEOF_CALL "1" + +// CoroSplit adds this attribute so that it can be accessed later +#define CORO_FRAMESIZE_ATTR "coroutine.framesize" + + #define CORO_DEVIRT_TRIGGER_FN "coro.devirt.trigger" namespace coro { Index: llvm/lib/Transforms/Coroutines/CoroSplit.cpp =================================================================== --- llvm/lib/Transforms/Coroutines/CoroSplit.cpp +++ llvm/lib/Transforms/Coroutines/CoroSplit.cpp @@ -349,7 +349,7 @@ } } -static void replaceFrameSize(coro::Shape &Shape) { +static void replaceFrameSize(Function &F, coro::Shape &Shape) { if (Shape.CoroSizes.empty()) return; @@ -360,6 +360,8 @@ auto Size = DL.getTypeAllocSize(Shape.FrameTy); auto *SizeConstant = ConstantInt::get(SizeIntrin->getType(), Size); + F.addFnAttr(CORO_FRAMESIZE_ATTR, std::to_string(Size)); + for (CoroSizeInst *CS : Shape.CoroSizes) { CS->replaceAllUsesWith(SizeConstant); CS->eraseFromParent(); @@ -708,6 +710,27 @@ } } +static void lowerSizeOf(Function &ParentF, CallGraph &CG, CallGraphSCC &SCC) { + SmallVector SizeOfInstrs; + for (auto IB = inst_begin(ParentF), IE = inst_end(ParentF); IB != IE;) { + Instruction &I = *IB++; + if (auto CS = CallSite(&I)) { + if (CS.getIntrinsicID() == Intrinsic::coro_sizeof) { + SizeOfInstrs.push_back(cast(&I)); + } + } + } + for (CoroSizeofInst *CS : SizeOfInstrs) { + auto CoroF = cast(CS->getArgOperand(0)->stripPointerCasts()); + Attribute Attr = CoroF->getFnAttribute(CORO_FRAMESIZE_ATTR); + auto Size = Attr.getValueAsInt(); + auto *SizeConstant = ConstantInt::get(CS->getType(), Size); + CS->replaceAllUsesWith(SizeConstant); + CS->eraseFromParent(); + } + ParentF.removeFnAttr(CORO_SIZEOF_ATTR); +} + static void splitCoroutine(Function &F, CallGraph &CG, CallGraphSCC &SCC) { coro::Shape Shape(F); if (!Shape.CoroBegin) @@ -716,7 +739,7 @@ simplifySuspendPoints(Shape); relocateInstructionBefore(Shape.CoroBegin, F); buildCoroutineFrame(F, Shape); - replaceFrameSize(Shape); + replaceFrameSize(F, Shape); // If there are no suspend points, no split required, just remove // the allocation and deallocation blocks, they are not needed. @@ -836,28 +859,39 @@ // Find coroutines for processing. SmallVector Coroutines; - for (CallGraphNode *CGN : SCC) - if (auto *F = CGN->getFunction()) + SmallVector FnsWithSizeOf; + for (CallGraphNode *CGN : SCC) { + if (auto *F = CGN->getFunction()) { if (F->hasFnAttribute(CORO_PRESPLIT_ATTR)) Coroutines.push_back(F); + if (F->hasFnAttribute(CORO_SIZEOF_ATTR)) + FnsWithSizeOf.push_back(F); + } + } - if (Coroutines.empty()) + if (Coroutines.empty() && FnsWithSizeOf.empty()) { return false; + } CallGraph &CG = getAnalysis().getCallGraph(); - createDevirtTriggerFunc(CG, SCC); - - for (Function *F : Coroutines) { - Attribute Attr = F->getFnAttribute(CORO_PRESPLIT_ATTR); - StringRef Value = Attr.getValueAsString(); - DEBUG(dbgs() << "CoroSplit: Processing coroutine '" << F->getName() - << "' state: " << Value << "\n"); - if (Value == UNPREPARED_FOR_SPLIT) { - prepareForSplit(*F, CG); - continue; + if (!Coroutines.empty()) { + createDevirtTriggerFunc(CG, SCC); + + for (Function *F : Coroutines) { + Attribute Attr = F->getFnAttribute(CORO_PRESPLIT_ATTR); + StringRef Value = Attr.getValueAsString(); + DEBUG(dbgs() << "CoroSplit: Processing coroutine '" << F->getName() + << "' state: " << Value << "\n"); + if (Value == UNPREPARED_FOR_SPLIT) { + prepareForSplit(*F, CG); + continue; + } + F->removeFnAttr(CORO_PRESPLIT_ATTR); + splitCoroutine(*F, CG, SCC); } - F->removeFnAttr(CORO_PRESPLIT_ATTR); - splitCoroutine(*F, CG, SCC); + } + for (Function *F : FnsWithSizeOf) { + lowerSizeOf(*F, CG, SCC); } return true; } Index: llvm/lib/Transforms/Coroutines/Coroutines.cpp =================================================================== --- llvm/lib/Transforms/Coroutines/Coroutines.cpp +++ llvm/lib/Transforms/Coroutines/Coroutines.cpp @@ -127,7 +127,8 @@ "llvm.coro.done", "llvm.coro.end", "llvm.coro.frame", "llvm.coro.free", "llvm.coro.id", "llvm.coro.param", "llvm.coro.promise", "llvm.coro.resume", "llvm.coro.save", - "llvm.coro.size", "llvm.coro.subfn.addr", "llvm.coro.suspend", + "llvm.coro.size", "llvm.coro.sizeof", "llvm.coro.subfn.addr", + "llvm.coro.suspend", }; return Intrinsic::lookupLLVMIntrinsicByName(CoroIntrinsics, Name) != -1; }