[coro] Add support for polymorphic return typed coro.suspend.async

[aschwaighofer]

Summary

This allows for suspend point specific resume function types.

Return values from a suspend point can therefore be modelled as
arguments to the resume function. Allowing for directly passed return
types.

[aschwaighofer]

Same patch. Lets restart testing. The previous failures all seemed very unrelated.

llvm/include/llvm/IR/Intrinsics.td

def int_coro_async_context_alloc : Intrinsic<[llvm_ptr_ty],
     [llvm_ptr_ty, llvm_ptr_ty],
     []>;
 def int_coro_async_context_dealloc : Intrinsic<[],
     [llvm_ptr_ty],
     []>;
 def int_coro_async_resume : Intrinsic<[llvm_ptr_ty],
     [],
     []>;
-def int_coro_suspend_async : Intrinsic<[llvm_ptr_ty, llvm_ptr_ty, llvm_ptr_ty],
-    [llvm_ptr_ty, llvm_ptr_ty, llvm_vararg_ty],
-    []>;
+def int_coro_suspend_async
+    : Intrinsic<[llvm_any_ty], [llvm_ptr_ty, llvm_ptr_ty, llvm_vararg_ty], []>;
 def int_coro_prepare_async : Intrinsic<[llvm_ptr_ty], [llvm_ptr_ty],
                                        [IntrNoMem]>;
 def int_coro_begin : Intrinsic<[llvm_ptr_ty], [llvm_token_ty, llvm_ptr_ty],
                                [WriteOnly<ArgIndex<1>>]>;
 
 def int_coro_free : Intrinsic<[llvm_ptr_ty], [llvm_token_ty, llvm_ptr_ty],
                               [IntrReadMem, IntrArgMemOnly,
                                ReadOnly<ArgIndex<1>>,

llvm/lib/Transforms/Coroutines/CoroInternal.h

struct RetconLoweringStorage {
     Function *Dealloc;
     BasicBlock *ReturnBlock;
     bool IsFrameInlineInStorage;
   };
 
   struct AsyncLoweringStorage {
     FunctionType *AsyncFuncTy;
     Value *Context;
+    CallingConv::ID AsyncCC;
     unsigned ContextArgNo;
     uint64_t ContextHeaderSize;
     uint64_t ContextAlignment;
     uint64_t FrameOffset; // Start of the frame.
     uint64_t ContextSize; // Includes frame size.
     GlobalVariable *AsyncFuncPointer;
 
     Align getContextAlignment() const { return Align(ContextAlignment); }
FunctionType *getResumeFunctionType() const {
     case coro::ABI::Switch: {
       auto *FnPtrTy = getSwitchResumePointerType();
       return cast<FunctionType>(FnPtrTy->getPointerElementType());
     }
     case coro::ABI::Retcon:
     case coro::ABI::RetconOnce:
       return RetconLowering.ResumePrototype->getFunctionType();
     case coro::ABI::Async:
-      return AsyncLowering.AsyncFuncTy;
+      // Not used. The function type depends on the active suspend.
+      return nullptr;
     }
 
     llvm_unreachable("Unknown coro::ABI enum");
   }
 
   ArrayRef<Type*> getRetconResultTypes() const {
     assert(ABI == coro::ABI::Retcon ||
            ABI == coro::ABI::RetconOnce);
CallingConv::ID getResumeFunctionCC() const {
     switch (ABI) {
     case coro::ABI::Switch:
       return CallingConv::Fast;
 
     case coro::ABI::Retcon:
     case coro::ABI::RetconOnce:
       return RetconLowering.ResumePrototype->getCallingConv();
     case coro::ABI::Async:
-      return CallingConv::Swift;
+      return AsyncLowering.AsyncCC;
     }
     llvm_unreachable("Unknown coro::ABI enum");
   }
 
   AllocaInst *getPromiseAlloca() const {
     if (ABI == coro::ABI::Switch)
       return SwitchLowering.PromiseAlloca;
     return nullptr;

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

if (isSwitchDestroyFunction()) {
     auto *Load = Builder.CreateLoad(Shape.getSwitchResumePointerType(),
                                     GepIndex);
     auto *Cond = Builder.CreateIsNull(Load);
     Builder.CreateCondBr(Cond, ResumeBB, NewSwitchBB);
     OldSwitchBB->getTerminator()->eraseFromParent();
   }
 }
 
+static FunctionType *
+getFunctionTypeFromAsyncSuspend(AnyCoroSuspendInst *Suspend) {
+  auto *AsyncSuspend = cast<CoroSuspendAsyncInst>(Suspend);
+  auto *StructTy = cast<StructType>(AsyncSuspend->getType());
+  auto &Context = Suspend->getParent()->getParent()->getContext();
+  auto *VoidTy = Type::getVoidTy(Context);
+  return FunctionType::get(VoidTy, StructTy->elements(), false);
+}
+
 static Function *createCloneDeclaration(Function &OrigF, coro::Shape &Shape,
                                         const Twine &Suffix,
-                                        Module::iterator InsertBefore) {
+                                        Module::iterator InsertBefore,
+                                        AnyCoroSuspendInst *ActiveSuspend) {
   Module *M = OrigF.getParent();
-  auto *FnTy = Shape.getResumeFunctionType();
+  auto *FnTy = (Shape.ABI != coro::ABI::Async)
+                   ? Shape.getResumeFunctionType()
+                   : getFunctionTypeFromAsyncSuspend(ActiveSuspend);
 
   Function *NewF =
       Function::Create(FnTy, GlobalValue::LinkageTypes::InternalLinkage,
                        OrigF.getName() + Suffix);
   NewF->addParamAttr(0, Attribute::NonNull);
 
   // For the async lowering ABI we can't guarantee that the context argument is
   // not access via a different pointer not based on the argument.

 }
 
 /// Clone the body of the original function into a resume function of
 /// some sort.
 void CoroCloner::create() {
   // Create the new function if we don't already have one.
   if (!NewF) {
     NewF = createCloneDeclaration(OrigF, Shape, Suffix,
-                                  OrigF.getParent()->end());
+                                  OrigF.getParent()->end(), ActiveSuspend);
   }
 
   // Replace all args with undefs. The buildCoroutineFrame algorithm already
   // rewritten access to the args that occurs after suspend points with loads
   // and stores to/from the coroutine frame.
   for (Argument &A : OrigF.args())
     VMap[&A] = UndefValue::get(A.getType());
 
static void splitAsyncCoroutine(Function &F, coro::Shape &Shape,
   auto NextF = std::next(F.getIterator());
 
   // Create a continuation function for each of the suspend points.
   Clones.reserve(Shape.CoroSuspends.size());
   for (size_t Idx = 0, End = Shape.CoroSuspends.size(); Idx != End; ++Idx) {
     auto *Suspend = cast<CoroSuspendAsyncInst>(Shape.CoroSuspends[Idx]);
 
     // Create the clone declaration.
-    auto *Continuation =
-        createCloneDeclaration(F, Shape, ".resume." + Twine(Idx), NextF);
+    auto *Continuation = createCloneDeclaration(
+        F, Shape, ".resume." + Twine(Idx), NextF, Suspend);
     Clones.push_back(Continuation);
 
     // Insert a branch to a new return block immediately before the suspend
     // point.
     auto *SuspendBB = Suspend->getParent();
     auto *NewSuspendBB = SuspendBB->splitBasicBlock(Suspend);
     auto *Branch = cast<BranchInst>(SuspendBB->getTerminator());
 
static void splitRetconCoroutine(Function &F, coro::Shape &Shape,
 
   // Create a continuation function for each of the suspend points.
   Clones.reserve(Shape.CoroSuspends.size());
   for (size_t i = 0, e = Shape.CoroSuspends.size(); i != e; ++i) {
     auto Suspend = cast<CoroSuspendRetconInst>(Shape.CoroSuspends[i]);
 
     // Create the clone declaration.
     auto Continuation =
-      createCloneDeclaration(F, Shape, ".resume." + Twine(i), NextF);
+        createCloneDeclaration(F, Shape, ".resume." + Twine(i), NextF, nullptr);
     Clones.push_back(Continuation);
 
     // Insert a branch to the unified return block immediately before
     // the suspend point.
     auto SuspendBB = Suspend->getParent();
     auto NewSuspendBB = SuspendBB->splitBasicBlock(Suspend);
     auto Branch = cast<BranchInst>(SuspendBB->getTerminator());
 

llvm/lib/Transforms/Coroutines/Coroutines.cpp

case Intrinsic::coro_id_async: {
     AsyncId->checkWellFormed();
     this->ABI = coro::ABI::Async;
     this->AsyncLowering.Context = AsyncId->getStorage();
     this->AsyncLowering.ContextArgNo = AsyncId->getStorageArgumentIndex();
     this->AsyncLowering.ContextHeaderSize = AsyncId->getStorageSize();
     this->AsyncLowering.ContextAlignment =
         AsyncId->getStorageAlignment().value();
     this->AsyncLowering.AsyncFuncPointer = AsyncId->getAsyncFunctionPointer();
-    auto &Context = F.getContext();
-    auto *Int8PtrTy = Type::getInt8PtrTy(Context);
-    auto *VoidTy = Type::getVoidTy(Context);
-    this->AsyncLowering.AsyncFuncTy =
-        FunctionType::get(VoidTy, {Int8PtrTy, Int8PtrTy, Int8PtrTy}, false);
+    this->AsyncLowering.AsyncCC = F.getCallingConv();
     break;
   };
   case Intrinsic::coro_id_retcon:
   case Intrinsic::coro_id_retcon_once: {
     auto ContinuationId = cast<AnyCoroIdRetconInst>(Id);
     ContinuationId->checkWellFormed();
     this->ABI = (IdIntrinsic == Intrinsic::coro_id_retcon
                   ? coro::ABI::Retcon

llvm/test/Transforms/Coroutines/coro-async.ll

call i1 (i8*, i1, ...) @llvm.coro.end.async(
                            i8* %async.ctxt, %async.task* %task.2, %async.actor* null)
   unreachable
 }
 
 ; CHECK-LABEL: define internal swiftcc void @multiple_coro_end_async.resume.0(
 ; CHECK: musttail call swiftcc void @asyncReturn(
 ; CHECK: ret void
 
+@polymorphic_suspend_return_fp = constant <{ i32, i32 }>
+  <{ i32 trunc ( ; Relative pointer to async function
+       i64 sub (
+         i64 ptrtoint (void (i8*, %async.task*, %async.actor*)* @polymorphic_suspend_return to i64),
+         i64 ptrtoint (i32* getelementptr inbounds (<{ i32, i32 }>, <{ i32, i32 }>* @polymorphic_suspend_return_fp, i32 0, i32 1) to i64)
+       )
+     to i32),
+     i32 64    ; Initial async context size without space for frame
+}>
+
+define swiftcc void @polymorphic_suspend_return(i8* %async.ctxt, %async.task* %task, %async.actor* %actor)  {
+entry:
+  %tmp = alloca { i64, i64 }, align 8
+  %proj.1 = getelementptr inbounds { i64, i64 }, { i64, i64 }* %tmp, i64 0, i32 0
+  %proj.2 = getelementptr inbounds { i64, i64 }, { i64, i64 }* %tmp, i64 0, i32 1
+
+  %id = call token @llvm.coro.id.async(i32 128, i32 16, i32 0,
+          i8* bitcast (<{i32, i32}>* @polymorphic_suspend_return_fp to i8*))
+  %hdl = call i8* @llvm.coro.begin(token %id, i8* null)
+  store i64 0, i64* %proj.1, align 8
+  store i64 1, i64* %proj.2, align 8
+  call void @some_may_write(i64* %proj.1)
+
+  ; Begin lowering: apply %my_other_async_function(%args...)
+
+  ; setup callee context
+  %arg0 = bitcast %async.task* %task to i8*
+  %arg1 = bitcast <{ i32, i32}>* @my_other_async_function_fp to i8*
+  %callee_context = call i8* @llvm.coro.async.context.alloc(i8* %arg0, i8* %arg1)
+  %callee_context.0 = bitcast i8* %callee_context to %async.ctxt*
+  ; store arguments ...
+  ; ... (omitted)
+
+  ; store the return continuation
+  %callee_context.return_to_caller.addr = getelementptr inbounds %async.ctxt, %async.ctxt* %callee_context.0, i32 0, i32 1
+  %return_to_caller.addr = bitcast void(i8*, %async.task*, %async.actor*)** %callee_context.return_to_caller.addr to i8**
+  %resume.func_ptr = call i8* @llvm.coro.async.resume()
+  store i8* %resume.func_ptr, i8** %return_to_caller.addr
+
+  ; store caller context into callee context
+  %callee_context.caller_context.addr = getelementptr inbounds %async.ctxt, %async.ctxt* %callee_context.0, i32 0, i32 0
+  store i8* %async.ctxt, i8** %callee_context.caller_context.addr
+  %resume_proj_fun = bitcast i8*(i8*)* @resume_context_projection to i8*
+  %callee = bitcast void(i8*, %async.task*, %async.actor*)* @asyncSuspend to i8*
+  %res = call {i8*, i8*, i8*, i8*} (i8*, i8*, ...) @llvm.coro.suspend.async.sl_p0i8p0i8p0i8p0i8s(
+                                                  i8* %resume.func_ptr,
+                                                  i8* %resume_proj_fun,
+                                                  void (i8*, i8*, %async.task*, %async.actor*)* @my_async_function.my_other_async_function_fp.apply,
+                                                  i8* %callee, i8* %callee_context, %async.task* %task, %async.actor *%actor)
+
+  call void @llvm.coro.async.context.dealloc(i8* %callee_context)
+  %continuation_task_arg = extractvalue {i8*, i8*, i8*, i8*} %res, 3
+  %task.2 =  bitcast i8* %continuation_task_arg to %async.task*
+  %val = load i64, i64* %proj.1
+  call void @some_user(i64 %val)
+  %val.2 = load i64, i64* %proj.2
+  call void @some_user(i64 %val.2)
+
+  tail call swiftcc void @asyncReturn(i8* %async.ctxt, %async.task* %task.2, %async.actor* %actor)
+  call i1 (i8*, i1, ...) @llvm.coro.end.async(i8* %hdl, i1 0)
+  unreachable
+}
+
+; CHECK-LABEL: define swiftcc void @polymorphic_suspend_return(i8* %async.ctxt, %async.task* %task, %async.actor* %actor)
+; CHECK-LABEL: define internal swiftcc void @polymorphic_suspend_return.resume.0(i8* {{.*}}%0, i8* {{.*}}%1, i8* {{.*}}%2, i8* {{.*}}%3)
+; CHECK: bitcast i8* %3 to %async.task*
+; CHECK: }
+
+declare { i8*, i8*, i8*, i8* } @llvm.coro.suspend.async.sl_p0i8p0i8p0i8p0i8s(i8*, i8*, ...)
 declare i8* @llvm.coro.prepare.async(i8*)
 declare token @llvm.coro.id.async(i32, i32, i32, i8*)
 declare i8* @llvm.coro.begin(token, i8*)
 declare i1 @llvm.coro.end.async(i8*, i1, ...)
 declare i1 @llvm.coro.end(i8*, i1)
 declare {i8*, i8*, i8*} @llvm.coro.suspend.async(i8*, i8*, ...)
 declare i8* @llvm.coro.async.context.alloc(i8*, i8*)
 declare void @llvm.coro.async.context.dealloc(i8*)
 declare swiftcc void @asyncReturn(i8*, %async.task*, %async.actor*)
 declare swiftcc void @asyncSuspend(i8*, %async.task*, %async.actor*)
 declare i8* @llvm.coro.async.resume()