Merge pull request swiftlang#34907 from eeckstein/concurrency-irgen2

[concurrency] IRGen: update task/executor/context on every suspend point

Author: eeckstein

lib/IRGen/GenCall.cpp

@@ -331,25 +331,60 @@ static Alignment getAsyncContextAlignment(IRGenModule &IGM) {
   return IGM.getPointerAlignment();
 }
 
+void IRGenFunction::setupAsync() {
+  llvm::Value *t = CurFn->getArg((unsigned)AsyncFunctionArgumentIndex::Task);
+  asyncTaskLocation = createAlloca(t->getType(), IGM.getPointerAlignment());
+  Builder.CreateStore(t, asyncTaskLocation);
+
+  llvm::Value *e = CurFn->getArg((unsigned)AsyncFunctionArgumentIndex::Executor);
+  asyncExecutorLocation = createAlloca(e->getType(), IGM.getPointerAlignment());
+  Builder.CreateStore(e, asyncExecutorLocation);
+
+  llvm::Value *c = CurFn->getArg((unsigned)AsyncFunctionArgumentIndex::Context);
+  asyncContextLocation = createAlloca(c->getType(), IGM.getPointerAlignment());
+  Builder.CreateStore(c, asyncContextLocation);
+}
+
 llvm::Value *IRGenFunction::getAsyncTask() {
   assert(isAsync());
-  auto *value = CurFn->getArg((unsigned)AsyncFunctionArgumentIndex::Task);
-  assert(value->getType() == IGM.SwiftTaskPtrTy);
-  return value;
+  return Builder.CreateLoad(asyncTaskLocation);
 }
 
 llvm::Value *IRGenFunction::getAsyncExecutor() {
   assert(isAsync());
-  auto *value = CurFn->getArg((unsigned)AsyncFunctionArgumentIndex::Executor);
-  assert(value->getType() == IGM.SwiftExecutorPtrTy);
-  return value;
+  return Builder.CreateLoad(asyncExecutorLocation);
 }
 
 llvm::Value *IRGenFunction::getAsyncContext() {
   assert(isAsync());
-  auto *value = CurFn->getArg((unsigned)AsyncFunctionArgumentIndex::Context);
-  assert(value->getType() == IGM.SwiftContextPtrTy);
-  return value;
+  return Builder.CreateLoad(asyncContextLocation);
+}
+
+llvm::CallInst *IRGenFunction::emitSuspendAsyncCall(ArrayRef<llvm::Value *> args) {
+  auto *id =
+    Builder.CreateIntrinsicCall(llvm::Intrinsic::coro_suspend_async, args);
+
+  // Update the current values of task, executor and context.
+
+  auto *rawTask = Builder.CreateExtractValue(id,
+      (unsigned)AsyncFunctionArgumentIndex::Task);
+  auto *task = Builder.CreateBitCast(rawTask, IGM.SwiftTaskPtrTy);
+  Builder.CreateStore(task, asyncTaskLocation);
+
+  auto *rawExecutor = Builder.CreateExtractValue(id,
+      (unsigned)AsyncFunctionArgumentIndex::Executor);
+  auto *executor = Builder.CreateBitCast(rawExecutor, IGM.SwiftExecutorPtrTy);
+  Builder.CreateStore(executor, asyncExecutorLocation);
+
+  auto *calleeContext = Builder.CreateExtractValue(id,
+      (unsigned)AsyncFunctionArgumentIndex::Context);
+  llvm::Constant *projectFn = cast<llvm::Constant>(args[1])->stripPointerCasts();
+  // Get the caller context from the calle context.
+  llvm::Value *context = Builder.CreateCall(projectFn, {calleeContext});
+  context = Builder.CreateBitCast(context, IGM.SwiftContextPtrTy);
+  Builder.CreateStore(context, asyncContextLocation);
+
+  return id;
 }
 
 llvm::Type *ExplosionSchema::getScalarResultType(IRGenModule &IGM) const {
@@ -2421,9 +2456,7 @@ class AsyncCallEmission final : public CallEmission {
         Builder.CreateBitOrPointerCast(fn.getRawPointer(), IGM.Int8PtrTy));
     for (auto arg: args)
       arguments.push_back(arg);
-    auto *id = Builder.CreateIntrinsicCall(llvm::Intrinsic::coro_suspend_async,
-                                           arguments);
-    return id;
+    return IGF.emitSuspendAsyncCall(arguments);
   }
 };
 

