Skip to content

Precompile the Metal operations that depend on the extension's tunings #200

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
maleadt wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Precompile the Metal operations that depend on the extension's tunings #200

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions Project.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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"
28 changes: 28 additions & 0 deletions ext/MetalExt.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Loading