Add TimerOutputs-based timers for internal kernels

Project.toml

@@ -21,6 +21,7 @@ ScopedValues = "7e506255-f358-4e82-b7e4-beb19740aa63"
 Strided = "5e0ebb24-38b0-5f93-81fe-25c709ecae67"
 TensorKitSectors = "13a9c161-d5da-41f0-bcbd-e1a08ae0647f"
 TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2"
+TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 TupleTools = "9d95972d-f1c8-5527-a6e0-b4b365fa01f6"
 VectorInterface = "409d34a3-91d5-4945-b6ec-7529ddf182d8"
 
@@ -67,6 +68,7 @@ ScopedValues = "1.3.0"
 Strided = "2.6.1"
 TensorKitSectors = "0.3.7"
 TensorOperations = "5.5.2, 5.6"
+TimerOutputs = "1"
 TupleTools = "1.5"
 VectorInterface = "0.6"
 julia = "1.10"

benchmark/README.md

@@ -29,3 +29,12 @@ benchpkgtable TensorKit \
     -i benchmark/results/ \
     -o benchmark/results/ \
 ```
+
+## Timing internal kernels
+
+To break down where time is spent *inside* the operations (symmetry work, bookkeeping,
+allocations, dense kernels), TensorKit ships built-in
+[TimerOutputs.jl](https://github.com/KristofferC/TimerOutputs.jl) instrumentation that is
+compiled away by default and can be enabled with `TensorKit.enable_timers!()`. See the
+"Profiling and timers" section of the documentation for the workflow and caveats.
+

docs/make.jl

@@ -33,6 +33,7 @@ pages = [
         "man/sectors.md", "man/gradedspaces.md",
         "man/fusiontrees.md",
         "Tensors" => TENSOR_PAGES,
+        "man/profiling.md",
     ],
     "Library" => [
         "lib/sectors.md", "lib/fusiontrees.md",

docs/src/man/profiling.md

@@ -0,0 +1,61 @@
+# [Profiling and timers](@id s_profiling)
+
+TensorKit's index manipulations, tensor contractions and factorizations are instrumented with [TimerOutputs.jl](https://github.com/KristofferC/TimerOutputs.jl) sections that are compiled away by default, so they incur no runtime cost.
+They can be enabled to obtain a detailed breakdown of where time is spent inside these operations, in particular the split between the different kinds of work involved in manipulating symmetric tensors:
+
+| category | contents |
+|:---|:---|
+| `symmetry` | fusion tree manipulations and recoupling coefficients (braiding, transposing, F- and R-symbols) |
+| `bookkeeping` | block structure computations, cache lookups, contraction planning |
+| `alloc` | allocation of output tensors and temporary buffers |
+| `dense` | dense tensor kernels (BLAS/LAPACK calls, strided permutations and additions) |
+| `other` | remaining time within an instrumented operation (dispatch, argument checking, uncovered overhead) |
+
+The canonical workflow looks as follows:
+
+```julia
+using TensorKit
+
+TensorKit.enable_timers!()      # triggers recompilation of the instrumented methods
+
+# warm up first, so that compilation does not pollute the timings
+V = SU2Space(0 => 4, 1//2 => 4, 1 => 2)
+t = rand(V ⊗ V ← V ⊗ V)
+permute(t, ((1, 3), (2, 4)))
+@tensor t2[a b; c d] := t[a x; c y] * t[y b; x d]
+svd_compact(t)
+
+TensorKit.reset_timers!()
+# ... run the workload of interest ...
+TensorKit.print_timers()        # full nested call tree
+TensorKit.timer_summary()       # symmetry / bookkeeping / alloc / dense / other totals
+
+TensorKit.disable_timers!()
+```
+
+[`TensorKit.print_timers`](@ref) displays the accumulated timings as a nested call tree, with sections for the top-level operations (`"permute!/braid!"`, `"contract!"`, `"svd_compact!"`, ...) and nested sections labeled by their category prefix (`"symmetry: ..."`, `"bookkeeping: ..."`, `"alloc: ..."`, `"dense: ..."`).
+[`TensorKit.timer_summary`](@ref) aggregates the *exclusive* time of each section (its own time minus that of its timed children) into per-category totals, such that every nanosecond is counted exactly once and the totals sum to the total measured time.
+
+A few caveats to keep in mind:
+
+*   Enabling or disabling the timers redefines internal functions, so instrumented methods recompile on first use afterwards.
+    This is a debug-session switch, not a runtime option.
+*   While timers are enabled, TensorKit-internal task parallelism is disabled, since the timer object may only be manipulated from a single task.
+    As a consequence, multi-threaded speedups are not measurable while timing, and TensorKit functions should not be called concurrently from multiple user tasks.
+*   Each section entry costs roughly 100–200 ns, which can distort measurements of workloads on very small tensors.
+*   The construction of fusion tree transformers and block structures is cached (see `empty_globalcaches!`), so their cost only shows up the first time a given structure is encountered.
+    Call `TensorKit.empty_globalcaches!()` before the measurement if you want the construction cost to be included, or after the warm-up if you want to measure the steady-state behavior with warm caches.
+*   For GPU tensors, the timings only reflect host-side dispatch of asynchronous kernels, unless the workload is explicitly synchronized.
+
+## Library documentation
+
+```@docs
+TensorKit.GLOBAL_TIMER
+TensorKit.enable_timers!
+TensorKit.disable_timers!
+TensorKit.reset_timers!
+TensorKit.print_timers
+TensorKit.timer_summary
+TensorKit.timer
+TensorKit.timers_enabled
+```

src/TensorKit.jl

@@ -123,6 +123,7 @@ using Dictionaries: Dictionaries, Dictionary, Indices, gettoken, gettokenvalue
 using LRUCache
 using OhMyThreads
 using ScopedValues
+using TimerOutputs: TimerOutputs, TimerOutput, @timeit_debug
 
 using TensorKitSectors
 import TensorKitSectors: dim, BraidingStyle, FusionStyle, ⊠, ⊗, ×
@@ -134,7 +135,7 @@ using Base: @boundscheck, @propagate_inbounds, @constprop,
     tuple_type_head, tuple_type_tail, tuple_type_cons,
     SizeUnknown, HasLength, HasShape, IsInfinite, EltypeUnknown, HasEltype
 using Base.Iterators: product, filter
-using Printf: @sprintf
+using Printf: @sprintf, @printf
 
 using LinearAlgebra: LinearAlgebra, BlasFloat
 using LinearAlgebra: norm, dot, normalize, normalize!, tr,
@@ -153,6 +154,7 @@ using Adapt: Adapt
 
 # Auxiliary files
 #-----------------
+include("auxiliary/timers.jl")
 include("auxiliary/auxiliary.jl")
 include("auxiliary/caches.jl")
 include("auxiliary/dicts.jl")
@@ -225,7 +227,8 @@ include("spaces/structure.jl")
 #-------------------------
 const TRANSFORMER_THREADS = Ref(1)
 
-get_num_transformer_threads() = TRANSFORMER_THREADS[]
+# while timing, force serial execution: timer sections may only be entered from one task
+get_num_transformer_threads() = timers_enabled() ? 1 : TRANSFORMER_THREADS[]
 
 function set_num_transformer_threads(n::Int)
     N = Base.Threads.nthreads()
@@ -238,7 +241,7 @@ end
 
 const TREEMANIPULATION_THREADS = Ref(1)
 
-get_num_manipulation_threads() = TREEMANIPULATION_THREADS[]
+get_num_manipulation_threads() = timers_enabled() ? 1 : TREEMANIPULATION_THREADS[]
 
 function set_num_manipulation_threads(n::Int)
     N = Base.Threads.nthreads()

src/auxiliary/caches.jl

@@ -48,6 +48,12 @@ function CacheStyle(args...)
     return GlobalLRUCache()
 end
 
+# category of the miss-path (construction) timer section of an `@cached` function
+function _cached_category(fname::Symbol)
+    return fname in (:fsbraid, :fstranspose, :treebraider, :treetransposer) ?
+        "symmetry" : "bookkeeping"
+end
+
 macro cached(ex)
     Meta.isexpr(ex, :function) ||
         error("cached macro can only be used on function definitions")
@@ -69,6 +75,9 @@ macro cached(ex)
     Meta.isexpr(fcall, :call) ||
         error("cached macro can only be used on function definitions")
     fname = fcall.args[1]
+    # timer labels for the cache lookup and the miss-path construction
+    lookuplabel = string("bookkeeping: cache ", fname)
+    misslabel = string(_cached_category(fname), ": compute ", fname)
     fargs = fcall.args[2:end]
     fargnames = map(fargs) do arg
         if Meta.isexpr(arg, :(::))
@@ -106,7 +115,7 @@ macro cached(ex)
     if hasparams
         fnocachecall = Expr(:where, fnocachecall, params...)
     end
-    fnocachebody = Expr(:call, _fname, fargnames...)
+    fnocachebody = :(@timeit_debug GLOBAL_TIMER $misslabel $(Expr(:call, _fname, fargnames...)))
     if typed
         T = gensym(:T)
         fnocachebody = Expr(:block, Expr(:(=), T, typeex), Expr(:(::), fnocachebody, T))
@@ -135,8 +144,8 @@ macro cached(ex)
         key = Expr(:tuple, fargnames...)
     end
     getvalex = :(
-        get!($cachevar, $key) do
-            return $_fname($(fargnames...))
+        @timeit_debug GLOBAL_TIMER $lookuplabel get!($cachevar, $key) do
+            return @timeit_debug GLOBAL_TIMER $misslabel $_fname($(fargnames...))
         end
     )
     if typed

src/auxiliary/timers.jl

@@ -0,0 +1,162 @@
+# This is the switch that every `@timeit_debug` section in this module checks, and that
+# `TimerOutputs.enable_debug_timings(TensorKit)` redefines to `true` (recompiling all
+# instrumented methods). TimerOutputs would define it automatically on the first
+# `@timeit_debug` expansion; defining it explicitly here lets ordinary code branch on it
+# as well, see `timers_enabled`.
+timeit_debug_enabled() = false
+
+"""
+    TensorKit.GLOBAL_TIMER
+
+The global `TimerOutput` object into which all timer sections of TensorKit accumulate.
+See [`enable_timers!`](@ref) for how to activate them, [`print_timers`](@ref) for
+displaying the resulting call tree and [`timer_summary`](@ref) for aggregating it into
+per-category totals.
+"""
+const GLOBAL_TIMER = TimerOutput("TensorKit")
+
+"""
+    TensorKit.timer() -> TimerOutputs.TimerOutput
+
+Return the global timer object [`GLOBAL_TIMER`](@ref) into which all TensorKit timer
+sections accumulate.
+"""
+timer() = GLOBAL_TIMER
+
+"""
+    TensorKit.timers_enabled() -> Bool
+
+Return whether the `@timeit_debug` timer sections of TensorKit are currently compiled in,
+i.e. whether [`enable_timers!`](@ref) has been called.
+
+This is a documented alias for `timeit_debug_enabled`, the switch that is redefined by
+`TimerOutputs.enable_debug_timings`. It is used internally to force serial execution of
+parallel regions while timing, since a `TimerOutput` may only be manipulated from a single
+task. When timers are disabled this check const-folds to `false`, so it has no runtime
+cost.
+"""
+timers_enabled() = timeit_debug_enabled()
+
+"""
+    TensorKit.enable_timers!()
+
+Enable all timer sections of TensorKit (including its submodules), which accumulate
+timings of the internal kernels into [`GLOBAL_TIMER`](@ref). Undone by
+[`disable_timers!`](@ref).
+
+!!! warning
+    Enabling or disabling timers redefines internal functions and therefore triggers
+    recompilation of the instrumented methods on first use. This is a debug-session
+    operation, not a runtime switch.
+
+!!! warning
+    While timers are enabled, TensorKit-internal task parallelism is disabled (threaded
+    regions run serially), so multi-threaded speedups are not measurable. Additionally,
+    TensorKit functions should not be called concurrently from multiple user tasks while
+    timing, as the timer object is not thread-safe.
+
+Note that each timer section adds an overhead of roughly 100-200 ns per entry, which can
+distort measurements of very small workloads. For GPU tensors, timings only reflect
+host-side dispatch of asynchronous kernels unless the workload is explicitly synchronized.
+"""
+function enable_timers!()
+    TimerOutputs.enable_debug_timings(TensorKit)
+    return nothing
+end
+
+"""
+    TensorKit.disable_timers!()
+
+Disable all timer sections of TensorKit again; the inverse of [`enable_timers!`](@ref).
+Also triggers recompilation of the instrumented methods on first use.
+"""
+function disable_timers!()
+    TimerOutputs.disable_debug_timings(TensorKit)
+    return nothing
+end
+
+"""
+    TensorKit.reset_timers!()
+
+Reset the accumulated timings in [`GLOBAL_TIMER`](@ref).
+"""
+function reset_timers!()
+    TimerOutputs.reset_timer!(GLOBAL_TIMER)
+    return nothing
+end
+
+"""
+    TensorKit.print_timers(io::IO = stdout; kwargs...)
+
+Print the accumulated timings in [`GLOBAL_TIMER`](@ref) as a nested table. Keyword
+arguments are forwarded to `TimerOutputs.print_timer`.
+"""
+print_timers(io::IO = stdout; kwargs...) = TimerOutputs.print_timer(io, GLOBAL_TIMER; kwargs...)
+
+const TIMER_CATEGORIES = (:symmetry, :bookkeeping, :alloc, :dense, :other)
+
+# "symmetry: recoupling" -> :symmetry; unprefixed or unknown prefix -> nothing
+function _timer_category(label::String)
+    i = findfirst(':', label)
+    i === nothing && return nothing
+    prefix = Symbol(label[1:prevind(label, i)])
+    return prefix in TIMER_CATEGORIES ? prefix : nothing
+end
+
+const TimerSummary = Dict{Symbol, @NamedTuple{time::Int64, allocated::Int64, ncalls::Int64}}
+
+"""
+    TensorKit.timer_summary([io::IO]; to = GLOBAL_TIMER)
+        -> Dict{Symbol, @NamedTuple{time::Int64, allocated::Int64, ncalls::Int64}}
+
+Aggregate the timer tree into per-category totals (time in ns, allocated bytes, number of
+section entries) for the categories `$(TIMER_CATEGORIES)`.
+
+Each section's *exclusive* time (its own time minus that of its timed children) is
+attributed to the category given by its label prefix or, for unprefixed labels, to the
+category of the nearest categorized ancestor (`:other` at the root). Every nanosecond is
+thus counted exactly once and the totals sum to the total measured time.
+
+When `io` is given (default `stdout`), a small table is printed; pass `nothing` to skip
+printing and only return the totals.
+"""
+function timer_summary(io::Union{IO, Nothing} = stdout; to::TimerOutput = GLOBAL_TIMER)
+    totals = TimerSummary(c => (time = 0, allocated = 0, ncalls = 0) for c in TIMER_CATEGORIES)
+    for child in to.root.children
+        _accumulate_summary!(totals, child, :other)
+    end
+    io === nothing || _print_timer_summary(io, totals)
+    return totals
+end
+
+function _accumulate_summary!(totals::TimerSummary, s, inherited::Symbol)
+    cat = something(_timer_category(s.name), inherited)
+    t = TimerOutputs.time(s)
+    b = TimerOutputs.allocated(s)
+    for child in s.children
+        t -= TimerOutputs.time(child)
+        b -= TimerOutputs.allocated(child)
+        _accumulate_summary!(totals, child, cat)
+    end
+    old = totals[cat]
+    totals[cat] = (
+        time = old.time + max(t, 0), allocated = old.allocated + max(b, 0),
+        ncalls = old.ncalls + TimerOutputs.ncalls(s),
+    )
+    return nothing
+end
+
+function _print_timer_summary(io::IO, totals::TimerSummary)
+    total_time = sum(x -> x.time, values(totals))
+    println(io, "TensorKit timer summary:")
+    for cat in TIMER_CATEGORIES
+        (; time, allocated, ncalls) = totals[cat]
+        percentage = total_time == 0 ? 0.0 : 100 * time / total_time
+        @printf(
+            io, "%12s: %s (%5.1f%%)  %s  %d sections\n",
+            cat, TimerOutputs.prettytime(time), percentage,
+            TimerOutputs.prettymemory(allocated), ncalls
+        )
+    end
+    return nothing
+end

src/factorizations/factorizations.jl

@@ -9,6 +9,8 @@ using ..TensorKit
 using ..TensorKit: AdjointTensorMap, SectorDict, SectorVector,
     blocktype, foreachblock, one!,
     similar_diagonal, similarstoragetype
+using ..TensorKit: GLOBAL_TIMER
+using TimerOutputs: @timeit_debug
 
 using LinearAlgebra: LinearAlgebra, BlasFloat, Diagonal,
     svdvals, svdvals!, eigen, eigen!,