From 65b126312e36146ced938ef8ff1b3b25afd7902f Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 28 Oct 2026 13:49:57 +0100 Subject: [PATCH] Precompile the Metal operations that depend on the extension's tunings Metal's package image can't hold AcceleratedKernels' operations as they run on Metal: this extension isn't loaded while Metal precompiles, so the kernels its workload compiles use the default tunings, and the first `sum`, `cumsum` or `any` of a session compiles its host code and kernels again. A `@compile_workload` in the extension caches them with the tunings that apply. It covers whole-array reductions (both passes), reductions of a matrix along either dimension, `cumsum`, `any` and `all` of Float32 arrays, and the sum of an Int array. On an M1 with Julia 1.13, the first call on 10^6 elements or a 1000x1000 matrix takes 50 ms instead of 450 for `sum`, 80 instead of 710 for `cumsum`, 100 instead of 730 for `any`, and 110 and 30 instead of 290 and 200 for `sum` along `dims=1` and `dims=2`. The extension's image grows by 6.5 MB, and loading Metal takes no measurable extra time. As in Metal's own workload, it runs on Julia 1.11 and later only, and only when Metal's workload launches kernels too, which it reports with `Metal.can_precompile_launches` (JuliaGPU/Metal.jl#1124). With older versions of Metal, which can't tell and implement some of these operations differently, the workload is skipped. --- Project.toml | 2 ++ ext/MetalExt.jl | 28 ++++++++++++++++++++++++++++ 2 files changed, 30 insertions(+) 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