Skip to content

Commit 3829214

Browse files
rjmccallrjmccall
committed
Generalize llvm.coro.suspend.retcon to allow an arbitrary number of arguments to be passed back to the continuation function.
llvm-svn: 368789
1 parent 94010b2 commit 3829214

File tree

10 files changed

+314
-64
lines changed

10 files changed

+314
-64
lines changed

‎llvm/include/llvm/IR/Intrinsics.td

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -982,7 +982,7 @@ def int_coro_size : Intrinsic<[llvm_anyint_ty], [], [IntrNoMem]>;
982982

983983
def int_coro_save : Intrinsic<[llvm_token_ty], [llvm_ptr_ty], []>;
984984
def int_coro_suspend : Intrinsic<[llvm_i8_ty], [llvm_token_ty, llvm_i1_ty], []>;
985-
def int_coro_suspend_retcon : Intrinsic<[llvm_i1_ty], [llvm_vararg_ty], []>;
985+
def int_coro_suspend_retcon : Intrinsic<[llvm_any_ty], [llvm_vararg_ty], []>;
986986
def int_coro_prepare_retcon : Intrinsic<[llvm_ptr_ty], [llvm_ptr_ty],
987987
[IntrNoMem]>;
988988

‎llvm/lib/Transforms/Coroutines/CoroInternal.h

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -171,14 +171,23 @@ struct LLVM_LIBRARY_VISIBILITY Shape {
171171
ABI == coro::ABI::RetconOnce);
172172
auto FTy = CoroBegin->getFunction()->getFunctionType();
173173

174-
// This is checked by AnyCoroIdRetconInst::isWellFormed().
174+
// The safety of all this is checked by checkWFRetconPrototype.
175175
if (auto STy = dyn_cast<StructType>(FTy->getReturnType())) {
176176
return STy->elements().slice(1);
177177
} else {
178178
return ArrayRef<Type*>();
179179
}
180180
}
181181

