From a1637c4aa085ae8c416825658298e504f4518d41 Mon Sep 17 00:00:00 2001 From: Simon Byrne Date: Thu, 11 Jan 2024 12:11:44 -0800 Subject: [PATCH] Add hook for inference (#44) * add hooks into compiler inference * add docs * update tests * bump tag --- Project.toml | 2 +- docs/src/api.md | 1 + docs/src/index.md | 2 +- src/NVTX.jl | 1 + src/julia.jl | 34 +++++++++++++++++++++++++++++++--- test/runtests.jl | 8 ++++---- 6 files changed, 39 insertions(+), 9 deletions(-) diff --git a/Project.toml b/Project.toml index 8a1fc2d..d249380 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "NVTX" uuid = "5da4648a-3479-48b8-97b9-01cb529c0a1f" authors = ["Simon Byrne "] -version = "0.3.3" +version = "0.3.4" [deps] Colors = "5ae59095-9a9b-59fe-a467-6f913c188581" diff --git a/docs/src/api.md b/docs/src/api.md index fc181d6..25b0d57 100644 --- a/docs/src/api.md +++ b/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() ``` diff --git a/docs/src/index.md b/docs/src/index.md index 620556c..251b465 100644 --- a/docs/src/index.md +++ b/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 diff --git a/src/NVTX.jl b/src/NVTX.jl index aa415b3..edccdf8 100644 --- a/src/NVTX.jl +++ b/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__() diff --git a/src/julia.jl b/src/julia.jl index d650777..47a306e 100644 --- a/src/julia.jl +++ b/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 diff --git a/test/runtests.jl b/test/runtests.jl index 43ff925..da0d958 100644 --- a/test/runtests.jl +++ b/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]