Merge pull request #80290 from nate-chandler/general-coro/20250325/1

nate-chandler · web-flow · commit e8d871ced5df · 2025-03-26T21:47:09.000-07:00
[CoroutineAccessors] Directly reference allocators.
diff --git a/include/swift/IRGen/Linking.h b/include/swift/IRGen/Linking.h
@@ -13,6 +13,7 @@
 #ifndef SWIFT_IRGEN_LINKING_H
 #define SWIFT_IRGEN_LINKING_H
 
+#include "swift/ABI/Coro.h"
 #include "swift/AST/Decl.h"
 #include "swift/AST/Module.h"
 #include "swift/AST/ProtocolAssociations.h"
@@ -24,8 +25,8 @@
 #include "swift/SIL/SILGlobalVariable.h"
 #include "swift/SIL/SILModule.h"
 #include "llvm/ADT/DenseMapInfo.h"
-#include "llvm/IR/GlobalValue.h"
 #include "llvm/IR/GlobalObject.h"
+#include "llvm/IR/GlobalValue.h"
 #include "llvm/IR/Module.h"
 
 namespace llvm {
@@ -165,6 +166,9 @@ class LinkEntity {
     ExtendedExistentialIsUniqueMask = 0x100,
     ExtendedExistentialIsSharedShift = 9,
     ExtendedExistentialIsSharedMask = 0x200,
+
+    // Used by CoroAllocator.  2 bits.
+    CoroAllocatorKindShift = 8, CoroAllocatorKindMask = 0x300,
   };
 #define LINKENTITY_SET_FIELD(field, value) (value << field##Shift)
 #define LINKENTITY_GET_FIELD(value, field) ((value & field##Mask) >> field##Shift)
@@ -590,10 +594,13 @@ class LinkEntity {
     /// The pointer is a const char* of the name.
     KnownCoroFunctionPointer,
 
-    /// An coro function pointer for a distributed accessor (method or
+    /// A coro function pointer for a distributed accessor (method or
     /// property).
     /// The pointer is a SILFunction*.
     DistributedAccessorCoroFunctionPointer,
+
+    /// An allocator to be passed to swift_coro_alloc and swift_coro_dealloc.
+    CoroAllocator,
   };
   friend struct llvm::DenseMapInfo<LinkEntity>;
 
@@ -1498,6 +1505,15 @@ class LinkEntity {
     return entity;
   }
 
+  static LinkEntity forCoroAllocator(CoroAllocatorKind kind) {
+    LinkEntity entity;
+    entity.Pointer = nullptr;
+    entity.SecondaryPointer = nullptr;
+    entity.Data = LINKENTITY_SET_FIELD(Kind, unsigned(Kind::CoroAllocator)) |
+                  LINKENTITY_SET_FIELD(CoroAllocatorKind, unsigned(kind));
+    return entity;
+  }
+
   static LinkEntity forCoroFunctionPointer(LinkEntity other) {
     LinkEntity entity;
     entity.Pointer = other.Pointer;
@@ -1709,6 +1725,11 @@ class LinkEntity {
         SecondaryPointer);
   }
 
+  CoroAllocatorKind getCoroAllocatorKind() const {
+    assert(getKind() == Kind::CoroAllocator);
+    return CoroAllocatorKind(LINKENTITY_GET_FIELD(Data, CoroAllocatorKind));
+  }
+
   CanGenericSignature getExtendedExistentialTypeShapeGenSig() const {
     assert(getKind() == Kind::ExtendedExistentialTypeShape);
     return CanGenericSignature(
diff --git a/include/swift/Runtime/Concurrency.h b/include/swift/Runtime/Concurrency.h
@@ -128,25 +128,11 @@ void swift_task_dealloc(void *ptr);
 SWIFT_EXPORT_FROM(swift_Concurrency)
 SWIFT_CC(swift) void swift_task_dealloc_through(void *ptr);
 
-// TODO: CoroutineAccessors: Eliminate this entry point and replace its uses
-//                           with direct references the globals.
-SWIFT_EXPORT_FROM(swift_Concurrency)
-SWIFT_CC(swift)
-CoroAllocator *swift_coro_getGlobalAllocator(CoroAllocatorFlags flags);
-
-// TODO: CoroutineAccessors: Mark the underlying struct const.
-SWIFT_EXPORT_FROM(swift_Concurrency)
-CoroAllocator *const _swift_coro_task_allocator;
-
 // 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: Mark the underlying struct const.
-SWIFT_EXPORT_FROM(swift_Concurrency)
-CoroAllocator *const _swift_coro_malloc_allocator;
 // }} TODO: CoroutineAccessors
 
 /// Cancel a task and all of its child tasks.
diff --git a/include/swift/Runtime/RuntimeFunctions.def b/include/swift/Runtime/RuntimeFunctions.def
@@ -3028,16 +3028,6 @@ FUNCTION(CoroDealloc,
          EFFECT(RuntimeEffect::Deallocating, RuntimeEffect::Concurrency),
          UNKNOWN_MEMEFFECTS)
 
-// const CoroAllocator *swift_coro_getGlobalAllocator(CoroAllocatorFlags flags);
-FUNCTION(CoroGetGlobalAllocator, 
-         _Concurrency, swift_coro_getGlobalAllocator, SwiftCC, 
-         CoroutineAccessorsAvailability, 
-         RETURNS(CoroAllocatorPtrTy), 
-         ARGS(CoroAllocatorFlagsTy), 
-         ATTRS(NoUnwind, WillReturn), 
-         EFFECT(RuntimeEffect::Allocating),
-         UNKNOWN_MEMEFFECTS)
-
 #undef RETURNS
 #undef ARGS
 #undef ATTRS
diff --git a/lib/IRGen/GenCall.cpp b/lib/IRGen/GenCall.cpp
@@ -44,6 +44,7 @@
 #include <optional>
 
 #include "CallEmission.h"
+#include "ConstantBuilder.h"
 #include "EntryPointArgumentEmission.h"
 #include "Explosion.h"
 #include "GenCall.h"
@@ -57,8 +58,8 @@
 #include "GenType.h"
 #include "IRGenDebugInfo.h"
 #include "IRGenFunction.h"
-#include "IRGenModule.h"
 #include "IRGenMangler.h"
+#include "IRGenModule.h"
 #include "LoadableTypeInfo.h"
 #include "NativeConventionSchema.h"
 #include "Signature.h"
@@ -5197,18 +5198,47 @@ Address irgen::emitAllocYieldManyCoroutineBuffer(IRGenFunction &IGF) {
   return createOpaqueBufferAlloca(IGF, getYieldManyCoroutineBufferSize(IGF.IGM),
                                  getYieldManyCoroutineBufferAlignment(IGF.IGM));
 }
+
+static llvm::Constant *getAddrOfGlobalCoroAllocator(IRGenModule &IGM,
+                                                    CoroAllocatorKind kind,
+                                                    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());
+        allocator.add(allocFn);
+        allocator.add(deallocFn);
+        return allocator.finishAndCreateFuture();
+      },
+      [&](llvm::GlobalVariable *var) { var->setConstant(true); });
+  return taskAllocator;
+}
+llvm::Constant *IRGenModule::getAddrOfGlobalCoroMallocAllocator() {
+  return getAddrOfGlobalCoroAllocator(*this, CoroAllocatorKind::Malloc,
+                                      getMallocFn(), getFreeFn());
+}
+llvm::Constant *IRGenModule::getAddrOfGlobalCoroAsyncTaskAllocator() {
+  return getAddrOfGlobalCoroAllocator(*this, CoroAllocatorKind::Async,
+                                      getTaskAllocFn(), getTaskDeallocFn());
+}
 llvm::Value *
 irgen::emitYieldOnce2CoroutineAllocator(IRGenFunction &IGF,
                                         std::optional<CoroAllocatorKind> kind) {
   if (!kind) {
     return IGF.getCoroutineAllocator();
   }
-  CoroAllocatorFlags flags(*kind);
-  auto *flagsValue = llvm::ConstantInt::get(IGF.IGM.CoroAllocatorFlagsTy,
-                                            flags.getOpaqueValue());
-
-  return IGF.Builder.CreateCall(
-      IGF.IGM.getCoroGetGlobalAllocatorFunctionPointer(), {flagsValue});
+  switch (*kind) {
+  case CoroAllocatorKind::Stack:
+    return llvm::ConstantPointerNull::get(IGF.IGM.CoroAllocatorPtrTy);
+  case CoroAllocatorKind::Async:
+    return IGF.IGM.getAddrOfGlobalCoroAsyncTaskAllocator();
+  case CoroAllocatorKind::Malloc:
+    return IGF.IGM.getAddrOfGlobalCoroMallocAllocator();
+  }
+  llvm_unreachable("unhandled case");
 }
 StackAddress irgen::emitAllocYieldOnce2CoroutineFrame(IRGenFunction &IGF,
                                                       llvm::Value *size) {
diff --git a/lib/IRGen/IRGenModule.cpp b/lib/IRGen/IRGenModule.cpp
@@ -832,8 +832,8 @@ IRGenModule::IRGenModule(IRGenerator &irgen,
   CoroAllocatorTy = createStructType(*this, "swift.coro_allocator",
                                      {
                                          Int32Ty, // CoroAllocator.Flags
-                                         CoroAllocateFnTy,
-                                         CoroDeallocateFnTy,
+                                         CoroAllocateFnTy->getPointerTo(),
+                                         CoroDeallocateFnTy->getPointerTo(),
                                      });
   CoroAllocatorPtrTy = CoroAllocatorTy->getPointerTo();
 }
diff --git a/lib/IRGen/IRGenModule.h b/lib/IRGen/IRGenModule.h
@@ -1671,6 +1671,9 @@ private:                                                                       \
                                             ConstantInit init);
   SILFunction *getSILFunctionForCoroFunctionPointer(llvm::Constant *cfp);
 
