Introduce llvm.coro.end.results intrinsics and use it to (optionally) return values

from `retcon.once` coroutines.

Author: asl

clang/lib/CodeGen/CGCoroutine.cpp

@@ -403,8 +403,11 @@ struct CallCoroEnd final : public EHScopeStack::Cleanup {
     llvm::Function *CoroEndFn = CGM.getIntrinsic(llvm::Intrinsic::coro_end);
     // See if we have a funclet bundle to associate coro.end with. (WinEH)
     auto Bundles = getBundlesForCoroEnd(CGF);
-    auto *CoroEnd = CGF.Builder.CreateCall(
-        CoroEndFn, {NullPtr, CGF.Builder.getTrue()}, Bundles);
+    auto *CoroEnd =
+      CGF.Builder.CreateCall(CoroEndFn,
+                             {NullPtr, CGF.Builder.getTrue(),
+                              llvm::ConstantTokenNone::get(CoroEndFn->getContext())},
+                             Bundles);
     if (Bundles.empty()) {
       // Otherwise, (landingpad model), create a conditional branch that leads
       // either to a cleanup block or a block with EH resume instruction.
@@ -755,7 +758,9 @@ void CodeGenFunction::EmitCoroutineBody(const CoroutineBodyStmt &S) {
   // Emit coro.end before getReturnStmt (and parameter destructors), since
   // resume and destroy parts of the coroutine should not include them.
   llvm::Function *CoroEnd = CGM.getIntrinsic(llvm::Intrinsic::coro_end);
-  Builder.CreateCall(CoroEnd, {NullPtr, Builder.getFalse()});
+  Builder.CreateCall(CoroEnd,
+                     {NullPtr, Builder.getFalse(),
+                      llvm::ConstantTokenNone::get(CoroEnd->getContext())});
 
   if (Stmt *Ret = S.getReturnStmt()) {
     // Since we already emitted the return value above, so we shouldn't
@@ -824,7 +829,11 @@ RValue CodeGenFunction::EmitCoroutineIntrinsic(const CallExpr *E,
   }
   for (const Expr *Arg : E->arguments())
     Args.push_back(EmitScalarExpr(Arg));
-
+  // @llvm.coro.end takes a token parameter. Add token 'none' as the last
+  // argument.
+  if (IID == llvm::Intrinsic::coro_end)
+    Args.push_back(llvm::ConstantTokenNone::get(getLLVMContext()));
+  
   llvm::Function *F = CGM.getIntrinsic(IID);
   llvm::CallInst *Call = Builder.CreateCall(F, Args);
 

clang/test/CodeGenCoroutines/coro-builtins.c

@@ -37,7 +37,7 @@ void f(int n) {
   // CHECK-NEXT: call ptr @llvm.coro.free(token %[[COROID]], ptr %[[FRAME]])
   __builtin_coro_free(__builtin_coro_frame());
 
-  // CHECK-NEXT: call i1 (ptr, i1, ...) @llvm.coro.end(ptr %[[FRAME]], i1 false)
+  // CHECK-NEXT: call i1  @llvm.coro.end(ptr %[[FRAME]], i1 false, token none)
   __builtin_coro_end(__builtin_coro_frame(), 0);
 
   // CHECK-NEXT: call i8 @llvm.coro.suspend(token none, i1 true)

clang/test/CodeGenCoroutines/coro-eh-cleanup.cpp

@@ -60,7 +60,7 @@ coro_t f() {
 
 // CHECK: [[COROENDBB]]:
 // CHECK-NEXT: %[[CLPAD:.+]] = cleanuppad within none
-// CHECK-NEXT: call i1 (ptr, i1, ...) @llvm.coro.end(ptr null, i1 true) [ "funclet"(token %[[CLPAD]]) ]
+// CHECK-NEXT: call i1 @llvm.coro.end(ptr null, i1 true, token none) [ "funclet"(token %[[CLPAD]]) ]
 // CHECK-NEXT: cleanupret from %[[CLPAD]] unwind label
 
 // CHECK-LPAD: @_Z1fv(
@@ -76,7 +76,7 @@ coro_t f() {
 // CHECK-LPAD:             to label %{{.+}} unwind label %[[UNWINDBB:.+]]
 
 // CHECK-LPAD: [[UNWINDBB]]:
-// CHECK-LPAD:   %[[I1RESUME:.+]] = call i1 (ptr, i1, ...) @llvm.coro.end(ptr null, i1 true)
+// CHECK-LPAD:   %[[I1RESUME:.+]] = call i1 @llvm.coro.end(ptr null, i1 true, token none)
 // CHECK-LPAD:   br i1  %[[I1RESUME]], label %[[EHRESUME:.+]], label
 // CHECK-LPAD: [[EHRESUME]]:
 // CHECK-LPAD-NEXT:  %[[exn:.+]] = load ptr, ptr %exn.slot, align 8

clang/test/CodeGenCoroutines/coro-lambda.cpp

@@ -55,4 +55,4 @@ void f() {
 //   CHECK: alloca %"struct.Task::promise_type"
 //   CHECK: call token @llvm.coro.id(
 //   CHECK: call i8 @llvm.coro.suspend(
-//   CHECK: call i1 (ptr, i1, ...) @llvm.coro.end(
+//   CHECK: call i1 @llvm.coro.end(

llvm/docs/Coroutines.rst

@@ -303,7 +303,7 @@ The LLVM IR for this coroutine looks like this:
     call void @free(ptr %mem)
     br label %suspend
   suspend:
-    %unused = call i1 (ptr, i1, ...) @llvm.coro.end(ptr %hdl, i1 false)
+    %unused = call i1 @llvm.coro.end(ptr %hdl, i1 false, token none)
     ret ptr %hdl
   }
 
@@ -630,7 +630,7 @@ store the current value produced by a coroutine.
     call void @free(ptr %mem)
     br label %suspend
   suspend:
-    %unused = call i1 (ptr, i1, ...) @llvm.coro.end(ptr %hdl, i1 false)
+    %unused = call i1 @llvm.coro.end(ptr %hdl, i1 false, token none)
     ret ptr %hdl
   }
 
@@ -1326,7 +1326,7 @@ A frontend should emit function attribute `presplitcoroutine` for the coroutine.
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 ::
 
-  declare i1 @llvm.coro.end(ptr <handle>, i1 <unwind>, ...)
+  declare i1 @llvm.coro.end(ptr <handle>, i1 <unwind>, token <result.token>)
 
 Overview:
 """""""""
@@ -1347,20 +1347,11 @@ The second argument should be `true` if this coro.end is in the block that is
 part of the unwind sequence leaving the coroutine body due to an exception and
 `false` otherwise.
 
-Other arguments can only be specified for unique-suspend returned-continuation
-coroutines where they will be normal returns of a coroutine continuation
-function. The number of arguments must match the return type of the continuation
-function:
+Non-trivial (non-none) token argument can only be specified for unique-suspend
+returned-continuation coroutines where it must be a token value produced by
+'``llvm.coro.end.results``' intrinsic.
 
-- if the return type of the continuation function is ``void`` there must be no
-  extra argumets
-
-- if the return type of the continuation function is a ``struct``, the arguments
-  will be element types of that ``struct`` in order;
-
-- otherwise, it is just the return value of the continuation function.
-
-No extra arguments are allowed for coro.end calls in unwind sections
+Only none token is allowed for coro.end calls in unwind sections
 
 Semantics:
 """"""""""
@@ -1393,7 +1384,7 @@ For landingpad based exception model, it is expected that frontend uses the
 .. code-block:: llvm
 
     ehcleanup:
-      %InResumePart = call i1 (ptr, i1, ...) @llvm.coro.end(ptr null, i1 true)
+      %InResumePart = call i1 @llvm.coro.end(ptr null, i1 true, token none)
       br i1 %InResumePart, label %eh.resume, label %cleanup.cont
 
     cleanup.cont:
@@ -1418,7 +1409,7 @@ referring to an enclosing cleanuppad as follows:
 
     ehcleanup:
       %tok = cleanuppad within none []
-      %unused = call i1 (ptr, i1, ...) @llvm.coro.end(ptr null, i1 true) [ "funclet"(token %tok) ]
+      %unused = call i1 @llvm.coro.end(ptr null, i1 true, token none) [ "funclet"(token %tok) ]
       cleanupret from %tok unwind label %RestOfTheCleanup
 
 The `CoroSplit` pass, if the funclet bundle is present, will insert
@@ -1443,6 +1434,53 @@ The following table summarizes the handling of `coro.end`_ intrinsic.
 |            | Landingpad  | mark coroutine as done | mark coroutine done             |
 +------------+-------------+------------------------+---------------------------------+
 
+.. _coro.end.results:
+
+'llvm.coro.end.results' Intrinsic
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+::
+
+  declare token @llvm.coro.end.results(...)
+
+Overview:
+"""""""""
+
+The '``llvm.coro.end.results``' intrinsic captures values to be returned from
+unique-suspend returned-continuation coroutines.
+
+Arguments:
+""""""""""
+
+The number of arguments must match the return type of the continuation function:
+
+- if the return type of the continuation function is ``void`` there must be no
+  arguments
+
+- if the return type of the continuation function is a ``struct``, the arguments
+  will be of element types of that ``struct`` in order;
+
+- otherwise, it is just the return value of the continuation function.
+
+.. code-block:: llvm
+
+  define {ptr, ptr} @g(ptr %buffer, ptr %ptr, i8 %val) presplitcoroutine {
+  entry:
+    %id = call token @llvm.coro.id.retcon.once(i32 8, i32 8, ptr %buffer,
+                                               ptr @prototype,
+                                               ptr @allocate, ptr @deallocate)
+    %hdl = call ptr @llvm.coro.begin(token %id, ptr null)
+
+  ...
+
+  cleanup:
+    %tok = call token (...) @llvm.coro.end.results(i8 %val)
+    call i1 @llvm.coro.end(ptr %hdl, i1 0, token %tok)
+    unreachable
+
+  ...
+
+  declare i8 @prototype(ptr, i1 zeroext)
+  
 
 'llvm.coro.end.async' Intrinsic
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

llvm/include/llvm/IR/Intrinsics.td

@@ -1643,7 +1643,8 @@ def int_coro_free : Intrinsic<[llvm_ptr_ty], [llvm_token_ty, llvm_ptr_ty],
                               [IntrReadMem, IntrArgMemOnly,
                                ReadOnly<ArgIndex<1>>,
                                NoCapture<ArgIndex<1>>]>;
-def int_coro_end : Intrinsic<[llvm_i1_ty], [llvm_ptr_ty, llvm_i1_ty, llvm_vararg_ty], []>;
+def int_coro_end : Intrinsic<[llvm_i1_ty], [llvm_ptr_ty, llvm_i1_ty, llvm_token_ty], []>;
+def int_coro_end_results : Intrinsic<[llvm_token_ty], [llvm_vararg_ty]>;
 def int_coro_end_async
     : Intrinsic<[llvm_i1_ty], [llvm_ptr_ty, llvm_i1_ty, llvm_vararg_ty], []>;
 

llvm/lib/IR/AutoUpgrade.cpp

@@ -951,7 +951,7 @@ static bool UpgradeIntrinsicFunction1(Function *F, Function *&NewFn) {
                                         F->arg_begin()->getType());
       return true;
     }
-    if (Name.equals("coro.end") && !F->isVarArg()) {
+    if (Name.equals("coro.end") && F->arg_size() == 2) {
       rename(F);
       NewFn = Intrinsic::getDeclaration(F->getParent(), Intrinsic::coro_end);
       return true;
@@ -4214,7 +4214,8 @@ void llvm::UpgradeIntrinsicCall(CallBase *CI, Function *NewFn) {
   }
 
   case Intrinsic::coro_end: {
-    SmallVector<Value *, 2> Args(CI->args());
+    SmallVector<Value *, 3> Args(CI->args());
+    Args.push_back(ConstantTokenNone::get(CI->getContext()));
     NewCall = Builder.CreateCall(NewFn, Args);
     break;
   }

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

@@ -3046,8 +3046,7 @@ void coro::buildCoroutineFrame(
   // Collect the spills for arguments and other not-materializable values.
   for (Argument &A : F.args())
     for (User *U : A.users())
-      if (Checker.isDefinitionAcrossSuspend(A, U) ||
-          isa<CoroEndInst>(U))
+      if (Checker.isDefinitionAcrossSuspend(A, U))
         FrameData.Spills[&A].push_back(cast<Instruction>(U));
 
   const DominatorTree DT(F);

llvm/lib/Transforms/Coroutines/CoroInstr.h

@@ -611,15 +611,53 @@ class LLVM_LIBRARY_VISIBILITY CoroAlignInst : public IntrinsicInst {
   }
 };
 
+/// This represents the llvm.end.results instruction.
+class LLVM_LIBRARY_VISIBILITY CoroEndResults : public IntrinsicInst {
+public:
+  op_iterator retval_begin() { return arg_begin(); }
+  const_op_iterator retval_begin() const { return arg_begin(); }
+
+  op_iterator retval_end() { return arg_end(); }
+  const_op_iterator retval_end() const { return arg_end(); }
+
+  iterator_range<op_iterator> return_values() {
+    return make_range(retval_begin(), retval_end());
+  }
+  iterator_range<const_op_iterator> return_values() const {
+    return make_range(retval_begin(), retval_end());
+  }
+
+  unsigned numReturns() const {
+    return std::distance(retval_begin(), retval_end());
+  }
+
+  // Methods to support type inquiry through isa, cast, and dyn_cast:
+  static bool classof(const IntrinsicInst *I) {
+    return I->getIntrinsicID() == Intrinsic::coro_end_results;
+  }
+  static bool classof(const Value *V) {
+    return isa<IntrinsicInst>(V) && classof(cast<IntrinsicInst>(V));
+  }
+};
+
 class LLVM_LIBRARY_VISIBILITY AnyCoroEndInst : public IntrinsicInst {
-  enum { FrameArg, UnwindArg };
+  enum { FrameArg, UnwindArg, TokenArg };
 
 public:
   bool isFallthrough() const { return !isUnwind(); }
   bool isUnwind() const {
     return cast<Constant>(getArgOperand(UnwindArg))->isOneValue();
   }
 
+  bool hasResults() const {
+    return !isa<ConstantTokenNone>(getArgOperand(TokenArg));
+  }
+
+  CoroEndResults *getResults() const {
+    assert(hasResults());
+    return cast<CoroEndResults>(getArgOperand(TokenArg));
+  }
+
   // Methods to support type inquiry through isa, cast, and dyn_cast:
   static bool classof(const IntrinsicInst *I) {
     auto ID = I->getIntrinsicID();
@@ -633,23 +671,6 @@ class LLVM_LIBRARY_VISIBILITY AnyCoroEndInst : public IntrinsicInst {
 /// This represents the llvm.coro.end instruction.
 class LLVM_LIBRARY_VISIBILITY CoroEndInst : public AnyCoroEndInst {
 public:
-  op_iterator retval_begin() { return std::next(arg_begin(), 2); }
-  const_op_iterator retval_begin() const { return std::next(arg_begin(), 2); }
-
-  op_iterator retval_end() { return arg_end(); }
-  const_op_iterator retval_end() const { return arg_end(); }
-
-  iterator_range<op_iterator> return_values() {
-    return make_range(retval_begin(), retval_end());
-  }
-  iterator_range<const_op_iterator> return_values() const {
-    return make_range(retval_begin(), retval_end());
-  }
-
-  unsigned numReturns() const {
-    return std::distance(retval_begin(), retval_end());
-  }
-
   // Methods to support type inquiry through isa, cast, and dyn_cast:
   static bool classof(const IntrinsicInst *I) {
     return I->getIntrinsicID() == Intrinsic::coro_end;

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

@@ -234,7 +234,7 @@ static void replaceFallthroughCoroEnd(AnyCoroEndInst *End,
   switch (Shape.ABI) {
   // The cloned functions in switch-lowering always return void.
   case coro::ABI::Switch:
-    assert(cast<CoroEndInst>(End)->numReturns() == 0 &&
+    assert(!cast<CoroEndInst>(End)->hasResults() &&
            "switch coroutine should not return any values");
     // coro.end doesn't immediately end the coroutine in the main function
     // in this lowering, because we need to deallocate the coroutine.
@@ -257,30 +257,40 @@ static void replaceFallthroughCoroEnd(AnyCoroEndInst *End,
     maybeFreeRetconStorage(Builder, Shape, FramePtr, CG);
     auto *CoroEnd = cast<CoroEndInst>(End);
     auto *RetTy = Shape.getResumeFunctionType()->getReturnType();
-    unsigned NumReturns = CoroEnd->numReturns();
+
+    if (!CoroEnd->hasResults()) {
+      assert(RetTy->isVoidTy());
+      Builder.CreateRetVoid();
+      break;
+    }
+
+    auto *CoroResults = CoroEnd->getResults();
+    unsigned NumReturns = CoroResults->numReturns();
 
     if (auto *RetStructTy = dyn_cast<StructType>(RetTy)) {
       assert(RetStructTy->getNumElements() == NumReturns &&
            "numbers of returns should match resume function singature");
       Value *ReturnValue = UndefValue::get(RetStructTy);
       unsigned Idx = 0;
-      for (Value *RetValEl : CoroEnd->return_values())
+      for (Value *RetValEl : CoroResults->return_values())
         ReturnValue = Builder.CreateInsertValue(ReturnValue, RetValEl, Idx++);
       Builder.CreateRet(ReturnValue);
     } else if (NumReturns == 0) {
       assert(RetTy->isVoidTy());
       Builder.CreateRetVoid();
     } else {
       assert(NumReturns == 1);
-      Builder.CreateRet(*CoroEnd->retval_begin());
+      Builder.CreateRet(*CoroResults->retval_begin());
     }
+    CoroResults->replaceAllUsesWith(ConstantTokenNone::get(CoroResults->getContext()));
+    CoroResults->eraseFromParent();
     break;
   }
 
   // In non-unique continuation lowering, we signal completion by returning
   // a null continuation.
   case coro::ABI::Retcon: {
-    assert(cast<CoroEndInst>(End)->numReturns() == 0 &&
+    assert(!cast<CoroEndInst>(End)->hasResults() &&
            "retcon coroutine should not return any values");
     maybeFreeRetconStorage(Builder, Shape, FramePtr, CG);
     auto RetTy = Shape.getResumeFunctionType()->getReturnType();