182+
ArrayRef<Type*> getRetconResumeTypes() const {
183+
assert(ABI == coro::ABI::Retcon ||
184+
ABI == coro::ABI::RetconOnce);
185+
186+
// The safety of all this is checked by checkWFRetconPrototype.
187+
auto FTy = RetconLowering.ResumePrototype->getFunctionType();
188+
return FTy->params().slice(1);
189+
}
190+
182191
CallingConv::ID getResumeFunctionCC() const {
183192
switch (ABI) {
184193
case coro::ABI::Switch:

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

Lines changed: 54 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -144,6 +144,7 @@ class CoroCloner {
144144
void createDeclaration();
145145
void replaceEntryBlock();
146146
Value *deriveNewFramePointer();
147+
void replaceRetconSuspendUses();
147148
void replaceCoroSuspends();
148149
void replaceCoroEnds();
149150
void handleFinalSuspend();
@@ -402,6 +403,52 @@ static Function *createCloneDeclaration(Function &OrigF, coro::Shape &Shape,
402403
return NewF;
403404
}
404405

406+
/// Replace uses of the active llvm.coro.suspend.retcon call with the
407+
/// arguments to the continuation function.
408+
///
409+
/// This assumes that the builder has a meaningful insertion point.
410+
void CoroCloner::replaceRetconSuspendUses() {
411+
assert(Shape.ABI == coro::ABI::Retcon ||
412+
Shape.ABI == coro::ABI::RetconOnce);
413+
414+
auto NewS = VMap[ActiveSuspend];
415+
if (NewS->use_empty()) return;
416+
417+
// Copy out all the continuation arguments after the buffer pointer into
418+
// an easily-indexed data structure for convenience.
419+
SmallVector<Value*, 8> Args;
420+
for (auto I = std::next(NewF->arg_begin()), E = NewF->arg_end(); I != E; ++I)
421+
Args.push_back(&*I);
422+
423+
// If the suspend returns a single scalar value, we can just do a simple
424+
// replacement.
425+
if (!isa<StructType>(NewS->getType())) {
426+
assert(Args.size() == 1);
427+
NewS->replaceAllUsesWith(Args.front());
428+
return;
429+
}
430+
431+
// Try to peephole extracts of an aggregate return.
432+
for (auto UI = NewS->use_begin(), UE = NewS->use_end(); UI != UE; ) {
433+
auto EVI = dyn_cast<ExtractValueInst>((UI++)->getUser());
434+
if (!EVI || EVI->getNumIndices() != 1)
435+
continue;
436+
437+
EVI->replaceAllUsesWith(Args[EVI->getIndices().front()]);
438+
EVI->eraseFromParent();
439+
}
440+
441+
// If we have no remaining uses, we're done.
442+
if (NewS->use_empty()) return;
443+
444+
// Otherwise, we need to create an aggregate.
445+
Value *Agg = UndefValue::get(NewS->getType());
446+
for (size_t I = 0, E = Args.size(); I != E; ++I)
447+
Agg = Builder.CreateInsertValue(Agg, Args[I], I);
448+
449+
NewS->replaceAllUsesWith(Agg);
450+
}
451+
405452
void CoroCloner::replaceCoroSuspends() {
406453
Value *SuspendResult;
407454

@@ -416,15 +463,12 @@ void CoroCloner::replaceCoroSuspends() {
416463
SuspendResult = Builder.getInt8(isSwitchDestroyFunction() ? 1 : 0);
417464
break;
418465

419-
// In continuation lowering, replace all of the suspend uses with false to
420-
// indicate that they're not unwinding resumes. We've already mapped the
421-
// active suspend to the appropriate argument, so any other suspend values
422-
// that are still being used must be from previous suspends. It's UB to try
423-
// to suspend during unwind, so they must be from regular resumes.
466+
// In returned-continuation lowering, the arguments from earlier
467+
// continuations are theoretically arbitrary, and they should have been
468+
// spilled.
424469
case coro::ABI::RetconOnce:
425470
case coro::ABI::Retcon:
426-
SuspendResult = Builder.getInt1(false);
427-
break;
471+
return;
428472
}
429473

430474
for (AnyCoroSuspendInst *CS : Shape.CoroSuspends) {
@@ -619,14 +663,11 @@ void CoroCloner::create() {
619663

620664
case coro::ABI::Retcon:
621665
case coro::ABI::RetconOnce:
622-
// Replace the active suspend with the should-unwind argument.
623-
// Coerce it to i1 if necessary.
666+
// Replace uses of the active suspend with the corresponding
667+
// continuation-function arguments.
624668
assert(ActiveSuspend != nullptr &&
625669
"no active suspend when lowering a continuation-style coroutine");
626-
Value *ShouldUnwind = &*std::next(NewF->arg_begin());
627-
if (!ShouldUnwind->getType()->isIntegerTy(1))
628-
ShouldUnwind = Builder.CreateIsNotNull(ShouldUnwind);
629-
VMap[ActiveSuspend]->replaceAllUsesWith(ShouldUnwind);
670+
replaceRetconSuspendUses();
630671
break;
631672
}
632673

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

Lines changed: 29 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -387,6 +387,7 @@ void coro::Shape::buildFrom(Function &F) {
387387
// Determine the result value types, and make sure they match up with
388388
// the values passed to the suspends.
389389
auto ResultTys = getRetconResultTypes();
390+
auto ResumeTys = getRetconResumeTypes();
390391

391392
for (auto AnySuspend : CoroSuspends) {
392393
auto Suspend = dyn_cast<CoroSuspendRetconInst>(AnySuspend);
@@ -396,6 +397,7 @@ void coro::Shape::buildFrom(Function &F) {
396397
"coro.suspend.retcon");
397398
}
398399

400+
// Check that the argument types of the suspend match the results.
399401
auto SI = Suspend->value_begin(), SE = Suspend->value_end();
400402
auto RI = ResultTys.begin(), RE = ResultTys.end();
401403
for (; SI != SE && RI != RE; ++SI, ++RI) {
@@ -411,6 +413,26 @@ void coro::Shape::buildFrom(Function &F) {
411413
Prototype->getFunctionType()->dump();
412414
report_fatal_error("wrong number of arguments to coro.suspend.retcon");
413415
}
416+
417+
// Check that the result type of the suspend matches the resume types.
418+
Type *SResultTy = Suspend->getType();
419+
ArrayRef<Type*> SuspendResultTys =
420+
(isa<StructType>(SResultTy)
421+
? cast<StructType>(SResultTy)->elements()
422+
: SResultTy); // forms an ArrayRef using SResultTy, be careful
423+
if (SuspendResultTys.size() != ResumeTys.size()) {
424+
Suspend->dump();
425+
Prototype->getFunctionType()->dump();
426+
report_fatal_error("wrong number of results from coro.suspend.retcon");
427+
}
428+
for (size_t I = 0, E = ResumeTys.size(); I != E; ++I) {
429+
if (SuspendResultTys[I] != ResumeTys[I]) {
430+
Suspend->dump();
431+
Prototype->getFunctionType()->dump();
432+
report_fatal_error("result from coro.suspend.retcon does not "
433+
"match corresponding prototype function param");
434+
}
435+
}
414436
}
415437
break;
416438
}
@@ -501,7 +523,7 @@ static void fail(const Instruction *I, const char *Reason, Value *V) {
501523
static void checkWFRetconPrototype(const AnyCoroIdRetconInst *I, Value *V) {
502524
auto F = dyn_cast<Function>(V->stripPointerCasts());
503525
if (!F)
504-
fail(I, "llvm.coro.retcon.* prototype not a Function", V);
526+
fail(I, "llvm.coro.id.retcon.* prototype not a Function", V);
505527

506528
auto FT = F->getFunctionType();
507529

@@ -517,23 +539,20 @@ static void checkWFRetconPrototype(const AnyCoroIdRetconInst *I, Value *V) {
517539
ResultOkay = false;
518540
}
519541
if (!ResultOkay)
520-
fail(I, "llvm.coro.retcon prototype must return pointer as first result",
521-
F);
542+
fail(I, "llvm.coro.id.retcon prototype must return pointer as first "
543+
"result", F);
522544

523545
if (FT->getReturnType() !=
524546
I->getFunction()->getFunctionType()->getReturnType())
525-
fail(I, "llvm.coro.retcon.* prototype return type must be same as"
547+
fail(I, "llvm.coro.id.retcon prototype return type must be same as"
526548
"current function return type", F);
527549
} else {
528550
// No meaningful validation to do here for llvm.coro.id.unique.once.
529551
}
530552

531-
if (FT->getNumParams() != 2)
532-
fail(I, "llvm.coro.retcon.* prototype must take exactly two parameters", F);
533-
if (!FT->getParamType(0)->isPointerTy())
534-
fail(I, "llvm.coro.retcon.* prototype must take pointer as 1st param", F);
535-
if (!FT->getParamType(1)->isIntegerTy()) // an i1, but not for abi purposes
536-
fail(I, "llvm.coro.retcon.* prototype must take integer as 2nd param", F);
553+
if (FT->getNumParams() == 0 || !FT->getParamType(0)->isPointerTy())
554+
fail(I, "llvm.coro.id.retcon.* prototype must take pointer as "
555+
"its first parameter", F);
537556
}
538557

539558
/// Check that the given value is a well-formed allocator.

‎llvm/test/Transforms/Coroutines/coro-retcon-once-value.ll

Lines changed: 13 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -4,22 +4,22 @@ target triple = "x86_64-apple-macosx10.12.0"
44

55
define {i8*, i32} @f(i8* %buffer, i32* %array) {
66
entry:
7-
%id = call token @llvm.coro.id.retcon.once(i32 8, i32 8, i8* %buffer, i8* bitcast (void (i8*, i8)* @prototype to i8*), i8* bitcast (i8* (i32)* @allocate to i8*), i8* bitcast (void (i8*)* @deallocate to i8*))
7+
%id = call token @llvm.coro.id.retcon.once(i32 8, i32 8, i8* %buffer, i8* bitcast (void (i8*, i1)* @prototype to i8*), i8* bitcast (i8* (i32)* @allocate to i8*), i8* bitcast (void (i8*)* @deallocate to i8*))
88
%hdl = call i8* @llvm.coro.begin(token %id, i8* null)
99
%load = load i32, i32* %array
1010
%load.pos = icmp sgt i32 %load, 0
1111
br i1 %load.pos, label %pos, label %neg
1212

1313
pos:
14-
%unwind0 = call i1 (...) @llvm.coro.suspend.retcon(i32 %load)
14+
%unwind0 = call i1 (...) @llvm.coro.suspend.retcon.i1(i32 %load)
1515
br i1 %unwind0, label %cleanup, label %pos.cont
1616

1717
pos.cont:
1818
store i32 0, i32* %array, align 4
1919
br label %cleanup
2020

2121
neg:
22-
%unwind1 = call i1 (...) @llvm.coro.suspend.retcon(i32 0)
22+
%unwind1 = call i1 (...) @llvm.coro.suspend.retcon.i1(i32 0)
2323
br i1 %unwind1, label %cleanup, label %neg.cont
2424

2525
neg.cont:
@@ -37,18 +37,17 @@ cleanup:
3737
; CHECK-NEXT: store i32* %array, i32** [[T0]], align 8
3838
; CHECK-NEXT: %load = load i32, i32* %array, align 4
3939
; CHECK-NEXT: %load.pos = icmp sgt i32 %load, 0
40-
; CHECK-NEXT: [[CONT:%.*]] = select i1 %load.pos, void (i8*, i8)* @f.resume.0, void (i8*, i8)* @f.resume.1
40+
; CHECK-NEXT: [[CONT:%.*]] = select i1 %load.pos, void (i8*, i1)* @f.resume.0, void (i8*, i1)* @f.resume.1
4141
; CHECK-NEXT: [[VAL:%.*]] = select i1 %load.pos, i32 %load, i32 0
42-
; CHECK-NEXT: [[CONT_CAST:%.*]] = bitcast void (i8*, i8)* [[CONT]] to i8*
42+
; CHECK-NEXT: [[CONT_CAST:%.*]] = bitcast void (i8*, i1)* [[CONT]] to i8*
4343
; CHECK-NEXT: [[T0:%.*]] = insertvalue { i8*, i32 } undef, i8* [[CONT_CAST]], 0
4444
; CHECK-NEXT: [[T1:%.*]] = insertvalue { i8*, i32 } [[T0]], i32 [[VAL]], 1
4545
; CHECK-NEXT: ret { i8*, i32 } [[T1]]
4646
; CHECK-NEXT: }
4747

4848
; CHECK-LABEL: define internal void @f.resume.0(i8* noalias nonnull %0, i1 zeroext %1)
4949
; CHECK-NEXT: :
50-
; CHECK-NEXT: [[T0:%.*]] = icmp eq i8 %1, 0
51-
; CHECK-NEXT: br i1 [[T0]],
50+
; CHECK-NEXT: br i1 %1,
5251
; CHECK: :
5352
; CHECK-NEXT: [[T0:%.*]] = bitcast i8* %0 to i32**
5453
; CHECK-NEXT: [[RELOAD:%.*]] = load i32*, i32** [[T0]], align 8
@@ -60,8 +59,7 @@ cleanup:
6059

6160
; CHECK-LABEL: define internal void @f.resume.1(i8* noalias nonnull %0, i1 zeroext %1)
6261
; CHECK-NEXT: :
63-
; CHECK-NEXT: [[T0:%.*]] = icmp eq i8 %1, 0
64-
; CHECK-NEXT: br i1 [[T0]],
62+
; CHECK-NEXT: br i1 %1,
6563
; CHECK: :
6664
; CHECK-NEXT: [[T0:%.*]] = bitcast i8* %0 to i32**
6765
; CHECK-NEXT: [[RELOAD:%.*]] = load i32*, i32** [[T0]], align 8
@@ -81,8 +79,8 @@ entry:
8179
%value = extractvalue {i8*, i32} %result, 1
8280
call void @print(i32 %value)
8381
%cont = extractvalue {i8*, i32} %result, 0
84-
%cont.cast = bitcast i8* %cont to void (i8*, i8)*
85-
call void %cont.cast(i8* %buffer, i8 zeroext 0)
82+
%cont.cast = bitcast i8* %cont to void (i8*, i1)*
83+
call void %cont.cast(i8* %buffer, i1 zeroext 0)
8684
ret void
8785
}
8886

@@ -95,19 +93,19 @@ entry:
9593
; CHECK-NEXT: store i32* %array, i32** [[BUFFER]], align 8
9694
; CHECK-NEXT: [[LOAD:%.*]] = load i32, i32* %array, align 4
9795
; CHECK-NEXT: [[LOAD_POS:%.*]] = icmp sgt i32 [[LOAD]], 0
98-
; CHECK-NEXT: [[CONT:%.*]] = select i1 [[LOAD_POS]], void (i8*, i8)* @f.resume.0, void (i8*, i8)* @f.resume.1
96+
; CHECK-NEXT: [[CONT:%.*]] = select i1 [[LOAD_POS]], void (i8*, i1)* @f.resume.0, void (i8*, i1)* @f.resume.1
9997
; CHECK-NEXT: [[VAL:%.*]] = select i1 [[LOAD_POS]], i32 [[LOAD]], i32 0
10098
; CHECK-NEXT: call void @print(i32 [[VAL]])
101-
; CHECK-NEXT: call void [[CONT]](i8* nonnull [[BUFFER_CAST]], i8 zeroext 0)
99+
; CHECK-NEXT: call void [[CONT]](i8* nonnull [[BUFFER_CAST]], i1 zeroext false)
102100
; CHECK-NEXT: ret void
103101

104102
declare token @llvm.coro.id.retcon.once(i32, i32, i8*, i8*, i8*, i8*)
105103
declare i8* @llvm.coro.begin(token, i8*)
106-
declare i1 @llvm.coro.suspend.retcon(...)
104+
declare i1 @llvm.coro.suspend.retcon.i1(...)
107105
declare i1 @llvm.coro.end(i8*, i1)
108106
declare i8* @llvm.coro.prepare.retcon(i8*)
109107

110-
declare void @prototype(i8*, i8 zeroext)
108+
declare void @prototype(i8*, i1 zeroext)
111109

112110
declare noalias i8* @allocate(i32 %size)
113111
declare void @deallocate(i8* %ptr)

‎llvm/test/Transforms/Coroutines/coro-retcon-once-value2.ll

Lines changed: 6 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -5,11 +5,11 @@ target triple = "x86_64-apple-macosx10.12.0"
55
define {i8*, i32*} @f(i8* %buffer, i32* %ptr) "coroutine.presplit"="1" {
66
entry:
77
%temp = alloca i32, align 4
8-
%id = call token @llvm.coro.id.retcon.once(i32 8, i32 8, i8* %buffer, i8* bitcast (void (i8*, i8)* @prototype to i8*), i8* bitcast (i8* (i32)* @allocate to i8*), i8* bitcast (void (i8*)* @deallocate to i8*))
8+
%id = call token @llvm.coro.id.retcon.once(i32 8, i32 8, i8* %buffer, i8* bitcast (void (i8*, i1)* @prototype to i8*), i8* bitcast (i8* (i32)* @allocate to i8*), i8* bitcast (void (i8*)* @deallocate to i8*))
99
%hdl = call i8* @llvm.coro.begin(token %id, i8* null)
1010
%oldvalue = load i32, i32* %ptr
1111
store i32 %oldvalue, i32* %temp
12-
%unwind = call i1 (...) @llvm.coro.suspend.retcon(i32* %temp)
12+
%unwind = call i1 (...) @llvm.coro.suspend.retcon.i1(i32* %temp)
1313
br i1 %unwind, label %cleanup, label %cont
1414

1515
cont:
@@ -33,7 +33,7 @@ cleanup:
3333
; CHECK-NEXT: store i32* %ptr, i32** [[SPILL]]
3434
; CHECK-NEXT: %oldvalue = load i32, i32* %ptr
3535
; CHECK-NEXT: store i32 %oldvalue, i32* %temp
36-
; CHECK-NEXT: [[T0:%.*]] = insertvalue { i8*, i32* } { i8* bitcast (void (i8*, i8)* @f.resume.0 to i8*), i32* undef }, i32* %temp, 1
36+
; CHECK-NEXT: [[T0:%.*]] = insertvalue { i8*, i32* } { i8* bitcast (void (i8*, i1)* @f.resume.0 to i8*), i32* undef }, i32* %temp, 1
3737
; CHECK-NEXT: ret { i8*, i32* } [[T0]]
3838
; CHECK-NEXT: }
3939

@@ -42,9 +42,8 @@ cleanup:
4242
; CHECK-NEXT: [[T0:%.*]] = bitcast i8* %0 to [[FRAME_T:%.*]]**
4343
; CHECK-NEXT: [[FRAME:%.*]] = load [[FRAME_T]]*, [[FRAME_T]]** [[T0]]
4444
; CHECK-NEXT: bitcast [[FRAME_T]]* [[FRAME]] to i8*
45-
; CHECK-NEXT: [[T0:%.*]] = icmp ne i8 %1, 0
4645
; CHECK-NEXT: %temp = getelementptr inbounds [[FRAME_T]], [[FRAME_T]]* [[FRAME]], i32 0, i32 1
47-
; CHECK-NEXT: br i1 [[T0]],
46+
; CHECK-NEXT: br i1 %1,
4847
; CHECK: :
4948
; CHECK-NEXT: [[TEMP_SLOT:%.*]] = getelementptr inbounds [[FRAME_T]], [[FRAME_T]]* [[FRAME]], i32 0, i32 1
5049
; CHECK-NEXT: [[PTR_SLOT:%.*]] = getelementptr inbounds [[FRAME_T]], [[FRAME_T]]* [[FRAME]], i32 0, i32 0
@@ -60,11 +59,10 @@ cleanup:
6059

6160
declare token @llvm.coro.id.retcon.once(i32, i32, i8*, i8*, i8*, i8*)
6261
declare i8* @llvm.coro.begin(token, i8*)
63-
declare i1 @llvm.coro.suspend.retcon(...)
62+
declare i1 @llvm.coro.suspend.retcon.i1(...)
6463
declare i1 @llvm.coro.end(i8*, i1)
65-
declare i8* @llvm.coro.prepare.retcon(i8*)
6664

67-
declare void @prototype(i8*, i8 zeroext)
65+
declare void @prototype(i8*, i1 zeroext)
6866

6967
declare noalias i8* @allocate(i32 %size)
7068
declare fastcc void @deallocate(i8* %ptr)

0 commit comments

Comments
 (0)