[CoroutineAccessors] Open code alloc fn.

Replace the runtime function with an open-coded version.

include/swift/Runtime/Coro.h

@@ -22,9 +22,6 @@
 #include <cstddef>
 
 namespace swift {
-SWIFT_RUNTIME_EXPORT
-SWIFT_CC(swift) void *swift_coro_alloc(CoroAllocator *allocator, size_t size);
-
 SWIFT_RUNTIME_EXPORT
 SWIFT_CC(swift) void swift_coro_dealloc(CoroAllocator *allocator, void *ptr);
 } // end namespace swift

include/swift/Runtime/RuntimeFunctions.def

@@ -3008,16 +3008,6 @@ FUNCTION(MemsetS, c, memset_s, C_CC, AlwaysAvailable,
          EFFECT(RuntimeEffect::NoEffect),
          UNKNOWN_MEMEFFECTS)
 
-// void *swift_coro_alloc(CoroAllocator *, size_t size);
-FUNCTION(CoroAlloc,
-         Swift, swift_coro_alloc, SwiftCC,
-         CoroutineAccessorsAvailability,
-         RETURNS(Int8PtrTy),
-         ARGS(CoroAllocatorPtrTy, SizeTy),
-         NO_ATTRS,
-         EFFECT(RuntimeEffect::Allocating, RuntimeEffect::Concurrency),
-         UNKNOWN_MEMEFFECTS)
-
 // void swift_coro_dealloc(CoroAllocator *, void *ptr);
 FUNCTION(CoroDealloc,
          Swift, swift_coro_dealloc, SwiftCC,

lib/IRGen/GenCall.cpp

@@ -5118,20 +5118,24 @@ void irgen::emitYieldManyCoroutineEntry(
                            allocFn, deallocFn, {});
 }
 
