diff --git a/llvm/include/llvm/IR/Intrinsics.td b/llvm/include/llvm/IR/Intrinsics.td --- a/llvm/include/llvm/IR/Intrinsics.td +++ b/llvm/include/llvm/IR/Intrinsics.td @@ -1216,6 +1216,7 @@ def int_coro_async_resume : Intrinsic<[llvm_ptr_ty], [], []>; +def int_coro_async_size_replace : Intrinsic<[], [llvm_ptr_ty, llvm_ptr_ty], []>; def int_coro_suspend_async : Intrinsic<[llvm_any_ty], [llvm_ptr_ty, llvm_ptr_ty, llvm_vararg_ty], []>; def int_coro_prepare_async : Intrinsic<[llvm_ptr_ty], [llvm_ptr_ty], diff --git a/llvm/lib/Transforms/Coroutines/CoroCleanup.cpp b/llvm/lib/Transforms/Coroutines/CoroCleanup.cpp --- a/llvm/lib/Transforms/Coroutines/CoroCleanup.cpp +++ b/llvm/lib/Transforms/Coroutines/CoroCleanup.cpp @@ -80,6 +80,23 @@ case Intrinsic::coro_subfn_addr: lowerSubFn(Builder, cast(II)); break; + case Intrinsic::coro_async_size_replace: + auto *Target = cast( + cast(II->getArgOperand(0)->stripPointerCasts()) + ->getInitializer()); + auto *Source = cast( + cast(II->getArgOperand(1)->stripPointerCasts()) + ->getInitializer()); + auto *TargetSize = Target->getOperand(1); + auto *SourceSize = Source->getOperand(1); + if (TargetSize->isElementWiseEqual(SourceSize)) { + break; + } + auto *TargetRelativeFunOffset = Target->getOperand(0); + auto *NewFuncPtrStruct = ConstantStruct::get( + Target->getType(), TargetRelativeFunOffset, SourceSize); + Target->replaceAllUsesWith(NewFuncPtrStruct); + break; } II->eraseFromParent(); Changed = true; @@ -95,10 +112,10 @@ } static bool declaresCoroCleanupIntrinsics(const Module &M) { - return coro::declaresIntrinsics(M, {"llvm.coro.alloc", "llvm.coro.begin", - "llvm.coro.subfn.addr", "llvm.coro.free", - "llvm.coro.id", "llvm.coro.id.retcon", - "llvm.coro.id.retcon.once"}); + return coro::declaresIntrinsics( + M, {"llvm.coro.alloc", "llvm.coro.begin", "llvm.coro.subfn.addr", + "llvm.coro.free", "llvm.coro.id", "llvm.coro.id.retcon", + "llvm.coro.id.retcon.once", "llvm.coro.async.size.replace"}); } PreservedAnalyses CoroCleanupPass::run(Function &F, diff --git a/llvm/lib/Transforms/Coroutines/CoroInstr.h b/llvm/lib/Transforms/Coroutines/CoroInstr.h --- a/llvm/lib/Transforms/Coroutines/CoroInstr.h +++ b/llvm/lib/Transforms/Coroutines/CoroInstr.h @@ -376,6 +376,18 @@ } }; +/// This represents the llvm.coro.async.size.replace instruction. +class LLVM_LIBRARY_VISIBILITY CoroAsyncSizeReplace : 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_async_size_replace; + } + static bool classof(const Value *V) { + return isa(V) && classof(cast(V)); + } +}; + /// This represents the llvm.coro.frame instruction. class LLVM_LIBRARY_VISIBILITY CoroFrameInst : public IntrinsicInst { public: diff --git a/llvm/lib/Transforms/Coroutines/Coroutines.cpp b/llvm/lib/Transforms/Coroutines/Coroutines.cpp --- a/llvm/lib/Transforms/Coroutines/Coroutines.cpp +++ b/llvm/lib/Transforms/Coroutines/Coroutines.cpp @@ -126,6 +126,7 @@ "llvm.coro.alloc", "llvm.coro.async.context.alloc", "llvm.coro.async.context.dealloc", + "llvm.coro.async.size.replace", "llvm.coro.async.store_resume", "llvm.coro.begin", "llvm.coro.destroy", diff --git a/llvm/test/Transforms/Coroutines/coro-async.ll b/llvm/test/Transforms/Coroutines/coro-async.ll --- a/llvm/test/Transforms/Coroutines/coro-async.ll +++ b/llvm/test/Transforms/Coroutines/coro-async.ll @@ -27,6 +27,15 @@ to i32), i32 128 ; Initial async context size without space for frame }> +@my_async_function_pa_fp = constant <{ i32, i32 }> + <{ i32 trunc ( + i64 sub ( + i64 ptrtoint (void (i8*, %async.task*, %async.actor*)* @my_async_function_pa to i64), + i64 ptrtoint (i32* getelementptr inbounds (<{ i32, i32 }>, <{ i32, i32 }>* @my_async_function_pa_fp, i32 0, i32 1) to i64) + ) + to i32), + i32 8 +}> ; Function that implements the dispatch to the callee function. define swiftcc void @my_async_function.my_other_async_function_fp.apply(i8* %fnPtr, i8* %async.ctxt, %async.task* %task, %async.actor* %actor) { @@ -99,8 +108,15 @@ unreachable } +define void @my_async_function_pa(i8* %ctxt, %async.task* %task, %async.actor* %actor) { + call void @llvm.coro.async.size.replace(i8* bitcast (<{i32, i32}>* @my_async_function_pa_fp to i8*), i8* bitcast (<{i32, i32}>* @my_async_function_fp to i8*)) + call swiftcc void @my_async_function(i8* %ctxt, %async.task* %task, %async.actor* %actor) + ret void +} + ; Make sure we update the async function pointer ; CHECK: @my_async_function_fp = constant <{ i32, i32 }> <{ {{.*}}, i32 176 } +; CHECK: @my_async_function_pa_fp = constant <{ i32, i32 }> <{ {{.*}}, i32 176 } ; CHECK: @my_async_function2_fp = constant <{ i32, i32 }> <{ {{.*}}, i32 176 } ; CHECK-LABEL: define swiftcc void @my_async_function(i8* %async.ctxt, %async.task* %task, %async.actor* %actor) { @@ -511,3 +527,4 @@ declare swiftcc void @asyncReturn(i8*, %async.task*, %async.actor*) declare swiftcc void @asyncSuspend(i8*, %async.task*, %async.actor*) declare i8* @llvm.coro.async.resume() +declare void @llvm.coro.async.size.replace(i8*, i8*)