Skip to content

Commit f8666b1

Browse files
nate-chandlernate-chandler
committed
[CoroutineAccessors] Pass cator to allocation fns.
Doing so enables allocators to contain additional context for use by allocation functions. Because the allocator is already passed to _swift_coro_alloc, on the fast path (no allocator, popless) no allocation function is used, and the allocator is passed in the swiftcoro register, this is cheap.
1 parent 4403625 commit f8666b1

File tree

5 files changed

+128
-56
lines changed

5 files changed

+128
-56
lines changed

‎include/swift/ABI/Coro.h‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -56,9 +56,11 @@ class CoroAllocatorFlags : public FlagSet<uint32_t> {
5656
setShouldDeallocateImmediately)
5757
};
5858

59+
struct CoroAllocator;
60+
5961
using CoroAllocation = void *;
60-
using CoroAllocateFn = CoroAllocation (*)(size_t);
61-
using CoroDealllocateFn = void (*)(CoroAllocation);
62+
using CoroAllocateFn = CoroAllocation (*)(CoroAllocator *, size_t);
63+
using CoroDealllocateFn = void (*)(CoroAllocator *, CoroAllocation);
6264

6365
struct CoroAllocator {
6466
CoroAllocatorFlags flags;

‎lib/IRGen/GenCoro.cpp‎

Lines changed: 96 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -671,6 +671,18 @@ struct Allocator {
671671
}
672672
}
673673

