Skip to content

Commit bf2be9e

Browse files
eecksteineeckstein
committed
[concurrency] IRGen: update task/executor/context on every suspend point
For this, store those 3 values on the stack at function entry and update them with the return values of coro_suspend_async intrinsic calls. This fixes a correctness issue, because the executor may be different after a resume. It also is more efficient, because this means that the 3 values don't have to preserved in the context over a suspension point.
1 parent 14350f8 commit bf2be9e

File tree

9 files changed

+79
-29
lines changed

9 files changed

+79
-29
lines changed

‎lib/IRGen/GenCall.cpp

Lines changed: 45 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -331,25 +331,60 @@ static Alignment getAsyncContextAlignment(IRGenModule &IGM) {
331331
return IGM.getPointerAlignment();
332332
}
333333

334+
void IRGenFunction::setupAsync() {
335+
llvm::Value *t = CurFn->getArg((unsigned)AsyncFunctionArgumentIndex::Task);
336+
asyncTaskLocation = createAlloca(t->getType(), IGM.getPointerAlignment());
337+
Builder.CreateStore(t, asyncTaskLocation);
338+
339+
llvm::Value *e = CurFn->getArg((unsigned)AsyncFunctionArgumentIndex::Executor);
340+
asyncExecutorLocation = createAlloca(e->getType(), IGM.getPointerAlignment());
341+
Builder.CreateStore(e, asyncExecutorLocation);
342+
343+
llvm::Value *c = CurFn->getArg((unsigned)AsyncFunctionArgumentIndex::Context);
344+
asyncContextLocation = createAlloca(c->getType(), IGM.getPointerAlignment());
345+
Builder.CreateStore(c, asyncContextLocation);
346+
}
347+
334348
llvm::Value *IRGenFunction::getAsyncTask() {
335349
assert(isAsync());
336-
auto *value = CurFn->getArg((unsigned)AsyncFunctionArgumentIndex::Task);
337-
assert(value->getType() == IGM.SwiftTaskPtrTy);
338-
return value;
350+
return Builder.CreateLoad(asyncTaskLocation);
339351
}
340352

341353
llvm::Value *IRGenFunction::getAsyncExecutor() {
342354
assert(isAsync());
343-
auto *value = CurFn->getArg((unsigned)AsyncFunctionArgumentIndex::Executor);
344-
assert(value->getType() == IGM.SwiftExecutorPtrTy);
345-
return value;
355+
return Builder.CreateLoad(asyncExecutorLocation);
346356
}
347357

348358
llvm::Value *IRGenFunction::getAsyncContext() {
349359
assert(isAsync());
350-
auto *value = CurFn->getArg((unsigned)AsyncFunctionArgumentIndex::Context);
351-
assert(value->getType() == IGM.SwiftContextPtrTy);
352-
return value;
360+
return Builder.CreateLoad(asyncContextLocation);
361+
}
362+
363+
llvm::CallInst *IRGenFunction::emitSuspendAsyncCall(ArrayRef<llvm::Value *> args) {
364+
auto *id =
365+
Builder.CreateIntrinsicCall(llvm::Intrinsic::coro_suspend_async, args);
366+
367+
// Update the current values of task, executor and context.
368+
369+
auto *rawTask = Builder.CreateExtractValue(id,
370+
(unsigned)AsyncFunctionArgumentIndex::Task);
371+
auto *task = Builder.CreateBitCast(rawTask, IGM.SwiftTaskPtrTy);
372+
Builder.CreateStore(task, asyncTaskLocation);
373+
374+
auto *rawExecutor = Builder.CreateExtractValue(id,
375+
(unsigned)AsyncFunctionArgumentIndex::Executor);
376+
auto *executor = Builder.CreateBitCast(rawExecutor, IGM.SwiftExecutorPtrTy);
377+
Builder.CreateStore(executor, asyncExecutorLocation);
378+
379+
auto *calleeContext = Builder.CreateExtractValue(id,
380+
(unsigned)AsyncFunctionArgumentIndex::Context);
381+
llvm::Constant *projectFn = cast<llvm::Constant>(args[1])->stripPointerCasts();
382+
// Get the caller context from the calle context.
383+
llvm::Value *context = Builder.CreateCall(projectFn, {calleeContext});
384+
context = Builder.CreateBitCast(context, IGM.SwiftContextPtrTy);
385+
Builder.CreateStore(context, asyncContextLocation);
386+
387+
return id;
353388
}
354389

