Skip to content

Commit e190b7c

Browse files
ChuanqiXu9ChuanqiXu9
committed
[Coroutines] Maintain the position of final suspend
Closing llvm/llvm-project#56329 The problem happens when we try to simplify the suspend points. We might break the assumption that the final suspend lives in the last slot of Shape.CoroSuspends. This patch tries to main the assumption and fixes the problem.
1 parent 8d19cfb commit e190b7c

File tree

3 files changed

+264
-0
lines changed

3 files changed

+264
-0
lines changed

‎clang/test/CodeGenCoroutines/pr56329.cpp‎

Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,117 @@
1+
// Test for PR56919. Tests the we won't contain the resumption of final suspend point.
2+
//
3+
// RUN: %clang_cc1 -triple %itanium_abi_triple -std=c++20 %s -O3 -S -emit-llvm -o - | FileCheck %s
4+
5+
#include "Inputs/coroutine.h"
6+
7+
void _exit(int status) __attribute__ ((__noreturn__));
8+
9+
class Promise;
10+
11+
// An object that can be co_awaited, but we always resume immediately from
12+
// await_suspend.
13+
struct ResumeFromAwaitSuspend{};
14+
15+
struct Task {
16+
using promise_type = Promise;
17+
Promise& promise;
18+
};
19+
20+
struct Promise {
21+
static std::coroutine_handle<Promise> GetHandle(Promise& promise) {
22+
return std::coroutine_handle<Promise>::from_promise(promise);
23+
}
24+
25+
void unhandled_exception() {}
26+
Task get_return_object() { return Task{*this}; }
27+
void return_void() {}
28+
29+
// Always suspend before starting the coroutine body. We actually run the body
30+
// when we are co_awaited.
31+
std::suspend_always initial_suspend() { return {}; }
32+
33+
// We support awaiting tasks. We do so by configuring them to resume us when
34+
// they are finished, and then resuming them from their initial suspend.
35+
auto await_transform(Task&& task) {
36+
struct Awaiter {
37+
bool await_ready() { return false; }
38+
39+
std::coroutine_handle<> await_suspend(
40+
const std::coroutine_handle<> handle) {
41+
// Tell the child to resume the parent once it finishes.
42+
child.resume_at_final_suspend = GetHandle(parent);
43+
44+
// Run the child.
45+
return GetHandle(child);
46+
}
47+
48+
void await_resume() {
49+
// The child is now at its final suspend point, and can be destroyed.
50+
return GetHandle(child).destroy();
51+
}
52+
53+
Promise& parent;
54+
Promise& child;
55+
};
56+
57+
return Awaiter{
58+
.parent = *this,
59+
.child = task.promise,
60+
};
61+
}
62+
63+
// Make evaluation of `co_await ResumeFromAwaitSuspend{}` go through the
64+
// await_suspend path, but cause it to resume immediately by returning our own
65+
// handle to resume.
66+
auto await_transform(ResumeFromAwaitSuspend) {
67+
struct Awaiter {
68+
bool await_ready() { return false; }
69+
70+
std::coroutine_handle<> await_suspend(const std::coroutine_handle<> h) {
71+
return h;
72+
}
73+
74+
void await_resume() {}
75+
};
76+
77+
return Awaiter{};
78+
}
79+
80+
// Always suspend at the final suspend point, transferring control back to our
81+
// caller. We expect never to be resumed from the final suspend.
82+
auto final_suspend() noexcept {
83+
struct FinalSuspendAwaitable final {
84+
bool await_ready() noexcept { return false; }
85+
86+
std::coroutine_handle<> await_suspend(std::coroutine_handle<>) noexcept {
87+
return promise.resume_at_final_suspend;
88+
}
89+
90+
void await_resume() noexcept {
91+
_exit(1);
92+
}
93+
94+
Promise& promise;
95+
};
96+
97+
return FinalSuspendAwaitable{.promise = *this};
98+
}
99+
100+
// The handle we will resume once we hit final suspend.
101+
std::coroutine_handle<> resume_at_final_suspend;
102+
};
103+
104+
Task Inner();
105+
106+
Task Outer() {
107+
co_await ResumeFromAwaitSuspend();
108+
co_await Inner();
109+
}
110+
111+
// CHECK: define{{.*}}@_Z5Outerv.resume(
112+
// CHECK-NOT: }
113+
// CHECK-NOT: _exit
114+
// CHECK: musttail call
115+
// CHECK: musttail call
116+
// CHECK-NEXT: ret void
117+
// CHEKC-NEXT: }