lib/IRGen/GenFunc.cpp

@@ -1332,7 +1332,8 @@ static llvm::Function *emitPartialApplicationForwarder(IRGenModule &IGM,
   fwd->addAttributes(llvm::AttributeList::FunctionIndex, b);
 
   IRGenFunction subIGF(IGM, fwd);
-  subIGF.setAsync(origType->isAsync());
+  if (origType->isAsync())
+    subIGF.setupAsync();
   if (IGM.DebugInfo)
     IGM.DebugInfo->emitArtificialFunction(subIGF, fwd);
 

lib/IRGen/GenThunk.cpp

@@ -144,7 +144,8 @@ void IRGenModule::emitDispatchThunk(SILDeclRef declRef) {
   }
 
   IRGenFunction IGF(*this, f);
-  IGF.setAsync(declRef.getAbstractFunctionDecl()->hasAsync());
+  if (declRef.getAbstractFunctionDecl()->hasAsync())
+    IGF.setupAsync();
 
   // Look up the method.
   auto fn = lookupMethod(IGF, declRef);

lib/IRGen/IRGenFunction.cpp

@@ -720,7 +720,7 @@ void IRGenFunction::emitAwaitAsyncContinuation(
         Builder.CreateBitOrPointerCast(getAsyncExecutor(), IGM.Int8PtrTy));
     arguments.push_back(Builder.CreateBitOrPointerCast(
         AsyncCoroutineCurrentContinuationContext, IGM.Int8PtrTy));
-    Builder.CreateIntrinsicCall(llvm::Intrinsic::coro_suspend_async, arguments);
+    emitSuspendAsyncCall(arguments);
 
     auto results = Builder.CreateAtomicCmpXchg(
         contAwaitSyncAddr, null, one,

lib/IRGen/IRGenFunction.h

@@ -132,6 +132,8 @@ class IRGenFunction {
   llvm::Value *getAsyncExecutor();
   llvm::Value *getAsyncContext();
 
+  llvm::CallInst *emitSuspendAsyncCall(ArrayRef<llvm::Value *> args);
+
   llvm::Function *getOrCreateResumePrjFn();
   llvm::Function *createAsyncDispatchFn(const FunctionPointer &fnPtr,
                                         ArrayRef<llvm::Value *> args);
@@ -162,7 +164,10 @@ class IRGenFunction {
   llvm::Value *CoroutineHandle = nullptr;
   llvm::Value *AsyncCoroutineCurrentResume = nullptr;
   llvm::Value *AsyncCoroutineCurrentContinuationContext = nullptr;
-  bool IsAsync = false;
+
+  Address asyncTaskLocation;
+  Address asyncExecutorLocation;
+  Address asyncContextLocation;
 
   /// The unique block that calls @llvm.coro.end.
   llvm::BasicBlock *CoroutineExitBlock = nullptr;
@@ -182,8 +187,8 @@ class IRGenFunction {
     return getEffectiveOptimizationMode() == OptimizationMode::ForSize;
   }
 
-  bool isAsync() const { return IsAsync; }
-  void setAsync(bool async = true) { IsAsync = async; }
+  void setupAsync();
+  bool isAsync() const { return asyncTaskLocation.isValid(); }
 
   Address createAlloca(llvm::Type *ty, Alignment align,
                        const llvm::Twine &name = "");

lib/IRGen/IRGenSIL.cpp

@@ -1504,7 +1504,8 @@ IRGenSILFunction::IRGenSILFunction(IRGenModule &IGM, SILFunction *f)
     IGM.createReplaceableProlog(*this, f);
   }
 
-  setAsync(f->getLoweredFunctionType()->isAsync());
+  if (f->getLoweredFunctionType()->isAsync())
+    setupAsync();
 }
 
 IRGenSILFunction::~IRGenSILFunction() {

test/IRGen/async.swift

@@ -15,7 +15,7 @@ public class SomeClass {}
 public func task_future_wait(_ task: __owned SomeClass) async throws -> Int
 
 // CHECK: define{{.*}} swiftcc void @"$s5async8testThisyyAA9SomeClassCnYF"(%swift.task* %0, %swift.executor* %1, %swift.context* %2)
-// CHECK-64: call swiftcc i8* @swift_task_alloc(%swift.task* %0, i64 64)
+// CHECK-64: call swiftcc i8* @swift_task_alloc(%swift.task* %{{[0-9]+}}, i64 64)
 // CHECK: tail call swiftcc void @swift_task_future_wait(
 public func testThis(_ task: __owned SomeClass) async {
   do {

test/IRGen/async/builtins.sil

@@ -11,9 +11,12 @@ import Swift
 // CHECK-LABEL: define hidden swiftcc void @get_task(%swift.task* %0, %swift.executor* %1, %swift.context* %2)
 sil hidden [ossa] @get_task : $@async @convention(thin) () -> @owned Builtin.NativeObject {
 bb0:
-  // CHECK: [[TASK:%.*]] = bitcast %swift.task* %0 to %swift.refcounted*
+  // CHECK: [[TASKLOC:%.*]] = alloca %swift.task*
+  // CHECK: store %swift.task* %0, %swift.task** [[TASKLOC]]
+  // CHECK: [[TASK:%.*]] = load %swift.task*, %swift.task** [[TASKLOC]]
+  // CHECK: [[TASKRC:%.*]] = bitcast %swift.task* [[TASK]] to %swift.refcounted*
   %0 = builtin "getCurrentAsyncTask"() : $Builtin.NativeObject
-  // CHECK-NEXT: [[TASK_COPY:%.*]] = call %swift.refcounted* @swift_retain(%swift.refcounted* returned [[TASK]])
+  // CHECK-NEXT: [[TASK_COPY:%.*]] = call %swift.refcounted* @swift_retain(%swift.refcounted* returned [[TASKRC]])
   %1 = copy_value %0 : $Builtin.NativeObject
   end_lifetime %0 : $Builtin.NativeObject
   return %1 : $Builtin.NativeObject

test/IRGen/async/get_async_continuation.sil

@@ -14,38 +14,44 @@ bb0:
 }
 
 // CHECK-LABEL: define{{.*}} @async_continuation(
+// CHECK: [[tsk_addr:%.*]] = alloca %swift.task*
+// CHECK: [[exe_addr:%.*]] = alloca %swift.executor*
+// CHECK: [[ctxt_addr:%.*]] = alloca %swift.context*
 // CHECK: [[cont_context:%.*]] = alloca %swift.async_continuation_context
 // CHECK: [[result_storage:%.*]] = alloca i32
 // CHECK: call token @llvm.coro.id.async
 // CHECK: call i8* @llvm.coro.begin(
 //   Create UnsafeContinuation<Int32>.
-// CHECK: [[tmp0:%.*]] = bitcast %swift.task* %0 to i8*
+// CHECK: [[tsk:%.*]] = load %swift.task*, %swift.task** [[tsk_addr]]
+// CHECK: [[tmp0:%.*]] = bitcast %swift.task* [[tsk]] to i8*
 // CHECK: [[tmp1:%.*]] = insertvalue %TSV undef, i8* [[tmp0]], 0
 // CHECK: [[continuation:%.*]] = insertvalue %T12_Concurrency18UnsafeContinuationV undef, %TSV [[tmp1]], 0
 //   Initialize the async  continuation context.
 // CHECK: [[context_addr:%.*]] = getelementptr inbounds %swift.async_continuation_context, %swift.async_continuation_context* [[cont_context]], i32 0, i32 0
-// CHECK: store %swift.context* %2, %swift.context** [[context_addr]]
+// CHECK: [[ctxt:%.*]] = load %swift.context*, %swift.context** [[ctxt_addr]]
+// CHECK: store %swift.context* [[ctxt]], %swift.context** [[context_addr]]
 // CHECK: [[error_addr:%.*]] = getelementptr inbounds %swift.async_continuation_context, %swift.async_continuation_context* [[cont_context]], i32 0, i32 2
 // CHECK: store %swift.error* null, %swift.error** [[error_addr]]
 // CHECK: [[result_addr:%.*]] = getelementptr inbounds %swift.async_continuation_context, %swift.async_continuation_context* [[cont_context]], i32 0, i32 3
 // CHECK: [[result_storage_as_opaque:%.*]] = bitcast i32* [[result_storage]] to %swift.opaque*
 // CHECK: store %swift.opaque* [[result_storage_as_opaque]], %swift.opaque** [[result_addr]]
 // CHECK:  [[exectuor_addr:%.*]] = getelementptr inbounds %swift.async_continuation_context, %swift.async_continuation_context* [[cont_context]], i32 0, i32 4
-// CHECK:  store %swift.executor* %1, %swift.executor** [[exectuor_addr]]
+// CHECK:  [[exe:%.*]] = load %swift.executor*, %swift.executor** [[exe_addr]]
+// CHECK:  store %swift.executor* [[exe]], %swift.executor** [[exectuor_addr]]
 //   Initialize the async task with the continuation function and async continuation context.
-// CHECK: [[task_continuation_fn_addr:%.*]] = getelementptr inbounds %swift.task, %swift.task* %0, i32 0, i32 3
+// CHECK: [[task_continuation_fn_addr:%.*]] = getelementptr inbounds %swift.task, %swift.task* [[tsk]], i32 0, i32 3
 // CHECK: [[continuation_fn:%.*]] = call i8* @llvm.coro.async.resume()
 // CHECK: store i8* [[continuation_fn]], i8** [[task_continuation_fn_addr]]
-// CHECK: [[task_resume_context_addr:%.*]] = getelementptr inbounds %swift.task, %swift.task* %0, i32 0, i32 4
+// CHECK: [[task_resume_context_addr:%.*]] = getelementptr inbounds %swift.task, %swift.task* [[tsk]], i32 0, i32 4
 // CHECK: [[cont_context2:%.*]] = bitcast %swift.async_continuation_context* [[cont_context]] to %swift.context*
 // CHECK: store %swift.context* [[cont_context2]], %swift.context** [[task_resume_context_addr]]
 //  Initialize the synchronization variable.
-// CHECK: [[synchronization_addr:%.*]] = getelementptr inbounds %swift.async_continuation_context, %swift.async_continuation_context* %3, i32 0, i32 1
+// CHECK: [[synchronization_addr:%.*]] = getelementptr inbounds %swift.async_continuation_context, %swift.async_continuation_context* [[cont_context]], i32 0, i32 1
 // CHECK: store atomic {{(i64|i32)}} 0, {{(i64|i32)}}* [[synchronization_addr]] release
 //  Do some stuff.
 // CHECK: call swiftcc void @not_async_test()
 //  Arrive at the await_async_continuation point.
-// CHECK: [[synchronization_addr_before_await:%.*]] = getelementptr inbounds %swift.async_continuation_context, %swift.async_continuation_context* %3, i32 0, i32 1
+// CHECK: [[synchronization_addr_before_await:%.*]] = getelementptr inbounds %swift.async_continuation_context, %swift.async_continuation_context* [[cont_context]], i32 0, i32 1
 // CHECK: [[first_at_sync_pt:%.*]] = cmpxchg {{(i64|i32)}}* [[synchronization_addr_before_await]], {{(i64|i32)}} 0, {{(i64|i32)}} 1 release acquire
 // CHECK: [[first_at_sync_pt_bool:%.*]] = extractvalue { {{(i64|i32)}}, i1 } [[first_at_sync_pt]], 1
 // CHECK: br i1 [[first_at_sync_pt_bool]], label %await.async.abort, label %await.async.maybe.resume