355390
llvm::Type *ExplosionSchema::getScalarResultType(IRGenModule &IGM) const {
@@ -2421,9 +2456,7 @@ class AsyncCallEmission final : public CallEmission {
24212456
Builder.CreateBitOrPointerCast(fn.getRawPointer(), IGM.Int8PtrTy));
24222457
for (auto arg: args)
24232458
arguments.push_back(arg);
2424-
auto *id = Builder.CreateIntrinsicCall(llvm::Intrinsic::coro_suspend_async,
2425-
arguments);
2426-
return id;
2459+
return IGF.emitSuspendAsyncCall(arguments);
24272460
}
24282461
};
24292462

‎lib/IRGen/GenFunc.cpp

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1332,7 +1332,8 @@ static llvm::Function *emitPartialApplicationForwarder(IRGenModule &IGM,
13321332
fwd->addAttributes(llvm::AttributeList::FunctionIndex, b);
13331333

13341334
IRGenFunction subIGF(IGM, fwd);
1335-
subIGF.setAsync(origType->isAsync());
1335+
if (origType->isAsync())
1336+
subIGF.setupAsync();
13361337
if (IGM.DebugInfo)
13371338
IGM.DebugInfo->emitArtificialFunction(subIGF, fwd);
13381339

‎lib/IRGen/GenThunk.cpp

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -144,7 +144,8 @@ void IRGenModule::emitDispatchThunk(SILDeclRef declRef) {
144144
}
145145

146146
IRGenFunction IGF(*this, f);
147-
IGF.setAsync(declRef.getAbstractFunctionDecl()->hasAsync());
147+
if (declRef.getAbstractFunctionDecl()->hasAsync())
148+
IGF.setupAsync();
148149

149150
// Look up the method.
150151
auto fn = lookupMethod(IGF, declRef);

‎lib/IRGen/IRGenFunction.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -720,7 +720,7 @@ void IRGenFunction::emitAwaitAsyncContinuation(
720720
Builder.CreateBitOrPointerCast(getAsyncExecutor(), IGM.Int8PtrTy));
721721
arguments.push_back(Builder.CreateBitOrPointerCast(
722722
AsyncCoroutineCurrentContinuationContext, IGM.Int8PtrTy));
723-
Builder.CreateIntrinsicCall(llvm::Intrinsic::coro_suspend_async, arguments);
723+
emitSuspendAsyncCall(arguments);
724724

725725
auto results = Builder.CreateAtomicCmpXchg(
726726
contAwaitSyncAddr, null, one,

‎lib/IRGen/IRGenFunction.h

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -132,6 +132,8 @@ class IRGenFunction {
132132
llvm::Value *getAsyncExecutor();
133133
llvm::Value *getAsyncContext();
134134

135+
llvm::CallInst *emitSuspendAsyncCall(ArrayRef<llvm::Value *> args);
136+
135137
llvm::Function *getOrCreateResumePrjFn();
136138
llvm::Function *createAsyncDispatchFn(const FunctionPointer &fnPtr,
137139
ArrayRef<llvm::Value *> args);
@@ -162,7 +164,10 @@ class IRGenFunction {
162164
llvm::Value *CoroutineHandle = nullptr;
163165
llvm::Value *AsyncCoroutineCurrentResume = nullptr;
164166
llvm::Value *AsyncCoroutineCurrentContinuationContext = nullptr;
165-
bool IsAsync = false;
167+
168+
Address asyncTaskLocation;
169+
Address asyncExecutorLocation;
170+
Address asyncContextLocation;
166171

167172
/// The unique block that calls @llvm.coro.end.
168173
llvm::BasicBlock *CoroutineExitBlock = nullptr;
@@ -182,8 +187,8 @@ class IRGenFunction {
182187
return getEffectiveOptimizationMode() == OptimizationMode::ForSize;
183188
}
184189

185-
bool isAsync() const { return IsAsync; }
186-
void setAsync(bool async = true) { IsAsync = async; }
190+
void setupAsync();
191+
bool isAsync() const { return asyncTaskLocation.isValid(); }
187192

188193
Address createAlloca(llvm::Type *ty, Alignment align,
189194
const llvm::Twine &name = "");

‎lib/IRGen/IRGenSIL.cpp

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1504,7 +1504,8 @@ IRGenSILFunction::IRGenSILFunction(IRGenModule &IGM, SILFunction *f)
15041504
IGM.createReplaceableProlog(*this, f);
15051505
}
15061506

1507-
setAsync(f->getLoweredFunctionType()->isAsync());
1507+
if (f->getLoweredFunctionType()->isAsync())
1508+
setupAsync();
15081509
}
15091510

15101511
IRGenSILFunction::~IRGenSILFunction() {

‎test/IRGen/async.swift

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ public class SomeClass {}
1515
public func task_future_wait(_ task: __owned SomeClass) async throws -> Int
1616

1717
// CHECK: define{{.*}} swiftcc void @"$s5async8testThisyyAA9SomeClassCnYF"(%swift.task* %0, %swift.executor* %1, %swift.context* %2)
18-
// CHECK-64: call swiftcc i8* @swift_task_alloc(%swift.task* %0, i64 64)
18+
// CHECK-64: call swiftcc i8* @swift_task_alloc(%swift.task* %{{[0-9]+}}, i64 64)
1919
// CHECK: tail call swiftcc void @swift_task_future_wait(
2020
public func testThis(_ task: __owned SomeClass) async {
2121
do {

‎test/IRGen/async/builtins.sil

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,9 +11,12 @@ import Swift
1111
// CHECK-LABEL: define hidden swiftcc void @get_task(%swift.task* %0, %swift.executor* %1, %swift.context* %2)
1212
sil hidden [ossa] @get_task : $@async @convention(thin) () -> @owned Builtin.NativeObject {
1313
bb0:
14-
// CHECK: [[TASK:%.*]] = bitcast %swift.task* %0 to %swift.refcounted*
14+
// CHECK: [[TASKLOC:%.*]] = alloca %swift.task*
15+
// CHECK: store %swift.task* %0, %swift.task** [[TASKLOC]]
16+
// CHECK: [[TASK:%.*]] = load %swift.task*, %swift.task** [[TASKLOC]]
17+
// CHECK: [[TASKRC:%.*]] = bitcast %swift.task* [[TASK]] to %swift.refcounted*
1518
%0 = builtin "getCurrentAsyncTask"() : $Builtin.NativeObject
16-
// CHECK-NEXT: [[TASK_COPY:%.*]] = call %swift.refcounted* @swift_retain(%swift.refcounted* returned [[TASK]])
19+
// CHECK-NEXT: [[TASK_COPY:%.*]] = call %swift.refcounted* @swift_retain(%swift.refcounted* returned [[TASKRC]])
1720
%1 = copy_value %0 : $Builtin.NativeObject
1821
end_lifetime %0 : $Builtin.NativeObject
1922
return %1 : $Builtin.NativeObject

‎test/IRGen/async/get_async_continuation.sil

Lines changed: 13 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -14,38 +14,44 @@ bb0:
1414
}
1515

1616
// CHECK-LABEL: define{{.*}} @async_continuation(
17+
// CHECK: [[tsk_addr:%.*]] = alloca %swift.task*
18+
// CHECK: [[exe_addr:%.*]] = alloca %swift.executor*
19+
// CHECK: [[ctxt_addr:%.*]] = alloca %swift.context*
1720
// CHECK: [[cont_context:%.*]] = alloca %swift.async_continuation_context
1821
// CHECK: [[result_storage:%.*]] = alloca i32
1922
// CHECK: call token @llvm.coro.id.async
2023
// CHECK: call i8* @llvm.coro.begin(
2124
// Create UnsafeContinuation<Int32>.
22-
// CHECK: [[tmp0:%.*]] = bitcast %swift.task* %0 to i8*
25+
// CHECK: [[tsk:%.*]] = load %swift.task*, %swift.task** [[tsk_addr]]
26+
// CHECK: [[tmp0:%.*]] = bitcast %swift.task* [[tsk]] to i8*
2327
// CHECK: [[tmp1:%.*]] = insertvalue %TSV undef, i8* [[tmp0]], 0
2428
// CHECK: [[continuation:%.*]] = insertvalue %T12_Concurrency18UnsafeContinuationV undef, %TSV [[tmp1]], 0
2529
// Initialize the async continuation context.
2630
// CHECK: [[context_addr:%.*]] = getelementptr inbounds %swift.async_continuation_context, %swift.async_continuation_context* [[cont_context]], i32 0, i32 0
27-
// CHECK: store %swift.context* %2, %swift.context** [[context_addr]]
31+
// CHECK: [[ctxt:%.*]] = load %swift.context*, %swift.context** [[ctxt_addr]]
32+
// CHECK: store %swift.context* [[ctxt]], %swift.context** [[context_addr]]
2833
// CHECK: [[error_addr:%.*]] = getelementptr inbounds %swift.async_continuation_context, %swift.async_continuation_context* [[cont_context]], i32 0, i32 2
2934
// CHECK: store %swift.error* null, %swift.error** [[error_addr]]
3035
// CHECK: [[result_addr:%.*]] = getelementptr inbounds %swift.async_continuation_context, %swift.async_continuation_context* [[cont_context]], i32 0, i32 3
3136
// CHECK: [[result_storage_as_opaque:%.*]] = bitcast i32* [[result_storage]] to %swift.opaque*
3237
// CHECK: store %swift.opaque* [[result_storage_as_opaque]], %swift.opaque** [[result_addr]]
3338
// CHECK: [[exectuor_addr:%.*]] = getelementptr inbounds %swift.async_continuation_context, %swift.async_continuation_context* [[cont_context]], i32 0, i32 4
34-
// CHECK: store %swift.executor* %1, %swift.executor** [[exectuor_addr]]
39+
// CHECK: [[exe:%.*]] = load %swift.executor*, %swift.executor** [[exe_addr]]
40+
// CHECK: store %swift.executor* [[exe]], %swift.executor** [[exectuor_addr]]
3541
// Initialize the async task with the continuation function and async continuation context.
36-
// CHECK: [[task_continuation_fn_addr:%.*]] = getelementptr inbounds %swift.task, %swift.task* %0, i32 0, i32 3
42+
// CHECK: [[task_continuation_fn_addr:%.*]] = getelementptr inbounds %swift.task, %swift.task* [[tsk]], i32 0, i32 3
3743
// CHECK: [[continuation_fn:%.*]] = call i8* @llvm.coro.async.resume()
3844
// CHECK: store i8* [[continuation_fn]], i8** [[task_continuation_fn_addr]]
39-
// CHECK: [[task_resume_context_addr:%.*]] = getelementptr inbounds %swift.task, %swift.task* %0, i32 0, i32 4
45+
// CHECK: [[task_resume_context_addr:%.*]] = getelementptr inbounds %swift.task, %swift.task* [[tsk]], i32 0, i32 4
4046
// CHECK: [[cont_context2:%.*]] = bitcast %swift.async_continuation_context* [[cont_context]] to %swift.context*
4147
// CHECK: store %swift.context* [[cont_context2]], %swift.context** [[task_resume_context_addr]]
4248
// Initialize the synchronization variable.
43-
// CHECK: [[synchronization_addr:%.*]] = getelementptr inbounds %swift.async_continuation_context, %swift.async_continuation_context* %3, i32 0, i32 1
49+
// CHECK: [[synchronization_addr:%.*]] = getelementptr inbounds %swift.async_continuation_context, %swift.async_continuation_context* [[cont_context]], i32 0, i32 1
4450
// CHECK: store atomic {{(i64|i32)}} 0, {{(i64|i32)}}* [[synchronization_addr]] release
4551
// Do some stuff.
4652
// CHECK: call swiftcc void @not_async_test()
4753
// Arrive at the await_async_continuation point.
48-
// CHECK: [[synchronization_addr_before_await:%.*]] = getelementptr inbounds %swift.async_continuation_context, %swift.async_continuation_context* %3, i32 0, i32 1
54+
// CHECK: [[synchronization_addr_before_await:%.*]] = getelementptr inbounds %swift.async_continuation_context, %swift.async_continuation_context* [[cont_context]], i32 0, i32 1
4955
// CHECK: [[first_at_sync_pt:%.*]] = cmpxchg {{(i64|i32)}}* [[synchronization_addr_before_await]], {{(i64|i32)}} 0, {{(i64|i32)}} 1 release acquire
5056
// CHECK: [[first_at_sync_pt_bool:%.*]] = extractvalue { {{(i64|i32)}}, i1 } [[first_at_sync_pt]], 1
5157
// CHECK: br i1 [[first_at_sync_pt_bool]], label %await.async.abort, label %await.async.maybe.resume

0 commit comments

Comments
 (0)