Extend retcon.once coroutines lowering to optionally produce a normal result (llvm#66333)

One of the main user of these kind of coroutines is swift. There yield-once (`retcon.once`) coroutines are used to temporary "expose" pointers to internal fields of various objects creating borrow scopes.

However, in some cases it might be useful also to allow these coroutines to produce a normal result, but there is no convenient way to represent this (as compared to switched-resume kind of coroutines where C++ `co_return`
is transformed to a member / callback call on promise object).

The extension is simple: we allow continuation function to have a non-void result and accept optional extra arguments via a special `llvm.coro.end.result` intrinsic that would essentially forward them as normal results.

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