Extend retcon.once coroutines lowering to optionally produce a normal result

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 @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 @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 @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

llvm/docs/Coroutines.rst

@@ -151,8 +151,8 @@ lowerings:
 - In yield-once returned-continuation lowering, the coroutine must
   suspend itself exactly once (or throw an exception).  The ramp
   function returns a continuation function pointer and yielded
-  values, but the continuation function simply returns `void`
-  when the coroutine has run to completion.
+  values, the continuation function may optionally return ordinary
+  results when the coroutine has run to completion.
 
 The coroutine frame is maintained in a fixed-size buffer that is
 passed to the `coro.id` intrinsic, which guarantees a certain size
@@ -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 @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 @llvm.coro.end(ptr %hdl, i1 false)
+    %unused = call i1 @llvm.coro.end(ptr %hdl, i1 false, token none)
     ret ptr %hdl
   }
 
@@ -1312,8 +1312,8 @@ Arguments:
 """"""""""
 
 As for ``llvm.core.id.retcon``, except that the return type of the
-continuation prototype must be `void` instead of matching the
-coroutine's return type.
+continuation prototype must represent the normal return type of the continuation
+(instead of matching the coroutine's return type).
 
 Semantics:
 """"""""""
@@ -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,6 +1347,12 @@ 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.
 
+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.
+
+Only none token is allowed for coro.end calls in unwind sections
+
 Semantics:
 """"""""""
 The purpose of this intrinsic is to allow frontends to mark the cleanup and
@@ -1378,7 +1384,7 @@ For landingpad based exception model, it is expected that frontend uses the
 .. code-block:: llvm
 
     ehcleanup:
-      %InResumePart = call 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:
@@ -1403,7 +1409,7 @@ referring to an enclosing cleanuppad as follows:
 
     ehcleanup:
       %tok = cleanuppad within none []
-      %unused = call 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
@@ -1428,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], []>;
+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,6 +951,12 @@ static bool UpgradeIntrinsicFunction1(Function *F, Function *&NewFn) {
                                         F->arg_begin()->getType());
       return true;
     }
+    if (Name.equals("coro.end") && F->arg_size() == 2) {
+      rename(F);
+      NewFn = Intrinsic::getDeclaration(F->getParent(), Intrinsic::coro_end);
+      return true;
+    }
+
     break;
   }
   case 'd':
@@ -4207,6 +4213,13 @@ void llvm::UpgradeIntrinsicCall(CallBase *CI, Function *NewFn) {
     break;
   }
 
+  case Intrinsic::coro_end: {
+    SmallVector<Value *, 3> Args(CI->args());
+    Args.push_back(ConstantTokenNone::get(CI->getContext()));
+    NewCall = Builder.CreateCall(NewFn, Args);
+    break;
+  }
+
   case Intrinsic::vector_extract: {
     StringRef Name = F->getName();
     Name = Name.substr(5); // Strip llvm

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();

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

@@ -234,6 +234,8 @@ 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)->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.
     if (!InResume)
@@ -251,14 +253,45 @@ static void replaceFallthroughCoroEnd(AnyCoroEndInst *End,
 
   // In unique continuation lowering, the continuations always return void.
   // But we may have implicitly allocated storage.
-  case coro::ABI::RetconOnce:
+  case coro::ABI::RetconOnce: {
     maybeFreeRetconStorage(Builder, Shape, FramePtr, CG);
-    Builder.CreateRetVoid();
+    auto *CoroEnd = cast<CoroEndInst>(End);
+    auto *RetTy = Shape.getResumeFunctionType()->getReturnType();
+
+    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 : 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(*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)->hasResults() &&
+           "retcon coroutine should not return any values");
     maybeFreeRetconStorage(Builder, Shape, FramePtr, CG);
     auto RetTy = Shape.getResumeFunctionType()->getReturnType();
     auto RetStructTy = dyn_cast<StructType>(RetTy);

llvm/test/Assembler/auto_upgrade_intrinsics.ll

@@ -47,6 +47,13 @@ entry:
   ret void
 }
 
+declare i1 @llvm.coro.end(ptr, i1)
+define void @test.coro.end(ptr %ptr) {
+; CHECK-LABEL: @test.coro.end(
+; CHECK: call i1 @llvm.coro.end(ptr %ptr, i1 false, token none)
+  call i1 @llvm.coro.end(ptr %ptr, i1 false)
+  ret void
+}
 
 @a = private global [60 x i8] zeroinitializer, align 1
 

llvm/test/Transforms/Coroutines/ArgAddr.ll

@@ -45,7 +45,7 @@ coro_Cleanup:
   br label %coro_Suspend
 
 coro_Suspend:
-  call i1 @llvm.coro.end(ptr null, i1 false)
+  call i1 @llvm.coro.end(ptr null, i1 false, token none)
   ret ptr %1
 }
 
@@ -69,7 +69,7 @@ declare i32 @llvm.coro.size.i32()
 declare ptr @llvm.coro.begin(token, ptr)
 declare i8 @llvm.coro.suspend(token, i1)
 declare ptr @llvm.coro.free(token, ptr)
-declare i1 @llvm.coro.end(ptr, i1)
+declare i1 @llvm.coro.end(ptr, i1, token)
 
 declare void @llvm.coro.resume(ptr)
 declare void @llvm.coro.destroy(ptr)

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

@@ -24,7 +24,7 @@ cleanup:
   br label %suspend
 
 suspend:
-  call i1 @llvm.coro.end(ptr %hdl, i1 0)
+  call i1 @llvm.coro.end(ptr %hdl, i1 0, token none)
   ret ptr %hdl
 }
 
@@ -44,7 +44,7 @@ declare void @llvm.coro.destroy(ptr)
 declare token @llvm.coro.id(i32, ptr, ptr, ptr)
 declare i1 @llvm.coro.alloc(token)
 declare ptr @llvm.coro.begin(token, ptr)
-declare i1 @llvm.coro.end(ptr, i1)
+declare i1 @llvm.coro.end(ptr, i1, token)
 
 declare void @capture_call(ptr)
 declare void @nocapture_call(ptr nocapture)