[IRGen] Use coro.end.async to return from an async function

The `coro.end.async` intrinsic allow specifying a function that is to be
tail-called as the last thing before returning.

LLVM lowering will inline the `must-tail-call` function argument to
`coro.end.async`. This `must-tail-call` function can contain a
`musttail` call.

```
define @my_must_tail_call_func(void (*)(i64) %fnptr, i64 %args) {
  musttail call void %fnptr(i64 %args)
  ret void
}

define @async_func() {
  ...
  coro.end.async(..., @my_must_tail_call_func, %return_continuation, i64 %args)
  unreachable
}
```

Author: aschwaighofer

lib/IRGen/GenCall.cpp

@@ -4797,8 +4797,24 @@ void irgen::emitAsyncReturn(IRGenFunction &IGF, AsyncContextLayout &asyncLayout,
   Args.push_back(IGF.getAsyncTask());
   Args.push_back(IGF.getAsyncExecutor());
   Args.push_back(IGF.getAsyncContext());
-  auto call = IGF.Builder.CreateCall(fnPtr, Args);
-  call->setTailCall();
+
+  // Setup the coro.end.async intrinsic call.
+  auto &Builder = IGF.Builder;
+  auto mustTailCallFn = IGF.createAsyncDispatchFn(fnPtr,Args);
+  auto handle = IGF.getCoroutineHandle();
+  auto rawFnPtr =
+      Builder.CreateBitOrPointerCast(fnPtr.getRawPointer(), IGF.IGM.Int8PtrTy);
+
+  SmallVector<llvm::Value*, 8> arguments;
+  arguments.push_back(handle);
+  arguments.push_back(/*is unwind*/Builder.getFalse());
+  arguments.push_back(mustTailCallFn);
+  arguments.push_back(rawFnPtr);
+  for (auto *arg: Args)
+    arguments.push_back(arg);
+
+  Builder.CreateIntrinsicCall(llvm::Intrinsic::coro_end_async, arguments);
+  Builder.CreateUnreachable();
 }
 
 void irgen::emitAsyncReturn(IRGenFunction &IGF, AsyncContextLayout &asyncLayout,

lib/IRGen/GenFunc.cpp

@@ -1286,6 +1286,7 @@ class AsyncPartialApplicationForwarderEmission
     return subIGF.Builder.CreateCall(fnPtr.getAsFunction(subIGF),
                                      asyncExplosion.claimAll());
   }
+  // [FIXME: swiftasynccc] This call should be marked musttail.
   void createReturn(llvm::CallInst *call) override {
     subIGF.Builder.CreateRetVoid();
   }
