From 3119c1befb3eba83fd802748722b6331bd89a590 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 28 Oct 2026 13:20:17 +0100 Subject: [PATCH 1/2] Compile a representative kernel in the precompilation workload The workload only compiled an empty kernel, so many LLVM.jl and GPUCompiler methods that every back-end calls while compiling a kernel (IR walkers, debug-info and metadata accessors, exception lowering) were only inferred into GPUCompiler's package image, without native code. Back-ends reach them through dynamic dispatch, and their own images cannot provide that code (JuliaLang/julia#58465), so the first kernel of every session compiled them. Compile a kernel with bounds-checked array accesses, a loop, Float16 arithmetic and math intrinsics instead. This removes about 20 methods from what Metal.jl and CUDA.jl compile for their first kernel. --- src/precompile.jl | 42 +++++++++++++++++++++++++++++++++++++++--- 1 file changed, 39 insertions(+), 3 deletions(-) diff --git a/src/precompile.jl b/src/precompile.jl index 9f35c128..bec3254f 100644 --- a/src/precompile.jl +++ b/src/precompile.jl @@ -22,17 +22,53 @@ import REPL const DummyCompilerJob = CompilerJob{NativeCompilerTarget, DummyCompilerParams} GPUCompiler.runtime_module(::DummyCompilerJob) = DummyRuntime - end - kernel() = nothing + # the runtime library isn't linked (see below), so don't reject calls to it + GPUCompiler.isintrinsic(::DummyCompilerJob, fn::String) = startswith(fn, "gpu_") + + # a minimal device array, so that kernels contain the bounds checks, exceptions and + # memory accesses that back-ends' array types generate. (back-ends override Base's + # `throw_boundserror`, which allocates the error, to throw without allocating.) + struct DeviceArray{T} <: DenseVector{T} + ptr::Ptr{T} + len::Int + end + Base.size(a::DeviceArray) = (a.len,) + @noinline throw_boundserror() = throw(nothing) + Base.@propagate_inbounds function Base.getindex(a::DeviceArray, i::Int) + @boundscheck 1 <= i <= a.len || throw_boundserror() + unsafe_load(a.ptr, i) + end + Base.@propagate_inbounds function Base.setindex!(a::DeviceArray, x, i::Int) + @boundscheck 1 <= i <= a.len || throw_boundserror() + unsafe_store!(a.ptr, x, i) + return + end + + # exercise the parts of the pipeline that back-ends' kernels typically need: bounds + # checks lowered to exceptions, loops, Float16 arithmetic and math intrinsics + function kernel(a::DeviceArray{Float32}, b::DeviceArray{Float16}, n::Int) + acc = 0f0 + for i in 1:n + x = a[i] + acc += fma(sqrt(abs(x)), x, floor(x)) + b[i] = b[i] * Float16(2) + Float16(x) + end + a[1] = min(acc, 1f0) + return + end + end @compile_workload begin - source = methodinstance(typeof(kernel), Tuple{}) target = NativeCompilerTarget() params = precompile_module.DummyCompilerParams() # XXX: on Windows, compiling the GPU runtime leaks GPU code in the native cache, # so prevent building the runtime library (see JuliaGPU/GPUCompiler.jl#601) config = CompilerConfig(target, params; libraries=false) + + tt = Tuple{precompile_module.DeviceArray{Float32}, + precompile_module.DeviceArray{Float16}, Int} + source = methodinstance(typeof(precompile_module.kernel), tt) job = CompilerJob(source, config) JuliaContext() do ctx From 482f804b3337ab062ac9b7ed38f9058ffcda8422 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 28 Oct 2026 13:20:17 +0100 Subject: [PATCH 2/2] Only load the tree-sitter grammars in the Highlights extension Loading the four grammar JLLs took about 10 ms of every `using GPUCompiler`, while only HighlightsExt uses them. They stay regular dependencies so that they're installed for the extension, which loads them instead. --- ext/HighlightsExt.jl | 14 ++++++++++++-- src/reflection.jl | 20 ++++++-------------- test/aqua.jl | 4 +++- 3 files changed, 21 insertions(+), 17 deletions(-) diff --git a/ext/HighlightsExt.jl b/ext/HighlightsExt.jl index 712e0646..4978036b 100644 --- a/ext/HighlightsExt.jl +++ b/ext/HighlightsExt.jl @@ -7,9 +7,19 @@ module HighlightsExt using GPUCompiler import Highlights +using tree_sitter_llvm_jll, tree_sitter_ptx_jll, tree_sitter_spirv_jll, + tree_sitter_gcn_jll -render(buf::IO, code, language, theme) = - Highlights.highlight(buf, MIME("text/ansi"), code, language, theme) +# tree-sitter grammars for `GPUCompiler.highlight_languages` +const languages = Dict( + "llvm" => tree_sitter_llvm_jll, + "ptx" => tree_sitter_ptx_jll, + "spirv" => tree_sitter_spirv_jll, + "gcn" => tree_sitter_gcn_jll, +) + +render(buf::IO, code, lexer, theme) = + Highlights.highlight(buf, MIME("text/ansi"), code, languages[lexer], theme) function __init__() GPUCompiler.highlighter[] = render diff --git a/src/reflection.jl b/src/reflection.jl index 9fbc7a60..d86170ad 100644 --- a/src/reflection.jl +++ b/src/reflection.jl @@ -11,17 +11,10 @@ import InteractiveUtils # needed interactively. The extension publishes its renderer through this hook. const highlighter = Ref{Union{Nothing,Function}}(nothing) -using tree_sitter_llvm_jll, tree_sitter_ptx_jll, tree_sitter_spirv_jll, - tree_sitter_gcn_jll - -# tree-sitter grammars for the assembly formats we emit, keyed by the -# `source_code` of the compiler target -const highlight_languages = Dict( - "llvm" => tree_sitter_llvm_jll, - "ptx" => tree_sitter_ptx_jll, - "spirv" => tree_sitter_spirv_jll, - "gcn" => tree_sitter_gcn_jll, -) +# the assembly formats we have tree-sitter grammars for, by the `source_code` of the +# compiler target. the grammars are dependencies of this package, so that they're installed +# along with it, but only the extension loads them. +const highlight_languages = ("llvm", "ptx", "spirv", "gcn") # Set to a Highlights.jl theme, or `nothing` to detect the terminal background. const highlight_theme = Ref{Any}(nothing) @@ -150,8 +143,7 @@ end function highlight(io::IO, code, lexer) have_color = get(io, :color, false) - language = get(highlight_languages, lexer, nothing) - if !have_color || language === nothing + if !have_color || !in(lexer, highlight_languages) print(io, code) return end @@ -169,7 +161,7 @@ function highlight(io::IO, code, lexer) # partial output buf = IOBuffer() try - Base.invokelatest(render, buf, code, language, selected_highlight_theme(io)) + Base.invokelatest(render, buf, code, lexer, selected_highlight_theme(io)) write(io, take!(buf)) catch err @warn "Failed to highlight $lexer code" exception=(err, catch_backtrace()) maxlog=1 diff --git a/test/aqua.jl b/test/aqua.jl index 54c6510d..927bc3a7 100644 --- a/test/aqua.jl +++ b/test/aqua.jl @@ -1,4 +1,6 @@ using Aqua +# the tree-sitter grammars are only loaded by the HighlightsExt extension Aqua.test_all(GPUCompiler; - stale_deps=(ignore=[:Tracy],),) + stale_deps=(ignore=[:Tracy, :tree_sitter_gcn_jll, :tree_sitter_llvm_jll, + :tree_sitter_ptx_jll, :tree_sitter_spirv_jll],),)