https://github.com/yuxuanchen1997 updated https://github.com/llvm/llvm-project/pull/207799
>From dbca16b70c1becb507741e8d7fb9f252b2754d19 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 When a switch coroutine is allocation-elided, resuming it can run through to a coro.end path and destroy the coroutine frame from inside the resume clone. That is fine for a heap-allocated frame, but it is wrong for the clang::coro_await_elidable CoroElide path: the caller may have supplied stack storage through the noalloc ramp. In that case, the frame's destroy slot is initialized to the cleanup clone, while the resume clone still emitted the normal coro.free/deallocation sequence and attempted to free storage owned by the caller. This is not limited to coroutines where CoroSplit fails to see a final suspend. Clang still emits a final-suspend coro.suspend for a C++ final_suspend returning suspend_never, but that suspend can flow through to coro.end and self-destroy when resumed. For switch resume clones with a generated CoroElide noalloc variant, the fallthrough coro.end path now loads the destroy/cleanup function pointer from the coroutine frame and tail-calls it. The frame slot is the runtime discriminator: elided frames dispatch to cleanup and heap frames dispatch to destroy. Add an IR regression for the split lowering and a C++ CodeGen regression for the original suspend_never final_suspend shape. The C++ test forces CoroAnnotationElide with -coro-elide-branch-ratio=0 because the current branch-frequency heuristic may otherwise skip the elision and hide the bug. 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 | 41 +++++++-- 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, 204 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..e8fcf8a501b84 --- /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: store i32 42, +// CHECK-NOT: call void @_Zdl +// CHECK: %[[DESTROY_ADDR:.+]] = getelementptr inbounds{{.*}} i8, ptr %{{.+}}, i64 8 +// CHECK-NEXT: %[[DESTROY:.+]] = load ptr, ptr %[[DESTROY_ADDR]] +// CHECK-NEXT: 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..d396d0c750618 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, + bool UseDestroySlot) { // Start inserting right before the coro.end. IRBuilder<> Builder(End); @@ -226,6 +236,14 @@ static void replaceFallthroughCoroEnd(AnyCoroEndInst *End, // in this lowering, because we need to deallocate the coroutine. if (InRamp) return; + if (UseDestroySlot) { + Value *DestroyAddr = createSwitchDestroyPtr(Shape, Builder, FramePtr); + Value *DestroyFn = Builder.CreateLoad(Shape.getSwitchResumePointerType(), + DestroyAddr, "destroy"); + auto *Call = Builder.CreateCall(Shape.getResumeFunctionType(), DestroyFn, + FramePtr); + Call->setTailCallKind(CallInst::TCK_Tail); + } Builder.CreateRetVoid(); break; @@ -384,11 +402,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, + bool UseDestroySlot = false) { if (End->isUnwind()) replaceUnwindCoroEnd(End, Shape, FramePtr, InRamp, CG); else - replaceFallthroughCoroEnd(End, Shape, FramePtr, InRamp, CG); + replaceFallthroughCoroEnd(End, Shape, FramePtr, InRamp, CG, UseDestroySlot); End->eraseFromParent(); } @@ -553,7 +572,9 @@ void coro::BaseCloner::replaceCoroEnds() { // 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, + FKind == CloneKind::SwitchResume && + Shape.SwitchLowering.HasCoroElideNoAllocVariant); } } @@ -1099,8 +1120,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 +2040,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..d8b36d2d423e0 100644 --- a/llvm/test/Transforms/Coroutines/coro-split-00.ll +++ b/llvm/test/Transforms/Coroutines/coro-split-00.ll @@ -55,7 +55,10 @@ entry: ; 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: %[[DESTROY_ADDR:.+]] = getelementptr inbounds i8, ptr %hdl, i64 8 +; CHECK-NEXT: %[[DESTROY:.+]] = load ptr, ptr %[[DESTROY_ADDR]] +; CHECK-NEXT: 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..4da8140b3de03 100644 --- a/llvm/test/Transforms/Coroutines/coro-split-addrspace.ll +++ b/llvm/test/Transforms/Coroutines/coro-split-addrspace.ll @@ -56,7 +56,10 @@ entry: ; 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: %[[DESTROY_ADDR:.+]] = getelementptr inbounds i8, ptr %hdl, i64 8 +; CHECK-NEXT: %[[DESTROY:.+]] = load ptr, ptr %[[DESTROY_ADDR]] +; CHECK-NEXT: 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..80dd6e3eebedc --- /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: call void @print(i32 1) +; CHECK-NOT: call void @free( +; CHECK: %[[DESTROY_ADDR:.+]] = getelementptr inbounds i8, ptr %hdl, i64 8 +; CHECK-NEXT: %[[DESTROY:.+]] = load ptr, ptr %[[DESTROY_ADDR]] +; CHECK-NEXT: 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