@@ -2478,7 +2479,7 @@ IRGenFunction::createAsyncDispatchFn(const FunctionPointer &fnPtr,
   llvm::Function *dispatch =
       llvm::Function::Create(dispatchFnTy, llvm::Function::InternalLinkage,
                              llvm::StringRef(name), &IGM.Module);
-  dispatch->setCallingConv(IGM.DefaultCC);
+  dispatch->setCallingConv(IGM.SwiftAsyncCC);
   dispatch->setDoesNotThrow();
   IRGenFunction dispatchIGF(IGM, dispatch);
   if (IGM.DebugInfo)
@@ -2499,7 +2500,7 @@ IRGenFunction::createAsyncDispatchFn(const FunctionPointer &fnPtr,
   auto callee = FunctionPointer(fnPtr.getKind(), fnPtrArg, newAuthInfo,
                                 fnPtr.getSignature());
   auto call = Builder.CreateCall(callee, callArgs);
-  call->setTailCall();
+  call->setTailCallKind(IGM.AsyncTailCallKind);
   Builder.CreateRetVoid();
   return dispatch;
 }
@@ -2566,7 +2567,7 @@ llvm::Function *IRGenFunction::createAsyncSuspendFn() {
   llvm::Function *suspendFn =
       llvm::Function::Create(suspendFnTy, llvm::Function::InternalLinkage,
                              name, &IGM.Module);
-  suspendFn->setCallingConv(IGM.DefaultCC);
+  suspendFn->setCallingConv(IGM.SwiftAsyncCC);
   suspendFn->setDoesNotThrow();
   IRGenFunction suspendIGF(IGM, suspendFn);
   if (IGM.DebugInfo)
@@ -2594,7 +2595,7 @@ llvm::Function *IRGenFunction::createAsyncSuspendFn() {
       { task, executor, targetExecutor });
   suspendCall->setDoesNotThrow();
   suspendCall->setCallingConv(IGM.SwiftAsyncCC);
-  suspendCall->setTailCall();
+  suspendCall->setTailCallKind(IGM.AsyncTailCallKind);
   Builder.CreateRetVoid();
   return suspendFn;
 }

lib/IRGen/GenThunk.cpp

@@ -354,7 +354,6 @@ void IRGenThunk::emit() {
 
   if (isAsync) {
     emitAsyncReturn(IGF, *asyncLayout, origTy, result);
-    IGF.emitCoroutineOrAsyncExit();
     return;
   }
 

lib/IRGen/IRGenModule.cpp

@@ -548,12 +548,16 @@ IRGenModule::IRGenModule(IRGenerator &irgen,
   DefaultCC = SWIFT_DEFAULT_LLVM_CC;
   SwiftCC = llvm::CallingConv::Swift;
 
-  bool isAsynCCSupported =
+  bool isAsyncCCSupported =
     clangASTContext.getTargetInfo().checkCallingConvention(clang::CC_SwiftAsync)
     == clang::TargetInfo::CCCR_OK;
-  SwiftAsyncCC = (opts.UseAsyncLowering && isAsynCCSupported)
-                     ? llvm::CallingConv::SwiftTail
-                     : SwiftCC;
+  if (opts.UseAsyncLowering && isAsyncCCSupported) {
+    SwiftAsyncCC = llvm::CallingConv::SwiftTail;
+    AsyncTailCallKind = llvm::CallInst::TCK_MustTail;
+  } else {
+    SwiftAsyncCC = SwiftCC;
+    AsyncTailCallKind = llvm::CallInst::TCK_Tail;
+  }
 
   if (opts.DebugInfoLevel > IRGenDebugInfoLevel::None)
     DebugInfo = IRGenDebugInfo::createIRGenDebugInfo(IRGen.Opts, *CI, *this,

lib/IRGen/IRGenModule.h

@@ -44,6 +44,7 @@
 #include "llvm/IR/Attributes.h"
 #include "llvm/IR/CallingConv.h"
 #include "llvm/IR/Constant.h"
+#include "llvm/IR/Instructions.h"
 #include "llvm/IR/ValueHandle.h"
 #include "llvm/Target/TargetMachine.h"
 
@@ -756,6 +757,9 @@ class IRGenModule {
   llvm::CallingConv::ID SwiftCC;       /// swift calling convention
   llvm::CallingConv::ID SwiftAsyncCC;  /// swift calling convention for async
 
+  /// What kind of tail call should be used for async->async calls.
+  llvm::CallInst::TailCallKind AsyncTailCallKind;
+
   Signature getAssociatedTypeWitnessTableAccessFunctionSignature();
 
   /// Get the bit width of an integer type for the target platform.

lib/IRGen/IRGenSIL.cpp

@@ -3434,9 +3434,14 @@ void IRGenFunction::emitCoroutineOrAsyncExit() {
   // Emit the block.
   Builder.emitBlock(coroEndBB);
   auto handle = getCoroutineHandle();
-  Builder.CreateIntrinsicCall(llvm::Intrinsic::coro_end,
-                              {handle,
-                               /*is unwind*/ Builder.getFalse()});
+  if (isAsync())
+    Builder.CreateIntrinsicCall(llvm::Intrinsic::coro_end_async,
+                                {handle,
+                                 /*is unwind*/ Builder.getFalse()});
+  else
+    Builder.CreateIntrinsicCall(llvm::Intrinsic::coro_end,
+                                {handle,
+                                 /*is unwind*/ Builder.getFalse()});
   Builder.CreateUnreachable();
 }
 
@@ -3467,10 +3472,8 @@ static void emitReturnInst(IRGenSILFunction &IGF,
     assert(!IGF.IndirectReturn.isValid() &&
            "Formally direct results should stay direct results for async "
            "functions");
-
     auto asyncLayout = getAsyncContextLayout(IGF);
     emitAsyncReturn(IGF, asyncLayout, fnType, result);
-    IGF.emitCoroutineOrAsyncExit();
   } else {
     auto funcLang = IGF.CurSILFn->getLoweredFunctionType()->getLanguage();
     auto swiftCCReturn = funcLang == SILFunctionLanguage::Swift;
@@ -3513,7 +3516,6 @@ void IRGenSILFunction::visitThrowInst(swift::ThrowInst *i) {
   if (isAsync()) {
     auto layout = getAsyncContextLayout(*this);
     emitAsyncReturn(*this, layout, i->getFunction()->getLoweredFunctionType());
-    emitCoroutineOrAsyncExit();
     return;
   }
 

test/IRGen/async.swift

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

test/IRGen/async/get_async_continuation.sil

@@ -84,7 +84,7 @@ bb0:
 // CHECK:  br label %coro.end
 
 // CHECK: coro.end:
-// CHECK:   call i1 @llvm.coro.end(
+// CHECK:   call i1 (i8*, i1, ...) @llvm.coro.end.async(
 // CHECK:   unreachable
 
 // CHECK: await.async.maybe.resume:

test/IRGen/async/hop_to_executor.sil

@@ -28,14 +28,18 @@ final actor MyActor {
 // CHECK: [[CAST_ACTOR:%[0-9]+]] = bitcast %T4test7MyActorC* [[ACTOR]] to %swift.executor*
 // CHECK-x86_64: call {{.*}} @llvm.coro.suspend.async{{.*}}(i32 2, i8* [[RESUME]], i8* bitcast (i8* (i8*)* @__swift_async_resume_get_context to i8*), i8* bitcast (void (i8*, %swift.executor*, %swift.task*, %swift.executor*, %swift.context*)* @__swift_suspend_point to i8*), i8* [[RESUME]], %swift.executor* [[CAST_ACTOR]], %swift.task* [[TASK]], %swift.executor* {{%[0-9]+}}, %swift.context* {{%[0-9]+}})
 // CHECK-arm64e: call {{.*}} @llvm.coro.suspend.async{{.*}}(i32 2, i8* [[RESUME]], i8* bitcast (i8* (i8*)* @__swift_async_resume_get_context to i8*), i8* bitcast (void (i8*, %swift.executor*, %swift.task*, %swift.executor*, %swift.context*)* @__swift_suspend_point to i8*), i8* [[SIGNED_RESUME]], %swift.executor* [[CAST_ACTOR]], %swift.task* [[TASK]], %swift.executor* {{%[0-9]+}}, %swift.context* {{%[0-9]+}})
+// CHECK: [[RET_CONTINUATION:%.*]] = bitcast void (%swift.task*, %swift.executor*, %swift.context*)* {{.*}} to i8*
+// CHECK:  call i1 (i8*, i1, ...) @llvm.coro.end.async(i8* {{.*}}, i1 false, void (i8*, %swift.task*, %swift.executor*, %swift.context*)* @[[TAIL_CALL_FUNC:.*]], i8* [[RET_CONTINUATION]]
+// CHECK: unreachable
+
 sil @test_simple : $@convention(method) @async (@guaranteed MyActor) -> () {
 bb0(%0 : $MyActor):
   hop_to_executor %0 : $MyActor
   %3 = tuple ()
   return %3 : $()
 }
 
-// CHECK-LABEL: define internal void @__swift_suspend_point
+// CHECK-LABEL: define internal swifttailcc void @__swift_suspend_point
 // CHECK-SAME:  (i8* %0, %swift.executor* %1, %swift.task* %2, %swift.executor* %3, %swift.context* [[CTXT:%[^,]+]])
 // CHECK:    [[RESUME_ADDR:%[0-9]+]] = getelementptr inbounds %swift.task, %swift.task* %2, i32 0, i32 4
 // CHECK:    store i8* %0, i8** [[RESUME_ADDR]]
@@ -46,8 +50,15 @@ bb0(%0 : $MyActor):
 // CHECK-arm64e: [[PTRAUTH_SIGN:%[^,]+]] = call i64 @llvm.ptrauth.sign.i64(i64 [[CTXT_INT]], i32 2, i64 [[PTRAUTH_BLEND]])
 // CHECK-arm64e: [[CTXT:%[^,]+]] = inttoptr i64 [[PTRAUTH_SIGN]] to %swift.context*
 // CHECK:    store %swift.context* [[CTXT]], %swift.context** [[CTXT_ADDR]]
-// CHECK:    tail call swift{{(tail)?}}cc void @swift_task_switch(%swift.task* %2, %swift.executor* %3, %swift.executor* %1)
+// CHECK:    {{(must)?}}tail call swift{{(tail)?}}cc void @swift_task_switch(%swift.task* %2, %swift.executor* %3, %swift.executor* %1)
 // CHECK:    ret void
 
+// CHECK: define{{.*}} void @[[TAIL_CALL_FUNC]](i8* %0, %swift.task* %1, %swift.executor* %2, %swift.context* %3)
+// CHECK:   %4 = bitcast i8* %0 to void (%swift.task*, %swift.executor*, %swift.context*)*
+// CHECK:   {{(must)?}}tail call swift{{(tail)?}}cc void %4(%swift.task* %1, %swift.executor* %2, %swift.context* swiftasync %3)
+// CHECK:   ret void
+// CHECK: }
+
+
 sil_vtable MyActor {
 }

test/IRGen/async/partial_apply.sil

@@ -405,7 +405,7 @@ bb0(%x : $*SwiftClassPair):
 sil public_external @use_closure2 : $@async @convention(thin) (@noescape @async @callee_guaranteed (Int) -> Int) -> ()
 
 // CHECK-LABEL: define{{( dllexport)?}}{{( protected)?}} swift{{(tail)?}}cc void @partial_apply_stack_callee_guaranteed_indirect_guaranteed_class_pair_param(%swift.task* {{%[0-9]+}}, %swift.executor* {{%[0-9]+}}, %swift.context* swiftasync {{%[0-9]+}}) {{#[0-9]+}} {
-// CHECK-LABEL: define internal swift{{(tail)?}}cc void @"$s45indirect_guaranteed_captured_class_pair_paramTA.70"(%swift.task* {{%[0-9]+}}, %swift.executor* {{%[0-9]+}}, %swift.context* swiftasync {{%[0-9]+}}) {{#[0-9]+}} {
+// CHECK-LABEL: define internal swift{{(tail)?}}cc void @"$s45indirect_guaranteed_captured_class_pair_paramTA.{{[0-9]+}}"(%swift.task* {{%[0-9]+}}, %swift.executor* {{%[0-9]+}}, %swift.context* swiftasync {{%[0-9]+}}) {{#[0-9]+}} {
 
 sil @partial_apply_stack_callee_guaranteed_indirect_guaranteed_class_pair_param : $@async @convention(thin) (@in_guaranteed SwiftClassPair) -> () {
 bb0(%x : $*SwiftClassPair):