Merge pull request #2101 from aschwaighofer/rdar_70097093_async_lowering

Cherry-pick of llvm.coro async lowering

Author: aschwaighofer

llvm/docs/Coroutines.rst

@@ -174,6 +174,88 @@ 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
+one of its arguments (the argument position is determined by
+`llvm.coro.id.async`). It is used to marshal arguments and return values of the
+coroutine. Therefore an async coroutine returns `void`.
+
+.. code-block:: llvm
+
+  define swiftcc void @async_coroutine(i8* %async.ctxt, i8*, i8*) {
+  }
+
+Values live accross a suspend point need to be stored in the coroutine frame to
+be available in the continuation function. This frame is stored as a tail to the
+`async context`.
+
+Every suspend point takes an `context projection function` argument which
+describes how-to obtain the continuations `async context` and 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 one of its arguments
+argument. The `resume function` can restore its (the caller's) `async context`
+by applying a `context projection function` that is provided by the frontend as
+a parameter to the `llvm.coro.suspend.async` intrinsic.
+
+.. code-block:: c
+
+  // For example:
+  struct async_context {
+    struct async_context *caller_context;
+    ...
+  }
+
+  char *context_projection_function(struct async_context *callee_ctxt) {
+     return callee_ctxt->caller_context;
+  }
+
+.. code-block:: llvm
+
+  %resume_func_ptr = call i8* @llvm.coro.async.resume()
+  call {i8*, i8*, i8*} (i8*, i8*, ...) @llvm.coro.suspend.async(
+                                              i8* %resume_func_ptr,
+                                              i8* %context_projection_function
+
+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 relative_function_pointer_to_async_impl;
+    uint32_t context_size;
+  }
+
+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.
+
+.. code-block:: llvm
+
+  call {i8*, i8*, i8*} (i8*, i8*, ...) @llvm.coro.suspend.async(
+                   i8* %resume_func_ptr,
+                   i8* %context_projection_function,
+                   i8* (bitcast void (i8*, i8*, i8*)* to i8*) %suspend_function,
+                   i8* %arg1, i8* %arg2, i8 %arg3)
+
 Coroutines by Example
 =====================
 
@@ -1093,6 +1175,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 +1501,68 @@ 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* <context projection function>,
+                             ... <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 `context projection function`. It should describe
+how-to restore the `async context` in the continuation function from the first
+argument of the continuation function. Its type is `i8* (i8*)`.
+
+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.prepare.async:
+
+'llvm.coro.prepare.async' Intrinsic
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+::
+
+  declare i8* @llvm.coro.prepare.async(i8* <coroutine function>)
+
+Overview:
+"""""""""
+
+The '``llvm.coro.prepare.async``' intrinsic is used to block inlining of the
+async coroutine until after coroutine splitting.
+
+Arguments:
+""""""""""
+
+The first argument should be an async coroutine of type `void (i8*, i8*, i8*)`.
+Lowering will replace this intrinsic with its coroutine function argument.
+
 .. _coro.suspend.retcon:
 
 'llvm.coro.suspend.retcon' Intrinsic

llvm/include/llvm/IR/Intrinsics.td

@@ -1139,6 +1139,23 @@ 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_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>>]>;
 

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/CoroElide.cpp

@@ -366,7 +366,7 @@ static bool replaceDevirtTrigger(Function &F) {
 }
 
 static bool declaresCoroElideIntrinsics(Module &M) {
-  return coro::declaresIntrinsics(M, {"llvm.coro.id"});
+  return coro::declaresIntrinsics(M, {"llvm.coro.id", "llvm.coro.id.async"});
 }
 
 PreservedAnalyses CoroElidePass::run(Function &F, FunctionAnalysisManager &AM) {

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

@@ -126,10 +126,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");
     }
@@ -617,6 +617,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;
@@ -901,7 +913,8 @@ static Instruction *insertSpills(const SpillInfo &Spills, coro::Shape &Shape) {
   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 (auto &P : Allocas) {
@@ -1514,7 +1527,8 @@ static void sinkSpillUsesAfterCoroBegin(Function &F, const SpillInfo &Spills,
     for (User *U : SpillDef->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);
@@ -1754,7 +1768,8 @@ void coro::buildCoroutineFrame(Function &F, Shape &Shape) {
     }
   }
   LLVM_DEBUG(dump("Spills", 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, Spills, Shape.CoroBegin);
   Shape.FrameTy = buildFrameType(F, Shape, Spills);
   Shape.FramePtr = insertSpills(Spills, Shape);