Compile a representative kernel in the precompilation workload

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

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

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

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],),)