diff --git a/Project.toml b/Project.toml index a821d19b..fdaca7da 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "GPUCompiler" uuid = "61eb1bfa-7361-4325-ad38-22787b887f55" -version = "2.5.4" +version = "2.6.0" authors = ["Tim Besard "] [workspace] diff --git a/src/GPUCompiler.jl b/src/GPUCompiler.jl index 5721c871..b74253ca 100644 --- a/src/GPUCompiler.jl +++ b/src/GPUCompiler.jl @@ -46,12 +46,27 @@ import CompilerCaching using ScopedValues: ScopedValue, with -# Hook used by the `@device_code_*` macros. Scope it to the current task and its -# children so concurrent reflection calls do not interfere. Defined here so the -# legacy `cached_compilation` in deprecated.jl can reference it. +""" + compile_hook + +Hook used by the `@device_code_*` macros to observe compilations in the current task +and its children. When non-`nothing`, call it with the compilation job using +`Base.invokelatest(compile_hook[], job)`. `GPUCompiler.compile` does this itself; +back-ends with their own pipeline must call it even when compilation results are cached. + +Jobs must have stable `hash` and `isequal` semantics for deduplication, support `show`, +and implement the requested GPUCompiler reflection functions: `code_lowered(job)`, +`code_typed(job; kwargs...)`, or `code_warntype`, `code_llvm`, and `code_native` with +`(io, job; kwargs...)`. Jobs need not be `CompilerJob`s or use LLVM. The all-stage +`@device_code` dump still requires a `CompilerJob` and the LLVM pipeline. + +Reflection runs with the hook disabled to avoid recursion. Wait for child tasks before +leaving a reflection macro so their output is included. +""" const compile_hook = ScopedValue{Union{Nothing,Function}}(nothing) include("utils.jl") +@public compile_hook, emit_hooked_compilation include("mangling.jl") # compiler interface and implementations diff --git a/src/reflection.jl b/src/reflection.jl index e84da2a9..1a422f33 100644 --- a/src/reflection.jl +++ b/src/reflection.jl @@ -278,6 +278,7 @@ function code_llvm(io::IO, @nospecialize(job::CompilerJob); optimize::Bool=job.c highlight(io, str, "llvm") end code_llvm(@nospecialize(job::CompilerJob); kwargs...) = code_llvm(stdout, job; kwargs...) +code_llvm(io::IO, @nospecialize(job); kwargs...) = throw(unsupported_reflection(:code_llvm, job)) """ code_native([io], f, types; cap::VersionNumber, kernel=false, raw=false) @@ -302,13 +303,35 @@ function code_native(io::IO, @nospecialize(job::CompilerJob); end code_native(@nospecialize(job::CompilerJob); kwargs...) = code_native(stdout, job; kwargs...) +code_native(io::IO, @nospecialize(job); kwargs...) = throw(unsupported_reflection(:code_native, job)) + +unsupported_reflection(f::Symbol, @nospecialize(job)) = + ArgumentError("$f is not supported for jobs of type $(typeof(job))") # # @device_code_* functions # -function emit_hooked_compilation(inner_hook, ex...) +""" + emit_hooked_compilation(hook, ex...; job_filter=Returns(true)) -> Expr + +Build the body of a `@device_code_*` macro: an expression that evaluates the user's +code (the last element of `ex`) with `hook` installed as the [`compile_hook`](@ref), +calling `hook(job; kwargs...)` once per distinct job with the remaining elements of +`ex` as keyword arguments. Back-ends define stage-specific macros with it, e.g. + +```julia +macro device_code_ptx(ex...) + hook = (job; io::IO=stdout) -> code_ptx(io, job) + GPUCompiler.emit_hooked_compilation(hook, ex...) +end +``` + +`job_filter(job)` selects which jobs to inspect. Jobs rejected by the filter do not +count toward the check that at least one kernel was observed. +""" +function emit_hooked_compilation(inner_hook, ex...; job_filter=Returns(true)) user_code = ex[end] user_kwargs = ex[1:end-1] quote @@ -316,6 +339,7 @@ function emit_hooked_compilation(inner_hook, ex...) jobs = Set() jobs_lock = ReentrantLock() function outer_hook(job) + $job_filter(job) || return Base.@lock jobs_lock begin job in jobs && return push!(jobs, job) @@ -350,7 +374,7 @@ See also: `InteractiveUtils.@code_lowered` macro device_code_lowered(ex...) quote buf = Any[] - function hook(job::CompilerJob) + function hook(job) append!(buf, code_lowered(job)) end $(emit_hooked_compilation(:hook, ex...)) @@ -368,8 +392,8 @@ See also: `InteractiveUtils.@code_typed` """ macro device_code_typed(ex...) quote - output = Dict{CompilerJob,Any}() - function hook(job::CompilerJob; kwargs...) + output = Dict{Any,Any}() + function hook(job; kwargs...) output[job] = code_typed(job; kwargs...) end $(emit_hooked_compilation(:hook, ex...)) @@ -386,7 +410,7 @@ Evaluates the expression `ex` and prints the result of See also: `InteractiveUtils.@code_warntype` """ macro device_code_warntype(ex...) - function hook(job::CompilerJob; io::IO=stdout, kwargs...) + function hook(@nospecialize(job); io::IO=stdout, kwargs...) println(io, "$job") println(io) code_warntype(io, job; kwargs...) @@ -404,7 +428,7 @@ to `io` for every compiled GPU kernel. For other supported keywords, see See also: InteractiveUtils.@code_llvm """ macro device_code_llvm(ex...) - function hook(job::CompilerJob; io::IO=stdout, kwargs...) + function hook(@nospecialize(job); io::IO=stdout, kwargs...) println(io, "; $job") code_llvm(io, job; kwargs...) end @@ -419,7 +443,7 @@ for every compiled GPU kernel. For other supported keywords, see [`GPUCompiler.code_native`](@ref). """ macro device_code_native(ex...) - function hook(job::CompilerJob; io::IO=stdout, kwargs...) + function hook(@nospecialize(job); io::IO=stdout, kwargs...) println(io, "// $job") println(io) code_native(io, job; kwargs...) @@ -431,11 +455,13 @@ end @device_code dir::AbstractString=... [...] ex Evaluates the expression `ex` and dumps all intermediate forms of code to the directory -`dir`. +`dir`. This dump requires a `CompilerJob` and includes LLVM IR; use the individual +`@device_code_*` macros for back-ends with other compilation stages. """ macro device_code(ex...) localUnique = 1 - function hook(job::CompilerJob; dir::AbstractString) + function hook(@nospecialize(job); dir::AbstractString) + job isa CompilerJob || throw(unsupported_reflection(:device_code, job)) name = job.source.def.name fn = "$(name)_$(localUnique)" mkpath(dir) diff --git a/src/reflection_compat.jl b/src/reflection_compat.jl index 92467a3d..0e00af8d 100644 --- a/src/reflection_compat.jl +++ b/src/reflection_compat.jl @@ -31,84 +31,96 @@ end function code_warntype_by_type(io::IO, @nospecialize(tt); debuginfo::Symbol=:default, optimize::Bool=false, kwargs...) - debuginfo = Base.IRShow.debuginfo(debuginfo) - lineprinter = Base.IRShow.__debuginfo[debuginfo] for (src, rettype) in Base.code_typed_by_type(tt; optimize, kwargs...) - if !(src isa Core.CodeInfo) + if src isa Core.CodeInfo + code_warntype(io, src, rettype; debuginfo) + else println(io, src) println(io, " failed to infer") - continue end - lambda_io::IOContext = io - p = src.parent - nargs::Int = 0 - if p isa Core.MethodInstance - println(io, p) - print(io, " from ") - println(io, p.def) - p.def isa Method && (nargs = p.def.nargs) - if !isempty(p.sparam_vals) - println(io, "Static Parameters") - sig = p.def.sig - warn_color = Base.warn_color() # more mild user notification - for i = 1:length(p.sparam_vals) - sig = sig::UnionAll - name = sig.var.name - val = p.sparam_vals[i] - print_highlighted(io::IO, v::String, color::Symbol) = - if highlighting[:warntype] - Base.printstyled(io, v; color) - else - Base.print(io, v) - end - if val isa TypeVar - if val.lb === Union{} - print(io, " ", name, " <: ") - print_highlighted(io, "$(val.ub)", warn_color) - elseif val.ub === Any - print(io, " ", sig.var.name, " >: ") - print_highlighted(io, "$(val.lb)", warn_color) - else - print(io, " ") - print_highlighted(io, "$(val.lb)", warn_color) - print(io, " <: ", sig.var.name, " <: ") - print_highlighted(io, "$(val.ub)", warn_color) - end - elseif val isa typeof(Vararg) - print(io, " ", name, "::") - print_highlighted(io, "Int", warn_color) + end + nothing +end + +""" + code_warntype(io::IO, src::Core.CodeInfo, rettype; debuginfo=:default) + +Print already-inferred code with type-instability highlighting. Back-ends can use +this method to display a particular specialization without running inference again. +""" +function code_warntype(io::IO, src::Core.CodeInfo, @nospecialize(rettype); + debuginfo::Symbol=:default) + debuginfo = Base.IRShow.debuginfo(debuginfo) + lineprinter = Base.IRShow.__debuginfo[debuginfo] + lambda_io::IOContext = io + p = src.parent + nargs::Int = 0 + if p isa Core.MethodInstance + println(io, p) + print(io, " from ") + println(io, p.def) + p.def isa Method && (nargs = p.def.nargs) + if !isempty(p.sparam_vals) + println(io, "Static Parameters") + sig = p.def.sig + warn_color = Base.warn_color() # more mild user notification + for i = 1:length(p.sparam_vals) + sig = sig::UnionAll + name = sig.var.name + val = p.sparam_vals[i] + print_highlighted(io::IO, v::String, color::Symbol) = + if highlighting[:warntype] + Base.printstyled(io, v; color) else - print(io, " ", sig.var.name, " = ") - print_highlighted(io, "$(val)", :cyan) # show the "good" type + Base.print(io, v) end - println(io) - sig = sig.body + if val isa TypeVar + if val.lb === Union{} + print(io, " ", name, " <: ") + print_highlighted(io, "$(val.ub)", warn_color) + elseif val.ub === Any + print(io, " ", sig.var.name, " >: ") + print_highlighted(io, "$(val.lb)", warn_color) + else + print(io, " ") + print_highlighted(io, "$(val.lb)", warn_color) + print(io, " <: ", sig.var.name, " <: ") + print_highlighted(io, "$(val.ub)", warn_color) + end + elseif val isa typeof(Vararg) + print(io, " ", name, "::") + print_highlighted(io, "Int", warn_color) + else + print(io, " ", sig.var.name, " = ") + print_highlighted(io, "$(val)", :cyan) # show the "good" type end + println(io) + sig = sig.body end end - if src.slotnames !== nothing - slotnames = Base.sourceinfo_slotnames(src) - lambda_io = IOContext(lambda_io, :SOURCE_SLOTNAMES => slotnames) - slottypes = src.slottypes - nargs > 0 && println(io, "Arguments") - for i = 1:length(slotnames) - if i == nargs + 1 - println(io, "Locals") - end - print(io, " ", slotnames[i]) - if isa(slottypes, Vector{Any}) - InteractiveUtils.warntype_type_printer(io; type=slottypes[i], used=true) - end - println(io) + end + if src.slotnames !== nothing + slotnames = Base.sourceinfo_slotnames(src) + lambda_io = IOContext(lambda_io, :SOURCE_SLOTNAMES => slotnames) + slottypes = src.slottypes + nargs > 0 && println(io, "Arguments") + for i = 1:length(slotnames) + if i == nargs + 1 + println(io, "Locals") + end + print(io, " ", slotnames[i]) + if isa(slottypes, Vector{Any}) + InteractiveUtils.warntype_type_printer(io; type=slottypes[i], used=true) end + println(io) end - print(io, "Body") - InteractiveUtils.warntype_type_printer(io; type=rettype, used=true) - println(io) - - irshow_config = Base.IRShow.IRShowConfig(lineprinter(src), InteractiveUtils.warntype_type_printer) - Base.IRShow.show_ir(lambda_io, src, irshow_config) - println(io) end + print(io, "Body") + InteractiveUtils.warntype_type_printer(io; type=rettype, used=true) + println(io) + + irshow_config = Base.IRShow.IRShowConfig(lineprinter(src), InteractiveUtils.warntype_type_printer) + Base.IRShow.show_ir(lambda_io, src, irshow_config) + println(io) nothing end diff --git a/test/native.jl b/test/native.jl index eadc3f1e..ceb780eb 100644 --- a/test/native.jl +++ b/test/native.jl @@ -90,6 +90,75 @@ end end @test inherited[] === hook @test fetch(Threads.@spawn GPUCompiler.compile_hook[]) === nothing + + # the macros accept jobs of back-ends that do not compile through LLVM, as long + # as they fire the hook and implement the reflection functions of their stages + mod3 = @eval module $(gensym()) + struct ForeignJob + name::Symbol + end + Base.show(io::IO, job::ForeignJob) = print(io, "ForeignJob(", job.name, ")") + $GPUCompiler.code_lowered(job::ForeignJob) = Any[job.name] + function $GPUCompiler.code_typed(job::ForeignJob; marker=job.name) + @assert $GPUCompiler.compile_hook[] === nothing + Any[marker] + end + $GPUCompiler.code_native(io::IO, job::ForeignJob; marker=job.name) = print(io, marker) + $GPUCompiler.code_warntype(io::IO, job::ForeignJob) = println(io, "typed ", job.name) + macro foreign_code(ex...) + hook = (job; io::IO=stdout) -> print(io, job.name) + $GPUCompiler.emit_hooked_compilation(hook, ex...; job_filter=job -> job isa ForeignJob) + end + report(job) = $GPUCompiler.compile_hook[] === nothing ? nothing : + Base.invokelatest($GPUCompiler.compile_hook[], job) + function filtered(io) + @foreign_code io=io begin + report(:other_backend) + report(ForeignJob(:foreign)) + report(ForeignJob(:foreign)) + end + end + rejected() = @foreign_code report(:other_backend) + end + typed = GPUCompiler.@device_code_typed begin + mod3.report(mod3.ForeignJob(:foreign)) + mod3.report(mod3.ForeignJob(:foreign)) + end + @test typed == Dict(mod3.ForeignJob(:foreign) => Any[:foreign]) + @test sprint(mod3.filtered) == "foreign" + @test_throws "no kernels executed" mod3.rejected() + warntype = sprint() do io + GPUCompiler.@device_code_warntype io=io mod3.report(mod3.ForeignJob(:foreign)) + end + @test occursin("ForeignJob(foreign)", warntype) && occursin("typed foreign", warntype) + @test GPUCompiler.@device_code_lowered(mod3.report(mod3.ForeignJob(:foreign))) == [:foreign] + typed = GPUCompiler.@device_code_typed marker=:forwarded begin + mod3.report(mod3.ForeignJob(:foreign)) + end + @test only(values(typed)) == [:forwarded] + native = sprint() do io + GPUCompiler.@device_code_native io=io marker=:assembly mod3.report(mod3.ForeignJob(:foreign)) + end + @test occursin("assembly", native) + @test_throws "code_llvm is not supported" GPUCompiler.@device_code_llvm( + io=devnull, mod3.report(mod3.ForeignJob(:foreign))) + @test_throws "code_native is not supported" GPUCompiler.code_native(devnull, :unsupported) + mktempdir() do dir + @test_throws "device_code is not supported" GPUCompiler.@device_code( + dir=dir, mod3.report(mod3.ForeignJob(:foreign))) + @test isempty(readdir(dir)) + end + @test GPUCompiler.compile_hook[] === nothing + + # One expression can report both CompilerJobs and a back-end's own job type. + mixed = GPUCompiler.@device_code_typed begin + Native.code_execution(mod.f, (Int,)) + mod3.report(mod3.ForeignJob(:foreign)) + mod3.report(mod3.ForeignJob(:foreign)) + end + @test length(mixed) == 2 + @test mixed[mod3.ForeignJob(:foreign)] == [:foreign] + @test count(job -> job isa CompilerJob, keys(mixed)) == 1 end @testset "method instances for type-valued callees and arguments" begin