Merge pull request #10179 from nate-chandler/coro-ret-popless

[Coro] Adopt ret.popless intrinsic.

Author: nate-chandler

llvm/include/llvm/IR/Intrinsics.td

@@ -838,6 +838,20 @@ def int_localrecover : DefaultAttrsIntrinsic<[llvm_ptr_ty],
                                  [llvm_ptr_ty, llvm_ptr_ty, llvm_i32_ty],
                                  [IntrNoMem, ImmArg<ArgIndex<2>>]>;
 
+// Marks the following ret instruction as a "popless" return, one that does not
+// not restore SP to its function-entry value (i.e., does not deallocate the
+// stack frame), allowing allocations made in the function to be accessible
+// by the caller.
+//
+// The function must be annotated with an appropriate target-specific calling
+// convention, so the caller can generate stack accesses accordingly, generally
+// by treating the call as a variably-sized alloca, so using FP-based addressing
+// for its own frame rather than relying on statically known SP offsets.
+//
+// Calls to this intrinsic need to be musttail, but don't follow the other ABI
+// requirements for musttail calls, since this is really annotating the ret.
+def int_ret_popless : DefaultAttrsIntrinsic<[], [], [IntrNoMem]>;
+
 // Given the frame pointer passed into an SEH filter function, returns a
 // pointer to the local variable area suitable for use with llvm.localrecover.
 def int_eh_recoverfp : DefaultAttrsIntrinsic<[llvm_ptr_ty],
@@ -1754,10 +1768,6 @@ def int_coro_await_suspend_handle : Intrinsic<[],
                                               [llvm_ptr_ty, llvm_ptr_ty, llvm_ptr_ty],
                                               [Throws]>;
 
-// FIXME: enforce musttail
-// XXX: attrs;  not throws, wb DefaultAttrsIntrinsic
-def int_coro_return : Intrinsic<[llvm_any_ty], [LLVMMatchType<0>], []>;
-
 // Coroutine Lowering Intrinsics. Used internally by coroutine passes.
 
 def int_coro_subfn_addr : DefaultAttrsIntrinsic<

llvm/lib/Analysis/InlineCost.cpp

