Skip to content

Commit 0e18f75

Browse files
GorNishanovGorNishanov
committedSep 5, 2016
[Coroutines] Part11: Add final suspend handling.
Summary: A frontend may designate a particular suspend to be final, by setting the second argument of the coro.suspend intrinsic to true. Such a suspend point has two properties: * it is possible to check whether a suspended coroutine is at the final suspend point via coro.done intrinsic; * a resumption of a coroutine stopped at the final suspend point leads to undefined behavior. The only possible action for a coroutine at a final suspend point is destroying it via coro.destroy intrinsic. This patch adds final suspend handling logic to CoroEarly and CoroSplit passes. Now, the final suspend point example from docs\Coroutines.rst compiles and produces expected result (see test/Transform/Coroutines/ex5.ll). Reviewers: majnemer Subscribers: mehdi_amini, llvm-commits Differential Revision: https://reviews.llvm.org/D24068 llvm-svn: 280646
1 parent 42b69dd commit 0e18f75

File tree

4 files changed

+166
-17
lines changed

4 files changed

+166
-17
lines changed
 

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

+29-4
Original file line numberDiff line numberDiff line change
@@ -26,10 +26,11 @@ namespace {
2626
// Created on demand if CoroEarly pass has work to do.
2727
class Lowerer : public coro::LowererBase {
2828
IRBuilder<> Builder;
29-
PointerType *AnyResumeFnPtrTy;
29+
PointerType *const AnyResumeFnPtrTy;
3030

3131
void lowerResumeOrDestroy(CallSite CS, CoroSubFnInst::ResumeKind);
3232
void lowerCoroPromise(CoroPromiseInst *Intrin);
33+
void lowerCoroDone(IntrinsicInst *II);
3334

3435
public:
3536
Lowerer(Module &M)
@@ -81,6 +82,27 @@ void Lowerer::lowerCoroPromise(CoroPromiseInst *Intrin) {
8182
Intrin->eraseFromParent();
8283
}
8384

85+
// When a coroutine reaches final suspend point, it zeros out ResumeFnAddr in
86+
// the coroutine frame (it is UB to resume from a final suspend point).
87+
// The llvm.coro.done intrinsic is used to check whether a coroutine is
88+
// suspended at the final suspend point or not.
89+
void Lowerer::lowerCoroDone(IntrinsicInst *II) {
90+
Value *Operand = II->getArgOperand(0);
91+
92+
// ResumeFnAddr is the first pointer sized element of the coroutine frame.
93+
auto *FrameTy = Int8Ptr;
94+
PointerType *FramePtrTy = FrameTy->getPointerTo();
95+
96+
Builder.SetInsertPoint(II);
97+
auto *BCI = Builder.CreateBitCast(Operand, FramePtrTy);
98+
auto *Gep = Builder.CreateConstInBoundsGEP1_32(FrameTy, BCI, 0);
99+
auto *Load = Builder.CreateLoad(Gep);
100+
auto *Cond = Builder.CreateICmpEQ(Load, NullPtr);
101+
102+
II->replaceAllUsesWith(Cond);
103+
II->eraseFromParent();
104+
}
105+
84106
// Prior to CoroSplit, calls to coro.begin needs to be marked as NoDuplicate,
85107
// as CoroSplit assumes there is exactly one coro.begin. After CoroSplit,
86108
// NoDuplicate attribute will be removed from coro.begin otherwise, it will
@@ -131,6 +153,9 @@ bool Lowerer::lowerEarlyIntrinsics(Function &F) {
131153
case Intrinsic::coro_promise:
132154
lowerCoroPromise(cast<CoroPromiseInst>(&I));
133155
break;
156+
case Intrinsic::coro_done:
157+
lowerCoroDone(cast<IntrinsicInst>(&I));
158+
break;
134159
}
135160
Changed = true;
136161
}
@@ -153,9 +178,9 @@ struct CoroEarly : public FunctionPass {
153178
// This pass has work to do only if we find intrinsics we are going to lower
154179
// in the module.
155180
bool doInitialization(Module &M) override {
156-
if (coro::declaresIntrinsics(M, {"llvm.coro.begin", "llvm.coro.resume",
157-
"llvm.coro.destroy", "llvm.coro.suspend",
158-
"llvm.coro.end"}))
181+
if (coro::declaresIntrinsics(M, {"llvm.coro.begin", "llvm.coro.end",
182+
"llvm.coro.resume", "llvm.coro.destroy",
183+
"llvm.coro.done", "llvm.coro.suspend"}))
159184
L = llvm::make_unique<Lowerer>(M);
160185
return false;
161186
}

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

+54-5
Original file line numberDiff line numberDiff line change
@@ -62,18 +62,27 @@ static BasicBlock *createResumeEntryBlock(Function &F, coro::Shape &Shape) {
6262
Builder.CreateSwitch(Index, UnreachBB, Shape.CoroSuspends.size());
6363
Shape.ResumeSwitch = Switch;
6464

65-
uint32_t SuspendIndex = 0;
66-
for (auto S : Shape.CoroSuspends) {
65+
size_t SuspendIndex = 0;
66+
for (CoroSuspendInst *S : Shape.CoroSuspends) {
6767
ConstantInt *IndexVal = Shape.getIndex(SuspendIndex);
6868

6969
// Replace CoroSave with a store to Index:
7070
// %index.addr = getelementptr %f.frame... (index field number)
7171
// store i32 0, i32* %index.addr1
7272
auto *Save = S->getCoroSave();
7373
Builder.SetInsertPoint(Save);
74-
auto *GepIndex = Builder.CreateConstInBoundsGEP2_32(
75-
FrameTy, FramePtr, 0, coro::Shape::IndexField, "index.addr");
76-
Builder.CreateStore(IndexVal, GepIndex);
74+
if (S->isFinal()) {
75+
// Final suspend point is represented by storing zero in ResumeFnAddr.
76+
auto *GepIndex = Builder.CreateConstInBoundsGEP2_32(FrameTy, FramePtr, 0,
77+
0, "ResumeFn.addr");
78+
auto *NullPtr = ConstantPointerNull::get(cast<PointerType>(
79+
cast<PointerType>(GepIndex->getType())->getElementType()));
80+
Builder.CreateStore(NullPtr, GepIndex);
81+
} else {
82+
auto *GepIndex = Builder.CreateConstInBoundsGEP2_32(
83+
FrameTy, FramePtr, 0, coro::Shape::IndexField, "index.addr");
84+
Builder.CreateStore(IndexVal, GepIndex);
85+
}
7786
Save->replaceAllUsesWith(ConstantTokenNone::get(C));
7887
Save->eraseFromParent();
7988

@@ -135,6 +144,37 @@ static void replaceFallthroughCoroEnd(IntrinsicInst *End,
135144
BB->getTerminator()->eraseFromParent();
136145
}
137146

147+
// Rewrite final suspend point handling. We do not use suspend index to
148+
// represent the final suspend point. Instead we zero-out ResumeFnAddr in the
149+
// coroutine frame, since it is undefined behavior to resume a coroutine
150+
// suspended at the final suspend point. Thus, in the resume function, we can
151+
// simply remove the last case (when coro::Shape is built, the final suspend
152+
// point (if present) is always the last element of CoroSuspends array).
153+
// In the destroy function, we add a code sequence to check if ResumeFnAddress
154+
// is Null, and if so, jump to the appropriate label to handle cleanup from the
155+
// final suspend point.
156+
static void handleFinalSuspend(IRBuilder<> &Builder, Value *FramePtr,
157+
coro::Shape &Shape, SwitchInst *Switch,
158+
bool IsDestroy) {
159+
assert(Shape.HasFinalSuspend);
160+
auto FinalCase = --Switch->case_end();
161+
BasicBlock *ResumeBB = FinalCase.getCaseSuccessor();
162+
Switch->removeCase(FinalCase);
163+
if (IsDestroy) {
164+
BasicBlock *OldSwitchBB = Switch->getParent();
165+
auto *NewSwitchBB = OldSwitchBB->splitBasicBlock(Switch, "Switch");
166+
Builder.SetInsertPoint(OldSwitchBB->getTerminator());
167+
auto *GepIndex = Builder.CreateConstInBoundsGEP2_32(Shape.FrameTy, FramePtr,
168+
0, 0, "ResumeFn.addr");
169+
auto *Load = Builder.CreateLoad(GepIndex);
170+
auto *NullPtr =
171+
ConstantPointerNull::get(cast<PointerType>(Load->getType()));
172+
auto *Cond = Builder.CreateICmpEQ(Load, NullPtr);
173+
Builder.CreateCondBr(Cond, ResumeBB, NewSwitchBB);
174+
OldSwitchBB->getTerminator()->eraseFromParent();
175+
}
176+
}
177+
138178
// Create a resume clone by cloning the body of the original function, setting
139179
// new entry block and replacing coro.suspend an appropriate value to force
140180
// resume or cleanup pass for every suspend point.
@@ -205,6 +245,15 @@ static Function *createClone(Function &F, Twine Suffix, coro::Shape &Shape,
205245
Value *OldVFrame = cast<Value>(VMap[Shape.CoroBegin]);
206246
OldVFrame->replaceAllUsesWith(NewVFrame);
207247

248+
// Rewrite final suspend handling as it is not done via switch (allows to
249+
// remove final case from the switch, since it is undefined behavior to resume
250+
// the coroutine suspended at the final suspend point.
251+
if (Shape.HasFinalSuspend) {
252+
auto *Switch = cast<SwitchInst>(VMap[Shape.ResumeSwitch]);
253+
bool IsDestroy = FnIndex != 0;
254+
handleFinalSuspend(Builder, NewFramePtr, Shape, Switch, IsDestroy);
255+
}
256+
208257
// Replace coro suspend with the appropriate resume index.
209258
// Replacing coro.suspend with (0) will result in control flow proceeding to
210259
// a resume label associated with a suspend point, replacing it with (1) will

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

+10-8
Original file line numberDiff line numberDiff line change
@@ -215,6 +215,7 @@ static CoroSaveInst *createCoroSave(CoroBeginInst *CoroBegin,
215215

216216
// Collect "interesting" coroutine intrinsics.
217217
void coro::Shape::buildFrom(Function &F) {
218+
size_t FinalSuspendIndex = 0;
218219
clear(*this);
219220
SmallVector<CoroFrameInst *, 8> CoroFrames;
220221
for (Instruction &I : instructions(F)) {
@@ -230,16 +231,12 @@ void coro::Shape::buildFrom(Function &F) {
230231
break;
231232
case Intrinsic::coro_suspend:
232233
CoroSuspends.push_back(cast<CoroSuspendInst>(II));
233-
// Make sure that the final suspend is the first suspend point in the
234-
// CoroSuspends vector.
235234
if (CoroSuspends.back()->isFinal()) {
235+
if (HasFinalSuspend)
236+
report_fatal_error(
237+
"Only one suspend point can be marked as final");
236238
HasFinalSuspend = true;
237-
if (CoroSuspends.size() > 1) {
238-
if (CoroSuspends.front()->isFinal())
239-
report_fatal_error(
240-
"Only one suspend point can be marked as final");
241-
std::swap(CoroSuspends.front(), CoroSuspends.back());
242-
}
239+
FinalSuspendIndex = CoroSuspends.size() - 1;
243240
}
244241
break;
245242
case Intrinsic::coro_begin: {
@@ -309,4 +306,9 @@ void coro::Shape::buildFrom(Function &F) {
309306
for (CoroSuspendInst *CS : CoroSuspends)
310307
if (!CS->getCoroSave())
311308
createCoroSave(CoroBegin, CS);
309+
310+
// Move final suspend to be the last element in the CoroSuspends vector.
311+
if (HasFinalSuspend &&
312+
FinalSuspendIndex != CoroSuspends.size() - 1)
313+
std::swap(CoroSuspends[FinalSuspendIndex], CoroSuspends.back());
312314
}

‎llvm/test/Transforms/Coroutines/ex5.ll

+73
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,73 @@
1+
; Fifth example from Doc/Coroutines.rst (final suspend)
2+
; RUN: opt < %s -O2 -enable-coroutines -S | FileCheck %s
3+
4+
define i8* @f(i32 %n) {
5+
entry:
6+
%id = call token @llvm.coro.id(i32 0, i8* null, i8* null, i8* null)
7+
%size = call i32 @llvm.coro.size.i32()
8+
%alloc = call i8* @malloc(i32 %size)
9+
%hdl = call noalias i8* @llvm.coro.begin(token %id, i8* %alloc)
10+
br label %while.cond
11+
while.cond:
12+
%n.val = phi i32 [ %n, %entry ], [ %dec, %while.body ]
13+
%cmp = icmp sgt i32 %n.val, 0
14+
br i1 %cmp, label %while.body, label %while.end
15+
16+
while.body:
17+
%dec = add nsw i32 %n.val, -1
18+
call void @print(i32 %n.val) #4
19+
%s = call i8 @llvm.coro.suspend(token none, i1 false)
20+
switch i8 %s, label %suspend [i8 0, label %while.cond
21+
i8 1, label %cleanup]
22+
while.end:
23+
%s.final = call i8 @llvm.coro.suspend(token none, i1 true)
24+
switch i8 %s.final, label %suspend [i8 0, label %trap
25+
i8 1, label %cleanup]
26+
trap:
27+
call void @llvm.trap()
28+
unreachable
29+
cleanup:
30+
%mem = call i8* @llvm.coro.free(token %id, i8* %hdl)
31+
call void @free(i8* %mem)
32+
br label %suspend
33+
suspend:
34+
call void @llvm.coro.end(i8* %hdl, i1 false)
35+
ret i8* %hdl
36+
}
37+
38+
declare noalias i8* @malloc(i32)
39+
declare void @print(i32)
40+
declare void @llvm.trap()
41+
declare void @free(i8* nocapture)
42+
43+
declare token @llvm.coro.id( i32, i8*, i8*, i8*)
44+
declare i32 @llvm.coro.size.i32()
45+
declare i8* @llvm.coro.begin(token, i8*)
46+
declare token @llvm.coro.save(i8*)
47+
declare i8 @llvm.coro.suspend(token, i1)
48+
declare i8* @llvm.coro.free(token, i8*)
49+
declare void @llvm.coro.end(i8*, i1)
50+
51+
; CHECK-LABEL: @main
52+
define i32 @main() {
53+
entry:
54+
%hdl = call i8* @f(i32 4)
55+
br label %while
56+
while:
57+
call void @llvm.coro.resume(i8* %hdl)
58+
%done = call i1 @llvm.coro.done(i8* %hdl)
59+
br i1 %done, label %end, label %while
60+
end:
61+
call void @llvm.coro.destroy(i8* %hdl)
62+
ret i32 0
63+
64+
; CHECK: call void @print(i32 4)
65+
; CHECK: call void @print(i32 3)
66+
; CHECK: call void @print(i32 2)
67+
; CHECK: call void @print(i32 1)
68+
; CHECK: ret i32 0
69+
}
70+
71+
declare i1 @llvm.coro.done(i8*)
72+
declare void @llvm.coro.resume(i8*)
73+
declare void @llvm.coro.destroy(i8*)

0 commit comments

Comments
 (0)
Please sign in to comment.