Start of an llvm.coro.async implementation

This patch adds the `async` lowering of coroutines.

This will be used by the Swift frontend to lower async functions. In
contrast to the `retcon` lowering the frontend needs to be in control
over control-flow at suspend points as execution might be suspended at
these points.

This is very much work in progress and the implementation will change as
it evolves with the frontend. As such the documentation is lacking
detail as some of it might change.

rdar://70097093

Differential Revision: https://reviews.llvm.org/D90612

Author: aschwaighofer

llvm/docs/Coroutines.rst

@@ -174,6 +174,61 @@ This may be acceptable if LLVM's coroutine support is primarily being
 used for low-level lowering and inlining is expected to be applied
 earlier in the pipeline.
 
+Async Lowering
+--------------
+
+In async-continuation lowering, signaled by the use of `llvm.coro.id.async`,
+handling of control-flow must be handled explicitly by the frontend.
+
+In this lowering, a coroutine is assumed to take the current `async context` as
+its first argument. It is used to marshal arguments and return values of the
+coroutine. Therefore a async coroutine returns `void`.
+
+.. code-block:: llvm
+  define swiftcc void @async_coroutine(i8* %async.ctxt, i8*, i8*) {
+  }
+
+
+Every suspend point takes an `async context` argument which provides the context
+and the coroutine frame of the callee function. Every
+suspend point has an associated `resume function` denoted by the
+`llvm.coro.async.resume` intrinsic. The coroutine is resumed by
+calling this `resume function` passing the `async context` as the first
+argument. It is assumed that the `resume function` can restore its (the
+caller's) `async context` by loading the first field in the `async context`.
+
+.. code-block:: c
+
+  struct async_context {
+    struct async_context *caller_context;
+    ...
+  }
+
+The frontend should provide a `async function pointer` struct associated with
+each async coroutine by `llvm.coro.id.async`'s argument. The initial size and
+alignment of the `async context` must be provided as arguments to the
+`llvm.coro.id.async` intrinsic. Lowering will update the size entry with the
+coroutine frame  requirements. The frontend is responsible for allocating the
+memory for the `async context` but can use the `async function pointer` struct
+to obtain the required size.
+
+.. code-block:: c
+  struct async_function_pointer {
+    uint32_t context_size;
+    uint32_t relative_function_pointer_to_async_impl;
+  }
+
+Lowering will split an async coroutine into a ramp function and one resume
+function per suspend point.
+
+How control-flow is passed between caller, suspension point, and back to
+resume function is left up to the frontend.
+
+The suspend point takes a function and its arguments. The function is intended
+to model the transfer to the callee function. It will be tail called by
+lowering and therefore must have the same signature and calling convention as
+the async coroutine.
+
 Coroutines by Example
 =====================
 
@@ -1093,6 +1148,45 @@ duplicated.
 
 A frontend should emit exactly one `coro.id` intrinsic per coroutine.
 
+.. _coro.id.async:
+
+'llvm.coro.id.async' Intrinsic
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+::
+
+  declare token @llvm.coro.id.async(i32 <context size>, i32 <align>,
+                                    i8* <context arg>,
+                                    i8* <async function pointer>)
+
+Overview:
+"""""""""
+
+The '``llvm.coro.id.async``' intrinsic returns a token identifying an async coroutine.
+
+Arguments:
+""""""""""
+
+The first argument provides the initial size of the `async context` as required
+from the frontend. Lowering will add to this size the size required by the frame
+storage and store that value to the `async function pointer`.
+
+The second argument, is the alignment guarantee of the memory of the
+`async context`. The frontend guarantees that the memory will be aligned by this
+value.
+
+The third argument is the `async context` argument in the current coroutine.
+
+The fourth argument is the address of the `async function pointer` struct.
+Lowering will update the context size requirement in this struct by adding the
+coroutine frame size requirement to the initial size requirement as specified by
+the first argument of this intrinisc.
+
+
+Semantics:
+""""""""""
+
+A frontend should emit exactly one `coro.id.async` intrinsic per coroutine.
+
 .. _coro.id.retcon:
 
 'llvm.coro.id.retcon' Intrinsic
@@ -1380,6 +1474,46 @@ to the coroutine:
     switch i8 %suspend1, label %suspend [i8 0, label %resume1
                                          i8 1, label %cleanup]
 
+.. _coro.suspend.async:
+
+'llvm.coro.suspend.async' Intrinsic
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+::
+
+  declare {i8*, i8*, i8*} @llvm.coro.suspend.async(i8* <resume function>,
+                                                   i8* <callee context>,
+                                                   ... <function to call>
+                                                   ... <arguments to function>)
+
+Overview:
+"""""""""
+
+The '``llvm.coro.suspend.async``' intrinsic marks the point where
+execution of a async coroutine is suspended and control is passed to a callee.
+
+Arguments:
+""""""""""
+
+The first argument should be the result of the `llvm.coro.async.resume` intrinsic.
+Lowering will replace this intrinsic with the resume function for this suspend
+point.
+
+The second argument is the `async context` allocation for the callee. It should
+provide storage the `async context` header and the coroutine frame.
+
+The third argument is the function that models tranfer to the callee at the
+suspend point. It should take 3 arguments. Lowering will `musttail` call this
+function.
+
+The fourth to six argument are the arguments for the third argument.
+
+Semantics:
+""""""""""
+
+The result of the intrinsic are mapped to the arguments of the resume function.
+Execution is suspended at this intrinsic and resumed when the resume function is
+called.
+
 .. _coro.suspend.retcon:
 
 'llvm.coro.suspend.retcon' Intrinsic

llvm/include/llvm/IR/Intrinsics.td

@@ -1187,6 +1187,21 @@ def int_coro_id_retcon_once : Intrinsic<[llvm_token_ty],
      llvm_ptr_ty, llvm_ptr_ty, llvm_ptr_ty],
     []>;
 def int_coro_alloc : Intrinsic<[llvm_i1_ty], [llvm_token_ty], []>;