@@ -2363,6 +2363,7 @@ bool CallAnalyzer::visitCallBase(CallBase &Call) {
       return false;
     case Intrinsic::icall_branch_funnel:
     case Intrinsic::localescape:
+    case Intrinsic::ret_popless:
       HasUninlineableIntrinsic = true;
       return false;
     case Intrinsic::vastart:

llvm/lib/IR/BasicBlock.cpp

@@ -306,6 +306,14 @@ const CallInst *BasicBlock::getTerminatingMustTailCall() const {
   if (!Prev)
     return nullptr;
 
+  // Some musttail intrinsic calls are special in being really simply ret
+  // annotations, and only need to be the last instruction before the ret.
+  // We don't need to look through the return value in those cases.
+  // FIXME: we should generalize getTerminatingDeoptimizeCall for this case.
+  if (auto *CI = dyn_cast<CallInst>(Prev))
+    if (CI->isMustTailCall() && CI->getIntrinsicID() == Intrinsic::ret_popless)
+      return CI;
+
   if (Value *RV = RI->getReturnValue()) {
     if (RV != Prev)
       return nullptr;

llvm/lib/IR/Verifier.cpp

@@ -3912,6 +3912,15 @@ void Verifier::verifyMustTailCall(CallInst &CI) {
             &CI);
     }
 #endif
+    Check(CI.getIntrinsicID() != Intrinsic::ret_popless,
+          "llvm.ret.popless call must be musttail", &CI);
+    return;
+  }
+
+  // Some musttail intrinsic calls are special, and don't have all the rules.
+  if (CI.getIntrinsicID() == Intrinsic::ret_popless) {
+    ReturnInst *Ret = dyn_cast_or_null<ReturnInst>(CI.getNextNode());
+    Check(Ret, "musttail intrinsic call must precede a ret", &CI);
     return;
   }
 

llvm/lib/Transforms/Coroutines/CoroCleanup.cpp

@@ -72,10 +72,6 @@ bool Lowerer::lower(Function &F) {
       case Intrinsic::coro_id_async:
         II->replaceAllUsesWith(ConstantTokenNone::get(Context));
         break;
-      case Intrinsic::coro_return:
-        // FIXME: Remove this case with backend support.
-        II->replaceAllUsesWith(II->getArgOperand(0));
-        break;
       case Intrinsic::coro_subfn_addr:
         lowerSubFn(Builder, cast<CoroSubFnInst>(II));
         break;

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

@@ -1987,9 +1987,12 @@ static void splitRetconCoroutine(Function &F, coro::Shape &Shape,
             F.getContext(), "coro.return.normal", &F, NewSuspendBB);
         Builder.CreateCondBr(NullAllocator, PoplessReturnBB, NormalReturnBB);
         IRBuilder<> PoplessBuilder(PoplessReturnBB);
-        auto *WrapRetV = PoplessBuilder.CreateIntrinsic(
-            RetV->getType(), Intrinsic::coro_return, {RetV});
-        PoplessBuilder.CreateRet(WrapRetV);
+        auto &Context = F.getContext();
+        auto *VoidTy = Type::getVoidTy(Context);
+        auto *RetPopless =
+            PoplessBuilder.CreateIntrinsic(VoidTy, Intrinsic::ret_popless, {});
+        RetPopless->setTailCallKind(CallInst::TailCallKind::TCK_MustTail);
+        PoplessBuilder.CreateRet(RetV);
         IRBuilder<> NormalBuilder(NormalReturnBB);
         NormalBuilder.CreateRet(RetV);
       } else {

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

@@ -1481,6 +1481,12 @@ static bool isSafeToHoistInstr(Instruction *I, unsigned Flags) {
     if (CB->getIntrinsicID() == Intrinsic::experimental_deoptimize)
       return false;
 
+  // Similarly for llvm.ret.popless (and likely generalizable to all musttail
+  // intrinsics).
+  if (auto *CB = dyn_cast<CallBase>(I))
+    if (CB->getIntrinsicID() == Intrinsic::ret_popless)
+      return false;
+
   // It's also unsafe/illegal to hoist an instruction above its instruction
   // operands
   BasicBlock *BB = I->getParent();

llvm/test/Transforms/Coroutines/coro-retcon-once-dynamic.ll

@@ -0,0 +1,103 @@
+; RUN: opt < %s -passes='module(coro-early),cgscc(coro-split),module(coro-cleanup)' -S | FileCheck %s
+
+target datalayout = "e-m:o-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "arm64-apple-macos99.99"
+
+; CHECK-LABEL: %func.Frame = type { ptr }
+
+; CHECK-LABEL: @func_cfp = constant <{ i32, i32 }> 
+; CHECK-SAME:  <{ 
+; CHECK-SAME:    i32 trunc
+; CHECK-SAME:    i32 16
+; CHECK-SAME:  }>
+@func_cfp = constant <{ i32, i32 }>
+  <{ i32 trunc ( ; offset to @func from @func_cfp
+       i64 sub (
+         i64 ptrtoint (ptr @func to i64),
+         i64 ptrtoint (ptr getelementptr inbounds (<{ i32, i32 }>, ptr @func_cfp, i32 0, i32 1) to i64)
+       )
+     to i32),
+     i32 64 ; frame size
+}>
+
+
+; CHECK-LABEL: @func(
+; CHECK-SAME:      ptr %buffer,
+; CHECK-SAME:      ptr %allocator
+; CHECK-SAME:      ptr %array
+; CHECK-SAME:  ) {
+; CHECK:           %array.spill.addr = getelementptr inbounds %func.Frame, ptr %buffer, i32 0, i32 0
+; CHECK:           store ptr %array, ptr %array.spill.addr
+; CHECK:           %load = load i32, ptr %array
+; CHECK:           %load.positive = icmp sgt i32 %load, 0
+; CHECK:           [[CONTINUATION:%.*]] = select i1 %load.positive
+; CHECK-SAME:          ptr @func.resume.0
+; CHECK-SAME:          ptr @func.resume.1
+; CHECK:           [[RETVAL_1:%.*]] = insertvalue { ptr, i32 } poison, ptr [[CONTINUATION:%.*]], 0
+; CHECK:           [[RETVAL_2:%.*]] = insertvalue { ptr, i32 } [[RETVAL_1:%.*]], i32 %load, 1
+; CHECK:           [[DONT_POP:%.*]] = icmp eq ptr %allocator, null
+; CHECK:           br i1 [[DONT_POP:%[^,]+]],
+; CHECK-SAME:          label %coro.return.popless
+; CHECK-SAME:          label %coro.return.normal
+; CHECK:         coro.return.popless:
+; CHECK:           musttail call void @llvm.ret.popless()
+; CHECK:           ret { ptr, i32 } [[RETVAL_2:%.*]]
+; CHECK:         coro.return.normal:
+; CHECK:           ret { ptr, i32 } [[RETVAL_2:%.*]]
+; CHECK:       }
+
+; CHECK-LABEL: @func.resume.0(
+; CHECK-SAME:      ptr [[BUFFER:[^,]+]]
+; CHECK-SAME:      ptr [[ALLOCATOR:%[^)]+]]
+; CHECK-SAME:  ) {
+; CHECK:           %array.reload.addr3 = getelementptr inbounds %func.Frame, ptr [[BUFFER:%.*]], i32 0, i32 0
+; CHECK:           %array.reload4 = load ptr, ptr %array.reload.addr3
+; CHECK:           store i32 0, ptr %array.reload4
+; CHECK:           ret void
+; CHECK:       }
+
+; CHECK-LABEL: @func.resume.1(
+; CHECK-SAME:      ptr [[BUFFER:[^,]+]]
+; CHECK-SAME:      ptr [[ALLOCATOR:%[^)]+]]
+; CHECK-SAME:  ) {
+; CHECK:           %array.reload.addr = getelementptr inbounds %func.Frame, ptr [[BUFFER:%.*]], i32 0, i32 0
+; CHECK:           %array.reload = load ptr, ptr %array.reload.addr
+; CHECK:           store i32 10, ptr %array.reload
+; CHECK:           ret void
+; CHECK:       }
+define swiftcorocc {ptr, i32} @func(ptr %buffer, ptr %allocator, ptr %array) {
+entry:
+  %id = call token @llvm.coro.id.retcon.once.dynamic(
+    i32 -1, 
+    i32 16, 
+    ptr @func_cfp, 
+    ptr %allocator, 
+    ptr %buffer, 
+    ptr @continuation_prototype, 
+    ptr @allocate, 
+    ptr @deallocate
+  )
+  %handle = call ptr @llvm.coro.begin(token %id, ptr null)
+  %load = load i32, ptr %array
+  %load.positive = icmp sgt i32 %load, 0
+  br i1 %load.positive, label %positive, label %negative
+
+positive:
+  call ptr (...) @llvm.coro.suspend.retcon.p0(i32 %load)
+  store i32 0, ptr %array, align 4
+  br label %cleanup
+
+negative:
+  call ptr (...) @llvm.coro.suspend.retcon.p0(i32 %load)
+  store i32 10, ptr %array, align 4
+  br label %cleanup
+
+cleanup:
+  call i1 @llvm.coro.end(ptr %handle, i1 0, token none)
+  unreachable
+}
+
+declare void @continuation_prototype(ptr, ptr)
+
+declare swiftcorocc noalias ptr @allocate(i32 %size)
+declare void @deallocate(ptr %ptr)

llvm/test/Verifier/ret_popless.ll

@@ -0,0 +1,14 @@
+; RUN: not opt -passes=verify < %s 2>&1 | FileCheck %s
+
+define void @test_ret_popless_not_musttail() {
+; CHECK: llvm.ret.popless call must be musttail
+  call void @llvm.ret.popless()
+  ret void
+}
+
+define i64 @test_ret_popless_not_returned(i64 %a) {
+; CHECK: musttail intrinsic call must precede a ret
+  musttail call void @llvm.ret.popless()
+  %res = bitcast i64 %a to i64
+  ret i64 %res
+}