https://github.com/yuxuanchen1997 updated https://github.com/llvm/llvm-project/pull/207799
>From f3528f86ab969030210998333741c0d6bbc2e902 Mon Sep 17 00:00:00 2001 From: Yuxuan Chen <[email protected]> Date: Mon, 6 Jul 2026 08:29:09 -0700 Subject: [PATCH] [Coroutines] Use destroy slot for CoroElide resume fallthrough Fixes https://github.com/llvm/llvm-project/issues/188230 A switch coroutine with a CoroElide noalloc variant uses the frame destroy slot as a runtime discriminator: heap-allocated frames store the destroy clone, while allocation-elided frames store the cleanup clone. Its resume clone can reach a fallthrough coro.end after the final suspend. The coroutine body can invoke a continuation before reaching that coro.end. A suspend_never continuation can complete and free its frame before control returns. When the callee frame was elided into that caller allocation, loading the destroy slot at coro.end reads freed storage, which AddressSanitizer reports as a use-after-free. Load and cache the destroy function in the switch-resume entry block before executing the resume body, then tail-call the cached value at fallthrough coro.end. The cache retains the frame slot as the allocation discriminator without dereferencing an elided frame after the continuation returns. Add a C++ CodeGen regression for the original suspend_never final-suspend shape and update CoroSplit tests to require the entry load. Assisted-By: Codex GPT 5.5 --- ...oro-await-elidable-suspend-never-final.cpp | 76 +++++++++++++++++ .../llvm/Transforms/Coroutines/CoroShape.h | 1 + llvm/lib/Transforms/Coroutines/CoroSplit.cpp | 52 ++++++++++-- llvm/lib/Transforms/Coroutines/Coroutines.cpp | 1 + .../Transforms/Coroutines/coro-split-00.ll | 5 +- .../Coroutines/coro-split-addrspace.ll | 5 +- ...o-split-resume-fallthrough-destroy-slot.ll | 83 +++++++++++++++++++ 7 files changed, 215 insertions(+), 8 deletions(-) create mode 100644 clang/test/CodeGenCoroutines/gh188230-coro-await-elidable-suspend-never-final.cpp create mode 100644 llvm/test/Transforms/Coroutines/coro-split-resume-fallthrough-destroy-slot.ll diff --git a/clang/test/CodeGenCoroutines/gh188230-coro-await-elidable-suspend-never-final.cpp b/clang/test/CodeGenCoroutines/gh188230-coro-await-elidable-suspend-never-final.cpp new file mode 100644 index 0000000000000..82c8255eb9af0 --- /dev/null +++ b/clang/test/CodeGenCoroutines/gh188230-coro-await-elidable-suspend-never-final.cpp @@ -0,0 +1,76 @@ +// Tests that a coro_await_elidable coroutine with a suspend_never final suspend +// destroys a resumed elided callee through the frame destroy slot. +// +// RUN: %clang_cc1 -triple x86_64-unknown-linux-gnu -std=c++20 -O2 \ +// RUN: -mllvm -coro-elide-branch-ratio=0 -emit-llvm %s -o - | FileCheck %s + +#include "Inputs/coroutine.h" + +struct gate { + std::coroutine_handle<> waiter = nullptr; + bool open = false; + + struct awaiter { + gate &g; + bool await_ready() noexcept { return g.open; } + void await_suspend(std::coroutine_handle<> h) noexcept { g.waiter = h; } + void await_resume() noexcept {} + }; + + awaiter operator co_await() noexcept { return {*this}; } +}; + +struct [[clang::coro_await_elidable]] task { + struct promise_type { + std::coroutine_handle<> continuation = nullptr; + + task get_return_object() noexcept { + return {std::coroutine_handle<promise_type>::from_promise(*this)}; + } + + std::suspend_never initial_suspend() noexcept { return {}; } + std::suspend_never final_suspend() noexcept { return {}; } + + void return_void() noexcept { + if (continuation) + continuation.resume(); + } + + void unhandled_exception() noexcept { __builtin_abort(); } + }; + + std::coroutine_handle<promise_type> handle; + + bool await_ready() noexcept { return false; } + + void await_suspend(std::coroutine_handle<> h) noexcept { + handle.promise().continuation = h; + } + + void await_resume() noexcept {} +}; + +task callee(gate &g, int &value) { + co_await g; + value = 42; +} + +task caller(gate &g, int &value, bool &finished) { + co_await callee(g, value); + finished = true; +} + +// CHECK-LABEL: define internal void @_Z6calleeR4gateRi.resume( +// CHECK: %[[DESTROY_ADDR:.+]] = getelementptr inbounds{{.*}} i8, ptr %{{.+}}, i64 8 +// CHECK-NEXT: %[[DESTROY:.+]] = load ptr, ptr %[[DESTROY_ADDR]] +// CHECK: store i32 42, +// CHECK-NOT: call void @_Zdl +// CHECK: tail call void %[[DESTROY]]( +// CHECK-NEXT: ret void + +// CHECK-LABEL: define internal {{.*}}void @_Z6calleeR4gateRi.destroy( +// CHECK: call void @_Zdl + +// CHECK-LABEL: define internal {{.*}}void @_Z6calleeR4gateRi.cleanup( +// CHECK-NOT: call void @_Zdl +// CHECK: ret void diff --git a/llvm/include/llvm/Transforms/Coroutines/CoroShape.h b/llvm/include/llvm/Transforms/Coroutines/CoroShape.h index 28931e3260e68..534e63f28ddc1 100644 --- a/llvm/include/llvm/Transforms/Coroutines/CoroShape.h +++ b/llvm/include/llvm/Transforms/Coroutines/CoroShape.h @@ -112,6 +112,7 @@ struct Shape { unsigned IndexOffset; bool HasFinalSuspend; bool HasUnwindCoroEnd; + bool HasCoroElideNoAllocVariant; }; struct RetconLoweringStorage { diff --git a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp index 7915fbf4cce05..4751e312676ff 100644 --- a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp +++ b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp @@ -164,6 +164,15 @@ static void maybeFreeRetconStorage(IRBuilder<> &Builder, Shape.emitDealloc(Builder, FramePtr, CG); } +/// Create a pointer to the switch destroy function field in the coroutine +/// frame. +static Value *createSwitchDestroyPtr(const coro::Shape &Shape, + IRBuilder<> &Builder, Value *FramePtr) { + auto *Offset = ConstantInt::get(Type::getInt64Ty(FramePtr->getContext()), + Shape.SwitchLowering.DestroyOffset); + return Builder.CreateInBoundsPtrAdd(FramePtr, Offset, "destroy.addr"); +} + /// Replace an llvm.coro.end.async. /// Will inline the must tail call function call if there is one. /// \returns true if cleanup of the coro.end block is needed, false otherwise. @@ -212,7 +221,8 @@ static bool replaceCoroEndAsync(AnyCoroEndInst *End) { /// Replace a non-unwind call to llvm.coro.end. static void replaceFallthroughCoroEnd(AnyCoroEndInst *End, const coro::Shape &Shape, Value *FramePtr, - bool InRamp, CallGraph *CG) { + bool InRamp, CallGraph *CG, + Value *DestroyFn) { // Start inserting right before the coro.end. IRBuilder<> Builder(End); @@ -226,6 +236,11 @@ static void replaceFallthroughCoroEnd(AnyCoroEndInst *End, // in this lowering, because we need to deallocate the coroutine. if (InRamp) return; + if (DestroyFn) { + auto *Call = Builder.CreateCall(Shape.getResumeFunctionType(), DestroyFn, + FramePtr); + Call->setTailCallKind(CallInst::TCK_Tail); + } Builder.CreateRetVoid(); break; @@ -384,11 +399,12 @@ static void replaceUnwindCoroEnd(AnyCoroEndInst *End, const coro::Shape &Shape, } static void replaceCoroEnd(AnyCoroEndInst *End, const coro::Shape &Shape, - Value *FramePtr, bool InRamp, CallGraph *CG) { + Value *FramePtr, bool InRamp, CallGraph *CG, + Value *DestroyFn = nullptr) { if (End->isUnwind()) replaceUnwindCoroEnd(End, Shape, FramePtr, InRamp, CG); else - replaceFallthroughCoroEnd(End, Shape, FramePtr, InRamp, CG); + replaceFallthroughCoroEnd(End, Shape, FramePtr, InRamp, CG, DestroyFn); End->eraseFromParent(); } @@ -549,11 +565,27 @@ void coro::BaseCloner::replaceCoroSuspends() { } void coro::BaseCloner::replaceCoroEnds() { + Value *DestroyFn = nullptr; + if (FKind == CloneKind::SwitchResume && + Shape.SwitchLowering.HasCoroElideNoAllocVariant) { + // A resumed coroutine can invoke user code that destroys the allocation + // containing an elided callee frame. Cache the function pointer before + // entering the resume body so the fallthrough coro.end does not read a + // freed frame. + IRBuilder<> EntryBuilder(&NewF->getEntryBlock(), + NewF->getEntryBlock().getFirstInsertionPt()); + Value *DestroyAddr = + createSwitchDestroyPtr(Shape, EntryBuilder, NewFramePtr); + DestroyFn = EntryBuilder.CreateLoad(Shape.getSwitchResumePointerType(), + DestroyAddr, "destroy"); + } + for (AnyCoroEndInst *CE : Shape.CoroEnds) { // We use a null call graph because there's no call graph node for // the cloned function yet. We'll just be rebuilding that later. auto *NewCE = cast<AnyCoroEndInst>(VMap[CE]); - replaceCoroEnd(NewCE, Shape, NewFramePtr, /*in ramp*/ false, nullptr); + replaceCoroEnd(NewCE, Shape, NewFramePtr, /*in ramp*/ false, nullptr, + DestroyFn); } } @@ -1099,8 +1131,13 @@ void coro::SwitchCloner::create() { // Clone the function coro::BaseCloner::create(); - // Replacing coro.free with 'null' in cleanup to suppress deallocation code. - if (FKind == coro::CloneKind::SwitchCleanup) + // Replacing coro.free with 'null' in cleanup suppresses deallocation code. + // Resume functions that fall through to coro.end delegate self-destruction + // through the frame destroy slot, which selects the cleanup clone for + // allocation-elided frames. + if (FKind == coro::CloneKind::SwitchCleanup || + (FKind == coro::CloneKind::SwitchResume && + Shape.SwitchLowering.HasCoroElideNoAllocVariant)) elideCoroFree(NewFramePtr); } @@ -2014,6 +2051,9 @@ static void doSplitCoroutine(Function &F, SmallVectorImpl<Function *> &Clones, bool shouldCreateNoAllocVariant = !isNoSuspendCoroutine && Shape.ABI == coro::ABI::Switch && hasSafeElideCaller(F) && !F.hasFnAttribute(llvm::Attribute::NoInline); + if (Shape.ABI == coro::ABI::Switch) + Shape.SwitchLowering.HasCoroElideNoAllocVariant = + shouldCreateNoAllocVariant; // If there are no suspend points, no split required, just remove // the allocation and deallocation blocks, they are not needed. diff --git a/llvm/lib/Transforms/Coroutines/Coroutines.cpp b/llvm/lib/Transforms/Coroutines/Coroutines.cpp index a922099a1f43f..2ecf66d18b5d3 100644 --- a/llvm/lib/Transforms/Coroutines/Coroutines.cpp +++ b/llvm/lib/Transforms/Coroutines/Coroutines.cpp @@ -294,6 +294,7 @@ void coro::Shape::analyze(Function &F, ABI = coro::ABI::Switch; SwitchLowering.HasFinalSuspend = HasFinalSuspend; SwitchLowering.HasUnwindCoroEnd = HasUnwindCoroEnd; + SwitchLowering.HasCoroElideNoAllocVariant = false; auto SwitchId = getSwitchCoroId(); SwitchLowering.ResumeSwitch = nullptr; diff --git a/llvm/test/Transforms/Coroutines/coro-split-00.ll b/llvm/test/Transforms/Coroutines/coro-split-00.ll index 727b9b2b9776e..0fd86d23ade39 100644 --- a/llvm/test/Transforms/Coroutines/coro-split-00.ll +++ b/llvm/test/Transforms/Coroutines/coro-split-00.ll @@ -51,11 +51,14 @@ entry: ; CHECK: ret ptr %hdl ; CHECK-LABEL: @f.resume({{.*}}) { +; CHECK: %[[DESTROY_ADDR:.+]] = getelementptr inbounds i8, ptr %hdl, i64 8 +; CHECK-NEXT: %[[DESTROY:.+]] = load ptr, ptr %[[DESTROY_ADDR]] ; CHECK-NOT: call ptr @malloc ; CHECK-NOT: call void @print(i32 0) ; CHECK: call void @print(i32 1) ; CHECK-NOT: call void @print(i32 0) -; CHECK: call void @free( +; CHECK-NOT: call void @free( +; CHECK: tail call void %[[DESTROY]](ptr %hdl) ; CHECK: ret void ; CHECK-LABEL: @f.destroy({{.*}}) { diff --git a/llvm/test/Transforms/Coroutines/coro-split-addrspace.ll b/llvm/test/Transforms/Coroutines/coro-split-addrspace.ll index 24db35d141b35..65a887a595230 100644 --- a/llvm/test/Transforms/Coroutines/coro-split-addrspace.ll +++ b/llvm/test/Transforms/Coroutines/coro-split-addrspace.ll @@ -52,11 +52,14 @@ entry: ; CHECK: ret ptr %hdl ; CHECK-LABEL: @f.resume({{.*}}) addrspace(200) { +; CHECK: %[[DESTROY_ADDR:.+]] = getelementptr inbounds i8, ptr %hdl, i64 8 +; CHECK-NEXT: %[[DESTROY:.+]] = load ptr, ptr %[[DESTROY_ADDR]] ; CHECK-NOT: call ptr @malloc ; CHECK-NOT: call void @print(i32 0) ; CHECK: call void @print(i32 1) ; CHECK-NOT: call void @print(i32 0) -; CHECK: call void @free( +; CHECK-NOT: call void @free( +; CHECK: tail call void %[[DESTROY]](ptr %hdl) ; CHECK: ret void ; CHECK-LABEL: @f.destroy({{.*}}) addrspace(200) { diff --git a/llvm/test/Transforms/Coroutines/coro-split-resume-fallthrough-destroy-slot.ll b/llvm/test/Transforms/Coroutines/coro-split-resume-fallthrough-destroy-slot.ll new file mode 100644 index 0000000000000..601c7694fa83e --- /dev/null +++ b/llvm/test/Transforms/Coroutines/coro-split-resume-fallthrough-destroy-slot.ll @@ -0,0 +1,83 @@ +; Tests that a switch coroutine resume clone which falls through to coro.end +; destroys itself through the frame destroy slot. This matters for allocation +; elision: the frame slot contains the cleanup clone for a stack-elided frame. +; +; RUN: opt < %s -passes='cgscc(coro-split),simplifycfg,early-cse' -S | FileCheck %s + +define ptr @f() presplitcoroutine { +entry: + %id = call token @llvm.coro.id(i32 0, ptr null, ptr @f, ptr null) + %need.alloc = call i1 @llvm.coro.alloc(token %id) + br i1 %need.alloc, label %dyn.alloc, label %begin + +dyn.alloc: + %size = call i32 @llvm.coro.size.i32() + %alloc = call ptr @malloc(i32 %size) + br label %begin + +begin: + %phi = phi ptr [ null, %entry ], [ %alloc, %dyn.alloc ] + %hdl = call ptr @llvm.coro.begin(token %id, ptr %phi) + call void @print(i32 0) + %0 = call i8 @llvm.coro.suspend(token none, i1 false) + switch i8 %0, label %suspend [i8 0, label %resume + i8 1, label %cleanup] + +resume: + call void @print(i32 1) + br label %cleanup + +cleanup: + %mem = call ptr @llvm.coro.free(token %id, ptr %hdl) + call void @free(ptr %mem) + br label %suspend + +suspend: + call void @llvm.coro.end(ptr %hdl, i1 false, token none) + ret ptr %hdl +} + +define void @caller() presplitcoroutine { +entry: + %ptr = call ptr @f() #0 + ret void +} + +; CHECK-LABEL: define ptr @f( +; CHECK: %[[NEED_ALLOC:.+]] = call i1 @llvm.coro.alloc( +; CHECK: %[[DESTROY_OR_CLEANUP:.+]] = select i1 %[[NEED_ALLOC]], ptr @f.destroy, ptr @f.cleanup +; CHECK: %[[DESTROY_ADDR:.+]] = getelementptr inbounds i8, ptr %hdl, i64 8 +; CHECK-NEXT: store ptr %[[DESTROY_OR_CLEANUP]], ptr %[[DESTROY_ADDR]] + +; CHECK-LABEL: define internal void @f.resume( +; CHECK: %[[DESTROY_ADDR:.+]] = getelementptr inbounds i8, ptr %hdl, i64 8 +; CHECK-NEXT: %[[DESTROY:.+]] = load ptr, ptr %[[DESTROY_ADDR]] +; CHECK: call void @print(i32 1) +; CHECK-NOT: call void @free( +; CHECK: tail call void %[[DESTROY]](ptr %hdl) +; CHECK-NEXT: ret void + +; CHECK-LABEL: define internal void @f.destroy( +; CHECK: call void @free( + +; CHECK-LABEL: define internal void @f.cleanup( +; CHECK-NOT: call void @free( +; CHECK: ret void + +; CHECK-LABEL: define internal ptr @f.noalloc( +; CHECK: %[[NOALLOC_DESTROY_ADDR:.+]] = getelementptr inbounds i8, ptr %{{.+}}, i64 8 +; CHECK-NEXT: store ptr @f.cleanup, ptr %[[NOALLOC_DESTROY_ADDR]] + +declare ptr @llvm.coro.free(token, ptr) +declare i32 @llvm.coro.size.i32() +declare i8 @llvm.coro.suspend(token, i1) +declare token @llvm.coro.id(i32, ptr, ptr, ptr) +declare i1 @llvm.coro.alloc(token) +declare ptr @llvm.coro.begin(token, ptr) +declare void @llvm.coro.end(ptr, i1, token) + +declare noalias ptr @malloc(i32) allockind("alloc,uninitialized") "alloc-family"="malloc" +declare void @print(i32) +declare void @free(ptr) willreturn allockind("free") "alloc-family"="malloc" + +attributes #0 = { coro_elide_safe } _______________________________________________ cfe-commits mailing list [email protected] https://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-commits