+def int_coro_id_async : Intrinsic<[llvm_token_ty],
+  [llvm_i32_ty, llvm_i32_ty, llvm_ptr_ty, llvm_ptr_ty],
+  []>;
+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_begin : Intrinsic<[llvm_ptr_ty], [llvm_token_ty, llvm_ptr_ty],
                                [WriteOnly<ArgIndex<1>>]>;
 

llvm/lib/Transforms/Coroutines/CoroCleanup.cpp

@@ -74,6 +74,7 @@ bool Lowerer::lowerRemainingCoroIntrinsics(Function &F) {
       case Intrinsic::coro_id:
       case Intrinsic::coro_id_retcon:
       case Intrinsic::coro_id_retcon_once:
+      case Intrinsic::coro_id_async:
         II->replaceAllUsesWith(ConstantTokenNone::get(Context));
         break;
       case Intrinsic::coro_subfn_addr:

llvm/lib/Transforms/Coroutines/CoroEarly.cpp

@@ -187,6 +187,7 @@ bool Lowerer::lowerEarlyIntrinsics(Function &F) {
         break;
       case Intrinsic::coro_id_retcon:
       case Intrinsic::coro_id_retcon_once:
+      case Intrinsic::coro_id_async:
         F.addFnAttr(CORO_PRESPLIT_ATTR, PREPARED_FOR_SPLIT);
         break;
       case Intrinsic::coro_resume:

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

@@ -135,10 +135,10 @@ struct SuspendCrossingInfo {
 
     BasicBlock *UseBB = I->getParent();
 
-    // As a special case, treat uses by an llvm.coro.suspend.retcon
-    // as if they were uses in the suspend's single predecessor: the
-    // uses conceptually occur before the suspend.
-    if (isa<CoroSuspendRetconInst>(I)) {
+    // As a special case, treat uses by an llvm.coro.suspend.retcon or an
+    // llvm.coro.suspend.async as if they were uses in the suspend's single
+    // predecessor: the uses conceptually occur before the suspend.
+    if (isa<CoroSuspendRetconInst>(I) || isa<CoroSuspendAsyncInst>(I)) {
       UseBB = UseBB->getSinglePredecessor();
       assert(UseBB && "should have split coro.suspend into its own block");
     }
@@ -788,6 +788,18 @@ static StructType *buildFrameType(Function &F, coro::Shape &Shape,
          B.getStructAlign() <= Id->getStorageAlignment());
     break;
   }
+  case coro::ABI::Async: {
+    Shape.AsyncLowering.FrameOffset =
+        alignTo(Shape.AsyncLowering.ContextHeaderSize, Shape.FrameAlign);
+    Shape.AsyncLowering.ContextSize =
+        Shape.AsyncLowering.FrameOffset + Shape.FrameSize;
+    if (Shape.AsyncLowering.getContextAlignment() < Shape.FrameAlign) {
+      report_fatal_error(
+          "The alignment requirment of frame variables cannot be higher than "
+          "the alignment of the async function context");
+    }
+    break;
+  }
   }
 
   return FrameTy;
@@ -1143,7 +1155,8 @@ static Instruction *insertSpills(const FrameDataInfo &FrameData,
   Shape.AllocaSpillBlock = SpillBlock;
 
   // retcon and retcon.once lowering assumes all uses have been sunk.
-  if (Shape.ABI == coro::ABI::Retcon || Shape.ABI == coro::ABI::RetconOnce) {
+  if (Shape.ABI == coro::ABI::Retcon || Shape.ABI == coro::ABI::RetconOnce ||
+      Shape.ABI == coro::ABI::Async) {
     // If we found any allocas, replace all of their remaining uses with Geps.
     Builder.SetInsertPoint(&SpillBlock->front());
     for (const auto &P : FrameData.Allocas) {
@@ -1866,7 +1879,8 @@ static void sinkSpillUsesAfterCoroBegin(Function &F,
     for (User *U : Def->users()) {
       auto Inst = cast<Instruction>(U);
       if (Inst->getParent() != CoroBegin->getParent() ||
-          Dom.dominates(CoroBegin, Inst))
+          Dom.dominates(CoroBegin, Inst) ||
+          isa<CoroIdAsyncInst>(Inst) /*'fake' use of async context argument*/)
         continue;
       if (ToMove.insert(Inst))
         Worklist.push_back(Inst);
@@ -2162,7 +2176,8 @@ void coro::buildCoroutineFrame(Function &F, Shape &Shape) {
       }
   }
   LLVM_DEBUG(dumpSpills("Spills", FrameData.Spills));
-  if (Shape.ABI == coro::ABI::Retcon || Shape.ABI == coro::ABI::RetconOnce)
+  if (Shape.ABI == coro::ABI::Retcon || Shape.ABI == coro::ABI::RetconOnce ||
+      Shape.ABI == coro::ABI::Async)
     sinkSpillUsesAfterCoroBegin(F, FrameData, Shape.CoroBegin);
   Shape.FrameTy = buildFrameType(F, Shape, FrameData);
   // Add PromiseAlloca to Allocas list so that it is processed in insertSpills.