‎llvm/lib/Transforms/Coroutines/CoroSplit.cpp‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1555,20 +1555,36 @@ static void simplifySuspendPoints(coro::Shape &Shape) {
15551555
size_t I = 0, N = S.size();
15561556
if (N == 0)
15571557
return;
1558+
1559+
size_t ChangedFinalIndex = std::numeric_limits<size_t>::max();
15581560
while (true) {
15591561
auto SI = cast<CoroSuspendInst>(S[I]);
15601562
// Leave final.suspend to handleFinalSuspend since it is undefined behavior
15611563
// to resume a coroutine suspended at the final suspend point.
15621564
if (!SI->isFinal() && simplifySuspendPoint(SI, Shape.CoroBegin)) {
15631565
if (--N == I)
15641566
break;
1567+
15651568
std::swap(S[I], S[N]);
1569+
1570+
if (cast<CoroSuspendInst>(S[I])->isFinal()) {
1571+
assert(Shape.SwitchLowering.HasFinalSuspend);
1572+
ChangedFinalIndex = I;
1573+
}
1574+
15661575
continue;
15671576
}
15681577
if (++I == N)
15691578
break;
15701579
}
15711580
S.resize(N);
1581+
1582+
// Maintain final.suspend in case final suspend was swapped.
1583+
// Due to we requrie the final suspend to be the last element of CoroSuspends.
1584+
if (ChangedFinalIndex < N) {
1585+
assert(cast<CoroSuspendInst>(S[ChangedFinalIndex])->isFinal());
1586+
std::swap(S[ChangedFinalIndex], S.back());
1587+
}
15721588
}
15731589

