[Coroutines] Documentation for plugin libraries and custom ABIs

llvm/docs/Coroutines.rst

@@ -312,6 +312,7 @@ lowered to a constant representing the size required for the coroutine frame.
 The `coro.begin`_ intrinsic initializes the coroutine frame and returns the
 coroutine handle. The second parameter of `coro.begin` is given a block of memory
 to be used if the coroutine frame needs to be allocated dynamically.
+
 The `coro.id`_ intrinsic serves as coroutine identity useful in cases when the
 `coro.begin`_ intrinsic get duplicated by optimization passes such as
 jump-threading.
@@ -749,6 +750,65 @@ and python iterator `__next__` would look like:
     return *(int*)coro.promise(hdl, 4, false);
   }
 
+Custom ABIs and Plugin Libraries
+--------------------------------
+
+Plugin libraries can extend coroutine lowering enabling a wide variety of users
+to utilize the coroutine transformation passes. An existing coroutine lowering
+is extended by: 1. defining custom ABIs that inherit from the existing ABIs,
+2. give a list of generators for the custom ABIs when constructing the
+`CoroSplit`_ pass, and 3. use `coro.begin.custom.abi` in place of `coro.begin`
+with an additional parameter for the index of the generator/ABI to be used for
+the coroutine.
+
+A custom ABI overriding the SwitchABI's materialization looks like:
+
+.. code-block:: c++
+
+  class CustomSwitchABI : public coro::SwitchABI {
+  public:
+    CustomSwitchABI(Function &F, coro::Shape &S)
+      : coro::SwitchABI(F, S, ExtraMaterializable) {}
+  };
+
+Giving a list of custom ABI generators while constructing the `CoroSplit`
+pass looks like:
+
+.. code-block:: c++
+
+  CoroSplitPass::BaseABITy GenCustomABI = [](Function &F, coro::Shape &S) {
+    return new CustomSwitchABI(F, S);
+  };
+
+  CGSCCPassManager CGPM;
+  CGPM.addPass(CoroSplitPass({GenCustomABI}));
+
+The LLVM IR for a coroutine using a Coroutine with a custom ABI looks like:
+
+.. code-block:: llvm
+
+  define ptr @f(i32 %n) presplitcoroutine_custom_abi {
+  entry:
+    %id = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr null)
+    %size = call i32 @llvm.coro.size.i32()
+    %alloc = call ptr @malloc(i32 %size)
+    %hdl = call noalias ptr @llvm.coro.begin.custom.abi(token %id, ptr %alloc, i32 0)
+    br label %loop
+  loop:
+    %n.val = phi i32 [ %n, %entry ], [ %inc, %loop ]
+    %inc = add nsw i32 %n.val, 1
+    call void @print(i32 %n.val)
+    %0 = call i8 @llvm.coro.suspend(token none, i1 false)
+    switch i8 %0, label %suspend [i8 0, label %loop
+                                  i8 1, label %cleanup]
+  cleanup:
+    %mem = call ptr @llvm.coro.free(token %id, ptr %hdl)
+    call void @free(ptr %mem)
+    br label %suspend
+  suspend:
+    %unused = call i1 @llvm.coro.end(ptr %hdl, i1 false, token none)
+    ret ptr %hdl
+  }
 
 Intrinsics
 ==========
@@ -1007,6 +1067,36 @@ with small positive and negative offsets).
 
 A frontend should emit exactly one `coro.begin` intrinsic per coroutine.
 
