diff --git a/Project.toml b/Project.toml index 1777ec22..52ff225d 100644 --- a/Project.toml +++ b/Project.toml @@ -9,6 +9,7 @@ ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197" Atomix = "a9b6321e-bd34-4604-b9c9-b65b8de01458" GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" +PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a" [weakdeps] CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" @@ -33,5 +34,6 @@ GPUArraysCore = "0.2.1" KernelAbstractions = "0.9.43, 0.10" Metal = "1.10.1" OpenCL = "0.10, 1" +PrecompileTools = "1" julia = "1.10" oneAPI = "1, 2" diff --git a/ext/CUDACoreExt.jl b/ext/CUDACoreExt.jl index b582963f..5bd5fdb5 100644 --- a/ext/CUDACoreExt.jl +++ b/ext/CUDACoreExt.jl @@ -2,6 +2,7 @@ module CUDACoreExt using CUDACore import AcceleratedKernels as AK +using PrecompileTools: @setup_workload, @compile_workload # Groups of threads read runs of up to 4096 Float32s: on an RTX 5080, summing a 4096×4096 matrix # along `dims=1` takes 17% less time than with a block per output. @@ -21,4 +22,31 @@ AK.sort_tuning(::CUDABackend, ::Type{T}) where {T} = radix_min_slice_len=1 << 19) : AK.SortTuning(bitonic_max_len=512, radix_min_len=1 << 23, bitonic_max_slice_len=256) + +# CUDACore's own workload can't cache these operations: it doesn't use the GPU, and this +# extension isn't loaded while CUDACore precompiles, so its kernels would use the default +# tunings rather than the ones above. The arrays are just large enough to select the kernels +# that a vector of 10^6 elements and a 1000×1000 matrix do. Without a functional GPU while +# precompiling, these operations are only compiled when first used. Like CUDACore's own +# workload, this needs 1.12: before, GPU code compiled while precompiling leaks into the host +# code of the image. +@static if VERSION >= v"1.12-" +@setup_workload begin + if CUDACore.functional() + @compile_workload begin + m = CuArray(ones(Float32, 256, 32)) + v = vec(m) + sum(v) + sum(m; dims=1) + sum(m; dims=2) + cumsum(v) + any(iszero, v) + all(iszero, v) + sum(CuArray(ones(Int, 8192))) + synchronize() + end + end +end +end + end