+  llvm::Constant *getAddrOfGlobalCoroMallocAllocator();
+  llvm::Constant *getAddrOfGlobalCoroAsyncTaskAllocator();
+
   llvm::Function *getAddrOfDispatchThunk(SILDeclRef declRef,
                                          ForDefinition_t forDefinition);
   void emitDispatchThunk(SILDeclRef declRef);
diff --git a/lib/IRGen/Linking.cpp b/lib/IRGen/Linking.cpp
@@ -577,6 +577,16 @@ std::string LinkEntity::mangleAsString(ASTContext &Ctx) const {
     Result.append("Twc");
     return Result;
   }
+  case Kind::CoroAllocator: {
+    switch (getCoroAllocatorKind()) {
+    case CoroAllocatorKind::Stack:
+      llvm::report_fatal_error("attempting to mangle the coro stack allocator");
+    case CoroAllocatorKind::Async:
+      return "_swift_coro_async_allocator";
+    case CoroAllocatorKind::Malloc:
+      return "_swift_coro_malloc_allocator";
+    }
+  }
   }
   llvm_unreachable("bad entity kind!");
 }
@@ -948,6 +958,8 @@ SILLinkage LinkEntity::getLinkage(ForDefinition_t forDefinition) const {
     return getUnderlyingEntityForCoroFunctionPointer().getLinkage(
         forDefinition);
     return getSILLinkage(getDeclLinkage(getDecl()), forDefinition);
+  case Kind::CoroAllocator:
+    return SILLinkage::Shared;
   }
   llvm_unreachable("bad link entity kind");
 }