-static llvm::Constant *getCoroAllocWrapperFn(IRGenModule &IGM) {
+static llvm::Constant *getCoroAllocFn(IRGenModule &IGM) {
+  auto isSwiftCoroCCAvailable = IGM.SwiftCoroCC == llvm::CallingConv::SwiftCoro;
   return IGM.getOrCreateHelperFunction(
-      "__swift_coro_alloc_", IGM.Int8PtrTy,
-      {IGM.CoroAllocatorPtrTy, IGM.SizeTy},
-      [](IRGenFunction &IGF) {
+      "_swift_coro_alloc", IGM.Int8PtrTy, {IGM.CoroAllocatorPtrTy, IGM.SizeTy},
+      [isSwiftCoroCCAvailable](IRGenFunction &IGF) {
         auto parameters = IGF.collectParameters();
         auto *allocator = parameters.claimNext();
         auto *size = parameters.claimNext();
+        if (isSwiftCoroCCAvailable) {
+          // swiftcorocc is available, so if there's no allocator pointer,
+          // allocate storage on the stack and return a pointer to it without
+          // popping the stack.
           auto *nullAllocator = IGF.Builder.CreateCmp(
               llvm::CmpInst::Predicate::ICMP_EQ, allocator,
               llvm::ConstantPointerNull::get(
                   cast<llvm::PointerType>(allocator->getType())));
-        auto *poplessReturn = IGF.createBasicBlock("coro.return.popless");
-        auto *normalReturn = IGF.createBasicBlock("coro.return.normal");
+          auto *poplessReturn = IGF.createBasicBlock("popless");
+          auto *normalReturn = IGF.createBasicBlock("normal");
           IGF.Builder.CreateCondBr(nullAllocator, poplessReturn, normalReturn);
           IGF.Builder.emitBlock(poplessReturn);
           // Emit the dynamic alloca.
@@ -5140,17 +5144,33 @@ static llvm::Constant *getCoroAllocWrapperFn(IRGenModule &IGM) {
           alloca->setAlignment(llvm::Align(MaximumAlignment));
           auto *retPopless = IGF.Builder.CreateIntrinsic(
               IGF.IGM.VoidTy, llvm::Intrinsic::ret_popless, {});
-        retPopless->setTailCallKind(llvm::CallInst::TailCallKind::TCK_MustTail);
+          retPopless->setTailCallKind(
+              llvm::CallInst::TailCallKind::TCK_MustTail);
           IGF.Builder.CreateRet(alloca);
+          // Start emitting the "normal" block.
           IGF.Builder.emitBlock(normalReturn);
-        auto *call = IGF.Builder.CreateCall(
-            IGF.IGM.getCoroAllocFunctionPointer(), {allocator, size});
+        }
+        auto *calleePtr = IGF.Builder.CreateInBoundsGEP(
+            IGF.IGM.CoroAllocatorTy, allocator,
+            {llvm::ConstantInt::get(IGF.IGM.Int32Ty, 0),
+             llvm::ConstantInt::get(IGF.IGM.Int32Ty, 1)});
+        auto *callee = IGF.Builder.CreateLoad(
+            Address(calleePtr, IGF.IGM.CoroAllocateFnTy->getPointerTo(),
+                    IGF.IGM.getPointerAlignment()),
+            "allocate_fn");
+        auto fnPtr = FunctionPointer::createUnsigned(
+            FunctionPointer::Kind::Function, callee,
+            Signature(cast<llvm::FunctionType>(IGF.IGM.CoroAllocateFnTy), {},
+                      IGF.IGM.SwiftCC));
+        auto *call = IGF.Builder.CreateCall(fnPtr, {size});
+        call->setDoesNotThrow();
+        call->setCallingConv(IGF.IGM.SwiftCC);
         IGF.Builder.CreateRet(call);
       },
-      /*setIsNoInline=*/false,
+      /*setIsNoInline=*/true,
       /*forPrologue=*/false,
       /*isPerformanceConstraint=*/false,
-      /*optionalLinkageOverride=*/nullptr, llvm::CallingConv::SwiftCoro);
+      /*optionalLinkageOverride=*/nullptr, IGM.SwiftCoroCC);
 }
 
 static llvm::Constant *getCoroDeallocWrapperFn(IRGenModule &IGM) {
@@ -5190,9 +5210,7 @@ void irgen::emitYieldOnce2CoroutineEntry(IRGenFunction &IGF,
   IGF.setCoroutineAllocator(allocator);
   auto isSwiftCoroCCAvailable =
       IGF.IGM.SwiftCoroCC == llvm::CallingConv::SwiftCoro;
-  auto allocFn = IGF.IGM.getOpaquePtr(isSwiftCoroCCAvailable
-                                          ? getCoroAllocWrapperFn(IGF.IGM)
-                                          : IGF.IGM.getCoroAllocFn());
+  auto allocFn = IGF.IGM.getOpaquePtr(getCoroAllocFn(IGF.IGM));
   auto deallocFn = IGF.IGM.getOpaquePtr(isSwiftCoroCCAvailable
                                             ? getCoroDeallocWrapperFn(IGF.IGM)
                                             : IGF.IGM.getCoroDeallocFn());

stdlib/public/runtime/Coro.cpp

@@ -16,10 +16,6 @@
 
 using namespace swift;
 
-void *swift::swift_coro_alloc(CoroAllocator *allocator, size_t size) {
-  return allocator->allocate(size);
-}
-
 void swift::swift_coro_dealloc(CoroAllocator *allocator, void *ptr) {
   assert(allocator);
   // Calls to swift_coro_dealloc are emitted in resume funclets for every

test/IRGen/coroutine_accessors.swift

@@ -24,6 +24,21 @@
 // CHECK-SAME:      swift_task_dealloc
 // CHECK-SAME:  }
 
+// CHECK-LABEL: @_swift_coro_alloc(
+// CHECK-SAME:      ptr [[ALLOCATOR:%[^,]+]]
+// CHECK-SAME:      [[INT]] [[SIZE:%[^)]+]]
+// CHECK-SAME:  )
+// CHECK-SAME:  {
+// CHECK:       entry:
+// CHECK:         [[ALLOCATE_FN_PTR:%[^,]+]] = getelementptr inbounds %swift.coro_allocator
+// CHECK-SAME:        ptr [[ALLOCATOR]]
+// CHECK-SAME:        i32 0
+// CHECK-SAME:        i32 1
+// CHECK:         [[ALLOCATE_FN:%[^,]+]] = load ptr, ptr [[ALLOCATE_FN_PTR]]
+// CHECK:         [[ALLOCATION:%[^,]+]] = call swiftcc ptr [[ALLOCATE_FN]]([[INT]] [[SIZE]])
+// CHECK:         ret ptr [[ALLOCATION]]
+// CHECK:       }
+
 @frozen
 public struct S {
 public var o: any AnyObject

test/IRGen/coroutine_accessors_popless.swift

@@ -28,7 +28,7 @@
 // CHECK-SAME:       free
 // CHECK-SAME:  }
 
-// CHECK-LABEL: @__swift_coro_alloc_(
+// CHECK-LABEL: @_swift_coro_alloc(
 // CHECK-SAME:      ptr [[ALLOCATOR:%[^,]+]]
 // CHECK-SAME:      i64 [[SIZE:%[^)]+]]
 // CHECK-SAME:  )
@@ -36,18 +36,20 @@
 // CHECK:       entry:
 // CHECK:         [[USE_POPLESS:%[^,]+]] = icmp eq ptr [[ALLOCATOR]], null
 // CHECK:         br i1 [[USE_POPLESS]],
-// CHECK-SAME:        label %coro.return.popless
-// CHECK-SAME:        label %coro.return.normal
-// CHECK:       coro.return.popless:
+// CHECK-SAME:        label %popless
+// CHECK-SAME:        label %normal
+// CHECK:       popless:
 // CHECK:         [[STACK_ALLOCATION:%[^,]+]] = alloca i8, i64 [[SIZE]]
 // CHECK:         musttail call void @llvm.ret.popless()
 // CHECK:         ret ptr [[STACK_ALLOCATION]]
-// CHECK:       coro.return.normal:
-// CHECK:         [[OTHER_ALLOCATION:%[^,]+]] = call swiftcc ptr @swift_coro_alloc(
+// CHECK:       normal:
+// CHECK:         [[ALLOCATE_FN_PTR:%[^,]+]] = getelementptr inbounds %swift.coro_allocator
 // CHECK-SAME:        ptr [[ALLOCATOR]]
-// CHECK-SAME:        i64 [[SIZE]]
-// CHECK-SAME:    )
-// CHECK:         ret ptr [[OTHER_ALLOCATION]]
+// CHECK-SAME:        i32 0
+// CHECK-SAME:        i32 1
+// CHECK:         [[ALLOCATE_FN:%[^,]+]] = load ptr, ptr [[ALLOCATE_FN_PTR]]
+// CHECK:         [[ALLOCATION:%[^,]+]] = call swiftcc ptr [[ALLOCATE_FN]](i64 [[SIZE]])
+// CHECK:         ret ptr [[ALLOCATION]]
 // CHECK:       }
 
 // CHECK-LABEL: @__swift_coro_dealloc_(