Merge pull request #80335 from nate-chandler/general-coro/20250326/2

[CoroutineAccessors] Move functions and add dealloc bit.

Author: nate-chandler

include/swift/ABI/Coro.h

@@ -39,6 +39,8 @@ class CoroAllocatorFlags : public FlagSet<uint32_t> {
     Kind           = 0,
     Kind_width     = 8,
 
+    ShouldDeallocateImmediately = 8,
+
     // 24 reserved bits
   };
   // clang-format on
@@ -49,6 +51,9 @@ class CoroAllocatorFlags : public FlagSet<uint32_t> {
 
   FLAGSET_DEFINE_FIELD_ACCESSORS(Kind, Kind_width, CoroAllocatorKind, getKind,
                                  setKind)
+  FLAGSET_DEFINE_FLAG_ACCESSORS(ShouldDeallocateImmediately,
+                                shouldDeallocateImmediately,
+                                setShouldDeallocateImmediately)
 };
 
 using CoroAllocation = void *;
@@ -62,16 +67,16 @@ struct CoroAllocator {
 
   /// Whether the allocator should deallocate memory on calls to
   /// swift_coro_dealloc.
-  constexpr bool shouldDeallocateImmediately() {
-    // Only the "mallocator" should immediately deallocate in
+  bool shouldDeallocateImmediately() {
+    // Currently, only the "mallocator" should immediately deallocate in
     // swift_coro_dealloc.  It must because the ABI does not provide a means for
     // the callee to return its allocations to the caller.
     //
     // Async allocations can be deferred until the first non-coroutine caller
     // from where swift_task_dealloc_through can be called and passed the
     // caller-allocated fixed-per-callee-sized-frame.
     // Stack allocations just pop the stack.
-    return flags.getKind() == CoroAllocatorKind::Malloc;
+    return flags.shouldDeallocateImmediately();
   }
 };
 

include/swift/Runtime/Concurrency.h

@@ -18,7 +18,6 @@
 #define SWIFT_RUNTIME_CONCURRENCY_H
 
 #include "swift/ABI/AsyncLet.h"
-#include "swift/ABI/Coro.h"
 #include "swift/ABI/Task.h"
 #include "swift/ABI/TaskGroup.h"
 
@@ -128,13 +127,6 @@ void swift_task_dealloc(void *ptr);
 SWIFT_EXPORT_FROM(swift_Concurrency)
 SWIFT_CC(swift) void swift_task_dealloc_through(void *ptr);
 
-// TODO: CoroutineAccessors: Move these declarations back to swiftCore {{
-SWIFT_EXPORT_FROM(swift_Concurrency)
-SWIFT_CC(swift) void *swift_coro_alloc(CoroAllocator *allocator, size_t size);
-SWIFT_EXPORT_FROM(swift_Concurrency)
-SWIFT_CC(swift) void swift_coro_dealloc(CoroAllocator *allocator, void *ptr);
-// }} TODO: CoroutineAccessors
-
 /// Cancel a task and all of its child tasks.
 ///
 /// This can be called from any thread.

include/swift/Runtime/Coro.h

@@ -21,6 +21,12 @@
 #include "swift/Runtime/Config.h"
 #include <cstddef>
 
-namespace swift {} // end namespace swift
+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
 
 #endif

include/swift/Runtime/RuntimeFunctions.def

@@ -3010,7 +3010,7 @@ FUNCTION(MemsetS, c, memset_s, C_CC, AlwaysAvailable,
 
 // void *swift_coro_alloc(CoroAllocator *, size_t size);
 FUNCTION(CoroAlloc,
-         _Concurrency, swift_coro_alloc, SwiftCC,
+         Swift, swift_coro_alloc, SwiftCC,
          CoroutineAccessorsAvailability,
          RETURNS(Int8PtrTy),
          ARGS(CoroAllocatorPtrTy, SizeTy),
@@ -3020,7 +3020,7 @@ FUNCTION(CoroAlloc,
 
 // void swift_coro_dealloc(CoroAllocator *, void *ptr);
 FUNCTION(CoroDealloc,
-         _Concurrency, swift_coro_dealloc, SwiftCC,
+         Swift, swift_coro_dealloc, SwiftCC,
          CoroutineAccessorsAvailability,
          RETURNS(VoidTy),
          ARGS(CoroAllocatorPtrTy, Int8PtrTy),

lib/IRGen/GenCall.cpp

@@ -5153,6 +5153,35 @@ static llvm::Constant *getCoroAllocWrapperFn(IRGenModule &IGM) {
       /*optionalLinkageOverride=*/nullptr, llvm::CallingConv::SwiftCoro);
 }
 
+static llvm::Constant *getCoroDeallocWrapperFn(IRGenModule &IGM) {
+  return IGM.getOrCreateHelperFunction(
+      "__swift_coro_dealloc_", IGM.VoidTy,
+      {IGM.CoroAllocatorPtrTy, IGM.Int8PtrTy},
+      [](IRGenFunction &IGF) {
+        auto parameters = IGF.collectParameters();
+        auto *allocator = parameters.claimNext();
+        auto *ptr = parameters.claimNext();
+        auto *nullAllocator = IGF.Builder.CreateCmp(
+            llvm::CmpInst::Predicate::ICMP_EQ, allocator,
+            llvm::ConstantPointerNull::get(
+                cast<llvm::PointerType>(allocator->getType())));
+        auto *bailBlock = IGF.createBasicBlock("bail");
+        auto *forwardBlock = IGF.createBasicBlock("forward");
+        IGF.Builder.CreateCondBr(nullAllocator, bailBlock, forwardBlock);
+        IGF.Builder.emitBlock(bailBlock);
+        // Emit the dynamic alloca.
+        IGF.Builder.CreateRetVoid();
+        IGF.Builder.emitBlock(forwardBlock);
+        IGF.Builder.CreateCall(
+            IGF.IGM.getCoroDeallocFunctionPointer(), {allocator, ptr});
+        IGF.Builder.CreateRetVoid();
+      },
+      /*setIsNoInline=*/false,
+      /*forPrologue=*/false,
+      /*isPerformanceConstraint=*/false,
+      /*optionalLinkageOverride=*/nullptr, llvm::CallingConv::SwiftCoro);
+}
+
 void irgen::emitYieldOnce2CoroutineEntry(IRGenFunction &IGF,
                                          CanSILFunctionType fnType,
                                          llvm::Value *buffer,
@@ -5164,7 +5193,9 @@ void irgen::emitYieldOnce2CoroutineEntry(IRGenFunction &IGF,
   auto allocFn = IGF.IGM.getOpaquePtr(isSwiftCoroCCAvailable
                                           ? getCoroAllocWrapperFn(IGF.IGM)
                                           : IGF.IGM.getCoroAllocFn());
-  auto deallocFn = IGF.IGM.getOpaquePtr(IGF.IGM.getCoroDeallocFn());
+  auto deallocFn = IGF.IGM.getOpaquePtr(isSwiftCoroCCAvailable
+                                            ? getCoroDeallocWrapperFn(IGF.IGM)
+                                            : IGF.IGM.getCoroDeallocFn());
   emitRetconCoroutineEntry(
       IGF, fnType, buffer, llvm::Intrinsic::coro_id_retcon_once_dynamic,
       Size(-1) /*dynamic-to-IRGen size*/, IGF.IGM.getCoroStaticFrameAlignment(),
@@ -5199,16 +5230,17 @@ Address irgen::emitAllocYieldManyCoroutineBuffer(IRGenFunction &IGF) {
                                  getYieldManyCoroutineBufferAlignment(IGF.IGM));
 }
 
-static llvm::Constant *getAddrOfGlobalCoroAllocator(IRGenModule &IGM,
-                                                    CoroAllocatorKind kind,
-                                                    llvm::Constant *allocFn,
-                                                    llvm::Constant *deallocFn) {
+static llvm::Constant *getAddrOfGlobalCoroAllocator(
+    IRGenModule &IGM, CoroAllocatorKind kind, bool shouldDeallocateImmediately,
+    llvm::Constant *allocFn, llvm::Constant *deallocFn) {
   auto entity = LinkEntity::forCoroAllocator(kind);
   auto taskAllocator = IGM.getOrCreateLazyGlobalVariable(
       entity,
       [&](ConstantInitBuilder &builder) -> ConstantInitFuture {
         auto allocator = builder.beginStruct(IGM.CoroAllocatorTy);
-        allocator.addInt32(CoroAllocatorFlags(kind).getOpaqueValue());
+        auto flags = CoroAllocatorFlags(kind);
+        flags.setShouldDeallocateImmediately(shouldDeallocateImmediately);
+        allocator.addInt32(flags.getOpaqueValue());
         allocator.add(allocFn);
         allocator.add(deallocFn);
         return allocator.finishAndCreateFuture();
@@ -5218,10 +5250,12 @@ static llvm::Constant *getAddrOfGlobalCoroAllocator(IRGenModule &IGM,
 }
 llvm::Constant *IRGenModule::getAddrOfGlobalCoroMallocAllocator() {
   return getAddrOfGlobalCoroAllocator(*this, CoroAllocatorKind::Malloc,
+                                      /*shouldDeallocateImmediately=*/true,
                                       getMallocFn(), getFreeFn());
 }
 llvm::Constant *IRGenModule::getAddrOfGlobalCoroAsyncTaskAllocator() {
   return getAddrOfGlobalCoroAllocator(*this, CoroAllocatorKind::Async,
+                                      /*shouldDeallocateImmediately=*/false,
                                       getTaskAllocFn(), getTaskDeallocFn());
 }
 llvm::Value *

stdlib/public/Concurrency/TaskAlloc.cpp

@@ -72,19 +72,3 @@ void swift::_swift_task_dealloc_specific(AsyncTask *task, void *ptr) {
 void swift::swift_task_dealloc_through(void *ptr) {
   allocator(swift_task_getCurrent()).deallocThrough(ptr);
 }
-
-void *swift::swift_coro_alloc(CoroAllocator *allocator, size_t size) {
-  return allocator->allocate(size);
-}
-
-void swift::swift_coro_dealloc(CoroAllocator *allocator, void *ptr) {
-  if (!allocator)
-    return;
-  // Calls to swift_coro_dealloc are emitted in resume funclets for every
-  // live-across dynamic allocation.  Whether such calls immediately deallocate
-  // memory depends on the allocator.
-  if (!allocator->shouldDeallocateImmediately()) {
-    return;
-  }
-  allocator->deallocate(ptr);
-}

stdlib/public/runtime/Coro.cpp

@@ -15,3 +15,18 @@
 #include "swift/Basic/FlagSet.h"
 
 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
+  // live-across dynamic allocation.  Whether such calls immediately deallocate
+  // memory depends on the allocator.
+  if (!allocator->shouldDeallocateImmediately()) {
+    return;
+  }
+  allocator->deallocate(ptr);
+}

test/IRGen/coroutine_accessors.swift

@@ -14,7 +14,7 @@
 // CHECK-SAME:  }>
 
 // CHECK-LABEL: _swift_coro_malloc_allocator = linkonce_odr hidden constant %swift.coro_allocator {
-// CHECK-SAME:       i32 2,
+// CHECK-SAME:       i32 258,
 // CHECK-SAME:       malloc,
 // CHECK-SAME:       free
 // CHECK-SAME:  }

test/IRGen/coroutine_accessors_popless.swift

@@ -23,7 +23,7 @@
 // CHECK-SAME:      swift_task_dealloc
 // CHECK-SAME:  }
 // CHECK-LABEL: _swift_coro_malloc_allocator = linkonce_odr hidden constant %swift.coro_allocator {
-// CHECK-SAME:       i32 2,
+// CHECK-SAME:       i32 258,
 // CHECK-SAME:       malloc,
 // CHECK-SAME:       free
 // CHECK-SAME:  }
@@ -50,6 +50,26 @@
 // CHECK:         ret ptr [[OTHER_ALLOCATION]]
 // CHECK:       }
 
+// CHECK-LABEL: @__swift_coro_dealloc_(
+// CHECK-SAME:      ptr [[ALLOCATOR:%[^,]+]]
+// CHECK-SAME:      ptr [[ADDRESS:%[^)]+]]
+// CHECK-SAME:  )
+// CHECK-SAME:  {
+// CHECK:       entry:
+// CHECK:         [[BAIL:%[^,]+]] = icmp eq ptr [[ALLOCATOR]], null
+// CHECK:         br i1 [[USE_POPLESS]],
+// CHECK-SAME:        label %bail
+// CHECK-SAME:        label %forward
+// CHECK:       bail:
+// CHECK:         ret void
+// CHECK:       forward:
+// CHECK:         call swiftcc void @swift_coro_dealloc(
+// CHECK-SAME:        ptr [[ALLOCATOR]]
+// CHECK-SAME:        ptr [[ADDRESS]]
+// CHECK-SAME:    )
+// CHECK:         ret void
+// CHECK:       }
+
 public var _i: Int = 0
 
 public var i: Int {

test/abi/macOS/arm64/concurrency.swift

@@ -408,10 +408,7 @@ Added: _$sScf25isIsolatingCurrentContextSbyFTj
 Added: _$sScf25isIsolatingCurrentContextSbyFTq
 Added: _$sScfsE25isIsolatingCurrentContextSbyF
 
-// add callee-allocated coro entrypoints
-// TODO: CoroutineAccessors: several of these symbols should be in swiftCore
-Added: _swift_coro_alloc
-Added: _swift_coro_dealloc
+// CoroutineAccessors
 Added: _swift_task_dealloc_through
 
 // SwiftSettings