Add hook for inference (#44)

* add hooks into compiler inference
* add docs
* update tests
* bump tag

Project.toml

@@ -1,7 +1,7 @@
 name = "NVTX"
 uuid = "5da4648a-3479-48b8-97b9-01cb529c0a1f"
 authors = ["Simon Byrne <simonbyrne@gmail.com>"]
-version = "0.3.3"
+version = "0.3.4"
 
 [deps]
 Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"

docs/src/api.md

@@ -19,6 +19,7 @@ The following macros provide the most convenient usage of this package
 
 ```@docs
 enable_gc_hooks()
+enable_inference_hook
 name_threads_julia()
 ```
 

docs/src/index.md

@@ -26,7 +26,7 @@ end
 
 ## Instrumenting Julia garbage collector
 
-Additionally, it is possible to annotate the Julia garbage collector by calling [`NVTX.enable_gc_hooks()`](@ref), or setting the `JULIA_NVTX_CALLBACKS` environment variable.
+Additionally, it is possible to annotate the Julia garbage collector by calling [`NVTX.enable_gc_hooks()`](@ref) and [`NVTX.enable_inference_hook()`](@ref), or setting the `JULIA_NVTX_CALLBACKS` environment variable.
 
 ## Running the profiler
 

src/NVTX.jl

@@ -31,6 +31,7 @@ function activate()
         alloc="alloc" in callbacks,
         free="free" in callbacks
     )
+    enable_inference_hook("inference" in callbacks)
 end
 
 function __init__()

src/julia.jl

@@ -43,6 +43,7 @@ const GC_FREE_MESSAGE = StringHandle(JULIA_DOMAIN, "free")
 const GC_COLOR = Ref{UInt32}(Colors.ARGB32(Colors.colorant"brown").color)
 const GC_ALLOC_COLOR = Ref{UInt32}(Colors.ARGB32(Colors.colorant"goldenrod1").color)
 const GC_FREE_COLOR = Ref{UInt32}(Colors.ARGB32(Colors.colorant"dodgerblue").color)
+const INFERENCE_COLOR = Ref{UInt32}(Colors.ARGB32(Colors.colorant"turquoise3").color)
 
 """
     NVTX.enable_gc_hooks(;gc=true, alloc=false, free=false)
@@ -71,9 +72,9 @@ function enable_gc_hooks(;gc::Bool=true, alloc::Bool=false, free::Bool=false)
         unsafe_store!(cglobal((:gc_message,libjulia_nvtx_callbacks),Ptr{Cvoid}), GC_MESSAGE.ptr)
         unsafe_store!(cglobal((:gc_color,libjulia_nvtx_callbacks),UInt32), GC_COLOR[])
         # https://github.com/JuliaLang/julia/blob/v1.8.3/src/julia.h#L879-L883
-        name_category(JULIA_DOMAIN, 1+0, "auto")
-        name_category(JULIA_DOMAIN, 1+1, "full")
-        name_category(JULIA_DOMAIN, 1+2, "incremental")
+        name_category(JULIA_DOMAIN, 1+0, "GC auto")
+        name_category(JULIA_DOMAIN, 1+1, "GC full")
+        name_category(JULIA_DOMAIN, 1+2, "GC incremental")
     end
     if alloc
         init!(GC_ALLOC_MESSAGE)
@@ -97,3 +98,30 @@ function enable_gc_hooks(;gc::Bool=true, alloc::Bool=false, free::Bool=false)
     return nothing
 end
 
+typeinf_ext_nvtx(mi::Base.Core.MethodInstance, world::UInt) = typeinf_ext_nvtx(Base.Core.Compiler.NativeInterpreter(world), mi)
+function typeinf_ext_nvtx(interp::Base.Core.Compiler.AbstractInterpreter, linfo::Base.Core.MethodInstance)
+    method = linfo.def
+    types = linfo.specTypes.parameters[2:end]
+    message = "$(method.name)($(join([string("::", t) for t in types], ", "))) @ $(method.module) $(method.file):$(method.line)"
+    id = range_start(JULIA_DOMAIN; message, color = INFERENCE_COLOR[], category = 11)
+    ret = Core.Compiler.typeinf_ext_toplevel(interp, linfo)
+    range_end(id)
+    return ret
+end
+precompile(typeinf_ext_nvtx, (Base.Core.Compiler.NativeInterpreter, Base.Core.MethodInstance))
+precompile(typeinf_ext_nvtx, (Base.Core.MethodInstance, UInt))
+
+"""
+    NVTX.enable_inference_hook(active::Bool=true)
+
+Add hooks for method inference. Can also be activated by adding `inference` to
+the `JULIA_NVTX_CALLBACKS` environment variable.
+"""
+function enable_inference_hook(enable::Bool=true)
+    if enable
+        name_category(JULIA_DOMAIN, 11, "compiler inference")
+        ccall(:jl_set_typeinf_func, Cvoid, (Any,), typeinf_ext_nvtx)
+    else
+        ccall(:jl_set_typeinf_func, Cvoid, (Any,), Core.Compiler.typeinf_ext_toplevel)
+    end
+end

test/runtests.jl

@@ -7,7 +7,7 @@ end
 
 nsys = get(ENV, "JULIA_NSYS", "nsys")
 
-run(`$nsys profile --env-var="JULIA_NVTX_CALLBACKS=gc|alloc|free" --output=$(joinpath(dirname, "basic")) --export=json,sqlite --trace=nvtx $(Base.julia_cmd()) --project=$(Base.active_project()) --threads=3 basic.jl`)
+run(`$nsys profile --env-var="JULIA_NVTX_CALLBACKS=gc|alloc|free|inference" --output=$(joinpath(dirname, "basic")) --export=json,sqlite --trace=nvtx $(Base.julia_cmd()) --project=$(Base.active_project()) --threads=3 basic.jl`)
 
 using DataFrames, SQLite, DBInterface, Colors, Test
 
@@ -84,14 +84,14 @@ julia_categories = DataFrame(DBInterface.execute(db, """
     WHERE eventType = $NvtxCategory AND domainId = $julia_domainId
     ORDER BY category
     """))
-@test julia_categories.category == [1, 2, 3]
-@test julia_categories.text == ["auto", "full", "incremental"]
+@test julia_categories.category == [1, 2, 3, 11]
+@test julia_categories.text == ["GC auto", "GC full", "GC incremental", "compiler inference"]
 
 julia_ranges = DataFrame(DBInterface.execute(db, """
     SELECT COALESCE(text, value) as text, category, color
     FROM NVTX_EVENTS
     LEFT JOIN StringIds on textId == id
-    WHERE eventType = $NvtxPushPopRange AND domainId = $julia_domainId AND category > 1
+    WHERE eventType = $NvtxPushPopRange AND domainId = $julia_domainId AND category > 1 AND category < 10
     ORDER BY start
     """)) # exclude auto GC
 @test julia_ranges.text == ["GC" for i = 1:2]