@@ -1051,6 +1063,7 @@ bool LinkEntity::isContextDescriptor() const {
   case Kind::CoroFunctionPointerAST:
   case Kind::DistributedThunkCoroFunctionPointer:
   case Kind::KnownCoroFunctionPointer:
+  case Kind::CoroAllocator:
     return false;
   }
   llvm_unreachable("invalid descriptor");
@@ -1188,6 +1201,8 @@ llvm::Type *LinkEntity::getDefaultDeclarationType(IRGenModule &IGM) const {
   case Kind::DistributedThunkCoroFunctionPointer:
   case Kind::KnownCoroFunctionPointer:
     return IGM.CoroFunctionPointerPtrTy;
+  case Kind::CoroAllocator:
+    return IGM.CoroAllocatorTy;
   default:
     llvm_unreachable("declaration LLVM type not specified");
   }
@@ -1220,6 +1235,7 @@ Alignment LinkEntity::getAlignment(IRGenModule &IGM) const {
   case Kind::OpaqueTypeDescriptorRecord:
   case Kind::AccessibleFunctionRecord:
   case Kind::ExtendedExistentialTypeShape:
+  case Kind::CoroAllocator:
     return Alignment(4);
   case Kind::AsyncFunctionPointer:
   case Kind::DispatchThunkAsyncFunctionPointer:
@@ -1359,6 +1375,7 @@ bool LinkEntity::isText() const {
   case Kind::DynamicallyReplaceableFunctionVariable:
   case Kind::CanonicalSpecializedGenericTypeMetadataAccessFunction:
   case Kind::ExtendedExistentialTypeShape:
+  case Kind::CoroAllocator:
     return true;
   case Kind::ObjCClass:
   case Kind::ObjCClassRef:
@@ -1520,6 +1537,7 @@ bool LinkEntity::isWeakImported(ModuleDecl *module) const {
   case Kind::DifferentiabilityWitness:
   case Kind::AccessibleFunctionRecord:
   case Kind::ExtendedExistentialTypeShape:
+  case Kind::CoroAllocator:
     return false;
 
   case Kind::AsyncFunctionPointer:
@@ -1656,6 +1674,7 @@ DeclContext *LinkEntity::getDeclContextForEmission() const {
   case Kind::TypeMetadataDemanglingCacheVariable:
   case Kind::NoncanonicalSpecializedGenericTypeMetadata:
   case Kind::NoncanonicalSpecializedGenericTypeMetadataCacheVariable:
+  case Kind::CoroAllocator:
     assert(isAlwaysSharedLinkage() && "kind should always be shared linkage");
     return nullptr;
 
@@ -1708,6 +1727,7 @@ bool LinkEntity::isAlwaysSharedLinkage() const {
   case Kind::TypeMetadataDemanglingCacheVariable:
   case Kind::NoncanonicalSpecializedGenericTypeMetadata:
   case Kind::NoncanonicalSpecializedGenericTypeMetadataCacheVariable:
+  case Kind::CoroAllocator:
     return true;
 
   default:
diff --git a/stdlib/public/Concurrency/TaskAlloc.cpp b/stdlib/public/Concurrency/TaskAlloc.cpp
@@ -69,39 +69,10 @@ void swift::_swift_task_dealloc_specific(AsyncTask *task, void *ptr) {
   allocator(task).dealloc(ptr);
 }
 
-static void *swift_task_alloc_thunk(size_t size) {
-  return swift_task_alloc(size);
-}
-
-static void swift_task_dealloc_thunk(void *ptr) { swift_task_dealloc(ptr); }
-
 void swift::swift_task_dealloc_through(void *ptr) {
   allocator(swift_task_getCurrent()).deallocThrough(ptr);
 }
 
-static CoroAllocator CoroTaskAllocatorImpl{
-    CoroAllocatorFlags(/*CoroAllocatorKind::Async*/ 1), swift_task_alloc_thunk,
-    swift_task_dealloc_thunk};
-
-CoroAllocator *const swift::_swift_coro_task_allocator = &CoroTaskAllocatorImpl;
-
-CoroAllocator *swift::swift_coro_getGlobalAllocator(CoroAllocatorFlags flags) {
-  switch (flags.getKind()) {
-  case CoroAllocatorKind::Stack:
-    return nullptr;
-  case CoroAllocatorKind::Async:
-    return _swift_coro_task_allocator;
-  case CoroAllocatorKind::Malloc:
-    return _swift_coro_malloc_allocator;
-  }
-}
-
-static CoroAllocator CoroMallocAllocatorImpl{
-    CoroAllocatorFlags(/*CoroAllocatorKind::Malloc*/ 2), malloc, free};
-
-CoroAllocator *const swift::_swift_coro_malloc_allocator =
-    &CoroMallocAllocatorImpl;
-
 void *swift::swift_coro_alloc(CoroAllocator *allocator, size_t size) {
   return allocator->allocate(size);
 }
diff --git a/test/IRGen/coroutine_accessors.swift b/test/IRGen/coroutine_accessors.swift
diff --git a/test/IRGen/coroutine_accessors_popless.swift b/test/IRGen/coroutine_accessors_popless.swift
diff --git a/test/abi/macOS/arm64/concurrency.swift b/test/abi/macOS/arm64/concurrency.swift
diff --git a/test/abi/macOS/x86_64/concurrency.swift b/test/abi/macOS/x86_64/concurrency.swift

Original file line number	Diff line number	Diff line change
`@@ -832,8 +832,8 @@ IRGenModule::IRGenModule(IRGenerator &irgen,`
`832`	`832`	`CoroAllocatorTy = createStructType(*this, "swift.coro_allocator",`
`833`	`833`	`{`
`834`	`834`	`Int32Ty, // CoroAllocator.Flags`
`835`		`- CoroAllocateFnTy,`
`836`		`- CoroDeallocateFnTy,`
	`835`	`+ CoroAllocateFnTy->getPointerTo(),`
	`836`	`+ CoroDeallocateFnTy->getPointerTo(),`
`837`	`837`	`});`
`838`	`838`	`CoroAllocatorPtrTy = CoroAllocatorTy->getPointerTo();`
`839`	`839`	`}`