diff --git a/Project.toml b/Project.toml index 920c06a0d..8a93b4358 100644 --- a/Project.toml +++ b/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" diff --git a/benchmark/README.md b/benchmark/README.md index c3415d26e..ccfeda717 100644 --- a/benchmark/README.md +++ b/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. + diff --git a/docs/make.jl b/docs/make.jl index 13d2e91fb..590f93bcc 100644 --- a/docs/make.jl +++ b/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", diff --git a/docs/src/man/profiling.md b/docs/src/man/profiling.md new file mode 100644 index 000000000..81b36abfc --- /dev/null +++ b/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 +``` diff --git a/src/TensorKit.jl b/src/TensorKit.jl index 7b2a9ae72..c418fe944 100644 --- a/src/TensorKit.jl +++ b/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() diff --git a/src/auxiliary/caches.jl b/src/auxiliary/caches.jl index 4c5db4652..f3bec7844 100644 --- a/src/auxiliary/caches.jl +++ b/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 diff --git a/src/auxiliary/timers.jl b/src/auxiliary/timers.jl new file mode 100644 index 000000000..2c3977458 --- /dev/null +++ b/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 diff --git a/src/factorizations/factorizations.jl b/src/factorizations/factorizations.jl index fbe87a63a..6ec8c7da2 100644 --- a/src/factorizations/factorizations.jl +++ b/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!, diff --git a/src/factorizations/matrixalgebrakit.jl b/src/factorizations/matrixalgebrakit.jl index 9bb8531b7..5974a60de 100644 --- a/src/factorizations/matrixalgebrakit.jl +++ b/src/factorizations/matrixalgebrakit.jl @@ -15,7 +15,9 @@ for f in return MAK.default_algorithm($f!, blocktype(T); kwargs...) end @eval function MAK.copy_input(::typeof($f), t::AbstractTensorMap) - return copy_oftype(t, factorisation_scalartype($f, t)) + return @timeit_debug GLOBAL_TIMER "alloc: copy_input" copy_oftype( + t, factorisation_scalartype($f, t) + ) end end @@ -38,13 +40,17 @@ for f! in ( ) @eval function MAK.$f!(t::AbstractTensorMap, F, alg::AbstractAlgorithm) $(f! in (:eig_full!, :eigh_full!) && :(LinearAlgebra.checksquare(t))) - foreachblock(t, F...) do _, (tblock, Fblocks...) - Fblocks′ = $f!(tblock, Fblocks, alg) - # deal with the case where the output is not in-place - for (b′, b) in zip(Fblocks′, Fblocks) - b === b′ || copy!(b, b′) + @timeit_debug GLOBAL_TIMER $(string(f!)) begin + foreachblock(t, F...) do _, (tblock, Fblocks...) + @timeit_debug GLOBAL_TIMER "dense: MatrixAlgebraKit" begin + Fblocks′ = $f!(tblock, Fblocks, alg) + # deal with the case where the output is not in-place + for (b′, b) in zip(Fblocks′, Fblocks) + b === b′ || copy!(b, b′) + end + end + return nothing end - return nothing end return F end @@ -59,11 +65,15 @@ for f! in ( ) @eval function MAK.$f!(t::AbstractTensorMap, N, alg::AbstractAlgorithm) $(f! in (:eig_vals!, :eigh_vals!, :project_hermitian!, :project_antihermitian!, :exponential!) && :(LinearAlgebra.checksquare(t))) - foreachblock(t, N) do _, (tblock, Nblock) - Nblock′ = $f!(tblock, Nblock, alg) - # deal with the case where the output is not the same as the input - Nblock === Nblock′ || copy!(Nblock, Nblock′) - return nothing + @timeit_debug GLOBAL_TIMER $(string(f!)) begin + foreachblock(t, N) do _, (tblock, Nblock) + @timeit_debug GLOBAL_TIMER "dense: MatrixAlgebraKit" begin + Nblock′ = $f!(tblock, Nblock, alg) + # deal with the case where the output is not the same as the input + Nblock === Nblock′ || copy!(Nblock, Nblock′) + end + return nothing + end end return N end @@ -72,11 +82,15 @@ end # Exponential with Tuple function MAK.exponential!((τ, t)::Tuple{E, T}, N, alg::AbstractAlgorithm) where {E <: Number, T <: AbstractTensorMap} LinearAlgebra.checksquare(t) - foreachblock(t, N) do _, (tblock, Nblock) - Nblock′ = exponential!((τ, tblock), Nblock, alg) - # deal with the case where the output is not the same as the input - Nblock === Nblock′ || copy!(Nblock, Nblock′) - return nothing + @timeit_debug GLOBAL_TIMER "exponential!" begin + foreachblock(t, N) do _, (tblock, Nblock) + @timeit_debug GLOBAL_TIMER "dense: MatrixAlgebraKit" begin + Nblock′ = exponential!((τ, tblock), Nblock, alg) + # deal with the case where the output is not the same as the input + Nblock === Nblock′ || copy!(Nblock, Nblock′) + end + return nothing + end end return N end @@ -126,118 +140,148 @@ MAK.exponential!((τ, t)::Tuple{E, T}, out, alg::DefaultAlgorithm) where {E <: N # Singular value decomposition # ---------------------------- function MAK.initialize_output(::typeof(svd_full!), t::AbstractTensorMap, ::AbstractAlgorithm) - V_cod = fuse(codomain(t)) - V_dom = fuse(domain(t)) - U = similar(t, codomain(t) ← V_cod) - S = similar(t, real(scalartype(t)), V_cod ← V_dom) - Vᴴ = similar(t, V_dom ← domain(t)) - return U, S, Vᴴ + @timeit_debug GLOBAL_TIMER "alloc: initialize_output" begin + V_cod = fuse(codomain(t)) + V_dom = fuse(domain(t)) + U = similar(t, codomain(t) ← V_cod) + S = similar(t, real(scalartype(t)), V_cod ← V_dom) + Vᴴ = similar(t, V_dom ← domain(t)) + return U, S, Vᴴ + end end function MAK.initialize_output(::typeof(svd_compact!), t::AbstractTensorMap, ::AbstractAlgorithm) - V_cod = V_dom = infimum(fuse(codomain(t)), fuse(domain(t))) - U = similar(t, codomain(t) ← V_cod) - S = similar_diagonal(t, real(scalartype(t)), V_cod) - Vᴴ = similar(t, V_dom ← domain(t)) - return U, S, Vᴴ + @timeit_debug GLOBAL_TIMER "alloc: initialize_output" begin + V_cod = V_dom = infimum(fuse(codomain(t)), fuse(domain(t))) + U = similar(t, codomain(t) ← V_cod) + S = similar_diagonal(t, real(scalartype(t)), V_cod) + Vᴴ = similar(t, V_dom ← domain(t)) + return U, S, Vᴴ + end end function MAK.initialize_output(::typeof(svd_vals!), t::AbstractTensorMap, alg::AbstractAlgorithm) - V_cod = infimum(fuse(codomain(t)), fuse(domain(t))) - T = real(scalartype(t)) - A = similarstoragetype(t, T) - return SectorVector{T, sectortype(t), A}(undef, V_cod) + @timeit_debug GLOBAL_TIMER "alloc: initialize_output" begin + V_cod = infimum(fuse(codomain(t)), fuse(domain(t))) + T = real(scalartype(t)) + A = similarstoragetype(t, T) + return SectorVector{T, sectortype(t), A}(undef, V_cod) + end end # Eigenvalue decomposition # ------------------------ function MAK.initialize_output(::typeof(eigh_full!), t::AbstractTensorMap, ::AbstractAlgorithm) - V_D = fuse(domain(t)) - D = similar_diagonal(t, real(scalartype(t)), V_D) - V = similar(t, codomain(t) ← V_D) - return D, V + @timeit_debug GLOBAL_TIMER "alloc: initialize_output" begin + V_D = fuse(domain(t)) + D = similar_diagonal(t, real(scalartype(t)), V_D) + V = similar(t, codomain(t) ← V_D) + return D, V + end end function MAK.initialize_output(::typeof(eig_full!), t::AbstractTensorMap, ::AbstractAlgorithm) - V_D = fuse(domain(t)) - Tc = complex(scalartype(t)) - D = similar_diagonal(t, Tc, V_D) - V = similar(t, Tc, codomain(t) ← V_D) - return D, V + @timeit_debug GLOBAL_TIMER "alloc: initialize_output" begin + V_D = fuse(domain(t)) + Tc = complex(scalartype(t)) + D = similar_diagonal(t, Tc, V_D) + V = similar(t, Tc, codomain(t) ← V_D) + return D, V + end end function MAK.initialize_output(::typeof(eigh_vals!), t::AbstractTensorMap, alg::AbstractAlgorithm) - V_D = fuse(domain(t)) - T = real(scalartype(t)) - A = similarstoragetype(t, T) - return SectorVector{T, sectortype(t), A}(undef, V_D) + @timeit_debug GLOBAL_TIMER "alloc: initialize_output" begin + V_D = fuse(domain(t)) + T = real(scalartype(t)) + A = similarstoragetype(t, T) + return SectorVector{T, sectortype(t), A}(undef, V_D) + end end function MAK.initialize_output(::typeof(eig_vals!), t::AbstractTensorMap, alg::AbstractAlgorithm) - V_D = fuse(domain(t)) - Tc = complex(scalartype(t)) - A = similarstoragetype(t, Tc) - return SectorVector{Tc, sectortype(t), A}(undef, V_D) + @timeit_debug GLOBAL_TIMER "alloc: initialize_output" begin + V_D = fuse(domain(t)) + Tc = complex(scalartype(t)) + A = similarstoragetype(t, Tc) + return SectorVector{Tc, sectortype(t), A}(undef, V_D) + end end # QR decomposition # ---------------- function MAK.initialize_output(::typeof(qr_full!), t::AbstractTensorMap, ::AbstractAlgorithm) - V_Q = fuse(codomain(t)) - Q = similar(t, codomain(t) ← V_Q) - R = similar(t, V_Q ← domain(t)) - return Q, R + @timeit_debug GLOBAL_TIMER "alloc: initialize_output" begin + V_Q = fuse(codomain(t)) + Q = similar(t, codomain(t) ← V_Q) + R = similar(t, V_Q ← domain(t)) + return Q, R + end end function MAK.initialize_output(::typeof(qr_compact!), t::AbstractTensorMap, ::AbstractAlgorithm) - V_Q = infimum(fuse(codomain(t)), fuse(domain(t))) - Q = similar(t, codomain(t) ← V_Q) - R = similar(t, V_Q ← domain(t)) - return Q, R + @timeit_debug GLOBAL_TIMER "alloc: initialize_output" begin + V_Q = infimum(fuse(codomain(t)), fuse(domain(t))) + Q = similar(t, codomain(t) ← V_Q) + R = similar(t, V_Q ← domain(t)) + return Q, R + end end function MAK.initialize_output(::typeof(qr_null!), t::AbstractTensorMap, ::AbstractAlgorithm) - V_Q = infimum(fuse(codomain(t)), fuse(domain(t))) - V_N = ⊖(fuse(codomain(t)), V_Q) - N = similar(t, codomain(t) ← V_N) - return N + @timeit_debug GLOBAL_TIMER "alloc: initialize_output" begin + V_Q = infimum(fuse(codomain(t)), fuse(domain(t))) + V_N = ⊖(fuse(codomain(t)), V_Q) + N = similar(t, codomain(t) ← V_N) + return N + end end # LQ decomposition # ---------------- function MAK.initialize_output(::typeof(lq_full!), t::AbstractTensorMap, ::AbstractAlgorithm) - V_Q = fuse(domain(t)) - L = similar(t, codomain(t) ← V_Q) - Q = similar(t, V_Q ← domain(t)) - return L, Q + @timeit_debug GLOBAL_TIMER "alloc: initialize_output" begin + V_Q = fuse(domain(t)) + L = similar(t, codomain(t) ← V_Q) + Q = similar(t, V_Q ← domain(t)) + return L, Q + end end function MAK.initialize_output(::typeof(lq_compact!), t::AbstractTensorMap, ::AbstractAlgorithm) - V_Q = infimum(fuse(codomain(t)), fuse(domain(t))) - L = similar(t, codomain(t) ← V_Q) - Q = similar(t, V_Q ← domain(t)) - return L, Q + @timeit_debug GLOBAL_TIMER "alloc: initialize_output" begin + V_Q = infimum(fuse(codomain(t)), fuse(domain(t))) + L = similar(t, codomain(t) ← V_Q) + Q = similar(t, V_Q ← domain(t)) + return L, Q + end end function MAK.initialize_output(::typeof(lq_null!), t::AbstractTensorMap, ::AbstractAlgorithm) - V_Q = infimum(fuse(codomain(t)), fuse(domain(t))) - V_N = ⊖(fuse(domain(t)), V_Q) - N = similar(t, V_N ← domain(t)) - return N + @timeit_debug GLOBAL_TIMER "alloc: initialize_output" begin + V_Q = infimum(fuse(codomain(t)), fuse(domain(t))) + V_N = ⊖(fuse(domain(t)), V_Q) + N = similar(t, V_N ← domain(t)) + return N + end end # Polar decomposition # ------------------- function MAK.initialize_output(::typeof(left_polar!), t::AbstractTensorMap, ::AbstractAlgorithm) - W = similar(t, space(t)) - P = similar(t, domain(t) ← domain(t)) - return W, P + @timeit_debug GLOBAL_TIMER "alloc: initialize_output" begin + W = similar(t, space(t)) + P = similar(t, domain(t) ← domain(t)) + return W, P + end end function MAK.initialize_output(::typeof(right_polar!), t::AbstractTensorMap, ::AbstractAlgorithm) - P = similar(t, codomain(t) ← codomain(t)) - Wᴴ = similar(t, space(t)) - return P, Wᴴ + @timeit_debug GLOBAL_TIMER "alloc: initialize_output" begin + P = similar(t, codomain(t) ← codomain(t)) + Wᴴ = similar(t, space(t)) + return P, Wᴴ + end end # Projections @@ -247,7 +291,7 @@ MAK.initialize_output(::typeof(project_hermitian!), tsrc::AbstractTensorMap, ::A MAK.initialize_output(::typeof(project_antihermitian!), tsrc::AbstractTensorMap, ::AbstractAlgorithm) = tsrc MAK.initialize_output(::typeof(project_isometric!), tsrc::AbstractTensorMap, ::AbstractAlgorithm) = - similar(tsrc) + @timeit_debug GLOBAL_TIMER "alloc: initialize_output" similar(tsrc) # Exponential # ---------------- diff --git a/src/planar/planaroperations.jl b/src/planar/planaroperations.jl index dce3d663c..59a968568 100644 --- a/src/planar/planaroperations.jl +++ b/src/planar/planaroperations.jl @@ -91,20 +91,22 @@ function planartrace!( q1 = $(q₁), q2 = $(q₂)")) end - if iszero(β) - fill!(C, β) - elseif !isone(β) - rmul!(C, β) - end - β′ = One() - for (f₁, f₂) in fusiontrees(A) - for ((f₁′, f₂′), coeff) in planar_trace((f₁, f₂), (p₁, p₂), (q₁, q₂)) - TO.tensortrace!( - C[f₁′, f₂′], - A[f₁, f₂], (p₁, p₂), (q₁, q₂), false, - α * coeff, β′, - backend, allocator - ) + @timeit_debug GLOBAL_TIMER "planartrace!" begin + if iszero(β) + fill!(C, β) + elseif !isone(β) + rmul!(C, β) + end + β′ = One() + for (f₁, f₂) in fusiontrees(A) + for ((f₁′, f₂′), coeff) in planar_trace((f₁, f₂), (p₁, p₂), (q₁, q₂)) + @timeit_debug GLOBAL_TIMER "dense: trace" TO.tensortrace!( + C[f₁′, f₂′], + A[f₁, f₂], (p₁, p₂), (q₁, q₂), false, + α * coeff, β′, + backend, allocator + ) + end end end return C @@ -159,35 +161,36 @@ function planarcontract!( return contract!(C, A, pA, B, pB, pAB, α, β, backend, allocator) end - codA, domA = codomainind(A), domainind(A) - codB, domB = codomainind(B), domainind(B) - oindA, cindA = pA - cindB, oindB = pB - oindA, cindA, oindB, cindB = reorder_indices( - codA, domA, codB, domB, oindA, cindA, oindB, cindB, pAB... - ) - - if oindA == codA && cindA == domA - A′ = A - else - A′ = TO.tensoralloc_add( - scalartype(A), A, (oindA, cindA), false, Val(true), allocator + @timeit_debug GLOBAL_TIMER "planarcontract!" begin + codA, domA = codomainind(A), domainind(A) + codB, domB = codomainind(B), domainind(B) + oindA, cindA = pA + cindB, oindB = pB + oindA, cindA, oindB, cindB = reorder_indices( + codA, domA, codB, domB, oindA, cindA, oindB, cindB, pAB... ) - transpose!(A′, A, (oindA, cindA), One(), Zero(), backend, allocator) - end - if cindB == codB && oindB == domB - B′ = B - else - B′ = TensorOperations.tensoralloc_add( - scalartype(B), B, (cindB, oindB), false, Val(true), allocator - ) - transpose!(B′, B, (cindB, oindB), One(), Zero(), backend, allocator) - end - mul!(C, A′, B′, α, β) - (oindA == codA && cindA == domA) || TO.tensorfree!(A′, allocator) - (cindB == codB && oindB == domB) || TO.tensorfree!(B′, allocator) + if oindA == codA && cindA == domA + A′ = A + else + A′ = @timeit_debug GLOBAL_TIMER "alloc: buffers" TO.tensoralloc_add( + scalartype(A), A, (oindA, cindA), false, Val(true), allocator + ) + transpose!(A′, A, (oindA, cindA), One(), Zero(), backend, allocator) + end + if cindB == codB && oindB == domB + B′ = B + else + B′ = @timeit_debug GLOBAL_TIMER "alloc: buffers" TensorOperations.tensoralloc_add( + scalartype(B), B, (cindB, oindB), false, Val(true), allocator + ) + transpose!(B′, B, (cindB, oindB), One(), Zero(), backend, allocator) + end + mul!(C, A′, B′, α, β) + (oindA == codA && cindA == domA) || TO.tensorfree!(A′, allocator) + (cindB == codB && oindB == domB) || TO.tensorfree!(B′, allocator) + end return C end diff --git a/src/tensors/indexmanipulations.jl b/src/tensors/indexmanipulations.jl index e788b9c1d..da1a62d84 100644 --- a/src/tensors/indexmanipulations.jl +++ b/src/tensors/indexmanipulations.jl @@ -307,14 +307,18 @@ See also [`braid`](@ref) for creating a new tensor. backend::AbstractBackend = TO.DefaultBackend(), allocator = TO.DefaultAllocator() ) @boundscheck spacecheck_transform(braid, tdst, tsrc, p, levels) - if has_array_view(tdst) && has_array_view(tsrc) - TO.tensoradd!(tdst[], tsrc[], p, false, α, β, backend, allocator) - return tdst + @timeit_debug GLOBAL_TIMER "permute!/braid!" begin + if has_array_view(tdst) && has_array_view(tsrc) + @timeit_debug GLOBAL_TIMER "dense: tensoradd" TO.tensoradd!( + tdst[], tsrc[], p, false, α, β, backend, allocator + ) + return tdst + end + levels1 = TupleTools.getindices(levels, codomainind(tsrc)) + levels2 = TupleTools.getindices(levels, domainind(tsrc)) + transformer = treebraider(tdst, tsrc, p, (levels1, levels2)) + return @inbounds add_transform!(tdst, tsrc, p, transformer, α, β, backend, allocator) end - levels1 = TupleTools.getindices(levels, codomainind(tsrc)) - levels2 = TupleTools.getindices(levels, domainind(tsrc)) - transformer = treebraider(tdst, tsrc, p, (levels1, levels2)) - return @inbounds add_transform!(tdst, tsrc, p, transformer, α, β, backend, allocator) end """ @@ -383,12 +387,16 @@ end backend::AbstractBackend = TO.DefaultBackend(), allocator = TO.DefaultAllocator() ) @boundscheck spacecheck_transform(transpose, tdst, tsrc, p) - if has_array_view(tdst) && has_array_view(tsrc) - TO.tensoradd!(tdst[], tsrc[], p, false, α, β, backend, allocator) - return tdst + @timeit_debug GLOBAL_TIMER "transpose!" begin + if has_array_view(tdst) && has_array_view(tsrc) + @timeit_debug GLOBAL_TIMER "dense: tensoradd" TO.tensoradd!( + tdst[], tsrc[], p, false, α, β, backend, allocator + ) + return tdst + end + transformer = treetransposer(tdst, tsrc, p) + return @inbounds add_transform!(tdst, tsrc, p, transformer, α, β, backend, allocator) end - transformer = treetransposer(tdst, tsrc, p) - return @inbounds add_transform!(tdst, tsrc, p, transformer, α, β, backend, allocator) end """ @@ -565,7 +573,9 @@ Base.@deprecate( else p2 = (linearize(p), ()) if has_array_view(tdst) && has_array_view(tsrc) - TO.tensoradd!(tdst[], tsrc[], p2, false, α, β, backend, allocator) + @timeit_debug GLOBAL_TIMER "dense: tensoradd" TO.tensoradd!( + tdst[], tsrc[], p2, false, α, β, backend, allocator + ) else ntasks = use_threaded_transform(tdst, transformer) ? get_num_transformer_threads() : 1 if tdst isa TensorMap && tsrc isa TensorMap # unpack data fields to avoid specializing @@ -591,23 +601,27 @@ function add_transform_kernel!( ) I = sectortype(tdst) if FusionStyle(I) === UniqueFusion() - taskforeach(fusiontrees(tsrc), ntasks) do (f₁, f₂) - (f₁′, f₂′), coeff = transformer((f₁, f₂)) - @inbounds TO.tensoradd!( - tdst[f₁′, f₂′], tsrc[f₁, f₂], p, false, α * coeff, β, backend, allocator - ) + @timeit_debug GLOBAL_TIMER "dense: tensoradd" begin + taskforeach(fusiontrees(tsrc), ntasks) do (f₁, f₂) + (f₁′, f₂′), coeff = transformer((f₁, f₂)) + @inbounds TO.tensoradd!( + tdst[f₁′, f₂′], tsrc[f₁, f₂], p, false, α * coeff, β, backend, allocator + ) + end end return nothing end - fblocks = fusionblocks(tsrc) - bufsize = buffersize(tsrc, fblocks) + @timeit_debug GLOBAL_TIMER "bookkeeping: fusionblocks" begin + fblocks = fusionblocks(tsrc) + bufsize = buffersize(tsrc, fblocks) + end # One max-sized workspace per task (a single one that is reused by all blocks when # serial), allocated on the calling thread before any task spawns, so that also # allocators that are not thread-safe can be used. cp = TO.allocator_checkpoint!(allocator) - buffers = [ + @timeit_debug GLOBAL_TIMER "alloc: buffers" buffers = [ TO.tensoralloc(storagetype(tdst), bufsize, Val(true), allocator) for _ in 1:clamp(length(fblocks), 1, ntasks) ] @@ -627,11 +641,13 @@ function add_transform_kernel!( data_dst::DenseVector, data_src::DenseVector, p, transformer::AbelianTreeTransformer, α, β, backend, allocator, ntasks::Int ) - taskforeach(transformer.data, ntasks) do (coeff, struct_dst, struct_src) - TO.tensoradd!( - StridedView(data_dst, struct_dst...), StridedView(data_src, struct_src...), - p, false, α * coeff, β, backend, allocator - ) + @timeit_debug GLOBAL_TIMER "dense: tensoradd" begin + taskforeach(transformer.data, ntasks) do (coeff, struct_dst, struct_src) + TO.tensoradd!( + StridedView(data_dst, struct_dst...), StridedView(data_src, struct_src...), + p, false, α * coeff, β, backend, allocator + ) + end end return nothing end @@ -645,7 +661,7 @@ function add_transform_kernel!( # serial), allocated on the calling thread before any task spawns, so that also # allocators that are not thread-safe can be used. cp = TO.allocator_checkpoint!(allocator) - buffers = [ + @timeit_debug GLOBAL_TIMER "alloc: buffers" buffers = [ TO.tensoralloc(typeof(data_dst), bufsize, Val(true), allocator) for _ in 1:clamp(length(transformer.data), 1, ntasks) ] @@ -668,7 +684,7 @@ function _add_transform_block!( if length(src) == 1 # Degenerate block with a single tree: no matmul needed. (f₁, f₂) = only(fusiontrees(src)) (f₁′, f₂′) = only(fusiontrees(dst)) - @inbounds TO.tensoradd!( + @timeit_debug GLOBAL_TIMER "dense: tensoradd" @inbounds TO.tensoradd!( tdst[f₁′, f₂′], tsrc[f₁, f₂], p, false, α * only(U), β, backend, allocator ) else # Multi-tree block: pack → recoupling matmul → unpack. @@ -683,7 +699,7 @@ function _add_transform_block!( # 1. Extract: copy each source block into column i of buffer_src as a flat vector, # using a trivial permutation so the layout is canonical before the matmul. - @inbounds for (i, (f₁, f₂)) in enumerate(fusiontrees(src)) + @timeit_debug GLOBAL_TIMER "dense: pack" @inbounds for (i, (f₁, f₂)) in enumerate(fusiontrees(src)) TO.tensoradd!( sreshape(view(buffer_src, :, i), sz_src), tsrc[f₁, f₂], ptriv, false, One(), Zero(), backend, allocator @@ -692,12 +708,14 @@ function _add_transform_block!( # 2. Recoupling: buffer_dst = α * buffer_src * U^T (each output tree is a linear # combination of input trees weighted by the recoupling coefficients). - U′ = _adapt_recoupling(storagetype(tdst), U) - mul!(buffer_dst, buffer_src, transpose(U′), α, Zero()) + @timeit_debug GLOBAL_TIMER "dense: recouple mul!" begin + U′ = _adapt_recoupling(storagetype(tdst), U) + mul!(buffer_dst, buffer_src, transpose(U′), α, Zero()) + end # 3. Insert: scatter column i of buffer_dst into the destination, applying the # actual index permutation p in the same tensoradd! call. - @inbounds for (i, (f₃, f₄)) in enumerate(fusiontrees(dst)) + @timeit_debug GLOBAL_TIMER "dense: unpack" @inbounds for (i, (f₃, f₄)) in enumerate(fusiontrees(dst)) TO.tensoradd!( tdst[f₃, f₄], sreshape(view(buffer_dst, :, i), sz_src), p, false, One(), β, backend, allocator @@ -714,7 +732,7 @@ function _add_transform_block!( ) if length(U) == 1 # Degenerate block with a single tree: no matmul needed. coeff = only(U) - TO.tensoradd!( + @timeit_debug GLOBAL_TIMER "dense: tensoradd" TO.tensoradd!( StridedView(data_dst, sz_dst, only(structs_dst)...), StridedView(data_src, sz_src, only(structs_src)...), p, false, α * coeff, β, backend, allocator @@ -728,7 +746,7 @@ function _add_transform_block!( # 1. Extract: copy each source block into column i of buffer_src as a flat vector, # using a trivial permutation so the layout is canonical before the matmul. - @inbounds for (i, struct_src_i) in enumerate(structs_src) + @timeit_debug GLOBAL_TIMER "dense: pack" @inbounds for (i, struct_src_i) in enumerate(structs_src) TO.tensoradd!( sreshape(view(buffer_src, :, i), sz_src), StridedView(data_src, sz_src, struct_src_i...), ptriv, false, One(), Zero(), backend, allocator @@ -737,12 +755,14 @@ function _add_transform_block!( # 2. Recoupling: buffer_dst = α * buffer_src * U^T (each output tree is a linear # combination of input trees weighted by the recoupling coefficients). - U′ = _adapt_recoupling(typeof(data_dst), U) - mul!(buffer_dst, buffer_src, transpose(U′), α, Zero()) + @timeit_debug GLOBAL_TIMER "dense: recouple mul!" begin + U′ = _adapt_recoupling(typeof(data_dst), U) + mul!(buffer_dst, buffer_src, transpose(U′), α, Zero()) + end # 3. Insert: scatter column i of buffer_dst into the destination, applying the # actual index permutation p in the same tensoradd! call. - @inbounds for (i, struct_dst_i) in enumerate(structs_dst) + @timeit_debug GLOBAL_TIMER "dense: unpack" @inbounds for (i, struct_dst_i) in enumerate(structs_dst) TO.tensoradd!( StridedView(data_dst, sz_dst, struct_dst_i...), sreshape(view(buffer_dst, :, i), sz_src), p, false, One(), β, backend, allocator diff --git a/src/tensors/linalg.jl b/src/tensors/linalg.jl index 4c885a6ed..b93f20cb3 100644 --- a/src/tensors/linalg.jl +++ b/src/tensors/linalg.jl @@ -333,37 +333,39 @@ function LinearAlgebra.mul!( compose(space(tA), space(tB)) == space(tC) || throw(SpaceMismatch(lazy"$(space(tC)) ≠ $(space(tA)) * $(space(tB))")) - iterC = blocks(tC) - iterA = blocks(tA) - iterB = blocks(tB) - nextA = iterate(iterA) - nextB = iterate(iterB) - nextC = iterate(iterC) - while !isnothing(nextC) - (cC, C), stateC = nextC - if !isnothing(nextA) && !isnothing(nextB) - (cA, A), stateA = nextA - (cB, B), stateB = nextB - if cA == cC && cB == cC - mul!(C, A, B, α, β) - nextA = iterate(iterA, stateA) - nextB = iterate(iterB, stateB) - nextC = iterate(iterC, stateC) - elseif cA < cC - nextA = iterate(iterA, stateA) - elseif cB < cC - nextB = iterate(iterB, stateB) + @timeit_debug GLOBAL_TIMER "dense: matmul" begin + iterC = blocks(tC) + iterA = blocks(tA) + iterB = blocks(tB) + nextA = iterate(iterA) + nextB = iterate(iterB) + nextC = iterate(iterC) + while !isnothing(nextC) + (cC, C), stateC = nextC + if !isnothing(nextA) && !isnothing(nextB) + (cA, A), stateA = nextA + (cB, B), stateB = nextB + if cA == cC && cB == cC + mul!(C, A, B, α, β) + nextA = iterate(iterA, stateA) + nextB = iterate(iterB, stateB) + nextC = iterate(iterC, stateC) + elseif cA < cC + nextA = iterate(iterA, stateA) + elseif cB < cC + nextB = iterate(iterB, stateB) + else + if β != one(β) + rmul!(C, β) + end + nextC = iterate(iterC, stateC) + end else if β != one(β) rmul!(C, β) end nextC = iterate(iterC, stateC) end - else - if β != one(β) - rmul!(C, β) - end - nextC = iterate(iterC, stateC) end end return tC diff --git a/src/tensors/tensoroperations.jl b/src/tensors/tensoroperations.jl index 147d27b16..4cb274cf3 100644 --- a/src/tensors/tensoroperations.jl +++ b/src/tensors/tensoroperations.jl @@ -234,10 +234,12 @@ function trace_permute!( q₁ = $(q₁), q₂ = $(q₂)")) end - if has_array_view(tdst) && has_array_view(tsrc) - TO.tensortrace!(tdst[], tsrc[], (p₁, p₂), (q₁, q₂), false, α, β, backend) - else - _trace_permute!(FusionStyle(I), tdst, tsrc, (p₁, p₂), (q₁, q₂), α, β, backend) + @timeit_debug GLOBAL_TIMER "trace_permute!" begin + if has_array_view(tdst) && has_array_view(tsrc) + TO.tensortrace!(tdst[], tsrc[], (p₁, p₂), (q₁, q₂), false, α, β, backend) + else + _trace_permute!(FusionStyle(I), tdst, tsrc, (p₁, p₂), (q₁, q₂), α, β, backend) + end end return tdst @@ -248,7 +250,7 @@ function _trace_permute!(::UniqueFusion, tdst, tsrc, (p₁, p₂), (q₁, q₂), r₁, r₂ = (p₁..., q₁...), (p₂..., q₂...) N₁, N₂ = length(p₁), length(p₂) - for (f₁, f₂) in fusiontrees(tsrc) + @timeit_debug GLOBAL_TIMER "dense: trace" for (f₁, f₂) in fusiontrees(tsrc) (f₁′, f₂′), coeff = permute((f₁, f₂), (r₁, r₂)) f₁′′, g₁ = split(f₁′, N₁) f₂′′, g₂ = split(f₂′, N₂) @@ -275,7 +277,7 @@ function _trace_permute!(::FusionStyle, tdst, tsrc, (p₁, p₂), (q₁, q₂), for src in fusionblocks(tsrc) dst, U = permute(src, (r₁, r₂)) - for (i, (f₁, f₂)) in enumerate(fusiontrees(src)) + @timeit_debug GLOBAL_TIMER "dense: trace" for (i, (f₁, f₂)) in enumerate(fusiontrees(src)) for (j, (f₁′, f₂′)) in enumerate(fusiontrees(dst)) coeff = U[j, i] iszero(coeff) && continue @@ -325,33 +327,35 @@ function contract!( length(pA[2]) == length(pB[1]) || throw(IndexError("number of contracted indices does not match")) - # find optimal contraction scheme by checking the following options: - # - sorting the contracted inds of A or B to avoid permutations - # - contracting B with A instead to avoid permutations - pA′, pB′, pA″, pB″, pAB′ = _contract_candidates(pA, pB, pAB) + @timeit_debug GLOBAL_TIMER "contract!" begin + # find optimal contraction scheme by checking the following options: + # - sorting the contracted inds of A or B to avoid permutations + # - contracting B with A instead to avoid permutations + pA′, pB′, pA″, pB″, pAB′ = _contract_candidates(pA, pB, pAB) - # dims are permutation-invariant, so compute them once here rather than in every memcost call - dA, dB, dC = dim(A), dim(B), dim(C) + # dims are permutation-invariant, so compute them once here rather than in every memcost call + dA, dB, dC = dim(A), dim(B), dim(C) - # keep order A en B, check possibilities for cind - memcost1 = _contract_memcost(dA, dB, dC, C, A, pA′, B, pB′, pAB) - memcost2 = _contract_memcost(dA, dB, dC, C, A, pA″, B, pB″, pAB) + # keep order A en B, check possibilities for cind + memcost1 = _contract_memcost(dA, dB, dC, C, A, pA′, B, pB′, pAB) + memcost2 = _contract_memcost(dA, dB, dC, C, A, pA″, B, pB″, pAB) - # reverse order A en B, check possibilities for cind - memcost3 = _contract_memcost(dB, dA, dC, C, B, reverse(pB′), A, reverse(pA′), pAB′) - memcost4 = _contract_memcost(dB, dA, dC, C, B, reverse(pB″), A, reverse(pA″), pAB′) + # reverse order A en B, check possibilities for cind + memcost3 = _contract_memcost(dB, dA, dC, C, B, reverse(pB′), A, reverse(pA′), pAB′) + memcost4 = _contract_memcost(dB, dA, dC, C, B, reverse(pB″), A, reverse(pA″), pAB′) - return if min(memcost1, memcost2) <= min(memcost3, memcost4) - if memcost1 <= memcost2 - return blas_contract!(C, A, pA′, B, pB′, pAB, α, β, backend, allocator) - else - return blas_contract!(C, A, pA″, B, pB″, pAB, α, β, backend, allocator) - end - else - if memcost3 <= memcost4 - return blas_contract!(C, B, reverse(pB′), A, reverse(pA′), pAB′, α, β, backend, allocator) + return if min(memcost1, memcost2) <= min(memcost3, memcost4) + if memcost1 <= memcost2 + return blas_contract!(C, A, pA′, B, pB′, pAB, α, β, backend, allocator) + else + return blas_contract!(C, A, pA″, B, pB″, pAB, α, β, backend, allocator) + end else - return blas_contract!(C, B, reverse(pB″), A, reverse(pA″), pAB′, α, β, backend, allocator) + if memcost3 <= memcost4 + return blas_contract!(C, B, reverse(pB′), A, reverse(pA′), pAB′, α, β, backend, allocator) + else + return blas_contract!(C, B, reverse(pB″), A, reverse(pA″), pAB′, α, β, backend, allocator) + end end end end @@ -423,7 +427,7 @@ function blas_contract!( # Bring A in the correct form for BLAS contraction if copyA - Anew = TO.tensoralloc_add(TC, A, pA, false, Val(true), allocator) + Anew = @timeit_debug GLOBAL_TIMER "alloc: buffers" TO.tensoralloc_add(TC, A, pA, false, Val(true), allocator) Anew = TO.tensoradd!(Anew, A, pA, false, One(), Zero(), backend, allocator) twistA && twist!(Anew, filter(!isdual ∘ Base.Fix1(space, Anew), domainind(Anew))) else @@ -433,7 +437,7 @@ function blas_contract!( # Bring B in the correct form for BLAS contraction if copyB - Bnew = TO.tensoralloc_add(TC, B, pB, false, Val(true), allocator) + Bnew = @timeit_debug GLOBAL_TIMER "alloc: buffers" TO.tensoralloc_add(TC, B, pB, false, Val(true), allocator) Bnew = TO.tensoradd!(Bnew, B, pB, false, One(), Zero(), backend, allocator) twistB && twist!(Bnew, filter(isdual ∘ Base.Fix1(space, Bnew), codomainind(Bnew))) else @@ -446,7 +450,7 @@ function blas_contract!( copyC = !TO.isblasdestination(C, ipAB) if copyC - Cnew = TO.tensoralloc_add(TC, C, ipAB, false, Val(true), allocator) + Cnew = @timeit_debug GLOBAL_TIMER "alloc: buffers" TO.tensoralloc_add(TC, C, ipAB, false, Val(true), allocator) mul!(Cnew, Anew, Bnew) TO.tensoradd!(C, Cnew, pAB, false, α, β, backend, allocator) TO.tensorfree!(Cnew, allocator) diff --git a/src/tensors/treetransformers.jl b/src/tensors/treetransformers.jl index 4a6401a36..00c0ec24b 100644 --- a/src/tensors/treetransformers.jl +++ b/src/tensors/treetransformers.jl @@ -23,7 +23,7 @@ function AbelianTreeTransformer(transform, p, Vdst, Vsrc) N = numind(Vsrc) data = Vector{Tuple{T, StridedStructure{N}, StridedStructure{N}}}(undef, L) - for (i, (f_src, stridestructure_src)) in enumerate(pairs(fts_src)) + @timeit_debug GLOBAL_TIMER "symmetry: tree transform" for (i, (f_src, stridestructure_src)) in enumerate(pairs(fts_src)) f_dst, coeff = transform(f_src) stridestructure_dst = fts_dst[f_dst] data[i] = (coeff, stridestructure_dst, stridestructure_src) @@ -83,27 +83,31 @@ function GenericTreeTransformer(transform, p, Vdst, Vsrc) N₁ = numout(Vsrc) N₂ = numin(Vsrc) - fblocks = fusionblocks(Vsrc) + fblocks = @timeit_debug GLOBAL_TIMER "bookkeeping: fusionblocks" fusionblocks(Vsrc) nblocks = length(fblocks) data = Vector{GenericTransformerData{T, N}}(undef, nblocks) nthreads = get_num_manipulation_threads() - taskforeach(1:nblocks, nthreads) do i - fs_src = fblocks[i] - fs_dst, U = transform(fs_src) - sz_src, newstructs_src = repack_transformer_structure(fusionstructure_src, fusiontrees(fs_src)) - sz_dst, newstructs_dst = repack_transformer_structure(fusionstructure_dst, fusiontrees(fs_dst)) - data[i] = U, (sz_dst, newstructs_dst), (sz_src, newstructs_src) - - @debug( - lazy"Created recoupling block for uncoupled: $(fs_src.uncoupled)", - sz = size(U), sparsity = count(!iszero, U) / length(U) - ) + @timeit_debug GLOBAL_TIMER "symmetry: recoupling matrices" begin + taskforeach(1:nblocks, nthreads) do i + fs_src = fblocks[i] + fs_dst, U = transform(fs_src) + @timeit_debug GLOBAL_TIMER "bookkeeping: repack" begin + sz_src, newstructs_src = repack_transformer_structure(fusionstructure_src, fusiontrees(fs_src)) + sz_dst, newstructs_dst = repack_transformer_structure(fusionstructure_dst, fusiontrees(fs_dst)) + end + data[i] = U, (sz_dst, newstructs_dst), (sz_src, newstructs_src) + + @debug( + lazy"Created recoupling block for uncoupled: $(fs_src.uncoupled)", + sz = size(U), sparsity = count(!iszero, U) / length(U) + ) + end end transformer = GenericTreeTransformer{T, N}(data) # sort by (approximate) weight to facilitate multi-threading strategies - sort!(transformer) + @timeit_debug GLOBAL_TIMER "bookkeeping: sort" sort!(transformer) Δt = Base.time() - t₀ diff --git a/test/Project.toml b/test/Project.toml index ca00312e2..ce51b70a3 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -24,6 +24,7 @@ TensorKitSectors = "13a9c161-d5da-41f0-bcbd-e1a08ae0647f" TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" TestExtras = "5ed8adda-3752-4e41-b88a-e8b09835ee3a" +TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f" TupleTools = "9d95972d-f1c8-5527-a6e0-b4b365fa01f6" VectorInterface = "409d34a3-91d5-4945-b6ec-7529ddf182d8" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" @@ -42,5 +43,6 @@ JET = "0.9, 0.10, 0.11" ParallelTestRunner = "2" Test = "1" TestExtras = "0.2,0.3" +TimerOutputs = "1" Zygote = "0.7" diff --git a/test/other/timers.jl b/test/other/timers.jl new file mode 100644 index 000000000..5da8e6f63 --- /dev/null +++ b/test/other/timers.jl @@ -0,0 +1,80 @@ +using Test, TestExtras +using TensorKit +using TensorOperations +using TimerOutputs +using TimerOutputs: @timeit + +@testset "timer API (no-op by default)" begin + @test TensorKit.timer() isa TimerOutput + @test TensorKit.timeit_debug_enabled() === false + @test TensorKit.timers_enabled() === false + + V = SU2Space(0 => 2, 1 // 2 => 2, 1 => 1) + t = rand(V ⊗ V ← V) + TensorKit.reset_timers!() + permute(t, ((2, 1), (3,))) + # nothing is recorded while timers are disabled + @test isempty(TensorKit.timer().root.children) +end + +@testset "timer_summary aggregation" begin + to = TimerOutput() + @timeit to "permute!" begin + @timeit to "bookkeeping: cache treebraider" begin + @timeit to "symmetry: compute treebraider" sleep(0.01) + end + @timeit to "dense: pack" sleep(0.01) + end + summary = TensorKit.timer_summary(nothing; to) + + @test keys(summary) == Set(TensorKit.TIMER_CATEGORIES) + @test summary[:bookkeeping].ncalls == 1 + @test summary[:symmetry].ncalls == 1 + @test summary[:dense].ncalls == 1 + @test summary[:alloc] == (time = 0, allocated = 0, ncalls = 0) + @test summary[:symmetry].time > 0 + @test summary[:dense].time > 0 + # the unprefixed top-level section contributes its exclusive time to :other + @test summary[:other].ncalls == 1 + + # exclusive-time attribution: category totals sum to the total measured time + total = sum(x -> x.time, values(summary)) + @test total ≈ TimerOutputs.time(only(to.root.children)) rtol = 0.01 + + # printing form does not error + @test sprint(io -> TensorKit.timer_summary(io; to)) isa String +end + +@testset "smoke test with timers enabled" begin + TensorKit.enable_timers!() + try + @test TensorKit.timeit_debug_enabled() === true + TensorKit.reset_timers!() + empty_globalcaches!() # ensure the (symmetry) construction work is not cached + + V = SU2Space(0 => 2, 1 // 2 => 2) + t = rand(V ⊗ V ← V ⊗ V) + permute(t, ((1, 3), (2, 4))) + @tensor t2[a; b] := t[a c; b c] + @tensor t3[a b; c d] := t[a x; c y] * t[y b; x d] + svd_compact(t) + + names = [child.name for child in TensorKit.timer().root.children] + @test "permute!/braid!" in names + @test "contract!" in names + # also verifies that `enable_timers!` reached the Factorizations submodule + @test "svd_compact!" in names + + summary = TensorKit.timer_summary(nothing) + @test summary[:dense].ncalls > 0 + @test summary[:symmetry].ncalls > 0 + @test summary[:bookkeeping].ncalls > 0 + @test summary[:alloc].ncalls > 0 + + # the full printed table renders without error + @test sprint(TensorKit.print_timers) isa String + finally + TensorKit.disable_timers!() + end + @test TensorKit.timeit_debug_enabled() === false +end