Support backend-defined jobs in reflection hooks

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 <tim.besard@gmail.com>"]
 
 [workspace]

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

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,20 +303,43 @@ 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
         # The job set and hook output are shared by child tasks, so update them together.
         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)

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

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