Index: include/llvm/IR/Intrinsics.td =================================================================== --- include/llvm/IR/Intrinsics.td +++ include/llvm/IR/Intrinsics.td @@ -840,6 +840,7 @@ def int_coro_frame : Intrinsic<[llvm_ptr_ty], [], [IntrNoMem]>; def int_coro_noop : Intrinsic<[llvm_ptr_ty], [], [IntrNoMem]>; def int_coro_size : Intrinsic<[llvm_anyint_ty], [], [IntrNoMem]>; +def int_coro_size_chk : Intrinsic<[], [llvm_anyint_ty], []>; 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: lib/Transforms/Coroutines/CoroCleanup.cpp =================================================================== --- lib/Transforms/Coroutines/CoroCleanup.cpp +++ lib/Transforms/Coroutines/CoroCleanup.cpp @@ -79,6 +79,9 @@ case Intrinsic::coro_subfn_addr: lowerSubFn(Builder, cast(II)); break; + case Intrinsic::coro_size_chk: + // Just remove it. + break; } II->eraseFromParent(); Changed = true; Index: lib/Transforms/Coroutines/CoroElide.cpp =================================================================== --- lib/Transforms/Coroutines/CoroElide.cpp +++ lib/Transforms/Coroutines/CoroElide.cpp @@ -196,6 +196,19 @@ ResumeAddr.clear(); DestroyAddr.clear(); + bool HasSizeChecksIntrinsics = false; + + // TODO: For now, just the mere presense of the coro_size_chk disables + // the Halo optimization. We need to investigate if less blunt approach + // is possible. + for (Instruction &I : instructions(CoroId->getFunction())) + if (isa(&I)) { + LLVM_DEBUG(dbgs() << "disabling Halo optimization " + " due to presence of coro_size_chk\n"); + HasSizeChecksIntrinsics = false; + break; + } + // Collect all coro.begin and coro.allocs associated with this coro.id. for (User *U : CoroId->users()) { if (auto *CB = dyn_cast(U)) @@ -243,7 +256,7 @@ replaceWithConstant(DestroyAddrConstant, DestroyAddr); - if (ShouldElide) { + if (ShouldElide && !HasSizeChecksIntrinsics) { auto *FrameTy = getFrameType(cast(ResumeAddrConstant)); elideHeapAllocations(CoroId->getFunction(), FrameTy, AA); coro::replaceCoroFree(CoroId, /*Elide=*/true); Index: lib/Transforms/Coroutines/CoroInstr.h =================================================================== --- lib/Transforms/Coroutines/CoroInstr.h +++ lib/Transforms/Coroutines/CoroInstr.h @@ -299,6 +299,32 @@ } }; + +/// This represents the llvm.coro.size instruction. +class LLVM_LIBRARY_VISIBILITY CoroSizeChkInst : public IntrinsicInst { + enum { SizeEst }; + + ConstantInt *getRawIndex() const { + return cast(getArgOperand(SizeEst)); + } +public: + bool isConstant() const { + return isa(getArgOperand(SizeEst)); + } + + int64_t getSizeEstimate() const { + return getRawIndex()->getValue().getSExtValue(); + } + + // Methods to support type inquiry through isa, cast, and dyn_cast: + static bool classof(const IntrinsicInst *I) { + return I->getIntrinsicID() == Intrinsic::coro_size_chk; + } + 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: lib/Transforms/Coroutines/CoroInternal.h =================================================================== --- lib/Transforms/Coroutines/CoroInternal.h +++ lib/Transforms/Coroutines/CoroInternal.h @@ -68,6 +68,7 @@ CoroBeginInst *CoroBegin; SmallVector CoroEnds; SmallVector CoroSizes; + SmallVector CoroSizeChks; SmallVector CoroSuspends; // Field Indexes for known coroutine frame fields. Index: lib/Transforms/Coroutines/CoroSplit.cpp =================================================================== --- lib/Transforms/Coroutines/CoroSplit.cpp +++ lib/Transforms/Coroutines/CoroSplit.cpp @@ -364,6 +364,20 @@ CS->replaceAllUsesWith(SizeConstant); CS->eraseFromParent(); } + + for (CoroSizeChkInst *CS : Shape.CoroSizeChks) { + if (!CS->isConstant()) + report_fatal_error("In function '" + CS->getFunction()->getName() + + "' with coroutine frame size " + Twine(Size) + + " front end size estimate is not a constant."); + else if ((int64_t)Size > CS->getSizeEstimate()) + report_fatal_error("In function '" + CS->getFunction()->getName() + + "' coroutine frame size " + Twine(Size) + " exceeds the frontend estimate of " + + Twine(CS->getSizeEstimate()) + "."); + + // Note that we do not remove the check intrinsic, so that it can be + // used to inhibit Halo optimization. + } } // Create a global constant array containing pointers to functions provided and Index: lib/Transforms/Coroutines/Coroutines.cpp =================================================================== --- lib/Transforms/Coroutines/Coroutines.cpp +++ lib/Transforms/Coroutines/Coroutines.cpp @@ -249,6 +249,9 @@ case Intrinsic::coro_size: CoroSizes.push_back(cast(II)); break; + case Intrinsic::coro_size_chk: + CoroSizeChks.push_back(cast(II)); + break; case Intrinsic::coro_frame: CoroFrames.push_back(cast(II)); break;