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/MetalExt.jl b/ext/MetalExt.jl index 2870a23d..7489c4af 100644 --- a/ext/MetalExt.jl +++ b/ext/MetalExt.jl @@ -2,6 +2,7 @@ module MetalExt using Metal import AcceleratedKernels as AK +using PrecompileTools: @setup_workload, @compile_workload # Scan tiles fit half of the 32 KiB of threadgroup memory. Metal's shader validation doubles every # kernel's threadgroup memory, so larger tiles can't run under it, and they aren't faster: on an @@ -31,4 +32,31 @@ AK.sort_tuning(::MetalBackend, ::Type{T}) where {T} = radix_min_slice_len=1 << 15) : AK.SortTuning(bitonic_max_len=512, radix_min_len=1 << 18, bitonic_max_slice_len=1024) + +# Metal's own workload can't cache AcceleratedKernels' operations as they run with the tunings +# above: this extension isn't loaded while Metal precompiles, so the kernels compiled there use +# the default tunings. Before Julia 1.11, GPU code cached in a package image shadows the host +# versions of the methods it shares with host code (see Metal's workload). +@static if VERSION >= v"1.11" +@setup_workload begin + # only when Metal's own workload launches kernels too (older versions of Metal can't tell, + # and implement some of these operations differently, so skip the workload there) + can_launch = isdefined(Metal, :can_precompile_launches) && Metal.can_precompile_launches() + @compile_workload if can_launch + # the vector spans several blocks, so that whole-array reductions take both passes, and + # the matrix is large enough for the kernels that reduce large matrices along dims + let m = MtlArray(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) + end + sum(MtlArray(ones(Int, 1024))) + synchronize() + end +end +end + end