674+
llvm::AttributeList getFunctionAttributeList(IRGenModule &IGM) {
675+
switch (kind) {
676+
case Flags:
677+
llvm_unreachable("not a function");
678+
case Field::Allocate:
679+
case Field::Deallocate:
680+
auto attrs = llvm::AttributeList();
681+
IGM.addSwiftCoroAttributes(attrs, 0);
682+
return attrs;
683+
}
684+
}
685+
674686
const PointerAuthSchema &getSchema(IRGenModule &IGM) {
675687
switch (kind) {
676688
case Flags:
@@ -732,7 +744,8 @@ struct Allocator {
732744
}
733745
return FunctionPointer::createUnsigned(
734746
FunctionPointer::Kind::Function, callee,
735-
Signature(field.getFunctionType(IGF.IGM), {}, IGF.IGM.SwiftCC));
747+
Signature(field.getFunctionType(IGF.IGM),
748+
field.getFunctionAttributeList(IGF.IGM), IGF.IGM.SwiftCC));
736749
}
737750
};
738751
} // end anonymous namespace
@@ -743,7 +756,8 @@ llvm::Constant *swift::irgen::getCoroAllocFn(IRGenModule &IGM) {
743756
"_swift_coro_alloc", IGM.Int8PtrTy, {IGM.CoroAllocatorPtrTy, IGM.SizeTy},
744757
[isSwiftCoroCCAvailable](IRGenFunction &IGF) {
745758
auto parameters = IGF.collectParameters();
746-
auto allocator = Allocator(parameters.claimNext(), IGF);
759+
auto *allocatorValue = parameters.claimNext();
760+
auto allocator = Allocator(allocatorValue, IGF);
747761
auto *size = parameters.claimNext();
748762
if (isSwiftCoroCCAvailable) {
749763
// swiftcorocc is available, so if there's no allocator pointer,
@@ -762,7 +776,7 @@ llvm::Constant *swift::irgen::getCoroAllocFn(IRGenModule &IGM) {
762776
});
763777
}
764778
auto fnPtr = allocator.getAllocate();
765-
auto *call = IGF.Builder.CreateCall(fnPtr, {size});
779+
auto *call = IGF.Builder.CreateCall(fnPtr, {allocatorValue, size});
766780
call->setDoesNotThrow();
767781
call->setCallingConv(IGF.IGM.SwiftCC);
768782
IGF.Builder.CreateRet(call);
@@ -784,7 +798,8 @@ llvm::Constant *swift::irgen::getCoroDeallocFn(IRGenModule &IGM) {
784798
{IGM.CoroAllocatorPtrTy, IGM.Int8PtrTy},
785799
[isSwiftCoroCCAvailable](IRGenFunction &IGF) {
786800
auto parameters = IGF.collectParameters();
787-
auto allocator = Allocator(parameters.claimNext(), IGF);
801+
auto *allocatorValue = parameters.claimNext();
802+
auto allocator = Allocator(allocatorValue, IGF);
788803
auto *ptr = parameters.claimNext();
789804
if (isSwiftCoroCCAvailable) {
790805
// swiftcorocc is available, so if there's no allocator pointer,
@@ -816,7 +831,7 @@ llvm::Constant *swift::irgen::getCoroDeallocFn(IRGenModule &IGM) {
816831
// Start emitting the "normal" block.
817832
IGF.Builder.emitBlock(normalBlock);
818833
auto fnPtr = allocator.getDeallocate();
819-
auto *call = IGF.Builder.CreateCall(fnPtr, {ptr});
834+
auto *call = IGF.Builder.CreateCall(fnPtr, {allocatorValue, ptr});
820835
call->setDoesNotThrow();
821836
call->setCallingConv(IGF.IGM.SwiftCC);
822837
IGF.Builder.CreateRetVoid();
@@ -854,51 +869,106 @@ static llvm::Constant *getAddrOfGlobalCoroAllocator(
854869
return taskAllocator;
855870
}
856871

857-
static llvm::Constant *getAddrOfSwiftCCMalloc(IRGenModule &IGM) {
858-
auto mallocFnPtr = IGM.getMallocFunctionPointer();
859-
auto sig = mallocFnPtr.getSignature();
860-
if (sig.getCallingConv() == IGM.SwiftCC) {
861-
return IGM.getMallocFn();
862-
}
872+
static llvm::Constant *getAddrOfSwiftCoroMalloc(IRGenModule &IGM) {
873+
auto *ty = IGM.CoroAllocateFnTy;
863874
return IGM.getOrCreateHelperFunction(
864-
"_swift_malloc", sig.getType()->getReturnType(), sig.getType()->params(),
875+
"_swift_coro_malloc", ty->getReturnType(), ty->params(),
865876
[](IRGenFunction &IGF) {
866877
auto parameters = IGF.collectParameters();
878+
parameters.claimNext(); // allocator
867879
auto *size = parameters.claimNext();
868-
auto malloc = IGF.IGM.getMallocFunctionPointer();
869-
auto *call = IGF.Builder.CreateCall(malloc, {size});
880+
auto alloc = IGF.IGM.getMallocFunctionPointer();
881+
auto *call = IGF.Builder.CreateCall(alloc, {size});
870882
IGF.Builder.CreateRet(call);
883+
},
884+
/*setIsNoInline=*/true, /*forPrologue=*/false,
885+
/*isPerformanceConstraint=*/false,
886+
/*optionalLinkage=*/nullptr,
887+
/*specialCallingConv=*/std::nullopt,
888+
/*transformAttributes=*/
889+
[&IGM](llvm::AttributeList &attrs) {
890+
IGM.addSwiftCoroAttributes(attrs, 0);
871891
});
872892
}
873893

874-
static llvm::Constant *getAddrOfSwiftCCFree(IRGenModule &IGM) {
875-
auto freeFnPtr = IGM.getFreeFunctionPointer();
876-
auto sig = freeFnPtr.getSignature();
877-
if (sig.getCallingConv() == IGM.SwiftCC) {
878-
return IGM.getFreeFn();
879-
}
894+
static llvm::Constant *getAddrOfSwiftCoroFree(IRGenModule &IGM) {
895+
auto *ty = IGM.CoroDeallocateFnTy;
880896
return IGM.getOrCreateHelperFunction(
881-
"_swift_free", sig.getType()->getReturnType(), sig.getType()->params(),
897+
"_swift_coro_free", ty->getReturnType(), ty->params(),
882898
[](IRGenFunction &IGF) {
883899
auto parameters = IGF.collectParameters();
900+
parameters.claimNext(); // allocator
884901
auto *ptr = parameters.claimNext();
885-
auto free = IGF.IGM.getFreeFunctionPointer();
886-
IGF.Builder.CreateCall(free, {ptr});
902+
auto dealloc = IGF.IGM.getFreeFunctionPointer();
903+
IGF.Builder.CreateCall(dealloc, {ptr});
887904
IGF.Builder.CreateRetVoid();
905+
},
906+
/*setIsNoInline=*/true, /*forPrologue=*/false,
907+
/*isPerformanceConstraint=*/false,
908+
/*optionalLinkage=*/nullptr,
909+
/*specialCallingConv=*/std::nullopt,
910+
/*transformAttributes=*/
911+
[&IGM](llvm::AttributeList &attrs) {
912+
IGM.addSwiftCoroAttributes(attrs, 0);
888913
});
889914
}
890915

891916
llvm::Constant *IRGenModule::getAddrOfGlobalCoroMallocAllocator() {
892917
return getAddrOfGlobalCoroAllocator(*this, CoroAllocatorKind::Malloc,
893918
/*shouldDeallocateImmediately=*/true,
894-
getAddrOfSwiftCCMalloc(*this),
895-
getAddrOfSwiftCCFree(*this));
919+
getAddrOfSwiftCoroMalloc(*this),
920+
getAddrOfSwiftCoroFree(*this));
921+
}
922+
923+
static llvm::Constant *getAddrOfSwiftCoroTaskAlloc(IRGenModule &IGM) {
924+
auto *ty = IGM.CoroAllocateFnTy;
925+
return IGM.getOrCreateHelperFunction(
926+
"_swift_coro_task_alloc", ty->getReturnType(), ty->params(),
927+
[](IRGenFunction &IGF) {
928+
auto parameters = IGF.collectParameters();
929+
parameters.claimNext(); // allocator
930+
auto *size = parameters.claimNext();
931+
auto alloc = IGF.IGM.getTaskAllocFunctionPointer();
932+
auto *call = IGF.Builder.CreateCall(alloc, {size});
933+
IGF.Builder.CreateRet(call);
934+
},
935+
/*setIsNoInline=*/true, /*forPrologue=*/false,
936+
/*isPerformanceConstraint=*/false,
937+
/*optionalLinkage=*/nullptr,
938+
/*specialCallingConv=*/std::nullopt,
939+
/*transformAttributes=*/
940+
[&IGM](llvm::AttributeList &attrs) {
941+
IGM.addSwiftCoroAttributes(attrs, 0);
942+
});
943+
}
944+
945+
static llvm::Constant *getAddrOfSwiftCoroTaskDealloc(IRGenModule &IGM) {
946+
auto *ty = IGM.CoroDeallocateFnTy;
947+
return IGM.getOrCreateHelperFunction(
948+
"_swift_coro_task_dealloc", ty->getReturnType(), ty->params(),
949+
[](IRGenFunction &IGF) {
950+
auto parameters = IGF.collectParameters();
951+
parameters.claimNext(); // allocator
952+
auto *ptr = parameters.claimNext();
953+
auto dealloc = IGF.IGM.getTaskDeallocFunctionPointer();
954+
IGF.Builder.CreateCall(dealloc, {ptr});
955+
IGF.Builder.CreateRetVoid();
956+
},
957+
/*setIsNoInline=*/true, /*forPrologue=*/false,
958+
/*isPerformanceConstraint=*/false,
959+
/*optionalLinkage=*/nullptr,
960+
/*specialCallingConv=*/std::nullopt,
961+
/*transformAttributes=*/
962+
[&IGM](llvm::AttributeList &attrs) {
963+
IGM.addSwiftCoroAttributes(attrs, 0);
964+
});
896965
}
897966

898967
llvm::Constant *IRGenModule::getAddrOfGlobalCoroAsyncTaskAllocator() {
899968
return getAddrOfGlobalCoroAllocator(*this, CoroAllocatorKind::Async,
900969
/*shouldDeallocateImmediately=*/false,
901-
getTaskAllocFn(), getTaskDeallocFn());
970+
getAddrOfSwiftCoroTaskAlloc(*this),
971+
getAddrOfSwiftCoroTaskDealloc(*this));
902972
}
903973

904974
llvm::Value *

‎lib/IRGen/IRGenModule.cpp‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -772,10 +772,10 @@ IRGenModule::IRGenModule(IRGenerator &irgen,
772772

773773
CoroFunctionPointerTy = createStructType(*this, "swift.coro_func_pointer",
774774
{RelativeAddressTy, Int32Ty}, true);
775-
CoroAllocateFnTy =
776-
llvm::FunctionType::get(CoroAllocationTy, SizeTy, /*isVarArg*/ false);
777-
CoroDeallocateFnTy =
778-
llvm::FunctionType::get(VoidTy, CoroAllocationTy, /*isVarArg*/ false);
775+
CoroAllocateFnTy = llvm::FunctionType::get(
776+
CoroAllocationTy, {CoroAllocatorPtrTy, SizeTy}, /*isVarArg*/ false);
777+
CoroDeallocateFnTy = llvm::FunctionType::get(
778+
VoidTy, {CoroAllocatorPtrTy, CoroAllocationTy}, /*isVarArg*/ false);
779779
CoroAllocatorFlagsTy = Int32Ty;
780780
// swift/ABI/Coro.h : CoroAllocator
781781
CoroAllocatorTy = createStructType(*this, "swift.coro_allocator",

‎test/IRGen/coroutine_accessors.swift‎

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -23,15 +23,15 @@
2323
// CHECK-SAME: i32 0
2424
// CHECK-SAME: }>
2525

26-
// CHECK-arm64e-LABEL: _swift_malloc.ptrauth = private constant {
27-
// CHECK-arm64e-SAME: ptr @_swift_malloc,
26+
// CHECK-arm64e-LABEL: _swift_coro_malloc.ptrauth = private constant {
27+
// CHECK-arm64e-SAME: ptr @_swift_coro_malloc,
2828
// CHECK-arm64e-SAME: i32 0,
2929
// CHECK-arm64e-SAME: i64 0,
3030
// CHECK-arm64e-SAME: i64 24469 }
3131
// CHECK-arm64e-SAME: section "llvm.ptrauth"
3232
// CHECK-arm64e-SAME: align 8
33-
// CHECK-arm64e-LABEL: _swift_free.ptrauth = private constant {
34-
// CHECK-arm64e-SAME: ptr @_swift_free,
33+
// CHECK-arm64e-LABEL: _swift_coro_free.ptrauth = private constant {
34+
// CHECK-arm64e-SAME: ptr @_swift_coro_free,
3535
// CHECK-arm64e-SAME: i32 0,
3636
// CHECK-arm64e-SAME: i64 0,
3737
// CHECK-arm64e-SAME: i64 40879 },
@@ -42,24 +42,24 @@
4242
// CHECK-SAME: malloc
4343
// CHECK-SAME: free
4444
// CHECK-SAME: }
45-
// CHECK-arm64e-LABEL: swift_task_alloc.ptrauth = private constant {
46-
// CHECK-arm64e-SAME: ptr @swift_task_alloc,
45+
// CHECK-arm64e-LABEL: _swift_coro_task_alloc.ptrauth = private constant {
46+
// CHECK-arm64e-SAME: ptr @_swift_coro_task_alloc,
4747
// CHECK-arm64e-SAME: i32 0,
4848
// CHECK-arm64e-SAME: i64 0,
4949
// CHECK-arm64e-SAME: i64 24469 }
5050
// CHECK-arm64e-SAME: section "llvm.ptrauth"
5151
// CHECK-arm64e-SAME: align 8
52-
// CHECK-arm64e-LABEL: @swift_task_dealloc.ptrauth = private constant {
53-
// CHECK-arm64e-SAME: ptr @swift_task_dealloc,
52+
// CHECK-arm64e-LABEL: @_swift_coro_task_dealloc.ptrauth = private constant {
53+
// CHECK-arm64e-SAME: ptr @_swift_coro_task_dealloc,
5454
// CHECK-arm64e-SAME: i32 0,
5555
// CHECK-arm64e-SAME: i64 0,
5656
// CHECK-arm64e-SAME: i64 40879 },
5757
// CHECK-arm64e-SAME: section "llvm.ptrauth",
5858
// CHECK-arm64e-SAME: align 8
5959
// CHECK-LABEL: _swift_coro_async_allocator = linkonce_odr hidden constant %swift.coro_allocator {
6060
// CHECK-SAME: i32 1,
61-
// CHECK-SAME: swift_task_alloc
62-
// CHECK-SAME: swift_task_dealloc
61+
// CHECK-SAME: _swift_coro_task_alloc
62+
// CHECK-SAME: _swift_coro_task_dealloc
6363
// CHECK-SAME: }
6464

6565
// CHECK-LABEL: @_swift_coro_alloc(
@@ -76,7 +76,7 @@
7676
// CHECK-arm64e: [[ALLOCATE_FN_BITS:%[^,]+]] = ptrtoint ptr [[ALLOCATE_FN]] to i64
7777
// CHECK-arm64e: [[ALLOCATE_FN_BITS_AUTHED:%[^,]+]] = call i64 @llvm.ptrauth.auth(i64 [[ALLOCATE_FN_BITS]], i32 0, i64 24469)
7878
// CHECK-arm64e: [[ALLOCATE_FN:%[^,]+]] = inttoptr i64 [[ALLOCATE_FN_BITS_AUTHED]]
79-
// CHECK: [[ALLOCATION:%[^,]+]] = call swiftcc ptr [[ALLOCATE_FN]]([[INT]] [[SIZE]])
79+
// CHECK: [[ALLOCATION:%[^,]+]] = call swiftcc ptr [[ALLOCATE_FN]](ptr swiftcoro [[ALLOCATOR]], [[INT]] [[SIZE]])
8080
// CHECK: ret ptr [[ALLOCATION]]
8181
// CHECK: }
8282

@@ -107,7 +107,7 @@
107107
// CHECK-arm64e: [[DEALLOCATE_FN_BITS:%[^,]+]] = ptrtoint ptr [[DEALLOCATE_FN]] to i64
108108
// CHECK-arm64e: [[DEALLOCATE_FN_BITS_AUTHED:%[^,]+]] = call i64 @llvm.ptrauth.auth(i64 [[DEALLOCATE_FN_BITS]], i32 0, i64 40879)
109109
// CHECK-arm64e: [[DEALLOCATE_FN:%[^,]+]] = inttoptr i64 [[DEALLOCATE_FN_BITS_AUTHED]]
110-
// CHECK: call swiftcc void [[DEALLOCATE_FN]](ptr [[ADDRESS]])
110+
// CHECK: call swiftcc void [[DEALLOCATE_FN]](ptr swiftcoro [[ALLOCATOR]], ptr [[ADDRESS]])
111111
// CHECK: ret void
112112
// CHECK: }
113113

‎test/IRGen/coroutine_accessors_popless.swift‎

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -25,34 +25,34 @@
2525
// CHECK-SAME: i32 0
2626
// CHECK-SAME: }>
2727

28-
// CHECK-arm64e-LABEL: swift_task_alloc.ptrauth = private constant {
29-
// CHECK-arm64e-SAME: ptr @swift_task_alloc,
28+
// CHECK-arm64e-LABEL: _swift_coro_task_alloc.ptrauth = private constant {
29+
// CHECK-arm64e-SAME: ptr @_swift_coro_task_alloc,
3030
// CHECK-arm64e-SAME: i32 0,
3131
// CHECK-arm64e-SAME: i64 0,
3232
// CHECK-arm64e-SAME: i64 24469 }
3333
// CHECK-arm64e-SAME: section "llvm.ptrauth"
3434
// CHECK-arm64e-SAME: align 8
35-
// CHECK-arm64e-LABEL: @swift_task_dealloc.ptrauth = private constant {
36-
// CHECK-arm64e-SAME: ptr @swift_task_dealloc,
35+
// CHECK-arm64e-LABEL: @_swift_coro_task_dealloc.ptrauth = private constant {
36+
// CHECK-arm64e-SAME: ptr @_swift_coro_task_dealloc,
3737
// CHECK-arm64e-SAME: i32 0,
3838
// CHECK-arm64e-SAME: i64 0,
3939
// CHECK-arm64e-SAME: i64 40879 },
4040
// CHECK-arm64e-SAME: section "llvm.ptrauth",
4141
// CHECK-arm64e-SAME: align 8
4242
// CHECK-LABEL: _swift_coro_async_allocator = linkonce_odr hidden constant %swift.coro_allocator {
4343
// CHECK-SAME: i32 1,
44-
// CHECK-SAME: swift_task_alloc
45-
// CHECK-SAME: swift_task_dealloc
44+
// CHECK-SAME: _swift_coro_task_alloc
45+
// CHECK-SAME: _swift_coro_task_dealloc
4646
// CHECK-SAME: }
47-
// CHECK-arm64e-LABEL: _swift_malloc.ptrauth = private constant {
48-
// CHECK-arm64e-SAME: ptr @_swift_malloc,
47+
// CHECK-arm64e-LABEL: _swift_coro_malloc.ptrauth = private constant {
48+
// CHECK-arm64e-SAME: ptr @_swift_coro_malloc,
4949
// CHECK-arm64e-SAME: i32 0,
5050
// CHECK-arm64e-SAME: i64 0,
5151
// CHECK-arm64e-SAME: i64 24469 }
5252
// CHECK-arm64e-SAME: section "llvm.ptrauth"
5353
// CHECK-arm64e-SAME: align 8
54-
// CHECK-arm64e-LABEL: _swift_free.ptrauth = private constant {
55-
// CHECK-arm64e-SAME: ptr @_swift_free,
54+
// CHECK-arm64e-LABEL: _swift_coro_free.ptrauth = private constant {
55+
// CHECK-arm64e-SAME: ptr @_swift_coro_free,
5656
// CHECK-arm64e-SAME: i32 0,
5757
// CHECK-arm64e-SAME: i64 0,
5858
// CHECK-arm64e-SAME: i64 40879 },
@@ -87,7 +87,7 @@
8787
// CHECK-arm64e: [[ALLOCATE_FN_BITS:%[^,]+]] = ptrtoint ptr [[ALLOCATE_FN]] to i64
8888
// CHECK-arm64e: [[ALLOCATE_FN_BITS_AUTHED:%[^,]+]] = call i64 @llvm.ptrauth.auth(i64 [[ALLOCATE_FN_BITS]], i32 0, i64 24469)
8989
// CHECK-arm64e: [[ALLOCATE_FN:%[^,]+]] = inttoptr i64 [[ALLOCATE_FN_BITS_AUTHED]]
90-
// CHECK: [[ALLOCATION:%[^,]+]] = call swiftcc ptr [[ALLOCATE_FN]]([[INT]] [[SIZE]])
90+
// CHECK: [[ALLOCATION:%[^,]+]] = call swiftcc ptr [[ALLOCATE_FN]](ptr swiftcoro [[ALLOCATOR]], [[INT]] [[SIZE]])
9191
// CHECK: ret ptr [[ALLOCATION]]
9292
// CHECK: }
9393

@@ -125,7 +125,7 @@
125125
// CHECK-arm64e: [[DEALLOCATE_FN_BITS:%[^,]+]] = ptrtoint ptr [[DEALLOCATE_FN]] to i64
126126
// CHECK-arm64e: [[DEALLOCATE_FN_BITS_AUTHED:%[^,]+]] = call i64 @llvm.ptrauth.auth(i64 [[DEALLOCATE_FN_BITS]], i32 0, i64 40879)
127127
// CHECK-arm64e: [[DEALLOCATE_FN:%[^,]+]] = inttoptr i64 [[DEALLOCATE_FN_BITS_AUTHED]]
128-
// CHECK: call swiftcc void [[DEALLOCATE_FN]](ptr [[ADDRESS]])
128+
// CHECK: call swiftcc void [[DEALLOCATE_FN]](ptr swiftcoro [[ALLOCATOR]], ptr [[ADDRESS]])
129129
// CHECK: ret void
130130
// CHECK: }
131131

0 commit comments

Comments
 (0)