+.. _coro.begin.custom.abi:
+
+'llvm.coro.begin.custom.abi' Intrinsic
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+::
+
+  declare ptr @llvm.coro.begin.custom.abi(token <id>, ptr <mem>, i32)
+
+Overview:
+"""""""""
+
+The '``llvm.coro.begin.custom.abi``' intrinsic is used in place of the
+`coro.begin` intrinsic with an additional parameter to specify the custom ABI
+for the coroutine. The return is identical to that of the `coro.begin`
+intrinsic.
+
+Arguments:
+""""""""""
+
+The first and second arguments are identical to those of the `coro.begin`
+intrinsic.
+
+The third argument is an i32 index of the generator list given to the
+`CoroSplit` pass specifying the custom ABI generator lor this coroutine.
+
+Semantics:
+""""""""""
+
+The semantics are identical to those of the `coro.begin` intrinsic.
+
 .. _coro.free:
 
 'llvm.coro.free' Intrinsic

llvm/include/llvm/Analysis/TargetTransformInfoImpl.h

@@ -769,6 +769,7 @@ class TargetTransformInfoImplBase {
     case Intrinsic::experimental_gc_relocate:
     case Intrinsic::coro_alloc:
     case Intrinsic::coro_begin:
+    case Intrinsic::coro_begin_custom_abi:
     case Intrinsic::coro_free:
     case Intrinsic::coro_end:
     case Intrinsic::coro_frame:

llvm/include/llvm/IR/Intrinsics.td

@@ -1716,7 +1716,8 @@ 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_begin_custom_abi : Intrinsic<[llvm_ptr_ty], [llvm_token_ty, llvm_ptr_ty, llvm_i32_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/ABI.h → llvm/include/llvm/Transforms/Coroutines/ABI.h

@@ -12,12 +12,13 @@
 // ABI enum and ABI class are used by the Coroutine passes when lowering.
 //===----------------------------------------------------------------------===//
 
-#ifndef LIB_TRANSFORMS_COROUTINES_ABI_H
-#define LIB_TRANSFORMS_COROUTINES_ABI_H
+#ifndef LLVM_TRANSFORMS_COROUTINES_ABI_H
+#define LLVM_TRANSFORMS_COROUTINES_ABI_H
 
-#include "CoroShape.h"
-#include "SuspendCrossingInfo.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
+#include "llvm/Transforms/Coroutines/CoroShape.h"
+#include "llvm/Transforms/Coroutines/MaterializationUtils.h"
+#include "llvm/Transforms/Coroutines/SuspendCrossingInfo.h"
 
 namespace llvm {
 
@@ -28,9 +29,15 @@ namespace coro {
 // This interface/API is to provide an object oriented way to implement ABI
 // functionality. This is intended to replace use of the ABI enum to perform
 // ABI operations. The ABIs (e.g. Switch, Async, Retcon{Once}) are the common
-// ABIs.
-
-class LLVM_LIBRARY_VISIBILITY BaseABI {
+// ABIs. However, specific users often need to modify the behavior of these,
+// such as for C++20 or Swift. This can be accomplished by inheriting one of
+// the common ABIs and overriding one or more of the methods to create a custom
+// ABI. The custom ABI is specified with the coro.begin.custom.abi intrinsic
+// instead of the coro.begin intrinsic by providing an i32 in the last arg.
+// This is used to lookup a generator for the custom ABI from a set of
+// generators provided to the CoroSplitPass constructor.
+
+class BaseABI {
 public:
   BaseABI(Function &F, coro::Shape &S,
           std::function<bool(Instruction &)> IsMaterializable)
@@ -41,7 +48,7 @@ class LLVM_LIBRARY_VISIBILITY BaseABI {
   virtual void init() = 0;
 
   // Allocate the coroutine frame and do spill/reload as needed.
-  virtual void buildCoroutineFrame();
+  virtual void buildCoroutineFrame(bool OptimizeFrame);
 
   // Perform the function splitting according to the ABI.
   virtual void splitCoroutine(Function &F, coro::Shape &Shape,
@@ -56,7 +63,7 @@ class LLVM_LIBRARY_VISIBILITY BaseABI {
   std::function<bool(Instruction &I)> IsMaterializable;
 };
 
-class LLVM_LIBRARY_VISIBILITY SwitchABI : public BaseABI {
+class SwitchABI : public BaseABI {
 public:
   SwitchABI(Function &F, coro::Shape &S,
             std::function<bool(Instruction &)> IsMaterializable)
@@ -69,7 +76,7 @@ class LLVM_LIBRARY_VISIBILITY SwitchABI : public BaseABI {
                       TargetTransformInfo &TTI) override;
 };
 
-class LLVM_LIBRARY_VISIBILITY AsyncABI : public BaseABI {
+class AsyncABI : public BaseABI {
 public:
   AsyncABI(Function &F, coro::Shape &S,
            std::function<bool(Instruction &)> IsMaterializable)
@@ -82,7 +89,7 @@ class LLVM_LIBRARY_VISIBILITY AsyncABI : public BaseABI {
                       TargetTransformInfo &TTI) override;
 };
 
-class LLVM_LIBRARY_VISIBILITY AnyRetconABI : public BaseABI {
+class AnyRetconABI : public BaseABI {
 public:
   AnyRetconABI(Function &F, coro::Shape &S,
                std::function<bool(Instruction &)> IsMaterializable)