From c47d0e6fbcf4ddaa0aafdd786c1ba7a6b4a22ab8 Mon Sep 17 00:00:00 2001 From: Valentin Churavy Date: Sun, 4 Oct 2026 09:18:07 +0200 Subject: [PATCH 1/9] Add a low-cost tracing subsystem `@profiling_range` and `profiling_mark` put named ranges and markers on the timeline of a tracing profiler, and kernel launches are annotated with the kernel's name. Profilers like NVTX, ITT and roctx annotate host threads and correlate device work themselves, so ranges go to every registered `Tracer` regardless of backend. With none registered, an annotation costs one atomic load. Tracers: NVTX (NVTXExt), roctx (ROCTXExt, via AMDGPU), Intel ITT (IntelITTExt), each registering only under its profiler, and a built-in NVTXT writer enabled with `JULIA_KA_NVTXT`. Supersedes #703 and #66. Assisted-by: Claude Code (Opus 5.5) --- Project.toml | 9 ++ docs/make.jl | 1 + docs/src/extras/profiling.md | 80 ++++++++++ ext/IntelITTExt.jl | 35 +++++ ext/NVTXExt.jl | 34 +++++ ext/ROCTXExt.jl | 76 ++++++++++ src/KernelAbstractions.jl | 7 + src/backend_launch.jl | 18 +++ src/profiling.jl | 282 +++++++++++++++++++++++++++++++++++ test/Project.toml | 2 + test/profiling.jl | 53 +++++++ test/runtests.jl | 156 +++++++++++++++++++ test/testsuite.jl | 5 + 13 files changed, 758 insertions(+) create mode 100644 docs/src/extras/profiling.md create mode 100644 ext/IntelITTExt.jl create mode 100644 ext/NVTXExt.jl create mode 100644 ext/ROCTXExt.jl create mode 100644 src/profiling.jl create mode 100644 test/profiling.jl diff --git a/Project.toml b/Project.toml index 4400dd267..091ce2340 100644 --- a/Project.toml +++ b/Project.toml @@ -25,7 +25,10 @@ SPIRV_Tools_jll = "6ac6d60f-d740-5983-97d7-a4482c0689f4" pocl_standalone_jll = "54f56a70-6062-5590-a942-1226658f6c83" [weakdeps] +AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e" +IntelITT = "c9b2f978-7543-4802-ae44-75068f23ee64" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +NVTX = "5da4648a-3479-48b8-97b9-01cb529c0a1f" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" @@ -33,19 +36,25 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" KernelInterface = {path = "lib/KernelInterface"} [extensions] +IntelITTExt = "IntelITT" LinearAlgebraExt = "LinearAlgebra" +NVTXExt = "NVTX" +ROCTXExt = "AMDGPU" SparseArraysExt = "SparseArrays" StaticArraysExt = "StaticArrays" [compat] +AMDGPU = "2" Adapt = "0.4, 1.0, 2.0, 3.0, 4" Atomix = "1.2.1" GPUCompiler = "2.10" GPUToolbox = "3.3.2" +IntelITT = "0.2" KernelInterface = "0.4" LLVM = "10" LinearAlgebra = "1.6" MacroTools = "0.5" +NVTX = "0.3, 1" PrecompileTools = "1" Printf = "<0.0.1, 1" Random = "1" diff --git a/docs/make.jl b/docs/make.jl index 2e0319f98..3d3c495db 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -41,6 +41,7 @@ function main() "Extras" => [ "extras/unrolling.md", "extras/pocl_debugging.md", + "extras/profiling.md", ], # Extras "Notes for implementations" => "implementations.md", ], # pages diff --git a/docs/src/extras/profiling.md b/docs/src/extras/profiling.md new file mode 100644 index 000000000..3240d5d71 --- /dev/null +++ b/docs/src/extras/profiling.md @@ -0,0 +1,80 @@ +# Profiling + +KernelAbstractions can put named ranges on the timeline of a tracing profiler, such as +NVIDIA Nsight Systems or Intel VTune, so that you can see which part of your program a +stretch of kernels belongs to. Annotations are cheap when no profiler is listening: a +single atomic load, and the label isn't even built. + +## Annotating code + +Wrap code in [`@profiling_range`](@ref): + +```julia +@profiling_range "volume integral" begin + volume_integral!(du, u, backend) +end +``` + +Ranges can be grouped with a `domain`, which maps to an NVTX or ITT domain: + +```julia +@profiling_range "time step $i" domain = "Trixi" begin + step!(integrator) +end +``` + +Kernel launches are annotated with the kernel's name automatically. For an instantaneous +event, use [`profiling_mark`](@ref), and for ranges that don't follow the structure of the +code, [`profiling_range_start`](@ref KernelAbstractions.profiling_range_start) and +[`profiling_range_end`](@ref KernelAbstractions.profiling_range_end). + +## Profilers + +Ranges are recorded on the host threads of the process, which is how NVTX, ITT and +roctx work: it is the profiler that attributes the device work launched within a range to +it. So which profiler records the ranges depends on what the process runs under, not on the +backend: running the CPU backend under Nsight Systems gives NVTX ranges, and a GPU backend +under VTune gives ITT tasks. + +Ranges go to every registered [`Tracer`](@ref KernelAbstractions.Tracer). These come with +KernelAbstractions, and only register themselves when their profiler is attached: + +- **Nsight Systems**: load [NVTX.jl](https://github.com/JuliaGPU/NVTX.jl) (CUDA.jl loads it + too) and run under `nsys profile --trace=nvtx,...`. +- **rocprof**: load [AMDGPU.jl](https://github.com/JuliaGPU/AMDGPU.jl) and run under + `rocprofv3 --marker-trace` (or the legacy `rocprof --roctx-trace`). Ranges are recorded + with roctx, which has no domains, so a `domain` other than `"KernelAbstractions"` prefixes + the label. +- **Intel VTune**: load [IntelITT.jl](https://github.com/JuliaPerf/IntelITT.jl) and run + under VTune. +- **NVTXT**, a text format that Nsight Systems imports, to trace without a profiler: set + `JULIA_KA_NVTXT=1` to write `ka-.nvtxt` to the working directory, or set it to a + path, in which `%p` is replaced by the process id. Then + ```sh + ImportNvtxt --cmd create --nvtxt ka-1234.nvtxt -o report.nsys-rep + ``` + To trace only part of a program, register an [`NVTXTTracer`](@ref + KernelAbstractions.NVTXTTracer) yourself: + ```julia + tracer = KernelAbstractions.register_tracer!(KernelAbstractions.NVTXTTracer("trace.nvtxt")) + run_simulation() + KernelAbstractions.unregister_tracer!(tracer) + close(tracer) + ``` + +Other profilers are supported by subtyping +[`Tracer`](@ref KernelAbstractions.Tracer). + +## API + +```@docs +@profiling_range +profiling_mark +KernelAbstractions.profiling_active +KernelAbstractions.profiling_range_start +KernelAbstractions.profiling_range_end +KernelAbstractions.Tracer +KernelAbstractions.register_tracer! +KernelAbstractions.unregister_tracer! +KernelAbstractions.NVTXTTracer +``` diff --git a/ext/IntelITTExt.jl b/ext/IntelITTExt.jl new file mode 100644 index 000000000..cd26b68b9 --- /dev/null +++ b/ext/IntelITTExt.jl @@ -0,0 +1,35 @@ +module IntelITTExt + +import KernelAbstractions as KA +import IntelITT + +# forwards ranges to Intel VTune as ITT tasks, one ITT domain per domain +struct ITTTracer <: KA.Tracer + domains::Dict{String, IntelITT.Domain} + lock::ReentrantLock +end +ITTTracer() = ITTTracer(Dict{String, IntelITT.Domain}(), ReentrantLock()) + +domain(tracer::ITTTracer, name::String) = + @lock tracer.lock get!(() -> IntelITT.Domain(name), tracer.domains, name) + +function KA.trace_range_start(tracer::ITTTracer, label, domain_name) + # overlapped tasks may end on another thread, and needn't nest + task = IntelITT.Task(domain(tracer, domain_name), label) + IntelITT.start(task) + return task +end + +KA.trace_range_end(::ITTTracer, task::IntelITT.Task) = (IntelITT.stop(task); nothing) + +const TRACER = Ref{ITTTracer}() + +function __init__() + # only under a collector, as otherwise every annotation would be wasted work + if IntelITT.isactive() + TRACER[] = KA.register_tracer!(ITTTracer()) + end + return +end + +end # module diff --git a/ext/NVTXExt.jl b/ext/NVTXExt.jl new file mode 100644 index 000000000..0fd9887a6 --- /dev/null +++ b/ext/NVTXExt.jl @@ -0,0 +1,34 @@ +module NVTXExt + +import KernelAbstractions as KA +import NVTX + +# forwards ranges to Nsight Systems as NVTX ranges, one NVTX domain per domain. NVTX ranges +# annotate host threads; Nsight Systems projects the GPU work launched within them itself. +struct NVTXTracer <: KA.Tracer + domains::Dict{String, NVTX.Domain} + lock::ReentrantLock +end +NVTXTracer() = NVTXTracer(Dict{String, NVTX.Domain}(), ReentrantLock()) + +domain(tracer::NVTXTracer, name::String) = + @lock tracer.lock get!(() -> NVTX.Domain(name), tracer.domains, name) + +# process ranges, rather than push/pop, since they may end on another thread +KA.trace_range_start(tracer::NVTXTracer, label, domain_name) = + NVTX.range_start(domain(tracer, domain_name); message = label) +KA.trace_range_end(::NVTXTracer, id::NVTX.RangeId) = (NVTX.range_end(id); nothing) +KA.trace_mark(tracer::NVTXTracer, label, domain_name) = + (NVTX.mark(domain(tracer, domain_name); message = label); nothing) + +const TRACER = Ref{NVTXTracer}() + +function __init__() + # only under Nsight, as otherwise every annotation would be wasted work + if NVTX.isactive() + TRACER[] = KA.register_tracer!(NVTXTracer()) + end + return +end + +end # module diff --git a/ext/ROCTXExt.jl b/ext/ROCTXExt.jl new file mode 100644 index 000000000..25c56de9b --- /dev/null +++ b/ext/ROCTXExt.jl @@ -0,0 +1,76 @@ +module ROCTXExt + +import KernelAbstractions as KA +import AMDGPU +using Base.Libc: Libdl + +# forwards ranges to rocprof as roctx ranges. like NVTX, roctx annotates host threads, and +# rocprof attributes the GPU work launched within them. roctx has no domains, so ranges in a +# domain other than "KernelAbstractions" are prefixed with it. +struct ROCTXTracer <: KA.Tracer + range_start::Ptr{Cvoid} + range_stop::Ptr{Cvoid} + mark::Ptr{Cvoid} +end + +""" + ROCTXTracer(library::AbstractString) + +A tracer that calls the roctx API in `library`, i.e. `librocprofiler-sdk-roctx` for +`rocprofv3`, or `libroctx64` for the legacy `rocprof`. +""" +function ROCTXTracer(library::AbstractString) + handle = Libdl.dlopen(library) + return ROCTXTracer( + Libdl.dlsym(handle, :roctxRangeStartA), Libdl.dlsym(handle, :roctxRangeStop), + Libdl.dlsym(handle, :roctxMarkA) + ) +end + +roctx_message(label, domain) = domain == "KernelAbstractions" ? label : string(domain, ": ", label) + +# process ranges, rather than push/pop, since they may end on another thread +KA.trace_range_start(tracer::ROCTXTracer, label, domain) = + ccall(tracer.range_start, UInt64, (Cstring,), roctx_message(label, domain)) +KA.trace_range_end(tracer::ROCTXTracer, id::UInt64) = + (ccall(tracer.range_stop, Cvoid, (UInt64,), id); nothing) +KA.trace_mark(tracer::ROCTXTracer, label, domain) = + (ccall(tracer.mark, Cvoid, (Cstring,), roctx_message(label, domain)); nothing) + +# rocprofv3 loads its tool through rocprofiler-register, and only intercepts the roctx of +# the rocprofiler-sdk; the legacy rocprof loads its tool into HSA, and intercepts libroctx64 +function roctx_libraries() + if haskey(ENV, "ROCP_TOOL_LIBRARIES") + return ["librocprofiler-sdk-roctx", "libroctx64"] + elseif haskey(ENV, "HSA_TOOLS_LIB") + return ["libroctx64", "librocprofiler-sdk-roctx"] + else + return String[] + end +end + +function rocm_libdir() + rocm_path = try + AMDGPU.ROCmDiscovery.find_roc_path() + catch + get(ENV, "ROCM_PATH", "/opt/rocm") + end + return joinpath(rocm_path, "lib") +end + +const TRACER = Ref{ROCTXTracer}() + +function __init__() + # only under rocprof, as otherwise every annotation would be wasted work + names = roctx_libraries() + isempty(names) && return + library = Libdl.find_library(names, [rocm_libdir()]) + if isempty(library) + @warn "Running under rocprof, but roctx wasn't found; KernelAbstractions' ranges won't be recorded" names + return + end + TRACER[] = KA.register_tracer!(ROCTXTracer(library)) + return +end + +end # module diff --git a/src/KernelAbstractions.jl b/src/KernelAbstractions.jl index 7bc531342..786cdaadc 100644 --- a/src/KernelAbstractions.jl +++ b/src/KernelAbstractions.jl @@ -6,6 +6,7 @@ export @index, @groupsize, @ndrange export @print export Backend, CPU export synchronize, get_backend, allocate +export @profiling_range, profiling_mark import PrecompileTools @@ -742,6 +743,12 @@ number as its compute units: `KernelAbstractions.POCL.device().max_compute_units """ const CPU = POCLBackend +include("profiling.jl") include("precompile.jl") +function __init__() + init_profiling() + return +end + end #module diff --git a/src/backend_launch.jl b/src/backend_launch.jl index fbd5cfd8d..f97710ce9 100644 --- a/src/backend_launch.jl +++ b/src/backend_launch.jl @@ -83,6 +83,24 @@ Core.kwcall(kwargs::NamedTuple, obj::Kernel{<:KI.Backend}, args::Vararg{Any, N}) launch_tuple(obj, args; kwargs...) function launch_tuple(obj::Kernel, args::Tuple; ndrange = nothing, workgroupsize = nothing) + if profiling_active() + return launch_traced(obj, args, ndrange, workgroupsize) + end + return launch_untraced(obj, args, ndrange, workgroupsize) +end + +# out of line, to keep the profiler out of the common path +@noinline function launch_traced(obj::Kernel, args::Tuple, ndrange, workgroupsize) + id = profiling_range_start(kernel_label(obj.f)) + try + launch_untraced(obj, args, ndrange, workgroupsize) + finally + profiling_range_end(id) + end + return nothing +end + +function launch_untraced(obj::Kernel, args::Tuple, ndrange, workgroupsize) ndrange, workgroupsize, iterspace, dynamic = launch_config(obj, ndrange, workgroupsize) # nothing to launch (or compile) for an empty ndrange any(iszero, size(blocks(iterspace))) && return nothing diff --git a/src/profiling.jl b/src/profiling.jl new file mode 100644 index 000000000..ba498fcdc --- /dev/null +++ b/src/profiling.jl @@ -0,0 +1,282 @@ +### +# Profiler integration +# +# Tracing profilers (Nsight Systems via NVTX, VTune via ITT, rocprof via roctx, ...) record +# ranges on the host threads of the process, and correlate the device work launched within +# them themselves. Which profiler records a range thus depends on what the process runs +# under, not on the backend: ranges go to every registered tracer, for every backend. With +# none registered, an annotation costs one atomic load. +### + +""" + Tracer + +Abstract supertype for profilers that record named ranges on the host. Register an +instance with [`register_tracer!`](@ref) to receive the ranges and markers of +[`@profiling_range`](@ref) and [`profiling_mark`](@ref), and of kernel launches. + +Subtypes implement + + trace_range_start(tracer, label::String, domain::String) -> id + trace_range_end(tracer, id) + trace_mark(tracer, label::String, domain::String) # optional + +Ranges may end on a different thread than they started on, and may overlap without nesting, +so implementations should use the profiler's start/end API (e.g. `nvtxRangeStartEx`) rather +than a thread-local push/pop stack. + +A tracer for a profiler that may not be attached should only be registered when it is, as +registering any tracer makes every annotation do work. +""" +abstract type Tracer end + +function trace_range_start end +function trace_range_end end +trace_mark(::Tracer, label, domain) = nothing + +# copy-on-write, so that checking for tracers is a single atomic load +mutable struct Tracers + @atomic tracers::Vector{Tracer} +end +const TRACERS = Tracers(Tracer[]) +const TRACERS_LOCK = ReentrantLock() + +tracers() = @atomic :acquire TRACERS.tracers + +""" + register_tracer!(tracer::Tracer) -> tracer + +Forward profiler ranges and markers to `tracer`, in addition to the tracers that are +registered already. +""" +function register_tracer!(tracer::Tracer) + @lock TRACERS_LOCK begin + current = tracers() + tracer in current || @atomic :release TRACERS.tracers = Tracer[current; tracer] + end + return tracer +end + +""" + unregister_tracer!(tracer::Tracer) + +Stop forwarding profiler ranges and markers to `tracer`. Ranges that are open keep going to +the tracers that were registered when they started. +""" +function unregister_tracer!(tracer::Tracer) + @lock TRACERS_LOCK begin + @atomic :release TRACERS.tracers = filter(t -> t !== tracer, tracers()) + end + return nothing +end + +""" + profiling_active()::Bool + +Whether a profiler is listening, i.e. whether a [`Tracer`](@ref) is registered. Check this +before doing work that only serves annotations. +""" +profiling_active() = !isempty(tracers()) + +struct ProfilingRange + tracers::Vector{Tracer} + ids::Vector{Any} +end + +""" + profiling_range_start(label::AbstractString; domain = "KernelAbstractions") + +Start a range named `label` in `domain`, and return a handle for +[`profiling_range_end`](@ref). Returns `nothing` if no profiler is listening. + +Prefer [`@profiling_range`](@ref), which ends the range even if an exception is thrown. Use +these for ranges that don't follow the structure of the code. +""" +function profiling_range_start(label; domain = "KernelAbstractions") + current = tracers() + isempty(current) && return nothing + label, domain = String(label), String(domain) + return ProfilingRange(current, Any[trace_range_start(t, label, domain) for t in current]) +end + +""" + profiling_range_end(range) + +End a range started with [`profiling_range_start`](@ref). +""" +function profiling_range_end(range::ProfilingRange) + for (tracer, id) in zip(range.tracers, range.ids) + trace_range_end(tracer, id) + end + return nothing +end +profiling_range_end(::Nothing) = nothing + +""" + profiling_mark(label::AbstractString; domain = "KernelAbstractions") + +Record an instantaneous marker named `label` in `domain`. +""" +function profiling_mark(label; domain = "KernelAbstractions") + current = tracers() + isempty(current) && return nothing + label, domain = String(label), String(domain) + for tracer in current + trace_mark(tracer, label, domain) + end + return nothing +end + +""" + @profiling_range label [domain = "KernelAbstractions"] expr + +Evaluate `expr` inside a profiler range named `label`, and return its value. `label` is +only evaluated if a profiler is listening (see [`profiling_active`](@ref)), so it can be +built with string interpolation at no cost to unprofiled runs. The range is ended if `expr` +throws. Assignments in `expr` are visible after the macro, as with `@time`. + +```julia +@profiling_range "volume integral" begin + volume_integral!(du, u, backend) +end + +@profiling_range "volume integral" domain = "Trixi" begin + volume_integral!(du, u, backend) +end +``` + +Ranges are recorded on the host, by whichever profiler the process runs under; device work +launched within a range is attributed to it by profilers that correlate the two, such as +Nsight Systems. Kernel launches are annotated with the name of the kernel automatically. +""" +macro profiling_range(label, args...) + isempty(args) && throw(ArgumentError("@profiling_range needs an expression to evaluate")) + expr = args[end] + domain = "KernelAbstractions" + for kw in args[1:(end - 1)] + if Meta.isexpr(kw, :(=)) && kw.args[1] === :domain + domain = kw.args[2] + else + throw(ArgumentError("@profiling_range: unexpected argument `$kw`; only `domain = ...` is accepted")) + end + end + id = gensym(:id) + return quote + local $id = if $profiling_active() + $profiling_range_start($(esc(label)); domain = $(esc(domain))) + else + nothing + end + # unlike `try`, `tryfinally` doesn't introduce a scope + $(Expr(:tryfinally, esc(expr), :($profiling_range_end($id)))) + end +end + +# the name kernel launches are annotated with: `@kernel function f` compiles to `gpu_f` +function kernel_label(f) + f isa Function || return string(typeof(f)) + name = string(nameof(f)) + return startswith(name, "gpu_") ? name[5:end] : name +end + + +## NVTXT + +""" + NVTXTTracer(path::AbstractString) + NVTXTTracer(io::IO) + +A [`Tracer`](@ref) that writes ranges and markers in the NVTXT text format, which NVIDIA +Nsight Systems imports with `ImportNvtxt`: + +``` +ImportNvtxt --cmd create --nvtxt ka-1234.nvtxt -o report.nsys-rep +``` + +This needs no profiler at run time, e.g. to trace the CPU backend on a machine without +Nsight Systems. Set `JULIA_KA_NVTXT` to start one when KernelAbstractions is loaded: to a +path, in which `%p` is replaced by the process id, or to `1` for `ka-%p.nvtxt` in the +working directory. Otherwise, register one with [`register_tracer!`](@ref), and `close` it +after unregistering it to flush the file. + +Records are written when a range ends, as a single line, so that threads don't interleave. +Ranges in a `domain` other than `"KernelAbstractions"` are prefixed with it. +""" +struct NVTXTTracer{IO_ <: IO} <: Tracer + io::IO_ + lock::ReentrantLock +end + +function NVTXTTracer(io::IO) + pid = getpid() + print( + io, """ + SetFileDisplayName, KernelAbstractions + @RangeStartEnd, Start, End, ThreadId, Message + ProcessId = $pid + CategoryId = 1 + Color = Blue + TimeBase = Manual + @Marker, Time, ThreadId, Message + ProcessId = $pid + CategoryId = 1 + Color = Blue + TimeBase = Manual + """ + ) + return NVTXTTracer(io, ReentrantLock()) +end +NVTXTTracer(path::AbstractString) = NVTXTTracer(open(path, "w")) + +Base.close(tracer::NVTXTTracer) = @lock tracer.lock close(tracer.io) + +struct NVTXTRange + start::UInt64 + thread::Int + message::String +end + +# the message is a quoted string, and a record a line +function nvtxt_message(label, domain) + message = domain == "KernelAbstractions" ? label : string(domain, ": ", label) + return replace(message, '"' => '\'', '\n' => ' ', '\r' => ' ') +end + +function nvtxt_record(tracer::NVTXTTracer, record::String) + @lock tracer.lock isopen(tracer.io) && write(tracer.io, record) + return nothing +end + +trace_range_start(::NVTXTTracer, label, domain) = + NVTXTRange(time_ns(), Threads.threadid(), nvtxt_message(label, domain)) + +function trace_range_end(tracer::NVTXTTracer, range::NVTXTRange) + stop = time_ns() + return nvtxt_record( + tracer, "RangeStartEnd, $(range.start), $stop, $(range.thread), \"$(range.message)\"\n" + ) +end + +function trace_mark(tracer::NVTXTTracer, label, domain) + time = time_ns() + return nvtxt_record( + tracer, "Marker, $time, $(Threads.threadid()), \"$(nvtxt_message(label, domain))\"\n" + ) +end + +function nvtxt_path(setting::AbstractString) + path = setting in ("1", "true", "yes") ? "ka-%p.nvtxt" : setting + return replace(path, "%p" => string(getpid())) +end + +function init_profiling() + setting = Base.get(ENV, "JULIA_KA_NVTXT", "") + if !isempty(setting) && !(setting in ("0", "false", "no")) + tracer = register_tracer!(NVTXTTracer(nvtxt_path(setting))) + atexit() do + unregister_tracer!(tracer) + close(tracer) + end + end + return +end diff --git a/test/Project.toml b/test/Project.toml index c191d6ced..547a12feb 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -2,10 +2,12 @@ Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" FileCheck = "4e644321-382b-4b05-b0b6-5d23c3d944fb" +IntelITT = "c9b2f978-7543-4802-ae44-75068f23ee64" InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" KernelInterface = "4ee993da-d684-4d17-a7dd-4e58e78d92bf" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +NVTX = "5da4648a-3479-48b8-97b9-01cb529c0a1f" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" diff --git a/test/profiling.jl b/test/profiling.jl new file mode 100644 index 000000000..97d8653ae --- /dev/null +++ b/test/profiling.jl @@ -0,0 +1,53 @@ +# records the ranges and markers it is given +struct RecordingTracer <: KernelAbstractions.Tracer + events::Vector{Any} + lock::ReentrantLock +end +RecordingTracer() = RecordingTracer([], ReentrantLock()) +function KernelAbstractions.trace_range_start(t::RecordingTracer, label, domain) + @lock t.lock push!(t.events, (:start, label, domain)) + return label +end +KernelAbstractions.trace_range_end(t::RecordingTracer, id) = + @lock t.lock push!(t.events, (:end, id)) +KernelAbstractions.trace_mark(t::RecordingTracer, label, domain) = + @lock t.lock push!(t.events, (:mark, label, domain)) + +function with_tracer(f, tracer = RecordingTracer()) + KernelAbstractions.register_tracer!(tracer) + try + f(tracer) + finally + KernelAbstractions.unregister_tracer!(tracer) + end + return tracer +end + +@kernel function profiling_fill!(A, x) + I = @index(Global) + @inbounds A[I] = x +end + +function profiling_testsuite(Backend, AT) + backend = Backend() + + # launches work whether or not a profiler listens + A = AT(zeros(Float32, 64)) + profiling_fill!(backend)(A, 1.0f0; ndrange = length(A)) + synchronize(backend) + @test all(Array(A) .== 1) + + # and are named after the kernel + tracer = with_tracer() do tracer + @profiling_range "step" profiling_fill!(backend)(A, 2.0f0; ndrange = length(A)) + synchronize(backend) + end + @test all(Array(A) .== 2) + @test tracer.events == [ + (:start, "step", "KernelAbstractions"), + (:start, "profiling_fill!", "KernelAbstractions"), (:end, "profiling_fill!"), + (:end, "step"), + ] + + return +end diff --git a/test/runtests.jl b/test/runtests.jl index 7e9c1b3ee..0a2ee9310 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -429,3 +429,159 @@ end @test CPU() isa KernelAbstractions.GPU @test NewBackend <: KernelAbstractions.GPU end + +@testset "Profiling" begin + RecordingTracer, with_tracer = Testsuite.RecordingTracer, Testsuite.with_tracer + # nothing is registered unless running under a profiler (or with `JULIA_KA_NVTXT`) + @test KernelAbstractions.profiling_active() == !isempty(KernelAbstractions.tracers()) + + if !KernelAbstractions.profiling_active() + @testset "inactive" begin + @test KernelAbstractions.profiling_range_start("label") === nothing + @test KernelAbstractions.profiling_range_end(nothing) === nothing + @test profiling_mark("label") === nothing + # the label isn't evaluated when nobody listens + @test (@profiling_range error("label") 1 + 2) == 3 + end + end + + @testset "macro" begin + @test (@profiling_range "label" 1 + 2) == 3 + @test (@profiling_range "label" domain = "Custom" 1 + 2) == 3 + @test_throws ErrorException @profiling_range "label" error("boom") + # assignments remain visible, as with `@time` + @profiling_range "assign" y = 42 + @test y == 42 + + @test_throws ArgumentError macroexpand(@__MODULE__, :(@profiling_range "label" foo = 1 2)) + end + + @testset "registration" begin + tracer = RecordingTracer() + with_tracer(tracer) do tracer + @test KernelAbstractions.profiling_active() + # registering twice doesn't duplicate + KernelAbstractions.register_tracer!(tracer) + @test count(t -> t === tracer, KernelAbstractions.tracers()) == 1 + end + @test !(tracer in KernelAbstractions.tracers()) + end + + @testset "ranges and markers" begin + tracer = with_tracer() do tracer + @test ( + @profiling_range "outer" domain = "Trixi" begin + profiling_mark("inside") + @profiling_range "inner $(1 + 1)" 7 + end + ) == 7 + id = KernelAbstractions.profiling_range_start("explicit"; domain = "X") + KernelAbstractions.profiling_range_end(id) + end + @test tracer.events == [ + (:start, "outer", "Trixi"), (:mark, "inside", "KernelAbstractions"), + (:start, "inner 2", "KernelAbstractions"), (:end, "inner 2"), (:end, "outer"), + (:start, "explicit", "X"), (:end, "explicit"), + ] + + # ranges end when the expression throws + tracer = with_tracer() do tracer + @test_throws ErrorException @profiling_range "throws" error("boom") + end + @test tracer.events == [(:start, "throws", "KernelAbstractions"), (:end, "throws")] + + # ranges end with the tracers they started with + tracer = KernelAbstractions.register_tracer!(RecordingTracer()) + id = KernelAbstractions.profiling_range_start("open") + KernelAbstractions.unregister_tracer!(tracer) + KernelAbstractions.profiling_range_end(id) + @test tracer.events == [(:start, "open", "KernelAbstractions"), (:end, "open")] + + # from many tasks at once + tracer = with_tracer() do tracer + @sync for i in 1:16 + Threads.@spawn @profiling_range "task $i" (yield(); i) + end + end + @test count(e -> e[1] === :start, tracer.events) == 16 + @test count(e -> e[1] === :end, tracer.events) == 16 + end + + @testset "multiple tracers" begin + a, b = RecordingTracer(), RecordingTracer() + with_tracer(a) do _ + with_tracer(b) do _ + @profiling_range "both" nothing + end + end + @test a.events == b.events == [(:start, "both", "KernelAbstractions"), (:end, "both")] + end + + # with a profiler listening + with_tracer() do _ + Testsuite.profiling_testsuite(CPU, Array) + end +end + +@testset "NVTXT" begin + @test KernelAbstractions.nvtxt_path("1") == "ka-$(getpid()).nvtxt" + @test KernelAbstractions.nvtxt_path("/tmp/trace-%p.nvtxt") == "/tmp/trace-$(getpid()).nvtxt" + + mktempdir() do dir + path = joinpath(dir, "trace.nvtxt") + tracer = KernelAbstractions.NVTXTTracer(path) + Testsuite.with_tracer(tracer) do _ + @profiling_range "range" nothing + @profiling_range "say \"hi\"\n" domain = "Trixi" nothing + profiling_mark("marker") + Testsuite.profiling_fill!(CPU())(zeros(Float32, 4), 1.0f0; ndrange = 4) + end + close(tracer) + # recording after closing is harmless + KernelAbstractions.trace_mark(tracer, "late", "KernelAbstractions") + + lines = readlines(path) + @test lines[1] == "SetFileDisplayName, KernelAbstractions" + @test "ProcessId = $(getpid())" in lines + records = filter(l -> startswith(l, "RangeStartEnd, ") || startswith(l, "Marker, "), lines) + @test length(records) == 4 + r = match(r"^RangeStartEnd, (\d+), (\d+), (\d+), \"range\"$", records[1]) + @test r !== nothing && parse(UInt64, r[1]) <= parse(UInt64, r[2]) + @test endswith(records[2], ", \"Trixi: say 'hi' \"") + @test match(r"^Marker, \d+, \d+, \"marker\"$", records[3]) !== nothing + @test endswith(records[4], ", \"profiling_fill!\"") + end + + # enabled with an environment variable + mktempdir() do dir + julia = Cmd(filter(arg -> !startswith(arg, "--code-coverage"), Base.julia_cmd().exec)) + script = """ + using KernelAbstractions + @profiling_range "from env" nothing + print(getpid()) + """ + cmd = `$julia --startup-file=no --project=$(Base.active_project()) -e $script` + env = ("JULIA_KA_NVTXT" => joinpath(dir, "env-%p.nvtxt"),) + pid = readchomp(setenv(cmd, copy(ENV)..., env...; dir)) + trace = read(joinpath(dir, "env-$pid.nvtxt"), String) + @test occursin("\"from env\"", trace) + end +end + +import IntelITT, NVTX +@testset "Profiler extensions" begin + # only registered under the profiler + itt = Base.get_extension(KernelAbstractions, :IntelITTExt) + @test isassigned(itt.TRACER) == IntelITT.isactive() + nvtx = Base.get_extension(KernelAbstractions, :NVTXExt) + @test isassigned(nvtx.TRACER) == NVTX.isactive() + + # but work without it + for tracer in (itt.ITTTracer(), nvtx.NVTXTracer()) + Testsuite.with_tracer(tracer) do _ + @test (@profiling_range "range" domain = "Ext" 1) == 1 + @test profiling_mark("mark") === nothing + Testsuite.profiling_fill!(CPU())(zeros(Float32, 4), 1.0f0; ndrange = 4) + end + end +end diff --git a/test/testsuite.jl b/test/testsuite.jl index cd75adbc3..06f43af83 100644 --- a/test/testsuite.jl +++ b/test/testsuite.jl @@ -44,6 +44,7 @@ include("convert.jl") include("specialfunctions.jl") include("random.jl") include("spawn.jl") +include("profiling.jl") function testsuite(backend, backend_str, backend_mod, AT, DAT; skip_tests = Set{String}()) @conditional_testset "Unittests" skip_tests begin @@ -118,6 +119,10 @@ function testsuite(backend, backend_str, backend_mod, AT, DAT; skip_tests = Set{ spawn_testsuite(backend, AT) end + @conditional_testset "Profiling" skip_tests begin + profiling_testsuite(backend, AT) + end + return end From 580a8e16c583d65bfb3176a8b667f40601334f5c Mon Sep 17 00:00:00 2001 From: Valentin Churavy Date: Sun, 4 Oct 2026 09:43:09 +0200 Subject: [PATCH 2/9] Add `KernelAbstractions.@profile` Records the ranges, markers and kernel launches of an expression with a temporary tracer, and summarizes them per name or, with `trace = true`, lists them in order. By default kernel launches synchronize their backend while profiling, so that their ranges measure execution rather than launch; tracers opt into this with `synchronizes_launches`. Assisted-by: Claude Code (Opus 5.5) --- docs/src/extras/profiling.md | 29 +++++ src/KernelAbstractions.jl | 1 + src/backend_launch.jl | 2 + src/profiler.jl | 229 +++++++++++++++++++++++++++++++++++ src/profiling.jl | 8 ++ test/runtests.jl | 59 +++++++++ 6 files changed, 328 insertions(+) create mode 100644 src/profiler.jl diff --git a/docs/src/extras/profiling.md b/docs/src/extras/profiling.md index 3240d5d71..7ceac8f9f 100644 --- a/docs/src/extras/profiling.md +++ b/docs/src/extras/profiling.md @@ -28,6 +28,33 @@ event, use [`profiling_mark`](@ref), and for ranges that don't follow the struct code, [`profiling_range_start`](@ref KernelAbstractions.profiling_range_start) and [`profiling_range_end`](@ref KernelAbstractions.profiling_range_end). +## Built-in profiler + +To see where time goes without an external profiler, run code under +[`KernelAbstractions.@profile`](@ref KernelAbstractions.@profile). It records the ranges +and kernel launches of an expression, and summarizes them: + +```julia-repl +julia> KernelAbstractions.@profile for i in 1:10 + @profiling_range "step" begin + mul2(backend)(A; ndrange = length(A)) + add(backend)(A, B; ndrange = length(A)) + end + end +Profiled 6.04 ms, recording 30 ranges. + + Time (%) Total time Calls Avg time Min time Max time Name + ──────── ────────── ───── ──────── ──────── ──────── ──── + 92.6 % 5.59 ms 10 559 µs 298 µs 2.9 ms step + 57.6 % 3.48 ms 10 348 µs 109 µs 2.48 ms mul2 + 37.6 % 2.27 ms 10 227 µs 179 µs 400 µs add +``` + +Kernel launches synchronize their backend while profiling, so that their ranges measure the +kernel rather than its launch; pass `synchronize = false` to measure launches. Pass +`trace = true` to list every range in order instead. The first call of a kernel includes +its compilation, so profile a warmed-up run. + ## Profilers Ranges are recorded on the host threads of the process, which is how NVTX, ITT and @@ -70,6 +97,8 @@ Other profilers are supported by subtyping ```@docs @profiling_range profiling_mark +KernelAbstractions.@profile +KernelAbstractions.ProfileResults KernelAbstractions.profiling_active KernelAbstractions.profiling_range_start KernelAbstractions.profiling_range_end diff --git a/src/KernelAbstractions.jl b/src/KernelAbstractions.jl index 786cdaadc..cbada2d55 100644 --- a/src/KernelAbstractions.jl +++ b/src/KernelAbstractions.jl @@ -744,6 +744,7 @@ number as its compute units: `KernelAbstractions.POCL.device().max_compute_units const CPU = POCLBackend include("profiling.jl") +include("profiler.jl") include("precompile.jl") function __init__() diff --git a/src/backend_launch.jl b/src/backend_launch.jl index f97710ce9..125588ff4 100644 --- a/src/backend_launch.jl +++ b/src/backend_launch.jl @@ -94,6 +94,8 @@ end id = profiling_range_start(kernel_label(obj.f)) try launch_untraced(obj, args, ndrange, workgroupsize) + # a profiler that measures kernels rather than launches + synchronizes_launches(id) && KI.synchronize(backend(obj)) finally profiling_range_end(id) end diff --git a/src/profiler.jl b/src/profiler.jl new file mode 100644 index 000000000..d6f523a87 --- /dev/null +++ b/src/profiler.jl @@ -0,0 +1,229 @@ +using Printf: @sprintf + +### +# Built-in profiler: records the ranges of an expression, and summarizes them +### + +struct ProfileRange + name::String + start::UInt64 + stop::UInt64 + thread::Int +end + +struct ProfileMarker + name::String + time::UInt64 + thread::Int +end + +struct ProfileTracer <: Tracer + synchronize::Bool + ranges::Vector{ProfileRange} + markers::Vector{ProfileMarker} + lock::ReentrantLock +end +ProfileTracer(synchronize::Bool) = + ProfileTracer(synchronize, ProfileRange[], ProfileMarker[], ReentrantLock()) + +synchronizes_launches(tracer::ProfileTracer) = tracer.synchronize + +# like NVTXT, without domains of their own +profile_name(label, domain) = domain == "KernelAbstractions" ? label : string(domain, ": ", label) + +trace_range_start(::ProfileTracer, label, domain) = + (profile_name(label, domain), time_ns(), Threads.threadid()) + +function trace_range_end(tracer::ProfileTracer, (name, start, thread)) + range = ProfileRange(name, start, time_ns(), thread) + @lock tracer.lock push!(tracer.ranges, range) + return nothing +end + +function trace_mark(tracer::ProfileTracer, label, domain) + marker = ProfileMarker(profile_name(label, domain), time_ns(), Threads.threadid()) + @lock tracer.lock push!(tracer.markers, marker) + return nothing +end + +""" + ProfileResults + +The ranges and markers recorded by [`@profile`](@ref KernelAbstractions.@profile). Shown, +it summarizes the time spent per range name or, with `trace = true`, lists the ranges in +the order they started. `results.ranges` and `results.markers` hold the raw records, with +times in nanoseconds from `time_ns()`. +""" +struct ProfileResults + start::UInt64 + stop::UInt64 + ranges::Vector{ProfileRange} + markers::Vector{ProfileMarker} + trace::Bool +end + +""" + KernelAbstractions.@profile [trace = false] [synchronize = true] expr + +Run `expr`, recording the ranges of [`@profiling_range`](@ref), the markers of +[`profiling_mark`](@ref) and the kernel launches within it, and return a +[`ProfileResults`](@ref KernelAbstractions.ProfileResults) that summarizes them: + +```julia-repl +julia> KernelAbstractions.@profile for i in 1:10 + @profiling_range "step" begin + mul2(backend)(A; ndrange = length(A)) + add(backend)(A, B; ndrange = length(A)) + end + end +Profiled 6.04 ms, recording 30 ranges. + + Time (%) Total time Calls Avg time Min time Max time Name + ──────── ────────── ───── ──────── ──────── ──────── ──── + 92.6 % 5.59 ms 10 559 µs 298 µs 2.9 ms step + 57.6 % 3.48 ms 10 348 µs 109 µs 2.48 ms mul2 + 37.6 % 2.27 ms 10 227 µs 179 µs 400 µs add +``` + +With `synchronize = true`, the default, kernel launches synchronize their backend before +their range ends, so that on GPU backends they measure the kernel's execution instead of +its launch. This serializes the host with the device, as `CUDA_LAUNCH_BLOCKING=1` does. + +With `trace = true`, the results list every range in the order it started instead, +indented by nesting. + +Other registered tracers, e.g. NVTX under Nsight Systems, record the ranges too. +""" +macro profile(args...) + isempty(args) && throw(ArgumentError("KernelAbstractions.@profile needs an expression to profile")) + expr = args[end] + trace, synchronize = false, true + for kw in args[1:(end - 1)] + if Meta.isexpr(kw, :(=)) && kw.args[1] === :trace + trace = kw.args[2] + elseif Meta.isexpr(kw, :(=)) && kw.args[1] === :synchronize + synchronize = kw.args[2] + else + throw(ArgumentError("KernelAbstractions.@profile: unexpected argument `$kw`; only `trace = ...` and `synchronize = ...` are accepted")) + end + end + return quote + $profile(() -> $(esc(expr)); trace = $(esc(trace)), synchronize = $(esc(synchronize))) + end +end + +function profile(f; trace::Bool = false, synchronize::Bool = true) + tracer = register_tracer!(ProfileTracer(synchronize)) + start = time_ns() + try + f() + finally + unregister_tracer!(tracer) + end + stop = time_ns() + return @lock tracer.lock ProfileResults(start, stop, copy(tracer.ranges), copy(tracer.markers), trace) +end + + +## report + +function format_time(ns::Real) + # switch units where three significant digits would round up to the next one + ns < 999.5 && return @sprintf("%.0f ns", ns) + ns < 999.5e3 && return @sprintf("%.3g µs", ns / 1.0e3) + ns < 999.5e6 && return @sprintf("%.3g ms", ns / 1.0e6) + return @sprintf("%.3g s", ns / 1.0e9) +end + +# columns of strings, the last one left-aligned +function print_table(io::IO, header::Vector{String}, rows::Vector{Vector{String}}) + widths = [maximum(textwidth, [h; getindex.(rows, i)]) for (i, h) in enumerate(header)] + function print_row(row) + print(io, " ") + for (i, (cell, width)) in enumerate(zip(row, widths)) + if i == length(row) + print(io, cell) + else + print(io, lpad(cell, width), " ") + end + end + return println(io) + end + print_row(header) + print_row(["─"^w for w in widths]) + foreach(print_row, rows) + return +end + +function Base.show(io::IO, ::MIME"text/plain", results::ProfileResults) + total = results.stop - results.start + nranges, nmarkers = length(results.ranges), length(results.markers) + print(io, "Profiled ", format_time(total), ", recording ", nranges, nranges == 1 ? " range" : " ranges") + nmarkers > 0 && print(io, " and ", nmarkers, nmarkers == 1 ? " marker" : " markers") + println(io, ".") + (nranges == 0 && nmarkers == 0) && return + println(io) + if results.trace + show_trace(io, results) + else + show_summary(io, results, total) + end + return +end + +function show_summary(io::IO, results::ProfileResults, total) + if !isempty(results.ranges) + durations = Dict{String, Vector{UInt64}}() + for range in results.ranges + push!(get!(durations, range.name, UInt64[]), range.stop - range.start) + end + rows = sort!(collect(durations); by = kv -> sum(kv[2]), rev = true) + print_table( + io, ["Time (%)", "Total time", "Calls", "Avg time", "Min time", "Max time", "Name"], + [ + [ + @sprintf("%.1f %%", 100 * sum(ds) / total), format_time(sum(ds)), string(length(ds)), + format_time(sum(ds) / length(ds)), format_time(minimum(ds)), format_time(maximum(ds)), + name, + ] for (name, ds) in rows + ] + ) + end + if !isempty(results.markers) + isempty(results.ranges) || println(io) + counts = Dict{String, Int}() + for marker in results.markers + counts[marker.name] = Base.get(counts, marker.name, 0) + 1 + end + rows = sort!(collect(counts); by = last, rev = true) + print_table(io, ["Count", "Marker"], [[string(n), name] for (name, n) in rows]) + end + return +end + +function show_trace(io::IO, results::ProfileResults) + # the nesting depth of each range among the ranges on its thread + events = sort!( + [ + [(r.start, r) for r in results.ranges]; + [(m.time, m) for m in results.markers] + ]; by = first + ) + open = Dict{Int, Vector{UInt64}}() # stop times of the open ranges, per thread + rows = Vector{String}[] + for (time, event) in events + stack = get!(open, event.thread, UInt64[]) + while !isempty(stack) && last(stack) <= time + pop!(stack) + end + indent = " "^length(stack) + if event isa ProfileRange + push!(stack, event.stop) + push!(rows, [format_time(time - results.start), format_time(event.stop - event.start), string(event.thread), indent * event.name]) + else + push!(rows, [format_time(time - results.start), "", string(event.thread), indent * "◆ " * event.name]) + end + end + print_table(io, ["Start", "Duration", "Thread", "Name"], rows) + return +end diff --git a/src/profiling.jl b/src/profiling.jl index ba498fcdc..a9a396d5a 100644 --- a/src/profiling.jl +++ b/src/profiling.jl @@ -20,6 +20,10 @@ Subtypes implement trace_range_start(tracer, label::String, domain::String) -> id trace_range_end(tracer, id) trace_mark(tracer, label::String, domain::String) # optional + synchronizes_launches(tracer)::Bool # optional, default `false` + +If `synchronizes_launches` is `true`, kernel launches synchronize their backend before their +range ends, so that the range measures the kernel's execution rather than its launch. Ranges may end on a different thread than they started on, and may overlap without nesting, so implementations should use the profiler's start/end API (e.g. `nvtxRangeStartEx`) rather @@ -33,6 +37,7 @@ abstract type Tracer end function trace_range_start end function trace_range_end end trace_mark(::Tracer, label, domain) = nothing +synchronizes_launches(::Tracer) = false # copy-on-write, so that checking for tracers is a single atomic load mutable struct Tracers @@ -112,6 +117,9 @@ function profiling_range_end(range::ProfilingRange) end profiling_range_end(::Nothing) = nothing +synchronizes_launches(range::ProfilingRange) = any(synchronizes_launches, range.tracers) +synchronizes_launches(::Nothing) = false + """ profiling_mark(label::AbstractString; domain = "KernelAbstractions") diff --git a/test/runtests.jl b/test/runtests.jl index 0a2ee9310..0ea973e88 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -523,6 +523,65 @@ end end end +@testset "@profile" begin + kfill! = Testsuite.profiling_fill! + A = zeros(Float32, 64) + + results = KernelAbstractions.@profile for i in 1:3 + @profiling_range "step" domain = "Demo" begin + kfill!(CPU())(A, Float32(i); ndrange = length(A)) + profiling_mark("half") + kfill!(CPU())(A, Float32(i); ndrange = length(A)) + end + end + @test !KernelAbstractions.profiling_active() + @test all(==(3), A) + @test count(r -> r.name == "Demo: step", results.ranges) == 3 + @test count(r -> r.name == "profiling_fill!", results.ranges) == 6 + @test length(results.markers) == 3 + @test all(r -> results.start <= r.start <= r.stop <= results.stop, results.ranges) + + summary = sprint(show, MIME"text/plain"(), results) + @test startswith(summary, "Profiled ") + @test occursin("recording 9 ranges and 3 markers.", summary) + lines = split(summary, '\n') + @test occursin("Total time", lines[3]) + # sorted by total time + @test endswith(lines[5], "Demo: step") && endswith(lines[6], "profiling_fill!") + @test any(l -> occursin(r"^ +3 half$", l), lines) + + trace = sprint( + show, MIME"text/plain"(), KernelAbstractions.@profile trace = true begin + @profiling_range "outer" begin + profiling_mark("mark") + @profiling_range "inner" nothing + end + end + ) + lines = split(trace, '\n') + @test occursin("Duration", lines[3]) + @test endswith(lines[5], " outer") && endswith(lines[6], " ◆ mark") && endswith(lines[7], " inner") + + @test occursin("recording 0 ranges.", sprint(show, MIME"text/plain"(), KernelAbstractions.@profile 1 + 1)) + + # launches synchronize their backend only if asked to + @test KernelAbstractions.synchronizes_launches(KernelAbstractions.ProfileTracer(true)) + @test !KernelAbstractions.synchronizes_launches(KernelAbstractions.ProfileTracer(false)) + @test !KernelAbstractions.synchronizes_launches(Testsuite.RecordingTracer()) + results = KernelAbstractions.@profile synchronize = false kfill!(CPU())(A, 1.0f0; ndrange = length(A)) + @test only(results.ranges).name == "profiling_fill!" + + # the profiler stops when the expression throws + @test_throws ErrorException KernelAbstractions.@profile error("boom") + @test !KernelAbstractions.profiling_active() + @test_throws ArgumentError macroexpand(@__MODULE__, :(KernelAbstractions.@profile foo = 1 2)) + + @test KernelAbstractions.format_time(5) == "5 ns" + @test KernelAbstractions.format_time(999.7) == "1 µs" + @test KernelAbstractions.format_time(1.234e6) == "1.23 ms" + @test KernelAbstractions.format_time(2.5e9) == "2.5 s" +end + @testset "NVTXT" begin @test KernelAbstractions.nvtxt_path("1") == "ka-$(getpid()).nvtxt" @test KernelAbstractions.nvtxt_path("/tmp/trace-%p.nvtxt") == "/tmp/trace-$(getpid()).nvtxt" From 36ef5dd2f75490d9060ca9fffb68562a2a818d7b Mon Sep 17 00:00:00 2001 From: Valentin Churavy Date: Sun, 4 Oct 2026 10:50:05 +0200 Subject: [PATCH 3/9] Scope `@profile` to the tasks of its expression Tracers are global, so `@profile` recorded every task in the process, including other profiles. It now records only the task running its expression and the tasks spawned from it, tracked with a ScopedValue, numbers those tasks, and nests the trace per task rather than per thread. It warns about ranges still open when it finishes, i.e. tasks it wasn't waited for. Assisted-by: Claude Code (Opus 5.5) --- Project.toml | 2 + docs/src/extras/profiling.md | 5 +++ src/profiler.jl | 71 ++++++++++++++++++++++++++--------- test/runtests.jl | 72 +++++++++++++++++++++++++++++++++++- 4 files changed, 132 insertions(+), 18 deletions(-) diff --git a/Project.toml b/Project.toml index 091ce2340..70b60739e 100644 --- a/Project.toml +++ b/Project.toml @@ -19,6 +19,7 @@ Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Random123 = "74087812-796a-5b5d-8853-05524746bad3" RandomNumbers = "e6cf234a-135c-5ec9-84dd-332b85af5143" +ScopedValues = "7e506255-f358-4e82-b7e4-beb19740aa63" SPIRVIntrinsics = "71d1d633-e7e8-4a92-83a1-de8814b09ba8" SPIRV_LLVM_Backend_jll = "4376b9bf-cff8-51b6-bb48-39421dff0d0c" SPIRV_Tools_jll = "6ac6d60f-d740-5983-97d7-a4482c0689f4" @@ -60,6 +61,7 @@ Printf = "<0.0.1, 1" Random = "1" Random123 = "1.7.1" RandomNumbers = "1.6.0" +ScopedValues = "1.3" SPIRVIntrinsics = "1.1.3" SPIRV_LLVM_Backend_jll = "23" SPIRV_Tools_jll = "2024.4, 2025.1" diff --git a/docs/src/extras/profiling.md b/docs/src/extras/profiling.md index 7ceac8f9f..3e3fbd406 100644 --- a/docs/src/extras/profiling.md +++ b/docs/src/extras/profiling.md @@ -55,6 +55,11 @@ kernel rather than its launch; pass `synchronize = false` to measure launches. P `trace = true` to list every range in order instead. The first call of a kernel includes its compilation, so profile a warmed-up run. +`@profile` records the task running the expression and the tasks it spawns, e.g. with +[`KernelAbstractions.@spawn`](@ref KernelAbstractions.@spawn), but not other tasks, so +profiles can run concurrently. Wait for spawned tasks within the expression, e.g. with +`@sync`, as what they record after it returns is lost. + ## Profilers Ranges are recorded on the host threads of the process, which is how NVTX, ITT and diff --git a/src/profiler.jl b/src/profiler.jl index d6f523a87..e04a67c0c 100644 --- a/src/profiler.jl +++ b/src/profiler.jl @@ -1,19 +1,24 @@ using Printf: @sprintf +using ScopedValues: ScopedValue, with ### # Built-in profiler: records the ranges of an expression, and summarizes them ### +# `task` numbers the tasks of a profile in the order they first recorded something, starting +# with 1 for the task that ran `@profile`; `thread` is the thread a range started on struct ProfileRange name::String start::UInt64 stop::UInt64 + task::Int thread::Int end struct ProfileMarker name::String time::UInt64 + task::Int thread::Int end @@ -21,27 +26,48 @@ struct ProfileTracer <: Tracer synchronize::Bool ranges::Vector{ProfileRange} markers::Vector{ProfileMarker} + tasks::IdDict{Task, Int} + open::Threads.Atomic{Int} # ranges started but not ended lock::ReentrantLock end -ProfileTracer(synchronize::Bool) = - ProfileTracer(synchronize, ProfileRange[], ProfileMarker[], ReentrantLock()) +ProfileTracer(synchronize::Bool) = ProfileTracer( + synchronize, ProfileRange[], ProfileMarker[], IdDict{Task, Int}(), Threads.Atomic{Int}(0), + ReentrantLock() +) -synchronizes_launches(tracer::ProfileTracer) = tracer.synchronize +# The profilers whose expression the current task is running, directly or in a task it +# spawned: tracers are global, but `@profile` only records the tasks of its expression. +const PROFILERS = ScopedValue{Vector{ProfileTracer}}(ProfileTracer[]) + +in_scope(tracer::ProfileTracer) = any(t -> t === tracer, PROFILERS[]) + +function task_number(tracer::ProfileTracer) + task = current_task() + return @lock tracer.lock get!(tracer.tasks, task, length(tracer.tasks) + 1) +end + +synchronizes_launches(tracer::ProfileTracer) = tracer.synchronize && in_scope(tracer) # like NVTXT, without domains of their own profile_name(label, domain) = domain == "KernelAbstractions" ? label : string(domain, ": ", label) -trace_range_start(::ProfileTracer, label, domain) = - (profile_name(label, domain), time_ns(), Threads.threadid()) +function trace_range_start(tracer::ProfileTracer, label, domain) + in_scope(tracer) || return nothing + Threads.atomic_add!(tracer.open, 1) + return (profile_name(label, domain), time_ns(), task_number(tracer), Threads.threadid()) +end -function trace_range_end(tracer::ProfileTracer, (name, start, thread)) - range = ProfileRange(name, start, time_ns(), thread) +trace_range_end(::ProfileTracer, ::Nothing) = nothing +function trace_range_end(tracer::ProfileTracer, (name, start, task, thread)) + range = ProfileRange(name, start, time_ns(), task, thread) @lock tracer.lock push!(tracer.ranges, range) + Threads.atomic_sub!(tracer.open, 1) return nothing end function trace_mark(tracer::ProfileTracer, label, domain) - marker = ProfileMarker(profile_name(label, domain), time_ns(), Threads.threadid()) + in_scope(tracer) || return nothing + marker = ProfileMarker(profile_name(label, domain), time_ns(), task_number(tracer), Threads.threadid()) @lock tracer.lock push!(tracer.markers, marker) return nothing end @@ -90,9 +116,15 @@ their range ends, so that on GPU backends they measure the kernel's execution in its launch. This serializes the host with the device, as `CUDA_LAUNCH_BLOCKING=1` does. With `trace = true`, the results list every range in the order it started instead, -indented by nesting. +indented by nesting on its task. + +Only the task running `expr` and the tasks it spawns (with `KernelAbstractions.@spawn`, +`Threads.@spawn` or `@async`) are recorded, so that other tasks, including other +`@profile`s, don't show up in the results. Wait for spawned tasks within `expr`, e.g. with +`@sync`: what a task records after `expr` has returned is lost, and `@profile` warns about +ranges that were still open. -Other registered tracers, e.g. NVTX under Nsight Systems, record the ranges too. +Other registered tracers, e.g. NVTX under Nsight Systems, record the ranges of all tasks. """ macro profile(args...) isempty(args) && throw(ArgumentError("KernelAbstractions.@profile needs an expression to profile")) @@ -114,13 +146,16 @@ end function profile(f; trace::Bool = false, synchronize::Bool = true) tracer = register_tracer!(ProfileTracer(synchronize)) + task_number(tracer) # the profiling task is task 1 start = time_ns() try - f() + with(f, PROFILERS => ProfileTracer[PROFILERS[]; tracer]) finally unregister_tracer!(tracer) end stop = time_ns() + open = tracer.open[] + open > 0 && @warn "$open profiled ranges were still open when `@profile` finished; wait for the tasks spawned within it, e.g. with `@sync`" return @lock tracer.lock ProfileResults(start, stop, copy(tracer.ranges), copy(tracer.markers), trace) end @@ -201,29 +236,31 @@ function show_summary(io::IO, results::ProfileResults, total) return end +location(event) = "task $(event.task) (thread $(event.thread))" + function show_trace(io::IO, results::ProfileResults) - # the nesting depth of each range among the ranges on its thread + # the nesting depth of each range among the ranges on its task events = sort!( [ [(r.start, r) for r in results.ranges]; [(m.time, m) for m in results.markers] ]; by = first ) - open = Dict{Int, Vector{UInt64}}() # stop times of the open ranges, per thread + open = Dict{Int, Vector{UInt64}}() # stop times of the open ranges, per task rows = Vector{String}[] for (time, event) in events - stack = get!(open, event.thread, UInt64[]) + stack = get!(open, event.task, UInt64[]) while !isempty(stack) && last(stack) <= time pop!(stack) end indent = " "^length(stack) if event isa ProfileRange push!(stack, event.stop) - push!(rows, [format_time(time - results.start), format_time(event.stop - event.start), string(event.thread), indent * event.name]) + push!(rows, [format_time(time - results.start), format_time(event.stop - event.start), location(event), indent * event.name]) else - push!(rows, [format_time(time - results.start), "", string(event.thread), indent * "◆ " * event.name]) + push!(rows, [format_time(time - results.start), "", location(event), indent * "◆ " * event.name]) end end - print_table(io, ["Start", "Duration", "Thread", "Name"], rows) + print_table(io, ["Start", "Duration", "On", "Name"], rows) return end diff --git a/test/runtests.jl b/test/runtests.jl index 0ea973e88..4fd9298ee 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -565,7 +565,12 @@ end @test occursin("recording 0 ranges.", sprint(show, MIME"text/plain"(), KernelAbstractions.@profile 1 + 1)) # launches synchronize their backend only if asked to - @test KernelAbstractions.synchronizes_launches(KernelAbstractions.ProfileTracer(true)) + tracer = KernelAbstractions.ProfileTracer(true) + # only for the tasks it profiles + @test !KernelAbstractions.synchronizes_launches(tracer) + @test KernelAbstractions.with(KernelAbstractions.PROFILERS => [tracer]) do + KernelAbstractions.synchronizes_launches(tracer) + end @test !KernelAbstractions.synchronizes_launches(KernelAbstractions.ProfileTracer(false)) @test !KernelAbstractions.synchronizes_launches(Testsuite.RecordingTracer()) results = KernelAbstractions.@profile synchronize = false kfill!(CPU())(A, 1.0f0; ndrange = length(A)) @@ -576,6 +581,71 @@ end @test !KernelAbstractions.profiling_active() @test_throws ArgumentError macroexpand(@__MODULE__, :(KernelAbstractions.@profile foo = 1 2)) + @testset "tasks" begin + As = [zeros(Float32, 64) for _ in 1:3] + work(i) = @profiling_range "task $i" kfill!(CPU())(As[i], 1.0f0; ndrange = length(As[i])) + + # spawned tasks are recorded, and numbered after the profiling task + results = KernelAbstractions.@profile @profiling_range "parent" begin + @sync for i in 1:3 + KernelAbstractions.@spawn CPU() work(i) + end + end + @test only(r.task for r in results.ranges if r.name == "parent") == 1 + @test sort([r.task for r in results.ranges if startswith(r.name, "task ")]) == 2:4 + # each kernel range is on the task that launched it + for i in 1:3 + task = only(r.task for r in results.ranges if r.name == "task $i") + @test count(r -> r.name == "profiling_fill!" && r.task == task, results.ranges) == 1 + end + trace = sprint(show, MIME"text/plain"(), KernelAbstractions.ProfileResults(results.start, results.stop, results.ranges, results.markers, true)) + @test occursin("task 1 (thread ", trace) + + # other tasks aren't + stop = Threads.Atomic{Bool}(false) + other = Threads.@spawn while !stop[] + @profiling_range "unrelated" yield() + end + results = KernelAbstractions.@profile for _ in 1:10 + @profiling_range "related" yield() + end + stop[] = true + wait(other) + @test all(r -> r.name == "related", results.ranges) + @test length(results.ranges) == 10 + + # nor are other profiles, at the same time or nested + t1 = Threads.@spawn KernelAbstractions.@profile for _ in 1:5 + @profiling_range "one" yield() + end + t2 = Threads.@spawn KernelAbstractions.@profile for _ in 1:7 + @profiling_range "two" yield() + end + r1, r2 = fetch(t1), fetch(t2) + @test all(r -> r.name == "one", r1.ranges) && length(r1.ranges) == 5 + @test all(r -> r.name == "two", r2.ranges) && length(r2.ranges) == 7 + local inner + outer = KernelAbstractions.@profile @profiling_range "outer" begin + inner = KernelAbstractions.@profile @profiling_range "inner" nothing + end + @test sort([r.name for r in outer.ranges]) == ["inner", "outer"] + @test [r.name for r in inner.ranges] == ["inner"] + + # ranges of tasks that outlive the profile are lost, with a warning + started, finish = Channel{Nothing}(1), Channel{Nothing}(1) + local task + results = @test_logs (:warn, r"1 profiled ranges were still open") KernelAbstractions.@profile begin + task = Threads.@spawn @profiling_range "outlives" begin + put!(started, nothing) + take!(finish) + end + take!(started) + end + put!(finish, nothing) + wait(task) + @test isempty(results.ranges) + end + @test KernelAbstractions.format_time(5) == "5 ns" @test KernelAbstractions.format_time(999.7) == "1 µs" @test KernelAbstractions.format_time(1.234e6) == "1.23 ms" From f1d1f3bbb1e5af3b38d1e255867864f681438ffd Mon Sep 17 00:00:00 2001 From: Valentin Churavy Date: Sun, 4 Oct 2026 13:34:43 +0200 Subject: [PATCH 4/9] Annotate `KernelAbstractions.@spawn` with a profiler range Each task is a range from `@spawn` until its queued work has completed, named after the call site or the new `name = ...` argument. The range starts in the spawning task, so `@profile` warns about a task it wasn't waited for even if the task hasn't run yet. `@profile` attributes a range to the task it ended on, which puts a spawn's range at the root of the spawned task. Assisted-by: Claude Code (Opus 5.5) --- src/profiler.jl | 14 +++++++---- src/spawn.jl | 62 +++++++++++++++++++++++++++++++++++------------- test/runtests.jl | 19 ++++++++++++++- test/spawn.jl | 6 +++++ 4 files changed, 78 insertions(+), 23 deletions(-) diff --git a/src/profiler.jl b/src/profiler.jl index e04a67c0c..8330a92f4 100644 --- a/src/profiler.jl +++ b/src/profiler.jl @@ -6,7 +6,9 @@ using ScopedValues: ScopedValue, with ### # `task` numbers the tasks of a profile in the order they first recorded something, starting -# with 1 for the task that ran `@profile`; `thread` is the thread a range started on +# with 1 for the task that ran `@profile`. A range belongs to the task, and thread, it ended +# on: the range of a `KernelAbstractions.@spawn` starts in the spawning task, and ends in +# the spawned one. struct ProfileRange name::String start::UInt64 @@ -54,12 +56,12 @@ profile_name(label, domain) = domain == "KernelAbstractions" ? label : string(do function trace_range_start(tracer::ProfileTracer, label, domain) in_scope(tracer) || return nothing Threads.atomic_add!(tracer.open, 1) - return (profile_name(label, domain), time_ns(), task_number(tracer), Threads.threadid()) + return (profile_name(label, domain), time_ns()) end trace_range_end(::ProfileTracer, ::Nothing) = nothing -function trace_range_end(tracer::ProfileTracer, (name, start, task, thread)) - range = ProfileRange(name, start, time_ns(), task, thread) +function trace_range_end(tracer::ProfileTracer, (name, start)) + range = ProfileRange(name, start, time_ns(), task_number(tracer), Threads.threadid()) @lock tracer.lock push!(tracer.ranges, range) Threads.atomic_sub!(tracer.open, 1) return nothing @@ -155,13 +157,15 @@ function profile(f; trace::Bool = false, synchronize::Bool = true) end stop = time_ns() open = tracer.open[] - open > 0 && @warn "$open profiled ranges were still open when `@profile` finished; wait for the tasks spawned within it, e.g. with `@sync`" + open > 0 && @warn "$(plural(open, "profiled range")) still open when `@profile` finished; wait for the tasks spawned within it, e.g. with `@sync`" return @lock tracer.lock ProfileResults(start, stop, copy(tracer.ranges), copy(tracer.markers), trace) end ## report +plural(n, what) = string(n, " ", what, n == 1 ? "" : "s") + function format_time(ns::Real) # switch units where three significant digits would round up to the next one ns < 999.5 && return @sprintf("%.0f ns", ns) diff --git a/src/spawn.jl b/src/spawn.jl index dc319a692..ac8bc0957 100644 --- a/src/spawn.jl +++ b/src/spawn.jl @@ -1,5 +1,5 @@ """ - @spawn [threadpool] backend [device=id] expr + @spawn [threadpool] backend [device=id] [name=label] expr Run `expr` on a new Julia task, like `Threads.@spawn`, and return the `Task`. Use it in place of `Threads.@spawn` to launch kernels from a task. It guarantees that @@ -30,6 +30,20 @@ end fetch(task) == 4 * length(A) ``` +# Profiling + +Each task is a profiler range (see [`@profiling_range`](@ref)), from `@spawn` until its +queued work has completed, i.e. what `wait(task)` waits for. It is named after the call +site, e.g. `"@spawn solver.jl:42"`, or after the `name` argument: + +```julia +task = KernelAbstractions.@spawn backend name = "halo exchange" exchange!(u) +``` + +The range starts in the spawning task, so [`KernelAbstractions.@profile`](@ref +KernelAbstractions.@profile) warns about a task it wasn't waited for, even if the task +hasn't run yet. + # Choosing the device Backends keep the active device in task-local state, and Julia does not copy that state @@ -68,25 +82,28 @@ for the protocol behind these guarantees, and for how to support it without a fu [`synchronize`](@ref). """ macro spawn(args...) - usage = "@spawn expects `@spawn [threadpool] backend [device=id] expr`" + usage = "@spawn expects `@spawn [threadpool] backend [device=id] [name=label] expr`" isempty(args) && throw(ArgumentError(usage)) - # `expr` is always last, so a top-level `=` anywhere before it is our `device=id` - # argument rather than part of the user's code. + # `expr` is always last, so a top-level `=` anywhere before it is our `device=id` or + # `name=label` argument rather than part of the user's code. expr = last(args) - is_device_arg(x) = Meta.isexpr(x, :(=), 2) && x.args[1] === :device - is_device_arg(expr) && throw(ArgumentError("$usage; `device=id` must be followed by the expression to run")) + is_kwarg(x) = Meta.isexpr(x, :(=), 2) && x.args[1] in (:device, :name) + is_kwarg(expr) && throw(ArgumentError("$usage; `$(expr.args[1])=...` must be followed by the expression to run")) - device = nothing + kwargs = Dict{Symbol, Any}() positional = Any[] for arg in args[1:(end - 1)] - if is_device_arg(arg) - device === nothing || throw(ArgumentError("@spawn accepts at most one `device=id` argument")) - device = arg.args[2] + if is_kwarg(arg) + key = arg.args[1] + haskey(kwargs, key) && throw(ArgumentError("@spawn accepts at most one `$key=...` argument")) + kwargs[key] = arg.args[2] else push!(positional, arg) end end + device = Base.get(kwargs, :device, nothing) + name = Base.get(kwargs, :name, "@spawn $(basename(string(__source__.file))):$(__source__.line)") if length(positional) == 1 threadpool = nothing @@ -101,15 +118,24 @@ macro spawn(args...) # scope: that is what lets an enclosing `@sync` see the task, and what makes `$x` # interpolation in `expr` work. Our own temporaries are gensyms so they cannot clash # with the user's variables. - b, dev, event, result = gensym(:backend), gensym(:dev), gensym(:event), gensym(:result) + b, dev, event, result, range = gensym(:backend), gensym(:dev), gensym(:event), gensym(:result), gensym(:range) # `device!` comes first because `wait_event` acts on the queue of the device that is # active when it is called: selecting the device afterwards would leave it unordered. + # The profiler range ends once the task's work has completed, also if `expr` throws. body = quote - $KI.device!($b, $dev) - $KI.wait_event($b, $event) - local $result = $expr - $KI.synchronize($b) - $result + $( + Expr( + :tryfinally, + quote + $KI.device!($b, $dev) + $KI.wait_event($b, $event) + local $result = $expr + $KI.synchronize($b) + $result + end, + :($profiling_range_end($range)) + ) + ) end task = if threadpool === nothing :(Threads.@spawn $body) @@ -118,12 +144,14 @@ macro spawn(args...) end # `device` and the event are both evaluated in the spawning task, so the event captures - # the work queued on the spawning task's device, not on `dev`. + # the work queued on the spawning task's device, not on `dev`. So is the start of the + # profiler range, so that `@profile` knows about the task even before it runs. return esc( quote local $b = $backend local $dev = $(device === nothing ? :($KI.device($b)) : device) local $event = $KI.record_event($b) + local $range = $profiling_active() ? $profiling_range_start($name) : nothing $task end ) diff --git a/test/runtests.jl b/test/runtests.jl index 4fd9298ee..d546eaaff 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -598,6 +598,15 @@ end task = only(r.task for r in results.ranges if r.name == "task $i") @test count(r -> r.name == "profiling_fill!" && r.task == task, results.ranges) == 1 end + # `@spawn` ranges are named after the call site, and belong to the spawned task + spawns = filter(r -> startswith(r.name, "@spawn runtests.jl:"), results.ranges) + @test sort([r.task for r in spawns]) == 2:4 + for r in spawns + child = only(c for c in results.ranges if c.task == r.task && startswith(c.name, "task ")) + @test r.start <= child.start <= child.stop <= r.stop + end + named = KernelAbstractions.@profile wait(KernelAbstractions.@spawn CPU() name = "named" nothing) + @test only(named.ranges).name == "named" trace = sprint(show, MIME"text/plain"(), KernelAbstractions.ProfileResults(results.start, results.stop, results.ranges, results.markers, true)) @test occursin("task 1 (thread ", trace) @@ -634,7 +643,7 @@ end # ranges of tasks that outlive the profile are lost, with a warning started, finish = Channel{Nothing}(1), Channel{Nothing}(1) local task - results = @test_logs (:warn, r"1 profiled ranges were still open") KernelAbstractions.@profile begin + results = @test_logs (:warn, r"1 profiled range still open") KernelAbstractions.@profile begin task = Threads.@spawn @profiling_range "outlives" begin put!(started, nothing) take!(finish) @@ -644,6 +653,14 @@ end put!(finish, nothing) wait(task) @test isempty(results.ranges) + + # also for a `@spawn` task that hasn't started yet, as its range starts at `@spawn` + go = Channel{Nothing}(1) + results = @test_logs (:warn, r"1 profiled range still open") KernelAbstractions.@profile begin + task = KernelAbstractions.@spawn CPU() take!(go) + end + put!(go, nothing) + wait(task) end @test KernelAbstractions.format_time(5) == "5 ns" diff --git a/test/spawn.jl b/test/spawn.jl index 88a2b1700..d337bff9f 100644 --- a/test/spawn.jl +++ b/test/spawn.jl @@ -89,6 +89,12 @@ function spawn_testsuite(Backend, AT) # A top-level assignment in the body is a body, not a `device=` argument. @test fetch(KernelAbstractions.@spawn backend y = 41 + 1) == 42 + + # `name=` labels the task's profiler range, and combines with `device=`. + @test fetch(KernelAbstractions.@spawn backend name = "labelled" 1) == 1 + @test fetch(KernelAbstractions.@spawn backend device = dev name = "label $dev" 2) == 2 + @test_throws ArgumentError macroexpand(@__MODULE__, :(KernelAbstractions.@spawn $backend name = "a" name = "b" 1)) + @test_throws ArgumentError macroexpand(@__MODULE__, :(KernelAbstractions.@spawn $backend name = "a")) end @testset "@sync" begin From e97d80bfd6d9ebb396006fb45f3dcbe75a779541 Mon Sep 17 00:00:00 2001 From: Valentin Churavy Date: Sun, 4 Oct 2026 14:06:42 +0200 Subject: [PATCH 5/9] Make annotations free when tracing is off - `@profiling_range` only enters the exception handler that ends the range when a profiler listens, compiling `expr` twice. This takes the cost of an annotated call from +6 ns to +0.3 ns when tracing is off. `@goto` and `@label` are therefore not supported in `expr`. - `profiling_mark` inlines the check for a profiler (+3.6 ns to +0 ns). - The traced launch path is inferred once with `@nospecializeinfer`, rather than along with the launch of every new kernel, which cost about 1.8 MiB of compiler allocations on each first launch. Assisted-by: Claude Code (Opus 5.5) --- src/backend_launch.jl | 7 +++++-- src/profiling.jl | 25 ++++++++++++++++++------- test/runtests.jl | 5 +++++ 3 files changed, 28 insertions(+), 9 deletions(-) diff --git a/src/backend_launch.jl b/src/backend_launch.jl index 125588ff4..94425e376 100644 --- a/src/backend_launch.jl +++ b/src/backend_launch.jl @@ -89,8 +89,11 @@ function launch_tuple(obj::Kernel, args::Tuple; ndrange = nothing, workgroupsize return launch_untraced(obj, args, ndrange, workgroupsize) end -# out of line, to keep the profiler out of the common path -@noinline function launch_traced(obj::Kernel, args::Tuple, ndrange, workgroupsize) +# Out of line, to keep the profiler out of the common path, and inferred only once rather +# than for every kernel, which would slow down every first launch. +Base.@nospecializeinfer @noinline function launch_traced( + @nospecialize(obj::Kernel), @nospecialize(args::Tuple), @nospecialize(ndrange), @nospecialize(workgroupsize) + ) id = profiling_range_start(kernel_label(obj.f)) try launch_untraced(obj, args, ndrange, workgroupsize) diff --git a/src/profiling.jl b/src/profiling.jl index a9a396d5a..61016301d 100644 --- a/src/profiling.jl +++ b/src/profiling.jl @@ -125,9 +125,12 @@ synchronizes_launches(::Nothing) = false Record an instantaneous marker named `label` in `domain`. """ -function profiling_mark(label; domain = "KernelAbstractions") +# the check is inlined into the caller, so that a marker costs nothing when nobody listens +@inline profiling_mark(label; domain = "KernelAbstractions") = + profiling_active() ? record_mark(label, domain) : nothing + +@noinline function record_mark(label, domain) current = tracers() - isempty(current) && return nothing label, domain = String(label), String(domain) for tracer in current trace_mark(tracer, label, domain) @@ -143,6 +146,10 @@ only evaluated if a profiler is listening (see [`profiling_active`](@ref)), so i built with string interpolation at no cost to unprofiled runs. The range is ended if `expr` throws. Assignments in `expr` are visible after the macro, as with `@time`. +`expr` is compiled twice, for when a profiler listens and for when none does, so that the +latter costs no more than a check. It therefore can't define labels: `@goto` and `@label` +are not supported in `expr`. + ```julia @profiling_range "volume integral" begin volume_integral!(du, u, backend) @@ -169,14 +176,18 @@ macro profiling_range(label, args...) end end id = gensym(:id) + # unlike `try`, `tryfinally` doesn't introduce a scope + traced = Expr(:tryfinally, esc(expr), :($profiling_range_end($id))) + # Entering the exception handler that ends the range costs more than checking for a + # profiler, so it is only entered when one listens, at the price of compiling `expr` + # twice. return quote - local $id = if $profiling_active() - $profiling_range_start($(esc(label)); domain = $(esc(domain))) + if $profiling_active() + local $id = $profiling_range_start($(esc(label)); domain = $(esc(domain))) + $traced else - nothing + $(esc(expr)) end - # unlike `try`, `tryfinally` doesn't introduce a scope - $(Expr(:tryfinally, esc(expr), :($profiling_range_end($id)))) end end diff --git a/test/runtests.jl b/test/runtests.jl index d546eaaff..d681050d0 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -453,6 +453,11 @@ end @profiling_range "assign" y = 42 @test y == 42 + # the expression is evaluated once + count = Ref(0) + @test (@profiling_range "once" (count[] += 1)) == 1 + @test count[] == 1 + @test_throws ArgumentError macroexpand(@__MODULE__, :(@profiling_range "label" foo = 1 2)) end From 056214243b517704d542e122f4d8aab026de3434 Mon Sep 17 00:00:00 2001 From: Valentin Churavy Date: Sun, 4 Oct 2026 16:10:03 +0200 Subject: [PATCH 6/9] Make tracing cheaper when it is on - Labels written as literals (in `@profiling_range` and `@spawn`), and kernel names, are passed to tracers as `Symbol`s. They are fixed in the code, so tracers can cache what they derive from them by identity; labels computed at run time stay `String`s. Domains are `Symbol`s. - The NVTX tracer registers `Symbol` labels with NVTX once, and caches domains by `Symbol` (a range from ~480 to ~410 ns under nsys). - The NVTXT tracer formats integers into a reused buffer instead of interpolating a string per record, and caches the messages of `Symbol` labels (a range from ~720 ns and 1.8 KB to ~400 ns and 160 B). - Kernel labels are computed once per kernel type, and the traced launch calls the same `launch_untraced` as the untraced one, rather than a dynamically dispatched copy. Assisted-by: Claude Code (Opus 5.5) --- ext/IntelITTExt.jl | 10 ++-- ext/NVTXExt.jl | 35 ++++++++++--- ext/ROCTXExt.jl | 5 +- src/backend_launch.jl | 19 +++---- src/profiler.jl | 2 +- src/profiling.jl | 112 +++++++++++++++++++++++++++++++----------- src/spawn.jl | 2 + test/profiling.jl | 15 ++++-- test/runtests.jl | 11 ++++- 9 files changed, 148 insertions(+), 63 deletions(-) diff --git a/ext/IntelITTExt.jl b/ext/IntelITTExt.jl index cd26b68b9..7142791ce 100644 --- a/ext/IntelITTExt.jl +++ b/ext/IntelITTExt.jl @@ -5,17 +5,17 @@ import IntelITT # forwards ranges to Intel VTune as ITT tasks, one ITT domain per domain struct ITTTracer <: KA.Tracer - domains::Dict{String, IntelITT.Domain} + domains::IdDict{Symbol, IntelITT.Domain} lock::ReentrantLock end -ITTTracer() = ITTTracer(Dict{String, IntelITT.Domain}(), ReentrantLock()) +ITTTracer() = ITTTracer(IdDict{Symbol, IntelITT.Domain}(), ReentrantLock()) -domain(tracer::ITTTracer, name::String) = - @lock tracer.lock get!(() -> IntelITT.Domain(name), tracer.domains, name) +domain(tracer::ITTTracer, name::Symbol) = + @lock tracer.lock get!(() -> IntelITT.Domain(String(name)), tracer.domains, name) function KA.trace_range_start(tracer::ITTTracer, label, domain_name) # overlapped tasks may end on another thread, and needn't nest - task = IntelITT.Task(domain(tracer, domain_name), label) + task = IntelITT.Task(domain(tracer, domain_name), String(label)) IntelITT.start(task) return task end diff --git a/ext/NVTXExt.jl b/ext/NVTXExt.jl index 0fd9887a6..7515d5e55 100644 --- a/ext/NVTXExt.jl +++ b/ext/NVTXExt.jl @@ -5,21 +5,40 @@ import NVTX # forwards ranges to Nsight Systems as NVTX ranges, one NVTX domain per domain. NVTX ranges # annotate host threads; Nsight Systems projects the GPU work launched within them itself. +struct NVTXDomain + domain::NVTX.Domain + # `Symbol` labels are fixed in the code, so they are registered with NVTX once, which + # makes recording them cheaper. Other labels are passed as they are. + strings::IdDict{Symbol, NVTX.StringHandle} +end + struct NVTXTracer <: KA.Tracer - domains::Dict{String, NVTX.Domain} + domains::IdDict{Symbol, NVTXDomain} lock::ReentrantLock end -NVTXTracer() = NVTXTracer(Dict{String, NVTX.Domain}(), ReentrantLock()) +NVTXTracer() = NVTXTracer(IdDict{Symbol, NVTXDomain}(), ReentrantLock()) -domain(tracer::NVTXTracer, name::String) = - @lock tracer.lock get!(() -> NVTX.Domain(name), tracer.domains, name) +function domain(tracer::NVTXTracer, name::Symbol) + return @lock tracer.lock get!(tracer.domains, name) do + NVTXDomain(NVTX.Domain(String(name)), IdDict{Symbol, NVTX.StringHandle}()) + end +end + +message(::NVTXTracer, d::NVTXDomain, label::String) = label +message(tracer::NVTXTracer, d::NVTXDomain, label::Symbol) = + @lock tracer.lock get!(() -> NVTX.StringHandle(d.domain, String(label)), d.strings, label) # process ranges, rather than push/pop, since they may end on another thread -KA.trace_range_start(tracer::NVTXTracer, label, domain_name) = - NVTX.range_start(domain(tracer, domain_name); message = label) +function KA.trace_range_start(tracer::NVTXTracer, label, domain_name) + d = domain(tracer, domain_name) + return NVTX.range_start(d.domain; message = message(tracer, d, label)) +end KA.trace_range_end(::NVTXTracer, id::NVTX.RangeId) = (NVTX.range_end(id); nothing) -KA.trace_mark(tracer::NVTXTracer, label, domain_name) = - (NVTX.mark(domain(tracer, domain_name); message = label); nothing) +function KA.trace_mark(tracer::NVTXTracer, label, domain_name) + d = domain(tracer, domain_name) + NVTX.mark(d.domain; message = message(tracer, d, label)) + return nothing +end const TRACER = Ref{NVTXTracer}() diff --git a/ext/ROCTXExt.jl b/ext/ROCTXExt.jl index 25c56de9b..a3f7366bc 100644 --- a/ext/ROCTXExt.jl +++ b/ext/ROCTXExt.jl @@ -6,7 +6,7 @@ using Base.Libc: Libdl # forwards ranges to rocprof as roctx ranges. like NVTX, roctx annotates host threads, and # rocprof attributes the GPU work launched within them. roctx has no domains, so ranges in a -# domain other than "KernelAbstractions" are prefixed with it. +# domain other than `:KernelAbstractions` are prefixed with it. struct ROCTXTracer <: KA.Tracer range_start::Ptr{Cvoid} range_stop::Ptr{Cvoid} @@ -27,7 +27,8 @@ function ROCTXTracer(library::AbstractString) ) end -roctx_message(label, domain) = domain == "KernelAbstractions" ? label : string(domain, ": ", label) +# a `Symbol` is passed to C as its name, without allocating +roctx_message(label, domain) = domain === KA.DEFAULT_DOMAIN ? label : string(domain, ": ", label) # process ranges, rather than push/pop, since they may end on another thread KA.trace_range_start(tracer::ROCTXTracer, label, domain) = diff --git a/src/backend_launch.jl b/src/backend_launch.jl index 94425e376..ac3a16c7d 100644 --- a/src/backend_launch.jl +++ b/src/backend_launch.jl @@ -83,18 +83,10 @@ Core.kwcall(kwargs::NamedTuple, obj::Kernel{<:KI.Backend}, args::Vararg{Any, N}) launch_tuple(obj, args; kwargs...) function launch_tuple(obj::Kernel, args::Tuple; ndrange = nothing, workgroupsize = nothing) - if profiling_active() - return launch_traced(obj, args, ndrange, workgroupsize) - end - return launch_untraced(obj, args, ndrange, workgroupsize) -end - -# Out of line, to keep the profiler out of the common path, and inferred only once rather -# than for every kernel, which would slow down every first launch. -Base.@nospecializeinfer @noinline function launch_traced( - @nospecialize(obj::Kernel), @nospecialize(args::Tuple), @nospecialize(ndrange), @nospecialize(workgroupsize) - ) - id = profiling_range_start(kernel_label(obj.f)) + profiling_active() || return launch_untraced(obj, args, ndrange, workgroupsize) + # The traced launch calls the same `launch_untraced`, so that it isn't inferred twice for + # every kernel. The helpers around it are only inferred once. + id = start_launch_range(obj) try launch_untraced(obj, args, ndrange, workgroupsize) # a profiler that measures kernels rather than launches @@ -105,6 +97,9 @@ Base.@nospecializeinfer @noinline function launch_traced( return nothing end +Base.@nospecializeinfer @noinline start_launch_range(@nospecialize(obj::Kernel)) = + profiling_range_start(kernel_label(obj.f)) + function launch_untraced(obj::Kernel, args::Tuple, ndrange, workgroupsize) ndrange, workgroupsize, iterspace, dynamic = launch_config(obj, ndrange, workgroupsize) # nothing to launch (or compile) for an empty ndrange diff --git a/src/profiler.jl b/src/profiler.jl index 8330a92f4..0ebee12d0 100644 --- a/src/profiler.jl +++ b/src/profiler.jl @@ -51,7 +51,7 @@ end synchronizes_launches(tracer::ProfileTracer) = tracer.synchronize && in_scope(tracer) # like NVTXT, without domains of their own -profile_name(label, domain) = domain == "KernelAbstractions" ? label : string(domain, ": ", label) +profile_name(label, domain) = domain === DEFAULT_DOMAIN ? String(label) : string(domain, ": ", label) function trace_range_start(tracer::ProfileTracer, label, domain) in_scope(tracer) || return nothing diff --git a/src/profiling.jl b/src/profiling.jl index 61016301d..96069deb7 100644 --- a/src/profiling.jl +++ b/src/profiling.jl @@ -17,14 +17,20 @@ instance with [`register_tracer!`](@ref) to receive the ranges and markers of Subtypes implement - trace_range_start(tracer, label::String, domain::String) -> id + trace_range_start(tracer, label::Label, domain::Symbol) -> id trace_range_end(tracer, id) - trace_mark(tracer, label::String, domain::String) # optional + trace_mark(tracer, label::Label, domain::Symbol) # optional synchronizes_launches(tracer)::Bool # optional, default `false` If `synchronizes_launches` is `true`, kernel launches synchronize their backend before their range ends, so that the range measures the kernel's execution rather than its launch. +A `Label` is a `Symbol` for labels that are fixed in the code: literals in +[`@profiling_range`](@ref), kernel names and `@spawn` call sites. As there are only so many +of those, tracers may cache what they derive from a `Symbol` label, e.g. a registered +string, keyed by its identity. Labels computed at run time are `String`s, which tracers +should not cache. The default domain is `:KernelAbstractions`. + Ranges may end on a different thread than they started on, and may overlap without nesting, so implementations should use the profiler's start/end API (e.g. `nvtxRangeStartEx`) rather than a thread-local push/pop stack. @@ -34,6 +40,16 @@ registering any tracer makes every annotation do work. """ abstract type Tracer end +const Label = Union{Symbol, String} +const DEFAULT_DOMAIN = :KernelAbstractions + +# literals are made `Symbol`s by the macros, anything else stays a `String` +as_label(label::Symbol) = label +as_label(label::AbstractString) = String(label) +as_label(label) = string(label) +as_domain(domain::Symbol) = domain +as_domain(domain) = Symbol(domain) + function trace_range_start end function trace_range_end end trace_mark(::Tracer, label, domain) = nothing @@ -89,7 +105,7 @@ struct ProfilingRange end """ - profiling_range_start(label::AbstractString; domain = "KernelAbstractions") + profiling_range_start(label; domain = :KernelAbstractions) Start a range named `label` in `domain`, and return a handle for [`profiling_range_end`](@ref). Returns `nothing` if no profiler is listening. @@ -97,10 +113,10 @@ Start a range named `label` in `domain`, and return a handle for Prefer [`@profiling_range`](@ref), which ends the range even if an exception is thrown. Use these for ranges that don't follow the structure of the code. """ -function profiling_range_start(label; domain = "KernelAbstractions") +function profiling_range_start(label; domain = DEFAULT_DOMAIN) current = tracers() isempty(current) && return nothing - label, domain = String(label), String(domain) + label, domain = as_label(label), as_domain(domain) return ProfilingRange(current, Any[trace_range_start(t, label, domain) for t in current]) end @@ -121,17 +137,17 @@ synchronizes_launches(range::ProfilingRange) = any(synchronizes_launches, range. synchronizes_launches(::Nothing) = false """ - profiling_mark(label::AbstractString; domain = "KernelAbstractions") + profiling_mark(label; domain = :KernelAbstractions) Record an instantaneous marker named `label` in `domain`. """ # the check is inlined into the caller, so that a marker costs nothing when nobody listens -@inline profiling_mark(label; domain = "KernelAbstractions") = +@inline profiling_mark(label; domain = DEFAULT_DOMAIN) = profiling_active() ? record_mark(label, domain) : nothing @noinline function record_mark(label, domain) current = tracers() - label, domain = String(label), String(domain) + label, domain = as_label(label), as_domain(domain) for tracer in current trace_mark(tracer, label, domain) end @@ -139,7 +155,7 @@ Record an instantaneous marker named `label` in `domain`. end """ - @profiling_range label [domain = "KernelAbstractions"] expr + @profiling_range label [domain = :KernelAbstractions] expr Evaluate `expr` inside a profiler range named `label`, and return its value. `label` is only evaluated if a profiler is listening (see [`profiling_active`](@ref)), so it can be @@ -167,7 +183,7 @@ Nsight Systems. Kernel launches are annotated with the name of the kernel automa macro profiling_range(label, args...) isempty(args) && throw(ArgumentError("@profiling_range needs an expression to evaluate")) expr = args[end] - domain = "KernelAbstractions" + domain = QuoteNode(DEFAULT_DOMAIN) for kw in args[1:(end - 1)] if Meta.isexpr(kw, :(=)) && kw.args[1] === :domain domain = kw.args[2] @@ -183,7 +199,7 @@ macro profiling_range(label, args...) # twice. return quote if $profiling_active() - local $id = $profiling_range_start($(esc(label)); domain = $(esc(domain))) + local $id = $profiling_range_start($(literal(label)); domain = $(literal(domain))) $traced else $(esc(expr)) @@ -191,11 +207,21 @@ macro profiling_range(label, args...) end end +# a string literal is a `Symbol` label, fixed in the code; anything else is evaluated +literal(x::String) = QuoteNode(Symbol(x)) +literal(x::QuoteNode) = x +literal(x) = esc(x) + # the name kernel launches are annotated with: `@kernel function f` compiles to `gpu_f` +# (computed once per kernel function type) +const KERNEL_LABELS = IdDict{Any, Symbol}() +const KERNEL_LABELS_LOCK = ReentrantLock() function kernel_label(f) - f isa Function || return string(typeof(f)) - name = string(nameof(f)) - return startswith(name, "gpu_") ? name[5:end] : name + return @lock KERNEL_LABELS_LOCK get!(KERNEL_LABELS, typeof(f)) do + f isa Function || return Symbol(typeof(f)) + name = string(nameof(f)) + return Symbol(startswith(name, "gpu_") ? name[5:end] : name) + end end @@ -219,11 +245,13 @@ working directory. Otherwise, register one with [`register_tracer!`](@ref), and after unregistering it to flush the file. Records are written when a range ends, as a single line, so that threads don't interleave. -Ranges in a `domain` other than `"KernelAbstractions"` are prefixed with it. +Ranges in a `domain` other than `:KernelAbstractions` are prefixed with it. """ struct NVTXTTracer{IO_ <: IO} <: Tracer io::IO_ lock::ReentrantLock + record::Vector{UInt8} # formatted under the lock + messages::IdDict{Tuple{Symbol, Symbol}, String} # of `Symbol` labels end function NVTXTTracer(io::IO) @@ -243,7 +271,7 @@ function NVTXTTracer(io::IO) TimeBase = Manual """ ) - return NVTXTTracer(io, ReentrantLock()) + return NVTXTTracer(io, ReentrantLock(), sizehint!(UInt8[], 256), IdDict{Tuple{Symbol, Symbol}, String}()) end NVTXTTracer(path::AbstractString) = NVTXTTracer(open(path, "w")) @@ -256,31 +284,57 @@ struct NVTXTRange end # the message is a quoted string, and a record a line -function nvtxt_message(label, domain) - message = domain == "KernelAbstractions" ? label : string(domain, ": ", label) +function nvtxt_message(label, domain::Symbol) + message = domain === DEFAULT_DOMAIN ? String(label) : string(domain, ": ", label) return replace(message, '"' => '\'', '\n' => ' ', '\r' => ' ') end - -function nvtxt_record(tracer::NVTXTTracer, record::String) - @lock tracer.lock isopen(tracer.io) && write(tracer.io, record) +nvtxt_message(tracer::NVTXTTracer, label::String, domain::Symbol) = nvtxt_message(label, domain) +nvtxt_message(tracer::NVTXTTracer, label::Symbol, domain::Symbol) = + @lock tracer.lock get!(() -> nvtxt_message(label, domain), tracer.messages, (label, domain)) + +# integers are formatted by hand, as `print` allocates a string for each +function append_decimal!(buffer::Vector{UInt8}, x::Unsigned) + first = length(buffer) + 1 + while true + push!(buffer, UInt8('0') + (x % 10) % UInt8) + x ÷= 10 + x == 0 && break + end + reverse!(buffer, first, length(buffer)) + return buffer +end +append_decimal!(buffer::Vector{UInt8}, x::Integer) = append_decimal!(buffer, unsigned(x)) + +function nvtxt_record(tracer::NVTXTTracer, kind::String, times, thread::Int, message::String) + @lock tracer.lock begin + isopen(tracer.io) || return nothing + buffer = empty!(tracer.record) + append!(buffer, codeunits(kind)) + for time in times + append!(buffer, codeunits(", ")) + append_decimal!(buffer, time) + end + append!(buffer, codeunits(", ")) + append_decimal!(buffer, thread) + append!(buffer, codeunits(", \"")) + append!(buffer, codeunits(message)) + append!(buffer, codeunits("\"\n")) + write(tracer.io, buffer) + end return nothing end -trace_range_start(::NVTXTTracer, label, domain) = - NVTXTRange(time_ns(), Threads.threadid(), nvtxt_message(label, domain)) +trace_range_start(tracer::NVTXTTracer, label, domain) = + NVTXTRange(time_ns(), Threads.threadid(), nvtxt_message(tracer, label, domain)) function trace_range_end(tracer::NVTXTTracer, range::NVTXTRange) stop = time_ns() - return nvtxt_record( - tracer, "RangeStartEnd, $(range.start), $stop, $(range.thread), \"$(range.message)\"\n" - ) + return nvtxt_record(tracer, "RangeStartEnd", (range.start, stop), range.thread, range.message) end function trace_mark(tracer::NVTXTTracer, label, domain) time = time_ns() - return nvtxt_record( - tracer, "Marker, $time, $(Threads.threadid()), \"$(nvtxt_message(label, domain))\"\n" - ) + return nvtxt_record(tracer, "Marker", (time,), Threads.threadid(), nvtxt_message(tracer, label, domain)) end function nvtxt_path(setting::AbstractString) diff --git a/src/spawn.jl b/src/spawn.jl index ac8bc0957..dac17a689 100644 --- a/src/spawn.jl +++ b/src/spawn.jl @@ -103,7 +103,9 @@ macro spawn(args...) end end device = Base.get(kwargs, :device, nothing) + # a literal name, like the default, is a `Symbol` label (see `Tracer`) name = Base.get(kwargs, :name, "@spawn $(basename(string(__source__.file))):$(__source__.line)") + name isa String && (name = QuoteNode(Symbol(name))) if length(positional) == 1 threadpool = nothing diff --git a/test/profiling.jl b/test/profiling.jl index 97d8653ae..788bfb21e 100644 --- a/test/profiling.jl +++ b/test/profiling.jl @@ -1,17 +1,22 @@ -# records the ranges and markers it is given +# records the ranges and markers it is given, with labels and domains as strings, and the +# types of the labels in `types` struct RecordingTracer <: KernelAbstractions.Tracer events::Vector{Any} + types::Vector{Any} lock::ReentrantLock end -RecordingTracer() = RecordingTracer([], ReentrantLock()) +RecordingTracer() = RecordingTracer([], [], ReentrantLock()) function KernelAbstractions.trace_range_start(t::RecordingTracer, label, domain) - @lock t.lock push!(t.events, (:start, label, domain)) - return label + @lock t.lock begin + push!(t.events, (:start, String(label), String(domain))) + push!(t.types, (typeof(label), typeof(domain))) + end + return String(label) end KernelAbstractions.trace_range_end(t::RecordingTracer, id) = @lock t.lock push!(t.events, (:end, id)) KernelAbstractions.trace_mark(t::RecordingTracer, label, domain) = - @lock t.lock push!(t.events, (:mark, label, domain)) + @lock t.lock push!(t.events, (:mark, String(label), String(domain))) function with_tracer(f, tracer = RecordingTracer()) KernelAbstractions.register_tracer!(tracer) diff --git a/test/runtests.jl b/test/runtests.jl index d681050d0..5d843d9e6 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -489,6 +489,15 @@ end (:start, "explicit", "X"), (:end, "explicit"), ] + # labels fixed in the code are `Symbol`s, which tracers may cache; others `String`s + @test tracer.types[1:3] == [(Symbol, Symbol), (String, Symbol), (String, Symbol)] + tracer = with_tracer() do tracer + Testsuite.profiling_fill!(CPU())(zeros(Float32, 4), 1.0f0; ndrange = 4) + wait(KernelAbstractions.@spawn CPU() nothing) + end + @test all(==((Symbol, Symbol)), tracer.types) + @test KernelAbstractions.kernel_label(Testsuite.gpu_profiling_fill!) === :profiling_fill! + # ranges end when the expression throws tracer = with_tracer() do tracer @test_throws ErrorException @profiling_range "throws" error("boom") @@ -689,7 +698,7 @@ end end close(tracer) # recording after closing is harmless - KernelAbstractions.trace_mark(tracer, "late", "KernelAbstractions") + KernelAbstractions.trace_mark(tracer, "late", :KernelAbstractions) lines = readlines(path) @test lines[1] == "SetFileDisplayName, KernelAbstractions" From aec082a487a1b034e278132ddfd1b6905a8a3c26 Mon Sep 17 00:00:00 2001 From: Valentin Churavy Date: Sun, 4 Oct 2026 16:22:26 +0200 Subject: [PATCH 7/9] Make the range of a traced kernel launch cheaper The label is a constant computed from the kernel's type, and the helper that starts the range is specialized again, now that the traced launch reuses the inferred `launch_untraced`. Starting and ending the range goes from ~310 to ~28 ns, i.e. a launch with a tracer registered costs ~70 ns more than one without, rather than ~500 ns. Assisted-by: Claude Code (Opus 5.5) --- src/backend_launch.jl | 5 ++--- src/profiling.jl | 15 +++++---------- 2 files changed, 7 insertions(+), 13 deletions(-) diff --git a/src/backend_launch.jl b/src/backend_launch.jl index ac3a16c7d..eef8c8984 100644 --- a/src/backend_launch.jl +++ b/src/backend_launch.jl @@ -85,7 +85,7 @@ Core.kwcall(kwargs::NamedTuple, obj::Kernel{<:KI.Backend}, args::Vararg{Any, N}) function launch_tuple(obj::Kernel, args::Tuple; ndrange = nothing, workgroupsize = nothing) profiling_active() || return launch_untraced(obj, args, ndrange, workgroupsize) # The traced launch calls the same `launch_untraced`, so that it isn't inferred twice for - # every kernel. The helpers around it are only inferred once. + # every kernel; the helpers around it are small. id = start_launch_range(obj) try launch_untraced(obj, args, ndrange, workgroupsize) @@ -97,8 +97,7 @@ function launch_tuple(obj::Kernel, args::Tuple; ndrange = nothing, workgroupsize return nothing end -Base.@nospecializeinfer @noinline start_launch_range(@nospecialize(obj::Kernel)) = - profiling_range_start(kernel_label(obj.f)) +@noinline start_launch_range(obj::Kernel) = profiling_range_start(kernel_label(obj.f)) function launch_untraced(obj::Kernel, args::Tuple, ndrange, workgroupsize) ndrange, workgroupsize, iterspace, dynamic = launch_config(obj, ndrange, workgroupsize) diff --git a/src/profiling.jl b/src/profiling.jl index 96069deb7..356651f6d 100644 --- a/src/profiling.jl +++ b/src/profiling.jl @@ -212,16 +212,11 @@ literal(x::String) = QuoteNode(Symbol(x)) literal(x::QuoteNode) = x literal(x) = esc(x) -# the name kernel launches are annotated with: `@kernel function f` compiles to `gpu_f` -# (computed once per kernel function type) -const KERNEL_LABELS = IdDict{Any, Symbol}() -const KERNEL_LABELS_LOCK = ReentrantLock() -function kernel_label(f) - return @lock KERNEL_LABELS_LOCK get!(KERNEL_LABELS, typeof(f)) do - f isa Function || return Symbol(typeof(f)) - name = string(nameof(f)) - return Symbol(startswith(name, "gpu_") ? name[5:end] : name) - end +# the name kernel launches are annotated with: `@kernel function f` compiles to `gpu_f`. It +# only depends on the type of the function, so it is a constant. +@generated function kernel_label(f) + name = string(f <: Function && isdefined(f, :instance) ? nameof(f.instance) : nameof(f)) + return QuoteNode(Symbol(startswith(name, "gpu_") ? name[5:end] : name)) end From e84c3ad972fb7b04f3bd08136b32e566a0fd7ba7 Mon Sep 17 00:00:00 2001 From: Valentin Churavy Date: Sun, 4 Oct 2026 17:29:46 +0200 Subject: [PATCH 8/9] Compile the traced launch's helpers once rather than per kernel They take the kernel's label, a constant, rather than the kernel, so only the call sites are specialized on the kernel. Assisted-by: Claude Code (Opus 5.5) --- src/backend_launch.jl | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/src/backend_launch.jl b/src/backend_launch.jl index eef8c8984..313dc9d73 100644 --- a/src/backend_launch.jl +++ b/src/backend_launch.jl @@ -85,19 +85,25 @@ Core.kwcall(kwargs::NamedTuple, obj::Kernel{<:KI.Backend}, args::Vararg{Any, N}) function launch_tuple(obj::Kernel, args::Tuple; ndrange = nothing, workgroupsize = nothing) profiling_active() || return launch_untraced(obj, args, ndrange, workgroupsize) # The traced launch calls the same `launch_untraced`, so that it isn't inferred twice for - # every kernel; the helpers around it are small. - id = start_launch_range(obj) + # every kernel. The label is a constant, and the helpers only take it and the range, so + # they are compiled once rather than for every kernel. + id = start_launch_range(kernel_label(obj.f)) try launch_untraced(obj, args, ndrange, workgroupsize) - # a profiler that measures kernels rather than launches - synchronizes_launches(id) && KI.synchronize(backend(obj)) + synchronize_launch(id, backend(obj)) finally profiling_range_end(id) end return nothing end -@noinline start_launch_range(obj::Kernel) = profiling_range_start(kernel_label(obj.f)) +@noinline start_launch_range(label::Symbol) = profiling_range_start(label) + +# for a profiler that measures kernels rather than launches +@noinline function synchronize_launch(id, backend) + synchronizes_launches(id) && KI.synchronize(backend) + return nothing +end function launch_untraced(obj::Kernel, args::Tuple, ndrange, workgroupsize) ndrange, workgroupsize, iterspace, dynamic = launch_config(obj, ndrange, workgroupsize) From d6c096731089007ea4bc48a7b2d74c223cf3743e Mon Sep 17 00:00:00 2001 From: Valentin Churavy Date: Sun, 4 Oct 2026 19:26:24 +0200 Subject: [PATCH 9/9] Trace kernel launches through a dynamic dispatch The traced launch is inferred and compiled once for all kernels, rather than as part of every kernel's launch, which added ~0.6 MiB and a few ms to the compilation of every new kernel. A traced launch pays for a dynamic dispatch to the untraced one instead (~50 ns). Assisted-by: Claude Code (Opus 5.5) --- src/backend_launch.jl | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/src/backend_launch.jl b/src/backend_launch.jl index 313dc9d73..68081a5e3 100644 --- a/src/backend_launch.jl +++ b/src/backend_launch.jl @@ -83,10 +83,16 @@ Core.kwcall(kwargs::NamedTuple, obj::Kernel{<:KI.Backend}, args::Vararg{Any, N}) launch_tuple(obj, args; kwargs...) function launch_tuple(obj::Kernel, args::Tuple; ndrange = nothing, workgroupsize = nothing) - profiling_active() || return launch_untraced(obj, args, ndrange, workgroupsize) - # The traced launch calls the same `launch_untraced`, so that it isn't inferred twice for - # every kernel. The label is a constant, and the helpers only take it and the range, so - # they are compiled once rather than for every kernel. + profiling_active() && return launch_traced(obj, args, ndrange, workgroupsize) + return launch_untraced(obj, args, ndrange, workgroupsize) +end + +# Inferred and compiled once for all kernels, rather than as part of every kernel's launch, +# which would add to the compilation of every new kernel. A traced launch pays for a dynamic +# dispatch to `launch_untraced` instead. +Base.@nospecializeinfer @noinline function launch_traced( + @nospecialize(obj::Kernel), @nospecialize(args::Tuple), @nospecialize(ndrange), @nospecialize(workgroupsize) + ) id = start_launch_range(kernel_label(obj.f)) try launch_untraced(obj, args, ndrange, workgroupsize)