Skip to content

Commit 5ad3373

Browse files
fpasserbyfpasserby
committed
Add await suspend lowering tests
1 parent 8dc3de9 commit 5ad3373

File tree

5 files changed

+328
-16
lines changed

5 files changed

+328
-16
lines changed

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

Lines changed: 12 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1412,10 +1412,10 @@ static bool hasCallsBetween(Instruction *Save, Instruction *ResumeOrDestroy) {
14121412
auto *ResumeOrDestroyBB = ResumeOrDestroy->getParent();
14131413

14141414
if (SaveBB == ResumeOrDestroyBB)
1415-
return hasCallsInBlockBetween(Save->getNextNode(), ResumeOrDestroy);
1415+
return hasCallsInBlockBetween(Save, ResumeOrDestroy);
14161416

14171417
// Any calls from Save to the end of the block?
1418-
if (hasCallsInBlockBetween(Save->getNextNode(), nullptr))
1418+
if (hasCallsInBlockBetween(Save, nullptr))
14191419
return true;
14201420

14211421
// Any calls from begging of the block up to ResumeOrDestroy?
@@ -1437,7 +1437,10 @@ static bool hasCallsBetween(Instruction *Save, Instruction *ResumeOrDestroy) {
14371437
// resume or destroy it
14381438
// FIXME: perform more sophisiticated analysis?
14391439
static bool isSimpleWrapper(CoroAwaitSuspendInst *AWS) {
1440-
auto Wrapper = AWS->getWrapperFunction();
1440+
auto *Wrapper = AWS->getWrapperFunction();
1441+
1442+
if (Wrapper->empty())
1443+
return false;
14411444

14421445
SmallVector<ReturnInst *, 4> Rets;
14431446

@@ -1575,17 +1578,12 @@ static void simplifySuspendPoints(coro::Shape &Shape) {
15751578
namespace {
15761579

15771580
struct SwitchCoroutineSplitter {
1578-
static void split(Module &M, Function &F, coro::Shape &Shape,
1581+
static void split(Function &F, coro::Shape &Shape,
15791582
SmallVectorImpl<Function *> &Clones,
15801583
TargetTransformInfo &TTI) {
15811584
assert(Shape.ABI == coro::ABI::Switch);
15821585

15831586
createResumeEntryBlock(F, Shape);
1584-
1585-
IRBuilder<> Builder(M.getContext());
1586-
for (auto *AWS : Shape.CoroAwaitSuspends)
1587-
lowerAwaitSuspend(Builder, AWS);
1588-
15891587
auto *ResumeClone =
15901588
createClone(F, ".resume", Shape, CoroCloner::Kind::SwitchResume);
15911589
auto *DestroyClone =
@@ -2122,14 +2120,18 @@ splitCoroutine(Module &M, Function &F, SmallVectorImpl<Function *> &Clones,
21222120
buildCoroutineFrame(F, Shape, TTI, MaterializableCallback);
21232121
replaceFrameSizeAndAlignment(Shape);
21242122

2123+
IRBuilder<> Builder(M.getContext());
2124+
for (auto *AWS : Shape.CoroAwaitSuspends)
2125+
lowerAwaitSuspend(Builder, AWS);
2126+
21252127
// If there are no suspend points, no split required, just remove
21262128
// the allocation and deallocation blocks, they are not needed.
21272129
if (Shape.CoroSuspends.empty()) {
21282130
handleNoSuspendCoroutine(Shape);
21292131
} else {
21302132
switch (Shape.ABI) {
21312133
case coro::ABI::Switch:
2132-
SwitchCoroutineSplitter::split(M, F, Shape, Clones, TTI);
2134+
SwitchCoroutineSplitter::split(F, Shape, Clones, TTI);
21332135
break;
21342136
case coro::ABI::Async:
21352137
splitAsyncCoroutine(F, Shape, Clones, TTI);

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

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -177,7 +177,11 @@ void coro::Shape::buildFrom(Function &F) {
177177
SmallVector<CoroSaveInst *, 2> UnusedCoroSaves;
178178

179179
for (Instruction &I : instructions(F)) {
180-
if (auto II = dyn_cast<IntrinsicInst>(&I)) {
180+
// FIXME: coro_await_suspend_* are not proper `IntrinisicInst`s
181+
// because they might be invoked
182+
if (auto AWS = dyn_cast<CoroAwaitSuspendInst>(&I)) {
183+
CoroAwaitSuspends.push_back(AWS);
184+
} else if (auto II = dyn_cast<IntrinsicInst>(&I)) {
181185
switch (II->getIntrinsicID()) {
182186
default:
183187
continue;
@@ -258,11 +262,6 @@ void coro::Shape::buildFrom(Function &F) {
258262
}
259263
}
260264
break;
261-
case Intrinsic::coro_await_suspend_void:
262-
case Intrinsic::coro_await_suspend_bool:
263-
case Intrinsic::coro_await_suspend_handle:
264-
CoroAwaitSuspends.push_back(cast<CoroAwaitSuspendInst>(II));
265-
break;
266265
}
267266
}
268267
}

‎llvm/test/Transforms/Coroutines/coro-await-suspend-lower-invoke.ll

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,123 @@
1+
; Tests that invoke <type> @llvm.coro.await.suspend lowers to invoke @helper
2+
; RUN: opt < %s -passes='module(coro-early),cgscc(coro-split),simplifycfg' -S | FileCheck %s
3+
4+
%Awaiter = type {}
5+
6+
; CHECK: define {{[^@]*}} @f.resume(ptr {{[^%]*}} %[[HDL:.+]])
7+
; CHECK: %[[AWAITER:.+]] = getelementptr inbounds %f.Frame, ptr %[[HDL]], i32 0, i32 0
8+
define void @f() presplitcoroutine personality i32 0 {
9+
entry:
10+
%awaiter = alloca %Awaiter
11+
%id = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr null)
12+
%size = call i32 @llvm.coro.size.i32()
13+
%alloc = call ptr @malloc(i32 %size)
14+
%hdl = call ptr @llvm.coro.begin(token %id, ptr %alloc)
15+
; Initial suspend so that all 3 await_suspend invocations are inside f.resume
16+
%suspend.init = call i8 @llvm.coro.suspend(token none, i1 false)
17+
switch i8 %suspend.init, label %ret [
18+
i8 0, label %step
19+
i8 1, label %cleanup
20+
]
21+
22+
; CHECK: invoke void @await_suspend_wrapper_void(ptr %[[AWAITER]], ptr %[[HDL]])
23+
; CHECK-NEXT: to label %[[STEP_CONT:[^ ]+]] unwind label %[[PAD:[^ ]+]]
24+
step:
25+
%save = call token @llvm.coro.save(ptr null)
26+
invoke void @llvm.coro.await.suspend.void(ptr %awaiter, ptr %hdl, ptr @await_suspend_wrapper_void)
27+
to label %step.continue unwind label %pad
28+
29+
; CHECK [[STEP_CONT]]:
30+
step.continue:
31+
%suspend = call i8 @llvm.coro.suspend(token %save, i1 false)
32+
switch i8 %suspend, label %ret [
33+
i8 0, label %step1
34+
i8 1, label %cleanup
35+
]
36+
37+
; CHECK: %[[RESUME:.+]] = invoke i1 @await_suspend_wrapper_bool(ptr %[[AWAITER]], ptr %[[HDL]])
38+
; CHECK-NEXT: to label %[[STEP1_CONT:[^ ]+]] unwind label %[[PAD]]
39+
step1:
40+
%save1 = call token @llvm.coro.save(ptr null)
41+
%resume.bool = invoke i1 @llvm.coro.await.suspend.bool(ptr %awaiter, ptr %hdl, ptr @await_suspend_wrapper_bool)
42+
to label %step1.continue unwind label %pad
43+
44+
; CHECK: [[STEP1_CONT]]:
45+
; CHECK-NEXT: br i1 %[[RESUME]], label %{{[^,]+}}, label %[[STEP2:.+]]
46+
step1.continue:
47+
br i1 %resume.bool, label %suspend.cond, label %step2
48+
49+
suspend.cond:
50+
%suspend1 = call i8 @llvm.coro.suspend(token %save1, i1 false)
51+
switch i8 %suspend1, label %ret [
52+
i8 0, label %step2
53+
i8 1, label %cleanup
54+
]
55+
56+
; CHECK: [[STEP2]]:
57+
; CHECK: %[[NEXT_HDL:.+]] = invoke ptr @await_suspend_wrapper_handle(ptr %[[AWAITER]], ptr %[[HDL]])
58+
; CHECK-NEXT: to label %[[STEP2_CONT:[^ ]+]] unwind label %[[PAD]]
59+
step2:
60+
%save2 = call token @llvm.coro.save(ptr null)
61+
%resume.handle = invoke ptr @llvm.coro.await.suspend.handle(ptr %awaiter, ptr %hdl, ptr @await_suspend_wrapper_handle)
62+
to label %step2.continue unwind label %pad
63+
64+
; CHECK: [[STEP2_CONT]]:
65+
; CHECK-NEXT: %[[NEXT_RESUME:.+]] = call ptr @llvm.coro.subfn.addr(ptr %[[NEXT_HDL]], i8 0)
66+
; CHECK-NEXT: musttail call {{.*}} void %[[NEXT_RESUME]](ptr %[[NEXT_HDL]])
67+
step2.continue:
68+
call void @llvm.coro.resume(ptr %resume.handle)
69+
%suspend2 = call i8 @llvm.coro.suspend(token %save2, i1 false)
70+
switch i8 %suspend2, label %ret [
71+
i8 0, label %step3
72+
i8 1, label %cleanup
73+
]
74+
75+
step3:
76+
br label %cleanup
77+
78+
pad:
79+
%lp = landingpad { ptr, i32 }
80+
catch ptr null
81+
%exn = extractvalue { ptr, i32 } %lp, 0
82+
call ptr @__cxa_begin_catch(ptr %exn)
83+
call void @__cxa_end_catch()
84+
br label %cleanup
85+
86+
cleanup:
87+
%mem = call ptr @llvm.coro.free(token %id, ptr %hdl)
88+
call void @free(ptr %mem)
89+
br label %ret
90+
91+
ret:
92+
call i1 @llvm.coro.end(ptr %hdl, i1 0, token none)
93+
ret void
94+
}
95+
96+
; check that we were haven't accidentally went out of @f.resume body
97+
; CHECK-LABEL: @f.destroy(
98+
; CHECK-LABEL: @f.cleanup(
99+
100+
declare void @await_suspend_wrapper_void(ptr, ptr)
101+
declare i1 @await_suspend_wrapper_bool(ptr, ptr)
102+
declare ptr @await_suspend_wrapper_handle(ptr, ptr)
103+
104+
declare ptr @llvm.coro.free(token, ptr)
105+
declare i32 @llvm.coro.size.i32()
106+
declare i8 @llvm.coro.suspend(token, i1)
107+
declare void @llvm.coro.resume(ptr)
108+
declare void @llvm.coro.destroy(ptr)
109+
110+
declare token @llvm.coro.id(i32, ptr, ptr, ptr)
111+
declare i1 @llvm.coro.alloc(token)
112+
declare ptr @llvm.coro.begin(token, ptr)
113+
declare void @llvm.coro.await.suspend.void(ptr, ptr, ptr)
114+
declare i1 @llvm.coro.await.suspend.bool(ptr, ptr, ptr)
115+
declare ptr @llvm.coro.await.suspend.handle(ptr, ptr, ptr)
116+
declare i1 @llvm.coro.end(ptr, i1, token)
117+
118+
declare ptr @__cxa_begin_catch(ptr)
119+
declare void @use_val(i32)
120+
declare void @__cxa_end_catch()
121+
122+
declare noalias ptr @malloc(i32)
123+
declare void @free(ptr)

‎llvm/test/Transforms/Coroutines/coro-await-suspend-lower.ll

Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,96 @@
1+
; Tests lowerings of different versions of coro.await.suspend
2+
; RUN: opt < %s -passes='module(coro-early),cgscc(coro-split),simplifycfg' -S | FileCheck %s
3+
4+
%Awaiter = type {}
5+
6+
; CHECK: define {{[^@]*}} @f.resume(ptr {{[^%]*}} %[[HDL:.+]])
7+
; CHECK: %[[AWAITER:.+]] = getelementptr inbounds %f.Frame, ptr %[[HDL]], i32 0, i32 0
8+
define void @f() presplitcoroutine {
9+
entry:
10+
%awaiter = alloca %Awaiter
11+
%id = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr null)
12+
%size = call i32 @llvm.coro.size.i32()
13+
%alloc = call ptr @malloc(i32 %size)
14+
%hdl = call ptr @llvm.coro.begin(token %id, ptr %alloc)
15+
%suspend.init = call i8 @llvm.coro.suspend(token none, i1 false)
16+
switch i8 %suspend.init, label %ret [
17+
i8 0, label %step
18+
i8 1, label %cleanup
19+
]
20+
21+
; CHECK: call void @await_suspend_wrapper_void(ptr %[[AWAITER]], ptr %[[HDL]])
22+
; CHECK-NEXT: br label %{{.*}}
23+
step:
24+
%save = call token @llvm.coro.save(ptr null)
25+
call void @llvm.coro.await.suspend.void(ptr %awaiter, ptr %hdl, ptr @await_suspend_wrapper_void)
26+
%suspend = call i8 @llvm.coro.suspend(token %save, i1 false)
27+
switch i8 %suspend, label %ret [
28+
i8 0, label %step1
29+
i8 1, label %cleanup
30+
]
31+
32+
; CHECK: %[[RESUME:.+]] = call i1 @await_suspend_wrapper_bool(ptr %[[AWAITER]], ptr %[[HDL]])
33+
; CHECK-NEXT: br i1 %[[RESUME]], label %{{[^,]+}}, label %[[STEP2:.+]]
34+
step1:
35+
%save1 = call token @llvm.coro.save(ptr null)
36+
%resume.bool = call i1 @llvm.coro.await.suspend.bool(ptr %awaiter, ptr %hdl, ptr @await_suspend_wrapper_bool)
37+
br i1 %resume.bool, label %suspend.cond, label %step2
38+
39+
suspend.cond:
40+
%suspend1 = call i8 @llvm.coro.suspend(token %save1, i1 false)
41+
switch i8 %suspend1, label %ret [
42+
i8 0, label %step2
43+
i8 1, label %cleanup
44+
]
45+
46+
; CHECK: [[STEP2]]:
47+
; CHECK: %[[NEXT_HDL:.+]] = call ptr @await_suspend_wrapper_handle(ptr %[[AWAITER]], ptr %[[HDL]])
48+
; CHECK-NEXT: %[[CONT:.+]] = call ptr @llvm.coro.subfn.addr(ptr %[[NEXT_HDL]], i8 0)
49+
; CHECK-NEXT: musttail call {{.*}} void %[[CONT]](ptr %[[NEXT_HDL]])
50+
step2:
51+
%save2 = call token @llvm.coro.save(ptr null)
52+
%resume.handle = call ptr @llvm.coro.await.suspend.handle(ptr %awaiter, ptr %hdl, ptr @await_suspend_wrapper_handle)
53+
call void @llvm.coro.resume(ptr %resume.handle)
54+
%suspend2 = call i8 @llvm.coro.suspend(token %save2, i1 false)
55+
switch i8 %suspend2, label %ret [
56+
i8 0, label %step3
57+
i8 1, label %cleanup
58+
]
59+
60+
step3:
61+
br label %cleanup
62+
63+
cleanup:
64+
%mem = call ptr @llvm.coro.free(token %id, ptr %hdl)
65+
call void @free(ptr %mem)
66+
br label %ret
67+
68+
ret:
69+
call i1 @llvm.coro.end(ptr %hdl, i1 0, token none)
70+
ret void
71+
}
72+
73+
; check that we were haven't accidentally went out of @f.resume body
74+
; CHECK-LABEL: @f.destroy(
75+
; CHECK-LABEL: @f.cleanup(
76+
77+
declare void @await_suspend_wrapper_void(ptr, ptr)
78+
declare i1 @await_suspend_wrapper_bool(ptr, ptr)
79+
declare ptr @await_suspend_wrapper_handle(ptr, ptr)
80+
81+
declare ptr @llvm.coro.free(token, ptr)
82+
declare i32 @llvm.coro.size.i32()
83+
declare i8 @llvm.coro.suspend(token, i1)
84+
declare void @llvm.coro.resume(ptr)
85+
declare void @llvm.coro.destroy(ptr)
86+
87+
declare token @llvm.coro.id(i32, ptr, ptr, ptr)
88+
declare i1 @llvm.coro.alloc(token)
89+
declare ptr @llvm.coro.begin(token, ptr)
90+
declare void @llvm.coro.await.suspend.void(ptr, ptr, ptr)
91+
declare i1 @llvm.coro.await.suspend.bool(ptr, ptr, ptr)
92+
declare ptr @llvm.coro.await.suspend.handle(ptr, ptr, ptr)
93+
declare i1 @llvm.coro.end(ptr, i1, token)
94+
95+
declare noalias ptr @malloc(i32)
96+
declare void @free(ptr)

0 commit comments

Comments
 (0)