[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.
This commit is contained in:
Nate Chandler
2025-10-16 09:23:24 -07:00
parent 4403625279
commit f8666b164b
5 changed files with 128 additions and 56 deletions

View File

@@ -56,9 +56,11 @@ public:
setShouldDeallocateImmediately)
};
struct CoroAllocator;
using CoroAllocation = void *;
using CoroAllocateFn = CoroAllocation (*)(size_t);
using CoroDealllocateFn = void (*)(CoroAllocation);
using CoroAllocateFn = CoroAllocation (*)(CoroAllocator *, size_t);
using CoroDealllocateFn = void (*)(CoroAllocator *, CoroAllocation);
struct CoroAllocator {
CoroAllocatorFlags flags;

View File

@@ -671,6 +671,18 @@ struct Allocator {
}
}
llvm::AttributeList getFunctionAttributeList(IRGenModule &IGM) {
switch (kind) {
case Flags:
llvm_unreachable("not a function");
case Field::Allocate:
case Field::Deallocate:
auto attrs = llvm::AttributeList();
IGM.addSwiftCoroAttributes(attrs, 0);
return attrs;
}
}
const PointerAuthSchema &getSchema(IRGenModule &IGM) {
switch (kind) {
case Flags:
@@ -732,7 +744,8 @@ private:
}
return FunctionPointer::createUnsigned(
FunctionPointer::Kind::Function, callee,
Signature(field.getFunctionType(IGF.IGM), {}, IGF.IGM.SwiftCC));
Signature(field.getFunctionType(IGF.IGM),
field.getFunctionAttributeList(IGF.IGM), IGF.IGM.SwiftCC));
}
};
} // end anonymous namespace
@@ -743,7 +756,8 @@ llvm::Constant *swift::irgen::getCoroAllocFn(IRGenModule &IGM) {
"_swift_coro_alloc", IGM.Int8PtrTy, {IGM.CoroAllocatorPtrTy, IGM.SizeTy},
[isSwiftCoroCCAvailable](IRGenFunction &IGF) {
auto parameters = IGF.collectParameters();
auto allocator = Allocator(parameters.claimNext(), IGF);
auto *allocatorValue = parameters.claimNext();
auto allocator = Allocator(allocatorValue, IGF);
auto *size = parameters.claimNext();
if (isSwiftCoroCCAvailable) {
// swiftcorocc is available, so if there's no allocator pointer,
@@ -762,7 +776,7 @@ llvm::Constant *swift::irgen::getCoroAllocFn(IRGenModule &IGM) {
});
}
auto fnPtr = allocator.getAllocate();
auto *call = IGF.Builder.CreateCall(fnPtr, {size});
auto *call = IGF.Builder.CreateCall(fnPtr, {allocatorValue, size});
call->setDoesNotThrow();
call->setCallingConv(IGF.IGM.SwiftCC);
IGF.Builder.CreateRet(call);
@@ -784,7 +798,8 @@ llvm::Constant *swift::irgen::getCoroDeallocFn(IRGenModule &IGM) {
{IGM.CoroAllocatorPtrTy, IGM.Int8PtrTy},
[isSwiftCoroCCAvailable](IRGenFunction &IGF) {
auto parameters = IGF.collectParameters();
auto allocator = Allocator(parameters.claimNext(), IGF);
auto *allocatorValue = parameters.claimNext();
auto allocator = Allocator(allocatorValue, IGF);
auto *ptr = parameters.claimNext();
if (isSwiftCoroCCAvailable) {
// swiftcorocc is available, so if there's no allocator pointer,
@@ -816,7 +831,7 @@ llvm::Constant *swift::irgen::getCoroDeallocFn(IRGenModule &IGM) {
// Start emitting the "normal" block.
IGF.Builder.emitBlock(normalBlock);
auto fnPtr = allocator.getDeallocate();
auto *call = IGF.Builder.CreateCall(fnPtr, {ptr});
auto *call = IGF.Builder.CreateCall(fnPtr, {allocatorValue, ptr});
call->setDoesNotThrow();
call->setCallingConv(IGF.IGM.SwiftCC);
IGF.Builder.CreateRetVoid();
@@ -854,51 +869,106 @@ static llvm::Constant *getAddrOfGlobalCoroAllocator(
return taskAllocator;
}
static llvm::Constant *getAddrOfSwiftCCMalloc(IRGenModule &IGM) {
auto mallocFnPtr = IGM.getMallocFunctionPointer();
auto sig = mallocFnPtr.getSignature();
if (sig.getCallingConv() == IGM.SwiftCC) {
return IGM.getMallocFn();
}
static llvm::Constant *getAddrOfSwiftCoroMalloc(IRGenModule &IGM) {
auto *ty = IGM.CoroAllocateFnTy;
return IGM.getOrCreateHelperFunction(
"_swift_malloc", sig.getType()->getReturnType(), sig.getType()->params(),
"_swift_coro_malloc", ty->getReturnType(), ty->params(),
[](IRGenFunction &IGF) {
auto parameters = IGF.collectParameters();
parameters.claimNext(); // allocator
auto *size = parameters.claimNext();
auto malloc = IGF.IGM.getMallocFunctionPointer();
auto *call = IGF.Builder.CreateCall(malloc, {size});
auto alloc = IGF.IGM.getMallocFunctionPointer();
auto *call = IGF.Builder.CreateCall(alloc, {size});
IGF.Builder.CreateRet(call);
},
/*setIsNoInline=*/true, /*forPrologue=*/false,
/*isPerformanceConstraint=*/false,
/*optionalLinkage=*/nullptr,
/*specialCallingConv=*/std::nullopt,
/*transformAttributes=*/
[&IGM](llvm::AttributeList &attrs) {
IGM.addSwiftCoroAttributes(attrs, 0);
});
}
static llvm::Constant *getAddrOfSwiftCCFree(IRGenModule &IGM) {
auto freeFnPtr = IGM.getFreeFunctionPointer();
auto sig = freeFnPtr.getSignature();
if (sig.getCallingConv() == IGM.SwiftCC) {
return IGM.getFreeFn();
}
static llvm::Constant *getAddrOfSwiftCoroFree(IRGenModule &IGM) {
auto *ty = IGM.CoroDeallocateFnTy;
return IGM.getOrCreateHelperFunction(
"_swift_free", sig.getType()->getReturnType(), sig.getType()->params(),
"_swift_coro_free", ty->getReturnType(), ty->params(),
[](IRGenFunction &IGF) {
auto parameters = IGF.collectParameters();
parameters.claimNext(); // allocator
auto *ptr = parameters.claimNext();
auto free = IGF.IGM.getFreeFunctionPointer();
IGF.Builder.CreateCall(free, {ptr});
auto dealloc = IGF.IGM.getFreeFunctionPointer();
IGF.Builder.CreateCall(dealloc, {ptr});
IGF.Builder.CreateRetVoid();
},
/*setIsNoInline=*/true, /*forPrologue=*/false,
/*isPerformanceConstraint=*/false,
/*optionalLinkage=*/nullptr,
/*specialCallingConv=*/std::nullopt,
/*transformAttributes=*/
[&IGM](llvm::AttributeList &attrs) {
IGM.addSwiftCoroAttributes(attrs, 0);
});
}
llvm::Constant *IRGenModule::getAddrOfGlobalCoroMallocAllocator() {
return getAddrOfGlobalCoroAllocator(*this, CoroAllocatorKind::Malloc,
/*shouldDeallocateImmediately=*/true,
getAddrOfSwiftCCMalloc(*this),
getAddrOfSwiftCCFree(*this));
getAddrOfSwiftCoroMalloc(*this),
getAddrOfSwiftCoroFree(*this));
}
static llvm::Constant *getAddrOfSwiftCoroTaskAlloc(IRGenModule &IGM) {
auto *ty = IGM.CoroAllocateFnTy;
return IGM.getOrCreateHelperFunction(
"_swift_coro_task_alloc", ty->getReturnType(), ty->params(),
[](IRGenFunction &IGF) {
auto parameters = IGF.collectParameters();
parameters.claimNext(); // allocator
auto *size = parameters.claimNext();
auto alloc = IGF.IGM.getTaskAllocFunctionPointer();
auto *call = IGF.Builder.CreateCall(alloc, {size});
IGF.Builder.CreateRet(call);
},
/*setIsNoInline=*/true, /*forPrologue=*/false,
/*isPerformanceConstraint=*/false,
/*optionalLinkage=*/nullptr,
/*specialCallingConv=*/std::nullopt,
/*transformAttributes=*/
[&IGM](llvm::AttributeList &attrs) {
IGM.addSwiftCoroAttributes(attrs, 0);
});
}
static llvm::Constant *getAddrOfSwiftCoroTaskDealloc(IRGenModule &IGM) {
auto *ty = IGM.CoroDeallocateFnTy;
return IGM.getOrCreateHelperFunction(
"_swift_coro_task_dealloc", ty->getReturnType(), ty->params(),
[](IRGenFunction &IGF) {
auto parameters = IGF.collectParameters();
parameters.claimNext(); // allocator
auto *ptr = parameters.claimNext();
auto dealloc = IGF.IGM.getTaskDeallocFunctionPointer();
IGF.Builder.CreateCall(dealloc, {ptr});
IGF.Builder.CreateRetVoid();
},
/*setIsNoInline=*/true, /*forPrologue=*/false,
/*isPerformanceConstraint=*/false,
/*optionalLinkage=*/nullptr,
/*specialCallingConv=*/std::nullopt,
/*transformAttributes=*/
[&IGM](llvm::AttributeList &attrs) {
IGM.addSwiftCoroAttributes(attrs, 0);
});
}
llvm::Constant *IRGenModule::getAddrOfGlobalCoroAsyncTaskAllocator() {
return getAddrOfGlobalCoroAllocator(*this, CoroAllocatorKind::Async,
/*shouldDeallocateImmediately=*/false,
getTaskAllocFn(), getTaskDeallocFn());
getAddrOfSwiftCoroTaskAlloc(*this),
getAddrOfSwiftCoroTaskDealloc(*this));
}
llvm::Value *

View File

@@ -772,10 +772,10 @@ IRGenModule::IRGenModule(IRGenerator &irgen,
CoroFunctionPointerTy = createStructType(*this, "swift.coro_func_pointer",
{RelativeAddressTy, Int32Ty}, true);
CoroAllocateFnTy =
llvm::FunctionType::get(CoroAllocationTy, SizeTy, /*isVarArg*/ false);
CoroDeallocateFnTy =
llvm::FunctionType::get(VoidTy, CoroAllocationTy, /*isVarArg*/ false);
CoroAllocateFnTy = llvm::FunctionType::get(
CoroAllocationTy, {CoroAllocatorPtrTy, SizeTy}, /*isVarArg*/ false);
CoroDeallocateFnTy = llvm::FunctionType::get(
VoidTy, {CoroAllocatorPtrTy, CoroAllocationTy}, /*isVarArg*/ false);
CoroAllocatorFlagsTy = Int32Ty;
// swift/ABI/Coro.h : CoroAllocator
CoroAllocatorTy = createStructType(*this, "swift.coro_allocator",

View File

@@ -23,15 +23,15 @@
// CHECK-SAME: i32 0
// CHECK-SAME: }>
// CHECK-arm64e-LABEL: _swift_malloc.ptrauth = private constant {
// CHECK-arm64e-SAME: ptr @_swift_malloc,
// CHECK-arm64e-LABEL: _swift_coro_malloc.ptrauth = private constant {
// CHECK-arm64e-SAME: ptr @_swift_coro_malloc,
// CHECK-arm64e-SAME: i32 0,
// CHECK-arm64e-SAME: i64 0,
// CHECK-arm64e-SAME: i64 24469 }
// CHECK-arm64e-SAME: section "llvm.ptrauth"
// CHECK-arm64e-SAME: align 8
// CHECK-arm64e-LABEL: _swift_free.ptrauth = private constant {
// CHECK-arm64e-SAME: ptr @_swift_free,
// CHECK-arm64e-LABEL: _swift_coro_free.ptrauth = private constant {
// CHECK-arm64e-SAME: ptr @_swift_coro_free,
// CHECK-arm64e-SAME: i32 0,
// CHECK-arm64e-SAME: i64 0,
// CHECK-arm64e-SAME: i64 40879 },
@@ -42,15 +42,15 @@
// CHECK-SAME: malloc
// CHECK-SAME: free
// CHECK-SAME: }
// CHECK-arm64e-LABEL: swift_task_alloc.ptrauth = private constant {
// CHECK-arm64e-SAME: ptr @swift_task_alloc,
// CHECK-arm64e-LABEL: _swift_coro_task_alloc.ptrauth = private constant {
// CHECK-arm64e-SAME: ptr @_swift_coro_task_alloc,
// CHECK-arm64e-SAME: i32 0,
// CHECK-arm64e-SAME: i64 0,
// CHECK-arm64e-SAME: i64 24469 }
// CHECK-arm64e-SAME: section "llvm.ptrauth"
// CHECK-arm64e-SAME: align 8
// CHECK-arm64e-LABEL: @swift_task_dealloc.ptrauth = private constant {
// CHECK-arm64e-SAME: ptr @swift_task_dealloc,
// CHECK-arm64e-LABEL: @_swift_coro_task_dealloc.ptrauth = private constant {
// CHECK-arm64e-SAME: ptr @_swift_coro_task_dealloc,
// CHECK-arm64e-SAME: i32 0,
// CHECK-arm64e-SAME: i64 0,
// CHECK-arm64e-SAME: i64 40879 },
@@ -58,8 +58,8 @@
// CHECK-arm64e-SAME: align 8
// CHECK-LABEL: _swift_coro_async_allocator = linkonce_odr hidden constant %swift.coro_allocator {
// CHECK-SAME: i32 1,
// CHECK-SAME: swift_task_alloc
// CHECK-SAME: swift_task_dealloc
// CHECK-SAME: _swift_coro_task_alloc
// CHECK-SAME: _swift_coro_task_dealloc
// CHECK-SAME: }
// CHECK-LABEL: @_swift_coro_alloc(
@@ -76,7 +76,7 @@
// CHECK-arm64e: [[ALLOCATE_FN_BITS:%[^,]+]] = ptrtoint ptr [[ALLOCATE_FN]] to i64
// CHECK-arm64e: [[ALLOCATE_FN_BITS_AUTHED:%[^,]+]] = call i64 @llvm.ptrauth.auth(i64 [[ALLOCATE_FN_BITS]], i32 0, i64 24469)
// CHECK-arm64e: [[ALLOCATE_FN:%[^,]+]] = inttoptr i64 [[ALLOCATE_FN_BITS_AUTHED]]
// CHECK: [[ALLOCATION:%[^,]+]] = call swiftcc ptr [[ALLOCATE_FN]]([[INT]] [[SIZE]])
// CHECK: [[ALLOCATION:%[^,]+]] = call swiftcc ptr [[ALLOCATE_FN]](ptr swiftcoro [[ALLOCATOR]], [[INT]] [[SIZE]])
// CHECK: ret ptr [[ALLOCATION]]
// CHECK: }
@@ -107,7 +107,7 @@
// CHECK-arm64e: [[DEALLOCATE_FN_BITS:%[^,]+]] = ptrtoint ptr [[DEALLOCATE_FN]] to i64
// CHECK-arm64e: [[DEALLOCATE_FN_BITS_AUTHED:%[^,]+]] = call i64 @llvm.ptrauth.auth(i64 [[DEALLOCATE_FN_BITS]], i32 0, i64 40879)
// CHECK-arm64e: [[DEALLOCATE_FN:%[^,]+]] = inttoptr i64 [[DEALLOCATE_FN_BITS_AUTHED]]
// CHECK: call swiftcc void [[DEALLOCATE_FN]](ptr [[ADDRESS]])
// CHECK: call swiftcc void [[DEALLOCATE_FN]](ptr swiftcoro [[ALLOCATOR]], ptr [[ADDRESS]])
// CHECK: ret void
// CHECK: }

View File

@@ -25,15 +25,15 @@
// CHECK-SAME: i32 0
// CHECK-SAME: }>
// CHECK-arm64e-LABEL: swift_task_alloc.ptrauth = private constant {
// CHECK-arm64e-SAME: ptr @swift_task_alloc,
// CHECK-arm64e-LABEL: _swift_coro_task_alloc.ptrauth = private constant {
// CHECK-arm64e-SAME: ptr @_swift_coro_task_alloc,
// CHECK-arm64e-SAME: i32 0,
// CHECK-arm64e-SAME: i64 0,
// CHECK-arm64e-SAME: i64 24469 }
// CHECK-arm64e-SAME: section "llvm.ptrauth"
// CHECK-arm64e-SAME: align 8
// CHECK-arm64e-LABEL: @swift_task_dealloc.ptrauth = private constant {
// CHECK-arm64e-SAME: ptr @swift_task_dealloc,
// CHECK-arm64e-LABEL: @_swift_coro_task_dealloc.ptrauth = private constant {
// CHECK-arm64e-SAME: ptr @_swift_coro_task_dealloc,
// CHECK-arm64e-SAME: i32 0,
// CHECK-arm64e-SAME: i64 0,
// CHECK-arm64e-SAME: i64 40879 },
@@ -41,18 +41,18 @@
// CHECK-arm64e-SAME: align 8
// CHECK-LABEL: _swift_coro_async_allocator = linkonce_odr hidden constant %swift.coro_allocator {
// CHECK-SAME: i32 1,
// CHECK-SAME: swift_task_alloc
// CHECK-SAME: swift_task_dealloc
// CHECK-SAME: _swift_coro_task_alloc
// CHECK-SAME: _swift_coro_task_dealloc
// CHECK-SAME: }
// CHECK-arm64e-LABEL: _swift_malloc.ptrauth = private constant {
// CHECK-arm64e-SAME: ptr @_swift_malloc,
// CHECK-arm64e-LABEL: _swift_coro_malloc.ptrauth = private constant {
// CHECK-arm64e-SAME: ptr @_swift_coro_malloc,
// CHECK-arm64e-SAME: i32 0,
// CHECK-arm64e-SAME: i64 0,
// CHECK-arm64e-SAME: i64 24469 }
// CHECK-arm64e-SAME: section "llvm.ptrauth"
// CHECK-arm64e-SAME: align 8
// CHECK-arm64e-LABEL: _swift_free.ptrauth = private constant {
// CHECK-arm64e-SAME: ptr @_swift_free,
// CHECK-arm64e-LABEL: _swift_coro_free.ptrauth = private constant {
// CHECK-arm64e-SAME: ptr @_swift_coro_free,
// CHECK-arm64e-SAME: i32 0,
// CHECK-arm64e-SAME: i64 0,
// CHECK-arm64e-SAME: i64 40879 },
@@ -87,7 +87,7 @@
// CHECK-arm64e: [[ALLOCATE_FN_BITS:%[^,]+]] = ptrtoint ptr [[ALLOCATE_FN]] to i64
// CHECK-arm64e: [[ALLOCATE_FN_BITS_AUTHED:%[^,]+]] = call i64 @llvm.ptrauth.auth(i64 [[ALLOCATE_FN_BITS]], i32 0, i64 24469)
// CHECK-arm64e: [[ALLOCATE_FN:%[^,]+]] = inttoptr i64 [[ALLOCATE_FN_BITS_AUTHED]]
// CHECK: [[ALLOCATION:%[^,]+]] = call swiftcc ptr [[ALLOCATE_FN]]([[INT]] [[SIZE]])
// CHECK: [[ALLOCATION:%[^,]+]] = call swiftcc ptr [[ALLOCATE_FN]](ptr swiftcoro [[ALLOCATOR]], [[INT]] [[SIZE]])
// CHECK: ret ptr [[ALLOCATION]]
// CHECK: }
@@ -125,7 +125,7 @@
// CHECK-arm64e: [[DEALLOCATE_FN_BITS:%[^,]+]] = ptrtoint ptr [[DEALLOCATE_FN]] to i64
// CHECK-arm64e: [[DEALLOCATE_FN_BITS_AUTHED:%[^,]+]] = call i64 @llvm.ptrauth.auth(i64 [[DEALLOCATE_FN_BITS]], i32 0, i64 40879)
// CHECK-arm64e: [[DEALLOCATE_FN:%[^,]+]] = inttoptr i64 [[DEALLOCATE_FN_BITS_AUTHED]]
// CHECK: call swiftcc void [[DEALLOCATE_FN]](ptr [[ADDRESS]])
// CHECK: call swiftcc void [[DEALLOCATE_FN]](ptr swiftcoro [[ALLOCATOR]], ptr [[ADDRESS]])
// CHECK: ret void
// CHECK: }