WIP: Tracing subsystem

Project.toml

@@ -22,11 +22,13 @@ pocl_jll = "627d6b7a-bbe6-5189-83e7-98cc0a5aeadd"
 
 [weakdeps]
 EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
+IntelITT = "c9b2f978-7543-4802-ae44-75068f23ee64"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
 [extensions]
 EnzymeExt = "EnzymeCore"
+IntelITTExt = "IntelITT"
 LinearAlgebraExt = "LinearAlgebra"
 SparseArraysExt = "SparseArrays"
 
@@ -35,6 +37,7 @@ Adapt = "0.4, 1.0, 2.0, 3.0, 4"
 Atomix = "0.1, 1"
 EnzymeCore = "0.7, 0.8.1"
 GPUCompiler = "1.13.3"
+IntelITT = "0.2"
 InteractiveUtils = "1.6"
 LLVM = "9.4.1"
 LinearAlgebra = "1.6"

docs/src/api.md

@@ -24,6 +24,20 @@ KernelAbstractions.zeros
 KernelAbstractions.supports_unified
 ```
 
+## Profiler integration
+
+KernelAbstractions exposes a small interface that backend or profiler packages
+can implement to forward named ranges to external tracing profilers such as
+NVIDIA Nsight Systems or Intel VTune. The defaults are no-ops, so calls are
+free when no integration is loaded.
+
+```@docs
+profiling_range_active
+profiling_range_start
+profiling_range_end
+@profiling_range
+```
+
 ## Internal
 
 ```@docs

ext/IntelITTExt.jl

@@ -0,0 +1,33 @@
+module IntelITTExt
+
+using KernelAbstractions: KernelAbstractions, CPU
+import KernelAbstractions: profiling_range_active, profiling_range_start, profiling_range_end
+
+import IntelITT
+
+const DOMAINS = Dict{String, IntelITT.Domain}()
+const DOMAINS_LOCK = ReentrantLock()
+
+function domain_for(name::AbstractString)
+    key = String(name)
+    return lock(DOMAINS_LOCK) do
+        get!(() -> IntelITT.Domain(key), DOMAINS, key)
+    end
+end
+
+function profiling_range_active(::CPU; domain = "KernelAbstractions")
+    return IntelITT.isactive()
+end
+
+function profiling_range_start(::CPU, label; domain = "KernelAbstractions")
+    task = IntelITT.Task(domain_for(domain), String(label))
+    IntelITT.start(task)
+    return task
+end
+
+function profiling_range_end(::CPU, id)
+    IntelITT.stop(id)
+    return nothing
+end
+
+end # module

src/KernelAbstractions.jl

@@ -6,6 +6,8 @@ export @index, @groupsize, @ndrange
 export @print
 export Backend, GPU, CPU
 export synchronize, get_backend, allocate
+export profiling_range_active, profiling_range_start, profiling_range_end
+export @profiling_range
 
 import PrecompileTools
 
@@ -194,6 +196,93 @@ function unsafe_free! end
 
 unsafe_free!(::AbstractArray) = return
 
