Skip to content

Commit 7f255fb

Browse files
fpasserbyfpasserby
committed
Add await suspend lowering tests
1 parent 9afcd59 commit 7f255fb

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
@@ -1404,10 +1404,10 @@ static bool hasCallsBetween(Instruction *Save, Instruction *ResumeOrDestroy) {
14041404
auto *ResumeOrDestroyBB = ResumeOrDestroy->getParent();
14051405

14061406
if (SaveBB == ResumeOrDestroyBB)
1407-
return hasCallsInBlockBetween(Save->getNextNode(), ResumeOrDestroy);
1407+
return hasCallsInBlockBetween(Save, ResumeOrDestroy);
14081408

14091409
// Any calls from Save to the end of the block?
1410-
if (hasCallsInBlockBetween(Save->getNextNode(), nullptr))
1410+
if (hasCallsInBlockBetween(Save, nullptr))
14111411
return true;
14121412

14131413
// Any calls from begging of the block up to ResumeOrDestroy?
@@ -1429,7 +1429,10 @@ static bool hasCallsBetween(Instruction *Save, Instruction *ResumeOrDestroy) {
14291429
// resume or destroy it
14301430
// FIXME: perform more sophisiticated analysis?
14311431
static bool isSimpleWrapper(CoroAwaitSuspendInst *AWS) {
1432-
auto Wrapper = AWS->getWrapperFunction();
1432+
auto *Wrapper = AWS->getWrapperFunction();
1433+
1434+
if (Wrapper->empty())
1435+
return false;
14331436

14341437
SmallVector<ReturnInst *, 4> Rets;
14351438

@@ -1567,17 +1570,12 @@ static void simplifySuspendPoints(coro::Shape &Shape) {
15671570
namespace {
15681571

15691572
struct SwitchCoroutineSplitter {
1570-
static void split(Module &M, Function &F, coro::Shape &Shape,
1573+
static void split(Function &F, coro::Shape &Shape,
15711574
SmallVectorImpl<Function *> &Clones,
15721575
TargetTransformInfo &TTI) {
15731576
assert(Shape.ABI == coro::ABI::Switch);
15741577

15751578
createResumeEntryBlock(F, Shape);
1576-
1577-
IRBuilder<> Builder(M.getContext());
1578-
for (auto *AWS : Shape.CoroAwaitSuspends)
1579-
lowerAwaitSuspend(Builder, AWS);
1580-
15811579
auto *ResumeClone =
15821580
createClone(F, ".resume", Shape, CoroCloner::Kind::SwitchResume);
15831581
auto *DestroyClone =
@@ -2106,14 +2104,18 @@ splitCoroutine(Module &M, Function &F, SmallVectorImpl<Function *> &Clones,
21062104
buildCoroutineFrame(F, Shape, MaterializableCallback);
21072105
replaceFrameSizeAndAlignment(Shape);
21082106

2107+
IRBuilder<> Builder(M.getContext());
2108+
for (auto *AWS : Shape.CoroAwaitSuspends)
2109+
lowerAwaitSuspend(Builder, AWS);
2110+
21092111
// If there are no suspend points, no split required, just remove
21102112
// the allocation and deallocation blocks, they are not needed.
21112113
if (Shape.CoroSuspends.empty()) {
21122114
handleNoSuspendCoroutine(Shape);
21132115
} else {
21142116
switch (Shape.ABI) {
21152117
case coro::ABI::Switch:
2116-
SwitchCoroutineSplitter::split(M, F, Shape, Clones, TTI);
2118+
SwitchCoroutineSplitter::split(F, Shape, Clones, TTI);
21172119
break;
21182120
case coro::ABI::Async:
21192121
splitAsyncCoroutine(F, Shape, Clones);

‎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)