Skip to content

Commit e8d871c

Browse files
nate-chandlernate-chandler
authored
Merge pull request #80290 from nate-chandler/general-coro/20250325/1
[CoroutineAccessors] Directly reference allocators.
2 parents f042eb8 + 0c2a38b commit e8d871c

File tree

12 files changed

+479
-72
lines changed

12 files changed

+479
-72
lines changed

‎include/swift/IRGen/Linking.h

Lines changed: 23 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
#ifndef SWIFT_IRGEN_LINKING_H
1414
#define SWIFT_IRGEN_LINKING_H
1515

16+
#include "swift/ABI/Coro.h"
1617
#include "swift/AST/Decl.h"
1718
#include "swift/AST/Module.h"
1819
#include "swift/AST/ProtocolAssociations.h"
@@ -24,8 +25,8 @@
2425
#include "swift/SIL/SILGlobalVariable.h"
2526
#include "swift/SIL/SILModule.h"
2627
#include "llvm/ADT/DenseMapInfo.h"
27-
#include "llvm/IR/GlobalValue.h"
2828
#include "llvm/IR/GlobalObject.h"
29+
#include "llvm/IR/GlobalValue.h"
2930
#include "llvm/IR/Module.h"
3031

3132
namespace llvm {
@@ -165,6 +166,9 @@ class LinkEntity {
165166
ExtendedExistentialIsUniqueMask = 0x100,
166167
ExtendedExistentialIsSharedShift = 9,
167168
ExtendedExistentialIsSharedMask = 0x200,
169+
170+
// Used by CoroAllocator. 2 bits.
171+
CoroAllocatorKindShift = 8, CoroAllocatorKindMask = 0x300,
168172
};
169173
#define LINKENTITY_SET_FIELD(field, value) (value << field##Shift)
170174
#define LINKENTITY_GET_FIELD(value, field) ((value & field##Mask) >> field##Shift)
@@ -590,10 +594,13 @@ class LinkEntity {
590594
/// The pointer is a const char* of the name.
591595
KnownCoroFunctionPointer,
592596

593-
/// An coro function pointer for a distributed accessor (method or
597+
/// A coro function pointer for a distributed accessor (method or
594598
/// property).
595599
/// The pointer is a SILFunction*.
596600
DistributedAccessorCoroFunctionPointer,
601+
602+
/// An allocator to be passed to swift_coro_alloc and swift_coro_dealloc.
603+
CoroAllocator,
597604
};
598605
friend struct llvm::DenseMapInfo<LinkEntity>;
599606

@@ -1498,6 +1505,15 @@ class LinkEntity {
14981505
return entity;
14991506
}
15001507

1508+
static LinkEntity forCoroAllocator(CoroAllocatorKind kind) {
1509+
LinkEntity entity;
1510+
entity.Pointer = nullptr;
1511+
entity.SecondaryPointer = nullptr;
1512+
entity.Data = LINKENTITY_SET_FIELD(Kind, unsigned(Kind::CoroAllocator)) |
1513+
LINKENTITY_SET_FIELD(CoroAllocatorKind, unsigned(kind));
1514+
return entity;
1515+
}
1516+
15011517
static LinkEntity forCoroFunctionPointer(LinkEntity other) {
15021518
LinkEntity entity;
15031519
entity.Pointer = other.Pointer;
@@ -1709,6 +1725,11 @@ class LinkEntity {
17091725
SecondaryPointer);
17101726
}
17111727

1728+
CoroAllocatorKind getCoroAllocatorKind() const {
1729+
assert(getKind() == Kind::CoroAllocator);
1730+
return CoroAllocatorKind(LINKENTITY_GET_FIELD(Data, CoroAllocatorKind));
1731+
}
1732+
17121733
CanGenericSignature getExtendedExistentialTypeShapeGenSig() const {
17131734
assert(getKind() == Kind::ExtendedExistentialTypeShape);
17141735
return CanGenericSignature(

‎include/swift/Runtime/Concurrency.h

Lines changed: 0 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -128,25 +128,11 @@ void swift_task_dealloc(void *ptr);
128128
SWIFT_EXPORT_FROM(swift_Concurrency)
129129
SWIFT_CC(swift) void swift_task_dealloc_through(void *ptr);
130130

131-
// TODO: CoroutineAccessors: Eliminate this entry point and replace its uses
132-
// with direct references the globals.
133-
SWIFT_EXPORT_FROM(swift_Concurrency)
134-
SWIFT_CC(swift)
135-
CoroAllocator *swift_coro_getGlobalAllocator(CoroAllocatorFlags flags);
136-
137-
// TODO: CoroutineAccessors: Mark the underlying struct const.
138-
SWIFT_EXPORT_FROM(swift_Concurrency)
139-
CoroAllocator *const _swift_coro_task_allocator;
140-
141131
// TODO: CoroutineAccessors: Move these declarations back to swiftCore {{
142132
SWIFT_EXPORT_FROM(swift_Concurrency)
143133
SWIFT_CC(swift) void *swift_coro_alloc(CoroAllocator *allocator, size_t size);
144134
SWIFT_EXPORT_FROM(swift_Concurrency)
145135
SWIFT_CC(swift) void swift_coro_dealloc(CoroAllocator *allocator, void *ptr);
146-
147-
// TODO: CoroutineAccessors: Mark the underlying struct const.
148-
SWIFT_EXPORT_FROM(swift_Concurrency)
149-
CoroAllocator *const _swift_coro_malloc_allocator;
150136
// }} TODO: CoroutineAccessors
151137

152138
/// Cancel a task and all of its child tasks.

‎include/swift/Runtime/RuntimeFunctions.def

Lines changed: 0 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -3028,16 +3028,6 @@ FUNCTION(CoroDealloc,
30283028
EFFECT(RuntimeEffect::Deallocating, RuntimeEffect::Concurrency),
30293029
UNKNOWN_MEMEFFECTS)
30303030

3031-
// const CoroAllocator *swift_coro_getGlobalAllocator(CoroAllocatorFlags flags);
3032-
FUNCTION(CoroGetGlobalAllocator,
3033-
_Concurrency, swift_coro_getGlobalAllocator, SwiftCC,
3034-
CoroutineAccessorsAvailability,
3035-
RETURNS(CoroAllocatorPtrTy),
3036-
ARGS(CoroAllocatorFlagsTy),
3037-
ATTRS(NoUnwind, WillReturn),
3038-
EFFECT(RuntimeEffect::Allocating),
3039-
UNKNOWN_MEMEFFECTS)
3040-
30413031
#undef RETURNS
30423032
#undef ARGS
30433033
#undef ATTRS

‎lib/IRGen/GenCall.cpp

Lines changed: 37 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@
4444
#include <optional>
4545

4646
#include "CallEmission.h"
47+
#include "ConstantBuilder.h"
4748
#include "EntryPointArgumentEmission.h"
4849
#include "Explosion.h"
4950
#include "GenCall.h"
@@ -57,8 +58,8 @@
5758
#include "GenType.h"
5859
#include "IRGenDebugInfo.h"
5960
#include "IRGenFunction.h"
60-
#include "IRGenModule.h"
6161
#include "IRGenMangler.h"
62+
#include "IRGenModule.h"
6263
#include "LoadableTypeInfo.h"
6364
#include "NativeConventionSchema.h"
6465
#include "Signature.h"
@@ -5197,18 +5198,47 @@ Address irgen::emitAllocYieldManyCoroutineBuffer(IRGenFunction &IGF) {
51975198
return createOpaqueBufferAlloca(IGF, getYieldManyCoroutineBufferSize(IGF.IGM),
51985199
getYieldManyCoroutineBufferAlignment(IGF.IGM));
51995200
}
5201+
5202+
static llvm::Constant *getAddrOfGlobalCoroAllocator(IRGenModule &IGM,
5203+
CoroAllocatorKind kind,
5204+
llvm::Constant *allocFn,
5205+
llvm::Constant *deallocFn) {
5206+
auto entity = LinkEntity::forCoroAllocator(kind);
5207+
auto taskAllocator = IGM.getOrCreateLazyGlobalVariable(
5208+
entity,
5209+
[&](ConstantInitBuilder &builder) -> ConstantInitFuture {
5210+
auto allocator = builder.beginStruct(IGM.CoroAllocatorTy);
5211+
allocator.addInt32(CoroAllocatorFlags(kind).getOpaqueValue());
5212+
allocator.add(allocFn);
5213+
allocator.add(deallocFn);
5214+
return allocator.finishAndCreateFuture();
5215+
},
5216+
[&](llvm::GlobalVariable *var) { var->setConstant(true); });
5217+
return taskAllocator;
5218+
}
5219+
llvm::Constant *IRGenModule::getAddrOfGlobalCoroMallocAllocator() {
5220+
return getAddrOfGlobalCoroAllocator(*this, CoroAllocatorKind::Malloc,
5221+
getMallocFn(), getFreeFn());
5222+
}
5223+
llvm::Constant *IRGenModule::getAddrOfGlobalCoroAsyncTaskAllocator() {
5224+
return getAddrOfGlobalCoroAllocator(*this, CoroAllocatorKind::Async,
5225+
getTaskAllocFn(), getTaskDeallocFn());
5226+
}
52005227
llvm::Value *
52015228
irgen::emitYieldOnce2CoroutineAllocator(IRGenFunction &IGF,
52025229
std::optional<CoroAllocatorKind> kind) {
52035230
if (!kind) {
52045231
return IGF.getCoroutineAllocator();
52055232
}
5206-
CoroAllocatorFlags flags(*kind);
5207-
auto *flagsValue = llvm::ConstantInt::get(IGF.IGM.CoroAllocatorFlagsTy,
5208-
flags.getOpaqueValue());
5209-
5210-
return IGF.Builder.CreateCall(
5211-
IGF.IGM.getCoroGetGlobalAllocatorFunctionPointer(), {flagsValue});
5233+
switch (*kind) {
5234+
case CoroAllocatorKind::Stack:
5235+
return llvm::ConstantPointerNull::get(IGF.IGM.CoroAllocatorPtrTy);
5236+
case CoroAllocatorKind::Async:
5237+
return IGF.IGM.getAddrOfGlobalCoroAsyncTaskAllocator();
5238+
case CoroAllocatorKind::Malloc:
5239+
return IGF.IGM.getAddrOfGlobalCoroMallocAllocator();
5240+
}
5241+
llvm_unreachable("unhandled case");
52125242
}
52135243
StackAddress irgen::emitAllocYieldOnce2CoroutineFrame(IRGenFunction &IGF,
52145244
llvm::Value *size) {

‎lib/IRGen/IRGenModule.cpp

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -832,8 +832,8 @@ IRGenModule::IRGenModule(IRGenerator &irgen,
832832
CoroAllocatorTy = createStructType(*this, "swift.coro_allocator",
833833
{
834834
Int32Ty, // CoroAllocator.Flags
835-
CoroAllocateFnTy,
836-
CoroDeallocateFnTy,
835+
CoroAllocateFnTy->getPointerTo(),
836+
CoroDeallocateFnTy->getPointerTo(),
837837
});
838838
CoroAllocatorPtrTy = CoroAllocatorTy->getPointerTo();
839839
}

‎lib/IRGen/IRGenModule.h

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1671,6 +1671,9 @@ private: \
16711671
ConstantInit init);
16721672
SILFunction *getSILFunctionForCoroFunctionPointer(llvm::Constant *cfp);
16731673

1674+
llvm::Constant *getAddrOfGlobalCoroMallocAllocator();
1675+
llvm::Constant *getAddrOfGlobalCoroAsyncTaskAllocator();
1676+
16741677
llvm::Function *getAddrOfDispatchThunk(SILDeclRef declRef,
16751678
ForDefinition_t forDefinition);
16761679
void emitDispatchThunk(SILDeclRef declRef);

‎lib/IRGen/Linking.cpp

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -577,6 +577,16 @@ std::string LinkEntity::mangleAsString(ASTContext &Ctx) const {
577577
Result.append("Twc");
578578
return Result;
579579
}
580+
case Kind::CoroAllocator: {
581+
switch (getCoroAllocatorKind()) {
582+
case CoroAllocatorKind::Stack:
583+
llvm::report_fatal_error("attempting to mangle the coro stack allocator");
584+
case CoroAllocatorKind::Async:
585+
return "_swift_coro_async_allocator";
586+
case CoroAllocatorKind::Malloc:
587+
return "_swift_coro_malloc_allocator";
588+
}
589+
}
580590
}
581591
llvm_unreachable("bad entity kind!");
582592
}
@@ -948,6 +958,8 @@ SILLinkage LinkEntity::getLinkage(ForDefinition_t forDefinition) const {
948958
return getUnderlyingEntityForCoroFunctionPointer().getLinkage(
949959
forDefinition);
950960
return getSILLinkage(getDeclLinkage(getDecl()), forDefinition);
961+
case Kind::CoroAllocator:
962+
return SILLinkage::Shared;
951963
}
952964
llvm_unreachable("bad link entity kind");
953965
}
@@ -1051,6 +1063,7 @@ bool LinkEntity::isContextDescriptor() const {
10511063
case Kind::CoroFunctionPointerAST:
10521064
case Kind::DistributedThunkCoroFunctionPointer:
10531065
case Kind::KnownCoroFunctionPointer:
1066+
case Kind::CoroAllocator:
10541067
return false;
10551068
}
10561069
llvm_unreachable("invalid descriptor");
@@ -1188,6 +1201,8 @@ llvm::Type *LinkEntity::getDefaultDeclarationType(IRGenModule &IGM) const {
11881201
case Kind::DistributedThunkCoroFunctionPointer:
11891202
case Kind::KnownCoroFunctionPointer:
11901203
return IGM.CoroFunctionPointerPtrTy;
1204+
case Kind::CoroAllocator:
1205+
return IGM.CoroAllocatorTy;
11911206
default:
11921207
llvm_unreachable("declaration LLVM type not specified");
11931208
}
@@ -1220,6 +1235,7 @@ Alignment LinkEntity::getAlignment(IRGenModule &IGM) const {
12201235
case Kind::OpaqueTypeDescriptorRecord:
12211236
case Kind::AccessibleFunctionRecord:
12221237
case Kind::ExtendedExistentialTypeShape:
1238+
case Kind::CoroAllocator:
12231239
return Alignment(4);
12241240
case Kind::AsyncFunctionPointer:
12251241
case Kind::DispatchThunkAsyncFunctionPointer:
@@ -1359,6 +1375,7 @@ bool LinkEntity::isText() const {
13591375
case Kind::DynamicallyReplaceableFunctionVariable:
13601376
case Kind::CanonicalSpecializedGenericTypeMetadataAccessFunction:
13611377
case Kind::ExtendedExistentialTypeShape:
1378+
case Kind::CoroAllocator:
13621379
return true;
13631380
case Kind::ObjCClass:
13641381
case Kind::ObjCClassRef:
@@ -1520,6 +1537,7 @@ bool LinkEntity::isWeakImported(ModuleDecl *module) const {
15201537
case Kind::DifferentiabilityWitness:
15211538
case Kind::AccessibleFunctionRecord:
15221539
case Kind::ExtendedExistentialTypeShape:
1540+
case Kind::CoroAllocator:
15231541
return false;
15241542

15251543
case Kind::AsyncFunctionPointer:
@@ -1656,6 +1674,7 @@ DeclContext *LinkEntity::getDeclContextForEmission() const {
16561674
case Kind::TypeMetadataDemanglingCacheVariable:
16571675
case Kind::NoncanonicalSpecializedGenericTypeMetadata:
16581676
case Kind::NoncanonicalSpecializedGenericTypeMetadataCacheVariable:
1677+
case Kind::CoroAllocator:
16591678
assert(isAlwaysSharedLinkage() && "kind should always be shared linkage");
16601679
return nullptr;
16611680

@@ -1708,6 +1727,7 @@ bool LinkEntity::isAlwaysSharedLinkage() const {
17081727
case Kind::TypeMetadataDemanglingCacheVariable:
17091728
case Kind::NoncanonicalSpecializedGenericTypeMetadata:
17101729
case Kind::NoncanonicalSpecializedGenericTypeMetadataCacheVariable:
1730+
case Kind::CoroAllocator:
17111731
return true;
17121732

17131733
default:

‎stdlib/public/Concurrency/TaskAlloc.cpp

Lines changed: 0 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -69,39 +69,10 @@ void swift::_swift_task_dealloc_specific(AsyncTask *task, void *ptr) {
6969
allocator(task).dealloc(ptr);
7070
}
7171

72-
static void *swift_task_alloc_thunk(size_t size) {
73-
return swift_task_alloc(size);
74-
}
75-
76-
static void swift_task_dealloc_thunk(void *ptr) { swift_task_dealloc(ptr); }
77-
7872
void swift::swift_task_dealloc_through(void *ptr) {
7973
allocator(swift_task_getCurrent()).deallocThrough(ptr);
8074
}
8175

82-
static CoroAllocator CoroTaskAllocatorImpl{
83-
CoroAllocatorFlags(/*CoroAllocatorKind::Async*/ 1), swift_task_alloc_thunk,
84-
swift_task_dealloc_thunk};
85-
86-
CoroAllocator *const swift::_swift_coro_task_allocator = &CoroTaskAllocatorImpl;
87-
88-
CoroAllocator *swift::swift_coro_getGlobalAllocator(CoroAllocatorFlags flags) {
89-
switch (flags.getKind()) {
90-
case CoroAllocatorKind::Stack:
91-
return nullptr;
92-
case CoroAllocatorKind::Async:
93-
return _swift_coro_task_allocator;
94-
case CoroAllocatorKind::Malloc:
95-
return _swift_coro_malloc_allocator;
96-
}
97-
}
98-
99-
static CoroAllocator CoroMallocAllocatorImpl{
100-
CoroAllocatorFlags(/*CoroAllocatorKind::Malloc*/ 2), malloc, free};
101-
102-
CoroAllocator *const swift::_swift_coro_malloc_allocator =
103-
&CoroMallocAllocatorImpl;
104-
10576
void *swift::swift_coro_alloc(CoroAllocator *allocator, size_t size) {
10677
return allocator->allocate(size);
10778
}

0 commit comments

Comments
 (0)