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