+"""
+    profiling_range_active(backend; domain = "KernelAbstractions") -> Bool
+
+Return `true` if `backend` supports profiler range annotations and a profiler
+is currently attached. Callers use this to skip building labels or handles when
+no profiler is listening.
+
+!!! note
+    Backend/profiler integrations **may** override this. The default returns `false`.
+"""
+profiling_range_active(::Any; domain = "KernelAbstractions") = false
+
+"""
+    profiling_range_start(backend, label; domain = "KernelAbstractions") -> id
+
+Start a named profiler range on `backend`. Returns an opaque handle that must
+be passed to [`profiling_range_end`](@ref).
+
+!!! note
+    Backend/profiler integrations **may** override this. The default is a no-op
+    that returns `nothing`.
+"""
+profiling_range_start(::Any, label; domain = "KernelAbstractions") = nothing
+
+"""
+    profiling_range_end(backend, id) -> nothing
+
+End a previously started profiler range identified by `id`.
+
+!!! note
+    Backend/profiler integrations **may** override this. The default is a no-op.
+"""
+profiling_range_end(::Any, id) = nothing
+
+"""
+    @profiling_range backend label [domain=...] expr
+
+Evaluate `expr` inside a profiler range named `label` on `backend`. The range
+is only started when [`profiling_range_active`](@ref) returns `true`, so the
+overhead is a single call when no profiler is attached. An optional `domain`
+keyword groups related ranges; it defaults to `"KernelAbstractions"`.
+
+```julia
+@profiling_range backend "volume integral" begin
+    volume_integral!(du, u, backend)
+end
+
+@profiling_range backend "volume integral" domain="Trixi" begin
+    volume_integral!(du, u, backend)
+end
+```
+"""
+macro profiling_range(args...)
+    length(args) >= 3 ||
+        error("@profiling_range requires at least `backend`, `label`, and an expression")
+    backend = args[1]
+    label = args[2]
+    expr = args[end]
+    domain = "KernelAbstractions"
+    for kw in args[3:(end - 1)]
+        if Meta.isexpr(kw, :(=)) && kw.args[1] === :domain
+            domain = kw.args[2]
+        else
+            error("@profiling_range: unexpected argument `$(kw)`; only `domain=...` is accepted")
+        end
+    end
+    backend_var = gensym(:backend)
+    domain_var = gensym(:domain)
+    active_var = gensym(:active)
+    id_var = gensym(:id)
+    return quote
+        local $backend_var = $(esc(backend))
+        local $domain_var = $(esc(domain))
+        local $active_var = $profiling_range_active($backend_var; domain = $domain_var)
+        local $id_var = $active_var ?
+            $profiling_range_start($backend_var, $(esc(label)); domain = $domain_var) :
+            nothing
+        try
+            $(esc(expr))
+        finally
+            if $active_var
+                $profiling_range_end($backend_var, $id_var)
+            end
+        end
+    end
+end
+
 """
 Abstract type for all KernelAbstractions backends.
 """

test/Project.toml

@@ -1,5 +1,6 @@
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+IntelITT = "c9b2f978-7543-4802-ae44-75068f23ee64"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"

test/runtests.jl

@@ -35,6 +35,44 @@ struct NewBackend <: KernelAbstractions.GPU end
     @test KernelAbstractions.functional(backend) === missing
 end
 
+@testset "profiling_range defaults" begin
+    backend = NewBackend()
+
+    @test KernelAbstractions.profiling_range_active(backend) == false
+    @test KernelAbstractions.profiling_range_active(nothing) == false
+
+    @test KernelAbstractions.profiling_range_start(backend, "label") === nothing
+    @test KernelAbstractions.profiling_range_start(backend, "label"; domain = "X") === nothing
+
+    @test KernelAbstractions.profiling_range_end(backend, nothing) === nothing
+
+    @test (@profiling_range backend "label" 1 + 2) == 3
+    @test (@profiling_range backend "label" domain = "Custom" 1 + 2) == 3
+
+    @test_throws ErrorException (@profiling_range backend "label" error("boom"))
+end
+
+@testset "profiling_range IntelITT extension" begin
+    using IntelITT
+
+    cpu = CPU()
+
+    @test KernelAbstractions.profiling_range_active(cpu) isa Bool
+
+    m = which(KernelAbstractions.profiling_range_active, Tuple{CPU})
+    @test parentmodule(m) === Base.get_extension(KernelAbstractions, :IntelITTExt)
+
+    task = KernelAbstractions.profiling_range_start(cpu, "label")
+    @test task !== nothing
+    @test KernelAbstractions.profiling_range_end(cpu, task) === nothing
+
+    ext = Base.get_extension(KernelAbstractions, :IntelITTExt)
+    @test ext.domain_for("Trixi") === ext.domain_for("Trixi")
+
+    @test (@profiling_range cpu "label" 1 + 2) == 3
+    @test (@profiling_range cpu "vi" domain = "Trixi" 7) == 7
+end
+
 
 # include("extensions/enzyme.jl")
 # @static if VERSION >= v"1.7.0"

test/testsuite.jl

@@ -39,6 +39,7 @@ include("reflection.jl")
 include("examples.jl")
 include("convert.jl")
 include("specialfunctions.jl")
+include("profiling_range.jl")
 
 function testsuite(backend, backend_str, backend_mod, AT, DAT; skip_tests = Set{String}())
     @conditional_testset "Unittests" skip_tests begin
@@ -93,6 +94,10 @@ function testsuite(backend, backend_str, backend_mod, AT, DAT; skip_tests = Set{
         examples_testsuite(backend, backend_str)
     end
 
+    @conditional_testset "ProfilingRange" skip_tests begin
+        profiling_range_testsuite(backend)
+    end
+
     return
 end