15741590
static void splitSwitchCoroutine(Function &F, coro::Shape &Shape,

‎llvm/test/Transforms/Coroutines/coro-preserve-final.ll‎

Lines changed: 131 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,131 @@
1+
; RUN: opt < %s -passes='cgscc(coro-split),simplifycfg,early-cse' -S | FileCheck %s
2+
3+
%"struct.std::__n4861::noop_coroutine_promise" = type { i8 }
4+
%struct.Promise = type { %"struct.std::__n4861::coroutine_handle" }
5+
%"struct.std::__n4861::coroutine_handle" = type { ptr }
6+
7+
define dso_local ptr @_Z5Outerv() #1 {
8+
entry:
9+
%__promise = alloca %struct.Promise, align 8
10+
%0 = call token @llvm.coro.id(i32 16, ptr nonnull %__promise, ptr nonnull @_Z5Outerv, ptr null)
11+
%1 = call i1 @llvm.coro.alloc(token %0)
12+
br i1 %1, label %coro.alloc, label %init.suspend
13+
14+
coro.alloc: ; preds = %entry
15+
%2 = tail call i64 @llvm.coro.size.i64()
16+
%call = call noalias noundef nonnull ptr @_Znwm(i64 noundef %2) #12
17+
br label %init.suspend
18+
19+
init.suspend: ; preds = %entry, %coro.alloc
20+
%3 = phi ptr [ null, %entry ], [ %call, %coro.alloc ]
21+
%4 = call ptr @llvm.coro.begin(token %0, ptr %3) #13
22+
call void @llvm.lifetime.start.p0(i64 8, ptr nonnull %__promise) #3
23+
store ptr null, ptr %__promise, align 8
24+
%5 = call token @llvm.coro.save(ptr null)
25+
%6 = call i8 @llvm.coro.suspend(token %5, i1 false)
26+
switch i8 %6, label %coro.ret [
27+
i8 0, label %await.suspend
28+
i8 1, label %cleanup62
29+
]
30+
31+
await.suspend: ; preds = %init.suspend
32+
%7 = call token @llvm.coro.save(ptr null)
33+
%8 = call ptr @llvm.coro.subfn.addr(ptr %4, i8 0)
34+
call fastcc void %8(ptr %4) #3
35+
%9 = call i8 @llvm.coro.suspend(token %7, i1 false)
36+
switch i8 %9, label %coro.ret [
37+
i8 0, label %await2.suspend
38+
i8 1, label %cleanup62
39+
]
40+
41+
await2.suspend: ; preds = %await.suspend
42+
%call27 = call ptr @_Z5Innerv() #3
43+
%10 = call token @llvm.coro.save(ptr null)
44+
%11 = getelementptr inbounds i8, ptr %__promise, i64 -16
45+
store ptr %11, ptr %call27, align 8
46+
%12 = getelementptr inbounds i8, ptr %call27, i64 -16
47+
%13 = call ptr @llvm.coro.subfn.addr(ptr nonnull %12, i8 0)
48+
call fastcc void %13(ptr nonnull %12) #3
49+
%14 = call i8 @llvm.coro.suspend(token %10, i1 false)
50+
switch i8 %14, label %coro.ret [
51+
i8 0, label %final.suspend
52+
i8 1, label %cleanup62
53+
]
54+
55+
final.suspend: ; preds = %await2.suspend
56+
%15 = call ptr @llvm.coro.subfn.addr(ptr nonnull %12, i8 1)
57+
call fastcc void %15(ptr nonnull %12) #3
58+
%16 = call token @llvm.coro.save(ptr null)
59+
%retval.sroa.0.0.copyload.i = load ptr, ptr %__promise, align 8
60+
%17 = call ptr @llvm.coro.subfn.addr(ptr %retval.sroa.0.0.copyload.i, i8 0)
61+
call fastcc void %17(ptr %retval.sroa.0.0.copyload.i) #3
62+
%18 = call i8 @llvm.coro.suspend(token %16, i1 true) #13
63+
switch i8 %18, label %coro.ret [
64+
i8 0, label %final.ready
65+
i8 1, label %cleanup62
66+
]
67+
68+
final.ready: ; preds = %final.suspend
69+
call void @_Z5_exiti(i32 noundef 1) #14
70+
unreachable
71+
72+
cleanup62: ; preds = %await2.suspend, %await.suspend, %init.suspend, %final.suspend
73+
call void @llvm.lifetime.end.p0(i64 8, ptr nonnull %__promise) #3
74+
%19 = call ptr @llvm.coro.free(token %0, ptr %4)
75+
%.not = icmp eq ptr %19, null
76+
br i1 %.not, label %coro.ret, label %coro.free
77+
78+
coro.free: ; preds = %cleanup62
79+
call void @_ZdlPv(ptr noundef nonnull %19) #3
80+
br label %coro.ret
81+
82+
coro.ret: ; preds = %coro.free, %cleanup62, %final.suspend, %await2.suspend, %await.suspend, %init.suspend
83+
%20 = call i1 @llvm.coro.end(ptr null, i1 false) #13
84+
ret ptr %__promise
85+
}
86+
87+
declare token @llvm.coro.id(i32, ptr readnone, ptr nocapture readonly, ptr) #2
88+
declare i1 @llvm.coro.alloc(token) #3
89+
declare dso_local noundef nonnull ptr @_Znwm(i64 noundef) local_unnamed_addr #4
90+
declare i64 @llvm.coro.size.i64() #5
91+
declare ptr @llvm.coro.begin(token, ptr writeonly) #3
92+
declare void @llvm.lifetime.start.p0(i64 immarg, ptr nocapture) #6
93+
declare token @llvm.coro.save(ptr) #7
94+
declare void @llvm.lifetime.end.p0(i64 immarg, ptr nocapture) #6
95+
declare i8 @llvm.coro.suspend(token, i1) #3
96+
declare dso_local ptr @_Z5Innerv() local_unnamed_addr #8
97+
declare dso_local void @_ZdlPv(ptr noundef) local_unnamed_addr #9
98+
declare ptr @llvm.coro.free(token, ptr nocapture readonly) #2
99+
declare i1 @llvm.coro.end(ptr, i1) #3
100+
declare dso_local void @_Z5_exiti(i32 noundef) local_unnamed_addr #10
101+
declare ptr @llvm.coro.subfn.addr(ptr nocapture readonly, i8) #11
102+
103+
attributes #0 = { mustprogress nounwind uwtable "frame-pointer"="none" "min-legal-vector-width"="0" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "tune-cpu"="generic" }
104+
attributes #1 = { nounwind presplitcoroutine uwtable "frame-pointer"="none" "min-legal-vector-width"="0" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "tune-cpu"="generic" }
105+
attributes #2 = { argmemonly nofree nounwind readonly }
106+
attributes #3 = { nounwind }
107+
attributes #4 = { nobuiltin allocsize(0) "frame-pointer"="none" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "tune-cpu"="generic" }
108+
attributes #5 = { nofree nosync nounwind readnone }
109+
attributes #6 = { argmemonly mustprogress nocallback nofree nosync nounwind willreturn }
110+
attributes #7 = { nomerge nounwind }
111+
attributes #8 = { "frame-pointer"="none" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "tune-cpu"="generic" }
112+
attributes #9 = { nobuiltin nounwind "frame-pointer"="none" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "tune-cpu"="generic" }
113+
attributes #10 = { noreturn "frame-pointer"="none" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "tune-cpu"="generic" }
114+
attributes #11 = { argmemonly nounwind readonly }
115+
attributes #12 = { nounwind allocsize(0) }
116+
attributes #13 = { noduplicate }
117+
attributes #14 = { noreturn nounwind }
118+
119+
; CHECK: define{{.*}}@_Z5Outerv.resume(
120+
; CHECK: entry.resume:
121+
; CHECK: switch i2 %index
122+
; CHECK-NEXT: i2 0, label %await2.suspend
123+
; CHECK-NEXT: i2 1, label %final.suspend
124+
;
125+
; CHECK: await2.suspend:
126+
; CHECK: musttail call
127+
; CHECK-NEXT: ret void
128+
;
129+
; CHECK: final.suspend:
130+
; CHECK: musttail call
131+
; CHECK-NEXT: ret void

0 commit comments

Comments
 (0)