diff --git a/Project.toml b/Project.toml index d1e0c834d..235b3e56c 100644 --- a/Project.toml +++ b/Project.toml @@ -9,9 +9,9 @@ projects = ["test", "docs"] [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" Dictionaries = "85a47980-9c8c-11e8-2b9f-f7ca1fa99fb4" -LRUCache = "8ac3fa9e-de4c-5943-b1dc-09c6b5f20637" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" MatrixAlgebraKit = "6c742aac-3347-4629-af66-fc926824e5e4" +MemoizationKit = "1b238080-9255-4fe9-b224-89eb24efe93b" OhMyThreads = "67456a42-1dca-4109-a031-0a68de7e3ad5" PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a" Preferences = "21216c6a-2e73-6563-6e65-726566657250" @@ -55,9 +55,9 @@ Enzyme = "0.13.207" EnzymeTestUtils = "0.2.8" FiniteDifferences = "0.12" GPUArrays = "11.4.1" -LRUCache = "1.6" LinearAlgebra = "1" MatrixAlgebraKit = "0.6.9" +MemoizationKit = "0.1" Mooncake = "0.5.27" OhMyThreads = "0.8.0" PrecompileTools = "1" diff --git a/docs/src/Changelog.md b/docs/src/Changelog.md index 2a853124a..58e06ac27 100644 --- a/docs/src/Changelog.md +++ b/docs/src/Changelog.md @@ -28,7 +28,8 @@ When releasing a new version, move the "Unreleased" changes to a new version sec - Extend VectorInterface compat to include v0.7 - Require TensorKitSectors v0.3.10 -- `@cached` can be used from other modules, such as package extensions, which then own the cache of the methods they add; these caches are shown with their module in `global_cache_info` after registration in `__init__` (internal) +- Caching is provided by [MemoizationKit.jl](https://github.com/QuantumKitHub/MemoizationKit.jl), which can inspect, resize and configure TensorKit's caches, also by memory size; this replaces the internal `@cached` machinery and the LRUCache.jl dependency +- Storage-specific transformers of package extensions are cached by adding `@cached` methods to `device_treebraider` and `device_treetransposer`, to which `treebraider` and `treetransposer` forward, such that their caches can be sized independently (internal) - The `TreeTransformer`s used in index manipulations are now constructed and cached per storagetype of the destination tensor, which is also a dispatch point for storage-specific transformers with their own cache (internal) ### Deprecated diff --git a/docs/src/man/fusiontrees.md b/docs/src/man/fusiontrees.md index 7797bc8bd..b6e9dd0f5 100644 --- a/docs/src/man/fusiontrees.md +++ b/docs/src/man/fusiontrees.md @@ -260,7 +260,29 @@ This is simply given by The `braid` and `permute` routines for double fusion trees will be the main access point for corresponding manipulations on tensors. As a consequence, results from these routines are memoized via a `@cached` mechanism. -The caching strategy is controlled by `CacheStyle`: for `FusionStyle(I) isa UniqueFusion` no caching is used (`NoCache()`), since the result is a single cheap scalar coefficient; for `MultipleFusion` or `GenericFusion` sectors, a global LRU cache (`GlobalLRUCache()`) is used, as computing the transformation matrix can be expensive and results are reused frequently. +The caching strategy is controlled by `MemoizationKit.CacheStyle`: for `FusionStyle(I) isa UniqueFusion` no caching is used (`NoCache()`), since the result is a single cheap scalar coefficient; for `MultipleFusion` or `GenericFusion` sectors, `GlobalCache()` uses MemoizationKit.jl's configured container (ClockCache by default), as computing the transformation matrix can be expensive and results are reused frequently. + +TensorKit also uses MemoizationKit.jl for sector structures, degeneracy structures and fusion tree transformers. +Sector structures share cache entries across spaces with the same sectors, even if their degeneracy dimensions differ, through `MemoizationKit.cachekey`. +The MemoizationKit APIs can inspect, resize and configure these caches: + +```julia +using TensorKit, MemoizationKit +MemoizationKit.cache_info(TensorKit) +MemoizationKit.set_cache_size!(TensorKit.treebraider, 50_000) +MemoizationKit.set_cache_size!(TensorKit.treebraider, 2^30; by = MemoizationKit.cachesize) # 1 GiB of cached values +MemoizationKit.set_cache_preferences!(TensorKit; maxsize = 50_000, measure = "bytes") +MemoizationKit.set_cache_preferences!(; container = "LRU") +``` + +Preferences are saved in `LocalPreferences.toml` and apply after restarting Julia; a runtime size change applies immediately. +`MemoizationKit.uncached(f, args...)` computes a result without consulting or filling the cache. +TensorKit supplies efficient `MemoizationKit.cachesize` methods for its structure, fusion-tree and transformer results. +These estimate bytes from the fixed-size fields and live array buffers; matrix coefficients and inline sector values are not traversed individually. +Reference-containing scalar or sector arrays fall back to `Base.summarysize`. +Buffers shared between cache entries count toward each entry, while identical source and destination structure buffers within a transformer count once. +The budget excludes cache keys, container and allocation overhead, and spare array capacity; it is a cached-value estimate rather than a limit on total process memory. +Further features, such as the live cache dashboard, disk caching and the full configuration options, are described in the [MemoizationKit.jl documentation](https://quantumkithub.github.io/MemoizationKit.jl/stable/). This caching implies that potential inefficiencies in the fusion tree manipulations (which we nonetheless try to avoid) will not seriously affect performance of tensor manipulations. diff --git a/docs/src/man/profiling.md b/docs/src/man/profiling.md index 81b36abfc..236ffe968 100644 --- a/docs/src/man/profiling.md +++ b/docs/src/man/profiling.md @@ -1,6 +1,7 @@ # [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. +Enabling TensorKit's timers also enables MemoizationKit.jl's cache lookup and computation timers, recorded in the same timer tree and categories. 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 | diff --git a/ext/TensorKitEnzymeExt/utility.jl b/ext/TensorKitEnzymeExt/utility.jl index e24c8fb0c..52a49579b 100644 --- a/ext/TensorKitEnzymeExt/utility.jl +++ b/ext/TensorKitEnzymeExt/utility.jl @@ -26,7 +26,7 @@ pullback_dC!(ΔC, β::Number) = scale!(ΔC, conj(β)) @inline EnzymeRules.inactive_type(::Type{<:TensorKit.GenericTreeTransformer}) = true @inline EnzymeRules.inactive_type(::Type{<:TensorKit.UniqueTreeTransformer}) = true @inline EnzymeRules.inactive_type(::Type{<:TensorKit.VectorSpace}) = true -@inline EnzymeRules.inactive_type(::Type{<:TensorKit.LRU}) = true +@inline EnzymeRules.inactive_type(::Type{<:TensorKit.MemoizationKit.AbstractCache}) = true function EnzymeRules.augmented_primal( config::EnzymeRules.RevConfigWidth{1}, @@ -142,13 +142,13 @@ end @inline EnzymeRules.inactive(::typeof(TensorKit.fusiontrees), ::Any) = nothing @inline EnzymeRules.inactive(::typeof(TensorKit.fusiontrees), ::Any, ::Any) = nothing @inline EnzymeRules.inactive(::typeof(TensorKit.fusiontrees), ::Any, ::Any, ::Any) = nothing -@inline EnzymeRules.inactive(::typeof(TensorKit.fsbraid), ::Any) = nothing -@inline EnzymeRules.inactive(::typeof(TensorKit.fsbraid), ::Any, ::Any) = nothing +@inline EnzymeRules.inactive(::typeof(TensorKit.fusiontreebraid), ::Any, ::Any, ::Any) = nothing @inline EnzymeRules.inactive(::typeof(TensorKit.artin_braid), ::Any, ::Any) = nothing -@inline EnzymeRules.inactive(::typeof(TensorKit.fstranspose), ::Any) = nothing -@inline EnzymeRules.inactive(::typeof(TensorKit.fstranspose), ::Any, ::Any) = nothing +@inline EnzymeRules.inactive(::typeof(TensorKit.fusiontreetranspose), ::Any, ::Any) = nothing @inline EnzymeRules.inactive(::typeof(TensorKit.treetransposer), ::Any...) = nothing @inline EnzymeRules.inactive(::typeof(TensorKit.treebraider), ::Any...) = nothing +@inline EnzymeRules.inactive(::typeof(TensorKit.device_treetransposer), ::Any...) = nothing +@inline EnzymeRules.inactive(::typeof(TensorKit.device_treebraider), ::Any...) = nothing @inline EnzymeRules.inactive(::typeof(TensorKit.insertleftunit), ::HomSpace, ::Any) = nothing @inline EnzymeRules.inactive(::typeof(TensorKit.insertrightunit), ::HomSpace, ::Any) = nothing @inline EnzymeRules.inactive(::typeof(TensorKit.removeunit), ::HomSpace, ::Any) = nothing diff --git a/ext/TensorKitGPUArraysExt.jl b/ext/TensorKitGPUArraysExt.jl index 948f9fa75..4ad4d75ca 100644 --- a/ext/TensorKitGPUArraysExt.jl +++ b/ext/TensorKitGPUArraysExt.jl @@ -196,6 +196,7 @@ end # strides of a dense array of shape `sz` _dense_strides(size::Dims) = (1, Base.front(cumprod(size))...) +_dense_strides(::Tuple{}) = () # `permute(Vsrc, p) == Vdst` is enforced when the transformer is built, so the permuted # source shape always matches `sz_dst` and the two views share Cartesian inds. @@ -271,8 +272,16 @@ function DeviceGenericTreeTransformer( ) end +TensorKit.treebraider( + A::Type{<:AnyGPUVector}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, + conjsrc::Bool, levels::IndexTuple +) = TensorKit.device_treebraider(A, Vdst, Vsrc, p, conjsrc, levels) +TensorKit.treetransposer( + A::Type{<:AnyGPUVector}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, conjsrc::Bool +) = TensorKit.device_treetransposer(A, Vdst, Vsrc, p, conjsrc) + # Cache the device transformer directly, constructing the cached host version only on a miss. -@cached function TensorKit.treebraider( +@cached function TensorKit.device_treebraider( A::Type{TA}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, conjsrc::Bool, levels::IndexTuple ) where {TA <: AnyGPUVector} @@ -280,31 +289,13 @@ end return device_transformer(A, host, linearize(p)) end -@cached function TensorKit.treetransposer( +@cached function TensorKit.device_treetransposer( A::Type{TA}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, conjsrc::Bool ) where {TA <: AnyGPUVector} host = TensorKit.treetransposer(Vector{scalartype(A)}, Vdst, Vsrc, p, conjsrc) return device_transformer(A, host, linearize(p)) end -const DEFAULT_DEVICECACHE_SIZE = 1000 - -function __init__() - # extensions can't modify GLOBAL_CACHES during precompilation, so register explicitly - append!( - TensorKit.GLOBAL_CACHES, ( - Symbol("TensorKitGPUArraysExt.GLOBAL_TREEBRAIDER_CACHE") => GLOBAL_TREEBRAIDER_CACHE, - Symbol("TensorKitGPUArraysExt.GLOBAL_TREETRANSPOSER_CACHE") => GLOBAL_TREETRANSPOSER_CACHE, - ) - ) - - # device-caches should probably have smaller sizes: - resize!(GLOBAL_TREEBRAIDER_CACHE; maxsize = DEFAULT_DEVICECACHE_SIZE) - resize!(GLOBAL_TREETRANSPOSER_CACHE; maxsize = DEFAULT_DEVICECACHE_SIZE) - - return nothing -end - # Allocations must outlive any GPUArrays allocation-cache scope around the operation. function device_transformer(A, transformer::UniqueTreeTransformer, p) packed = DeviceUniqueTreeTransformer(transformer, p) diff --git a/src/TensorKit.jl b/src/TensorKit.jl index 6bdbebfb5..4b4494ef4 100644 --- a/src/TensorKit.jl +++ b/src/TensorKit.jl @@ -119,7 +119,6 @@ const TO = TensorOperations using MatrixAlgebraKit using Dictionaries: Dictionaries, Dictionary, Indices, gettoken, gettokenvalue -using LRUCache using OhMyThreads using ScopedValues using TimerOutputs: TimerOutputs, TimerOutput, @timeit_debug diff --git a/src/auxiliary/caches.jl b/src/auxiliary/caches.jl index bdc07ca08..c327f34ba 100644 --- a/src/auxiliary/caches.jl +++ b/src/auxiliary/caches.jl @@ -1,16 +1,23 @@ -""" - const GLOBAL_CACHES +# Caching of the fusion tree manipulations and tensor structures is provided by MemoizationKit.jl. +# `CacheStyle` is imported (not just used) so that the `CacheStyle(::typeof(f), ...)` +# specializations in TensorKit, and those of users, extend `MemoizationKit.CacheStyle`. +import MemoizationKit: CacheStyle +using MemoizationKit: MemoizationKit, @cached, NoCache, GlobalCache, Hashed, LRU -Registry of the global caches TensorKit maintains, as `name => cache` pairs. +# Buffer payloads can be measured without visiting their elements when stored inline. +# Reference-containing sector/scalar types need the recursive fallback to include their data. +_cache_payload_size(a::Array{T}) where {T} = isbitstype(T) ? sizeof(a) : Base.summarysize(a) -See also [`empty_globalcaches!`](@ref) and [`global_cache_info`](@ref). -""" -const GLOBAL_CACHES = Pair{Symbol, Any}[] +# Include the index tables as well as the values (Dictionaries 0.4's Indices layout). +function _cache_payload_size(inds::Indices) + return sizeof(getfield(inds, :slots)) + sizeof(getfield(inds, :hashes)) + + _cache_payload_size(getfield(inds, :values)) +end """ empty_globalcaches!() -Empty every global cache in [`GLOBAL_CACHES`](@ref). +Empty the global caches of TensorKit. These mostly contain bookkeeping for various different index manipulations and tensor structures, so clearing this out can free up some memory whenever you have a workflow that involves a large variety of structures. For example, you may want to clear the cache when working with different symmetries, or in algorithms that dynamically alter the sizes of tensors. @@ -19,218 +26,20 @@ Since everything is recomputed on demand, this is purely a memory measure. See also [`global_cache_info`](@ref) to display the status. """ function empty_globalcaches!() - foreach(empty! ∘ last, GLOBAL_CACHES) + MemoizationKit.empty_caches!(TensorKit) return nothing end """ global_cache_info([io::IO = stdout]) -Print the hit/miss statistics and current size of every global cache in [`GLOBAL_CACHES`](@ref). +Print the hit/miss statistics and current size of every global cache of TensorKit. See also [`empty_globalcaches!`](@ref). """ function global_cache_info(io::IO = stdout) - for (name, cache) in GLOBAL_CACHES - println(io, name, ":\t", LRUCache.cache_info(cache)) + for (f, cache) in MemoizationKit.cache_info(TensorKit) + println(io, nameof(f), ":\t", cache) end return end - -abstract type CacheStyle end -struct NoCache <: CacheStyle end -struct TaskLocalCache{D <: AbstractDict} <: CacheStyle end -struct GlobalLRUCache <: CacheStyle end - -const DEFAULT_GLOBALCACHE_SIZE = Ref(10^4) - -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 - -# timer sections are controlled by the `timeit_debug_enabled` switch of the module using `@cached` -_timeit_expr(label, ex) = - Expr(:macrocall, GlobalRef(TimerOutputs, Symbol("@timeit_debug")), LineNumberNode(@__LINE__, @__FILE__), GlobalRef(@__MODULE__, :GLOBAL_TIMER), label, ex) - -macro cached(ex) - return _cached(__module__, ex) -end - -function _cached_error(msg, ex) - ex isa Expr && (ex = Base.remove_linenums!(copy(ex))) - Meta.isexpr(ex, :kw) && (ex = Expr(:(=), ex.args...)) # prints as `x = 1` - throw(ArgumentError("@cached: $msg, got `$ex`")) -end - -# the arguments form the cache key, so they must be positional; unnamed ones get a name -function _cached_canonicalize_arg(arg) - Meta.isexpr(arg, :parameters) && _cached_error("keyword arguments are not supported", arg) - Meta.isexpr(arg, :kw) && _cached_error("default values are not supported", arg) - Meta.isexpr(arg, :...) && _cached_error("varargs are not supported", arg) - arg isa Symbol && return arg - Meta.isexpr(arg, :(::)) || _cached_error("unsupported argument", arg) - length(arg.args) == 1 && return Expr(:(::), gensym(:arg), arg.args[1]) - return arg -end - -function _cached(mod::Module, ex) - if !Meta.isexpr(ex, :function) - Meta.isexpr(ex, :(=)) && Meta.isexpr(ex.args[1], (:call, :where, :(::))) && - _cached_error("short-form function definitions are not supported, use `function ... end`", ex.args[1]) - _cached_error("expected a function definition", ex) - end - fcall = ex.args[1] - if Meta.isexpr(fcall, :where) - hasparams = true - params = fcall.args[2:end] - fcall = fcall.args[1] - else - hasparams = false - end - if Meta.isexpr(fcall, :(::)) - typed = true - typeex = fcall.args[2] - fcall = fcall.args[1] - else - typed = false - end - Meta.isexpr(fcall, :call) || _cached_error("expected a function signature", fcall) - fname = fcall.args[1] - # qualified names such as `TensorKit.treebraider` add methods to a function of another module - basename = Meta.isexpr(fname, :.) ? fname.args[end].value : fname - basename isa Symbol || - _cached_error("expected a function name or a qualified name such as `Module.f`", fname) - fargs = map(_cached_canonicalize_arg, fcall.args[2:end]) - fcall = Expr(:call, fname, fargs...) - fargnames = map(arg -> arg isa Symbol ? arg : arg.args[1], fargs) - # timer labels for the cache lookup and the miss-path construction - lookuplabel = string("bookkeeping: cache ", basename) - misslabel = string(_cached_category(basename), ": compute ", basename) - _fbody = ex.args[2] - - # actual implenetation, with underscore name - _fname = Symbol(:_, basename) - _fcall = Expr(:call, _fname, fargs...) - if hasparams - _fcall = Expr(:where, _fcall, params...) - end - _fex = Expr(:function, _fcall, _fbody) - - # implementation that chooses the cache style - newfcall = fcall - if hasparams - newfcall = Expr(:where, newfcall, params...) - end - cachestylevar = gensym(:cachestyle) - cachestyleex = Expr( - :(=), cachestylevar, Expr(:call, CacheStyle, fname, fargnames...) - ) - newfbody = Expr( - :block, cachestyleex, Expr(:call, fname, fargnames..., cachestylevar) - ) - newfex = Expr(:function, newfcall, newfbody) - - # nocache implementation - fnocachecall = Expr(:call, fname, fargs..., :(::$NoCache)) - if hasparams - fnocachecall = Expr(:where, fnocachecall, params...) - end - fnocachebody = _timeit_expr(misslabel, Expr(:call, _fname, fargnames...)) - if typed - T = gensym(:T) - fnocachebody = Expr(:block, Expr(:(=), T, typeex), Expr(:(::), fnocachebody, T)) - end - fnocacheex = Expr(:function, fnocachecall, fnocachebody) - - # tasklocal cache implementation - Dvar = gensym(:D) - flocalcachecall = Expr(:call, fname, fargs..., :(::$TaskLocalCache{$Dvar})) - if hasparams - flocalcachecall = Expr(:where, flocalcachecall, params..., Dvar) - else - flocalcachecall = Expr(:where, flocalcachecall, Dvar) - end - globalcachename = Symbol(:GLOBAL_, uppercase(string(basename)), :_CACHE) - localcachename = Symbol(:_tasklocal_, globalcachename) - cachevar = gensym(:cache) - getlocalcacheex = :( - $cachevar::$Dvar = get!(task_local_storage(), $(QuoteNode(localcachename))) do - return $Dvar() - end - ) - valvar = gensym(:val) - if length(fargnames) == 1 - key = fargnames[1] - else - key = Expr(:tuple, fargnames...) - end - missex = _timeit_expr(misslabel, Expr(:call, _fname, fargnames...)) - getvalex = _timeit_expr(lookuplabel, :(get!(() -> $missex, $cachevar, $key))) - if typed - T = gensym(:T) - flocalcachebody = Expr( - :block, - getlocalcacheex, - Expr(:(=), T, typeex), - Expr(:(=), Expr(:(::), valvar, T), getvalex), - Expr(:return, valvar) - ) - else - flocalcachebody = Expr( - :block, - getlocalcacheex, - Expr(:(=), valvar, getvalex), - Expr(:return, valvar) - ) - end - flocalcacheex = Expr(:function, flocalcachecall, flocalcachebody) - - # # global cache implementation - fglobalcachecall = Expr(:call, fname, fargs..., :(::$GlobalLRUCache)) - if hasparams - fglobalcachecall = Expr(:where, fglobalcachecall, params...) - end - getglobalcachex = Expr(:(=), cachevar, globalcachename) - if typed - T = gensym(:T) - fglobalcachebody = Expr( - :block, - getglobalcachex, - Expr(:(=), T, typeex), - Expr(:(=), Expr(:(::), valvar, T), getvalex), - Expr(:return, valvar) - ) - else - fglobalcachebody = Expr( - :block, - getglobalcachex, - Expr(:(=), valvar, getvalex), - Expr(:return, valvar) - ) - end - fglobalcacheex = Expr(:function, fglobalcachecall, fglobalcachebody) - fglobalcachedef = Expr( - :const, - Expr(:(=), globalcachename, :($LRU{Any, Any}(; maxsize = $DEFAULT_GLOBALCACHE_SIZE[]))) - ) - # Other modules register their caches explicitly from `__init__`, since - # mutations of this registry during their precompilation do not survive loading. - fglobalcacheregister = mod === (@__MODULE__) ? Expr( - :call, :push!, GlobalRef(@__MODULE__, :GLOBAL_CACHES), - :($(QuoteNode(globalcachename)) => $globalcachename) - ) : nothing - - # # total expression - return esc( - Expr( - :block, _fex, newfex, fnocacheex, flocalcacheex, - fglobalcachedef, fglobalcacheregister, fglobalcacheex - ) - ) -end diff --git a/src/auxiliary/dicts.jl b/src/auxiliary/dicts.jl index 642373e19..b5923e117 100644 --- a/src/auxiliary/dicts.jl +++ b/src/auxiliary/dicts.jl @@ -242,25 +242,3 @@ function Base.mergewith(combine, d1::SortedVectorDict{K, V}, d2::SortedVectorDic unmatched = combine == min ? Returns(nothing) : identity return _sortedmerge(combine, unmatched, unmatched, d1, d2) end - -""" - Hashed(value, hashfunction = Base.hash, isequal = Base.isequal) - -Wrapper struct to alter the `hash` and `isequal` implementations of a given value. -This is useful in the contexts of dictionaries, where you either want to customize the hashfunction, -or consider various values as equal with a different notion of equality. -""" -struct Hashed{T, H <: Function, E <: Function} - val::T - hashf::H - eqf::E -end - -Hashed(val, hashf = Base.hash, eqf = Base.isequal) = - Hashed{typeof(val), typeof(hashf), typeof(eqf)}(val, hashf, eqf) - -Base.parent(h::Hashed) = h.val -Base.hash(h::Hashed, seed::UInt) = h.hashf(parent(h), seed) -# Note: requires the equality functions to be equal to avoid asymmetric results -Base.isequal(h1::Hashed{<:Any, <:Any, E}, h2::Hashed{<:Any, <:Any, E}) where {E} = - h1.eqf(parent(h1), parent(h2)) diff --git a/src/auxiliary/timers.jl b/src/auxiliary/timers.jl index 2c3977458..77bde1c60 100644 --- a/src/auxiliary/timers.jl +++ b/src/auxiliary/timers.jl @@ -61,6 +61,7 @@ host-side dispatch of asynchronous kernels unless the workload is explicitly syn """ function enable_timers!() TimerOutputs.enable_debug_timings(TensorKit) + MemoizationKit.enable_cache_timers!(TensorKit, GLOBAL_TIMER) return nothing end @@ -72,6 +73,7 @@ Also triggers recompilation of the instrumented methods on first use. """ function disable_timers!() TimerOutputs.disable_debug_timings(TensorKit) + MemoizationKit.disable_cache_timers!(TensorKit) return nothing end diff --git a/src/fusiontrees/braiding_manipulations.jl b/src/fusiontrees/braiding_manipulations.jl index ce08f510c..f2a702684 100644 --- a/src/fusiontrees/braiding_manipulations.jl +++ b/src/fusiontrees/braiding_manipulations.jl @@ -284,42 +284,45 @@ function braid(src::Union{FusionTreePair, FusionTreeBlock}, p::Index2Tuple, leve @assert numind(src) == length(p[1]) + length(p[2]) @assert numout(src) == length(levels[1]) && numin(src) == length(levels[2]) @assert TupleTools.isperm((p[1]..., p[2]...)) - return fsbraid((src, p, levels)) + return fusiontreebraid(src, p, levels) end -const FSPBraidKey{I, N₁, N₂} = Tuple{FusionTreePair{I}, Index2Tuple{N₁, N₂}, Index2Tuple} -const FSBBraidKey{I, N₁, N₂} = Tuple{FusionTreeBlock{I}, Index2Tuple{N₁, N₂}, Index2Tuple} - -Base.@assume_effects :foldable function _fsdicttype(::Type{T}) where {I, N₁, N₂, T <: FSPBraidKey{I, N₁, N₂}} - E = sectorscalartype(I) - return Pair{fusiontreetype(I, N₁, N₂), E} -end -Base.@assume_effects :foldable function _fsdicttype(::Type{T}) where {I, N₁, N₂, T <: FSBBraidKey{I, N₁, N₂}} - F₁ = fusiontreetype(I, N₁) - F₂ = fusiontreetype(I, N₂) - E = sectorscalartype(I) - return Pair{FusionTreeBlock{I, N₁, N₂, Tuple{F₁, F₂}}, Matrix{E}} +# Braiding can promote real fusion coefficients to complex values. Keep the result +# type explicit so all braid paths remain inferred, including anyonic sectors. +Base.@assume_effects :foldable function fusiontreebraidtype( + src::Union{FusionTreePair{I}, FusionTreeBlock{I}}, ::Index2Tuple{N₁, N₂} + ) where {I, N₁, N₂} + T = sectorscalartype(I) + if src isa FusionTreePair + return Pair{fusiontreetype(I, N₁, N₂), T} + else + F = Tuple{fusiontreetype(I, N₁), fusiontreetype(I, N₂)} + return Pair{FusionTreeBlock{I, N₁, N₂, F}, Matrix{T}} + end end -@cached function fsbraid(key::K)::_fsdicttype(K) where {I, N₁, N₂, K <: Union{FSPBraidKey{I, N₁, N₂}, FSBBraidKey{I, N₁, N₂}}} - if K <: FSPBraidKey - ((f₁, f₂), (p1, p2), (l1, l2)) = key - p = linearizepermutation(p1, p2, length(f₁), length(f₂)) - levels = (l1..., reverse(l2)...) - (f, f0), coeff1 = repartition((f₁, f₂), N₁ + N₂) - f′, coeff2 = braid(f, p, levels) - (f₁′, f₂′), coeff3 = repartition((f′, f0), N₁) - return (f₁′, f₂′) => coeff1 * coeff2 * coeff3 +@cached function fusiontreebraid( + src::Union{FusionTreePair{I}, FusionTreeBlock{I}}, + p::Index2Tuple{N₁, N₂}, levels::Index2Tuple + )::fusiontreebraidtype(src, p) where {I, N₁, N₂} + p₁, p₂ = p + l₁, l₂ = levels + if src isa FusionTreePair + f₁, f₂ = src + p′ = linearizepermutation(p₁, p₂, length(f₁), length(f₂)) + levels = (l₁..., reverse(l₂)...) + (f, f₀), coeff₁ = repartition((f₁, f₂), N₁ + N₂) + f′, coeff₂ = braid(f, p′, levels) + (f₁′, f₂′), coeff₃ = repartition((f′, f₀), N₁) + return (f₁′, f₂′) => coeff₁ * coeff₂ * coeff₃ else - src, (p1, p2), (l1, l2) = key - - p = linearizepermutation(p1, p2, numout(src), numin(src)) - levels = (l1..., reverse(l2)...) + p′ = linearizepermutation(p₁, p₂, numout(src), numin(src)) + levels = (l₁..., reverse(l₂)...) dst, U = repartition(src, numind(src)) braid_style = BraidingStyle(I) - for s in permutation2swaps(p) + for s in permutation2swaps(p′) _check_levels(braid_style, levels, s) inv = levels[s] > levels[s + 1] dst, U_tmp = artin_braid(dst, s; inv) @@ -339,8 +342,8 @@ end end end -CacheStyle(::typeof(fsbraid), k::Union{FSPBraidKey{I}, FSBBraidKey{I}}) where {I} = - FusionStyle(I) isa UniqueFusion ? NoCache() : GlobalLRUCache() +CacheStyle(::typeof(fusiontreebraid), src::Union{FusionTreePair{I}, FusionTreeBlock{I}}, args...) where {I} = + FusionStyle(I) isa UniqueFusion ? NoCache() : GlobalCache() """ permute((f₁, f₂)::FusionTreePair, (p1, p2)::Index2Tuple) diff --git a/src/fusiontrees/duality_manipulations.jl b/src/fusiontrees/duality_manipulations.jl index c595afdc2..b53267989 100644 --- a/src/fusiontrees/duality_manipulations.jl +++ b/src/fusiontrees/duality_manipulations.jl @@ -524,28 +524,28 @@ function Base.transpose(src::Union{FusionTreePair, FusionTreeBlock}, p::Index2Tu N == length(p[1]) + length(p[2]) || throw(ArgumentError("invalid permutation p = $p of length N = $N")) p′ = linearizepermutation(p..., numout(src), numin(src)) iscyclicpermutation(p′) || throw(ArgumentError("invalid cyclic or planar permutation p = $p")) - return fstranspose((src, p)) + return fusiontreetranspose(src, p) end -const FSPTransposeKey{I, N₁, N₂} = Tuple{FusionTreePair{I}, Index2Tuple{N₁, N₂}} -const FSBTransposeKey{I, N₁, N₂} = Tuple{FusionTreeBlock{I}, Index2Tuple{N₁, N₂}} - -Base.@assume_effects :foldable function _fsdicttype(::Type{T}) where {I, N₁, N₂, T <: FSPTransposeKey{I, N₁, N₂}} - E = fusionscalartype(I) - return Pair{fusiontreetype(I, N₁, N₂), E} -end -Base.@assume_effects :foldable function _fsdicttype(::Type{T}) where {I, N₁, N₂, T <: FSBTransposeKey{I, N₁, N₂}} - F₁ = fusiontreetype(I, N₁) - F₂ = fusiontreetype(I, N₂) - E = fusionscalartype(I) - return Pair{FusionTreeBlock{I, N₁, N₂, Tuple{F₁, F₂}}, Matrix{E}} +Base.@assume_effects :foldable function fusiontreetransposetype( + src::Union{FusionTreePair{I}, FusionTreeBlock{I}}, ::Index2Tuple{N₁, N₂} + ) where {I, N₁, N₂} + T = fusionscalartype(I) + if src isa FusionTreePair + return Pair{fusiontreetype(I, N₁, N₂), T} + else + F = Tuple{fusiontreetype(I, N₁), fusiontreetype(I, N₂)} + return Pair{FusionTreeBlock{I, N₁, N₂, F}, Matrix{T}} + end end -@cached function fstranspose(key::K)::_fsdicttype(K) where {I, N₁, N₂, K <: Union{FSPTransposeKey{I, N₁, N₂}, FSBTransposeKey{I, N₁, N₂}}} - src, (p1, p2) = key +@cached function fusiontreetranspose( + src::Union{FusionTreePair{I}, FusionTreeBlock{I}}, p::Index2Tuple{N₁, N₂} + )::fusiontreetransposetype(src, p) where {I, N₁, N₂} + p₁, p₂ = p N = N₁ + N₂ - p = linearizepermutation(p1, p2, numout(src), numin(src)) + p = linearizepermutation(p₁, p₂, numout(src), numin(src)) dst, U = repartition(src, N₁) length(p) == 0 && return dst => U @@ -567,10 +567,8 @@ end return dst => U end -CacheStyle(::typeof(fstranspose), k::FSPTransposeKey{I}) where {I} = - FusionStyle(I) isa UniqueFusion ? NoCache() : GlobalLRUCache() -CacheStyle(::typeof(fstranspose), k::FSBTransposeKey{I}) where {I} = - FusionStyle(I) isa UniqueFusion ? NoCache() : GlobalLRUCache() +CacheStyle(::typeof(fusiontreetranspose), src::Union{FusionTreePair{I}, FusionTreeBlock{I}}, args...) where {I} = + FusionStyle(I) isa UniqueFusion ? NoCache() : GlobalCache() # COMPOSITE DUALITY MANIPULATIONS PART 2: Planar traces #------------------------------------------------------------------- diff --git a/src/fusiontrees/fusiontrees.jl b/src/fusiontrees/fusiontrees.jl index 7b97223ec..fccad1596 100644 --- a/src/fusiontrees/fusiontrees.jl +++ b/src/fusiontrees/fusiontrees.jl @@ -122,6 +122,15 @@ struct FusionTreeBlock{I, N₁, N₂, F <: FusionTreePair{I, N₁, N₂}} trees::Vector{F} end +MemoizationKit.cachesize(b::FusionTreeBlock) = sizeof(b) + _cache_payload_size(b.trees) + +# The scalar result is inline for ordinary sectors. Matrix results include the output +# trees and the recoupling matrix, without walking every coefficient. +MemoizationKit.cachesize(p::Pair{<:FusionTreePair, <:Number}) = + isbitstype(typeof(p)) ? sizeof(p) : Base.summarysize(p) +MemoizationKit.cachesize(p::Pair{<:FusionTreeBlock, <:Matrix}) = + sizeof(p) + MemoizationKit.cachesize(first(p)) + _cache_payload_size(last(p)) + function FusionTreeBlock{I}( uncoupled::Tuple{NTuple{N₁, I}, NTuple{N₂, I}}, isdual::Tuple{NTuple{N₁, Bool}, NTuple{N₂, Bool}}; diff --git a/src/spaces/structure.jl b/src/spaces/structure.jl index 9a4ddf5c0..1485e4691 100644 --- a/src/spaces/structure.jl +++ b/src/spaces/structure.jl @@ -21,7 +21,10 @@ struct SectorStructure{I <: Sector, F <: FusionTreePair{I}} fusiontrees::Indices{F} end -Base.@assume_effects :foldable function sectorstructuretype(key::Hashed{S}) where {S <: HomSpace} +MemoizationKit.cachesize(s::SectorStructure) = + sizeof(s) + _cache_payload_size(s.blocksectors) + _cache_payload_size(s.fusiontrees) + +Base.@assume_effects :foldable function sectorstructuretype(::S) where {S <: HomSpace} I = sectortype(S) F = fusiontreetype(I, numout(S), numin(S)) return SectorStructure{I, F} @@ -35,11 +38,8 @@ pairs as `Indices`. The result is cached based on the sector structure of `W` (i degeneracy dimensions). See also [`degeneracystructure`](@ref), [`fusiontrees`](@ref), [`blocksectors`](@ref). -""" sectorstructure(::HomSpace) -sectorstructure(W::HomSpace) = sectorstructure(Hashed(W, sectorhash, sectorequal)) - -@cached function sectorstructure(key::Hashed{S})::sectorstructuretype(key) where {S <: HomSpace} - W = parent(key) +""" +@cached function sectorstructure(W::S)::sectorstructuretype(W) where {S <: HomSpace} codom, dom = codomain(W), domain(W) I = sectortype(S) @@ -71,7 +71,7 @@ sectorstructure(W::HomSpace) = sectorstructure(Hashed(W, sectorhash, sectorequal return SectorStructure{I, F}(Indices(bs), Indices(trees)) end -CacheStyle(::typeof(sectorstructure), ::Hashed{<:HomSpace}) = GlobalLRUCache() +MemoizationKit.cachekey(::typeof(sectorstructure), W::HomSpace) = (Hashed(W, sectorhash, sectorequal),) # DegeneracyStructure: degeneracy-dependent characterization of HomSpaces # ----------------------------------------------------------------------- @@ -97,6 +97,8 @@ struct DegeneracyStructure{N} subblockstructure::Vector{StridedStructure{N}} end +MemoizationKit.cachesize(s::DegeneracyStructure) = sizeof(s) + sizeof(s.blockstructure) + sizeof(s.subblockstructure) + function degeneracystructuretype(W::HomSpace) N = length(codomain(W)) + length(domain(W)) return DegeneracyStructure{N} @@ -194,5 +196,3 @@ function _subblock_strides(subsz, sz, str) throw(ArgumentError("unexpected error in computing subblock strides")) return strides end - -CacheStyle(::typeof(degeneracystructure), ::HomSpace) = GlobalLRUCache() diff --git a/src/tensors/treetransformers.jl b/src/tensors/treetransformers.jl index 990e50d2d..27ebe98ae 100644 --- a/src/tensors/treetransformers.jl +++ b/src/tensors/treetransformers.jl @@ -68,6 +68,22 @@ struct GenericTreeTransformer{T, N} <: TreeTransformer structure_src::Vector{StridedStructure{N}} end +# Both structure arrays are retained by the transformer, even when another cache also +# holds them. Within an entry, identical source/destination arrays are counted only once. +_cache_structure_size(t::TreeTransformer) = sizeof(t.structure_src) + + (t.structure_src === t.structure_dst ? 0 : sizeof(t.structure_dst)) + +MemoizationKit.cachesize(t::UniqueTreeTransformer) = sizeof(t) + _cache_payload_size(t.data) + _cache_structure_size(t) + +function MemoizationKit.cachesize(t::GenericTreeTransformer) + bytes = sizeof(t) + sizeof(t.data) + _cache_structure_size(t) + for (U, inds_dst, inds_src) in t.data + bytes += _cache_payload_size(U) + sizeof(inds_src) + + (inds_src === inds_dst ? 0 : sizeof(inds_dst)) + end + return bytes +end + function UniqueTreeTransformer(::Type{A}, transform, p, Vdst, Vsrc, conjsrc::Bool) where {A} t₀ = Base.time() @@ -210,4 +226,23 @@ end return TreeTransformer(A, fusiontreetransform, p, Vdst, Vsrc, conjsrc) end -# default cachestyle is GlobalLRUCache +""" + device_treebraider(A, Vdst, Vsrc, p, conjsrc, levels) + device_treetransposer(A, Vdst, Vsrc, p, conjsrc) + +Tree transformers for device storage types `A`, to which `treebraider` and `treetransposer` +forward. Their methods are added by package extensions. Being separate functions, the device +transformers have their own cache, which can be configured independently of the host ones. +""" +function device_treebraider end +function device_treetransposer end +@doc (@doc device_treebraider) device_treetransposer + +# Keep TensorKit's profiling categories for the phases instrumented by MemoizationKit. +const TensorKitCachedFunction = Union{ + typeof(fusiontreebraid), typeof(fusiontreetranspose), typeof(treebraider), typeof(treetransposer), + typeof(device_treebraider), typeof(device_treetransposer), + typeof(sectorstructure), typeof(degeneracystructure), +} +MemoizationKit.instrument_label(f::TensorKitCachedFunction, ::Val{:lookup}) = "bookkeeping: cache $(nameof(f))" +MemoizationKit.instrument_label(f::TensorKitCachedFunction, ::Val{:compute}) = "symmetry: compute $(nameof(f))" diff --git a/test/cuda/tensors.jl b/test/cuda/tensors.jl index 973118cac..d859981d9 100644 --- a/test/cuda/tensors.jl +++ b/test/cuda/tensors.jl @@ -573,3 +573,17 @@ end end end end + +@timedtestset "Device transformers are cached separately" begin + TensorKit.empty_globalcaches!() + V = SU2Space(0 => 2, 1 // 2 => 2) + t = CUDA.rand(Float64, V ⊗ V ← V ⊗ V) + p = ((2, 1), (4, 3)) + @test permute(t, p) ≈ adapt(CuArray, permute(ad(t), p)) + devicecache = only(TensorKit.MemoizationKit.cache_info(TensorKit.device_treebraider)).second + hostcache = only(TensorKit.MemoizationKit.cache_info(TensorKit.treebraider)).second + @test length(devicecache) == 1 + @test all(key -> !(first(key) <: CuArray), keys(hostcache)) + TensorKit.empty_globalcaches!() + @test isempty(devicecache) +end diff --git a/test/other/caches.jl b/test/other/caches.jl new file mode 100644 index 000000000..ec750d48e --- /dev/null +++ b/test/other/caches.jl @@ -0,0 +1,89 @@ +using Test, TestExtras +using TensorKit +using TensorKit: MemoizationKit, sectorstructure, degeneracystructure + +@testset "Cache structure keys" begin + V1 = SU2Space(0 => 2, 1 // 2 => 3) + V2 = SU2Space(0 => 4, 1 // 2 => 1) + W1, W2 = V1 ⊗ V1 ← V1, V2 ⊗ V2 ← V2 + empty_globalcaches!() + s1 = @constinferred sectorstructure(W1) + @test s1 === @constinferred sectorstructure(W2) + fresh = MemoizationKit.uncached(sectorstructure, W1) + @test s1.blocksectors == fresh.blocksectors + @test s1.fusiontrees == fresh.fusiontrees + @test parent(only(MemoizationKit.cachekey(sectorstructure, W1))) === W1 + @test length(only(MemoizationKit.cache_info(sectorstructure)).second) == 1 + @test (@constinferred degeneracystructure(W1)).totaldim != + (@constinferred degeneracystructure(W2)).totaldim + @test length(only(MemoizationKit.cache_info(degeneracystructure)).second) == 2 +end + +@testset "TensorKit cached-value sizes" begin + for I in (Z2Irrep, SU2Irrep, A4Irrep, FibonacciAnyon) + labels = collect(Iterators.take(values(I), 2)) + V = GradedSpace{I}(labels[1] => 2, labels[2] => 2) + W = V ⊗ V ← V ⊗ V + p, levels = ((2, 1), (4, 3)), (1, 2, 3, 4) + transformer = TensorKit.treebraider(Vector{Float64}, W, W, p, false, levels) + for value in (sectorstructure(W), degeneracystructure(W), transformer) + bytes = @constinferred MemoizationKit.cachesize(value) + @test 0 < bytes <= Base.summarysize(value) + end + for src in TensorKit.fusionblocks(W) + result = TensorKit.fusiontreebraid(src, p, ((1, 2), (3, 4))) + @test 0 < (@constinferred MemoizationKit.cachesize(result)) <= Base.summarysize(result) + end + + # Two distinct transformer entries share structure buffers. Each entry must + # retain its own size budget even if the degeneracy cache has been cleared. + MemoizationKit.set_cache_size!(TensorKit.treebraider, MemoizationKit.cachesize(transformer); by = MemoizationKit.cachesize) + try + TensorKit.treebraider(Vector{Float64}, W, W, p, false, levels) + TensorKit.treebraider(Vector{Float64}, W, W, ((1, 2), (3, 4)), false, levels) + MemoizationKit.empty_caches!(degeneracystructure) + cache = only(MemoizationKit.cache_info(TensorKit.treebraider)).second + stats = MemoizationKit.cache_stats(cache) + @test stats.by === MemoizationKit.cachesize + @test stats.currentsize <= stats.maxsize + @test stats.length <= 1 + @test stats.currentsize == sum(MemoizationKit.cachesize, values(cache); init = 0) + + MemoizationKit.set_cache_size!(TensorKit.treebraider, 0; by = MemoizationKit.cachesize) + @test isempty(cache) + @test TensorKit.treebraider(Vector{Float64}, W, W, p, false, levels).data == transformer.data + @test isempty(cache) + finally + MemoizationKit.set_cache_size!(TensorKit.treebraider, 10_000) + end + end + + # An empty transformer still costs space, and arbitrary-precision coefficients + # include their referenced storage instead of just the array's pointer slots. + strides = TensorKit.StridedStructure{2}[] + empty_transformer = TensorKit.UniqueTreeTransformer(Tuple{Float64, Int, Int}[], strides, strides) + @test MemoizationKit.cachesize(empty_transformer) > 0 + big_transformer = TensorKit.UniqueTreeTransformer([(BigFloat(1), 1, 1)], strides, strides) + @test MemoizationKit.cachesize(big_transformer) >= Base.summarysize(big_transformer.data) +end + +# A cache owned by another module must survive TensorKit's cache clearing. +MemoizationKit.@cached unrelated_cache(x) = [x] + +@testset "TensorKit cache management" begin + retained = unrelated_cache(1) + try + V = SU2Space(0 => 2, 1 // 2 => 2) + t = rand(V ⊗ V ← V ⊗ V) + @test (@constinferred transpose(transpose(t))) ≈ t + p = ((2, 1), (4, 3)) + @test (@constinferred permute(permute(t, p), p)) ≈ t + info = sprint(TensorKit.global_cache_info) + @test occursin("treebraider", info) + @test empty_globalcaches!() === nothing + @test all(isempty(cache) for (_, cache) in MemoizationKit.cache_info(TensorKit)) + @test unrelated_cache(1) === retained + finally + MemoizationKit.empty_caches!(unrelated_cache) + end +end diff --git a/test/other/hashed.jl b/test/other/hashed.jl deleted file mode 100644 index ae0c18915..000000000 --- a/test/other/hashed.jl +++ /dev/null @@ -1,48 +0,0 @@ -using Test, TestExtras -using TensorKit -using TensorKit: Hashed - -@testset "Hashed" begin - @testset "default constructor" begin - h1 = @constinferred Hashed(42) - h2 = Hashed(42) - @test isequal(h1, h2) - @test hash(h1) == hash(h2) - @test parent(h1) == 42 - end - - @testset "custom hash function" begin - # hash only the length, ignoring contents - lenhash = (v, seed) -> hash(length(v), seed) - h1 = Hashed([1, 2, 3], lenhash) - h2 = Hashed([4, 5, 6], lenhash) - @test hash(h1) == hash(h2) - h3 = Hashed([1, 2], lenhash) - @test hash(h1) != hash(h3) - end - - @testset "custom isequal" begin - # consider vectors equal if they have the same length - lenequal = (a, b) -> length(a) == length(b) - h1 = Hashed([1, 2, 3], Base.hash, lenequal) - h2 = Hashed([4, 5, 6], Base.hash, lenequal) - h3 = Hashed([1, 2], Base.hash, lenequal) - @test isequal(h1, h2) - @test !isequal(h1, h3) - end - - @testset "Dict key usage" begin - d = Dict(Hashed(1) => "one", Hashed(2) => "two") - @test d[Hashed(1)] == "one" - @test d[Hashed(2)] == "two" - @test length(d) == 2 - end - - @testset "Dict with custom hash and isequal" begin - lenhash = (v, seed) -> hash(length(v), seed) - lenequal = (a, b) -> length(a) == length(b) - d = Dict(Hashed([1, 2, 3], lenhash, lenequal) => "length3") - # lookup with different contents but same length should succeed - @test d[Hashed([7, 8, 9], lenhash, lenequal)] == "length3" - end -end diff --git a/test/other/timers.jl b/test/other/timers.jl index 5da8e6f63..809917669 100644 --- a/test/other/timers.jl +++ b/test/other/timers.jl @@ -65,6 +65,13 @@ end # also verifies that `enable_timers!` reached the Factorizations submodule @test "svd_compact!" in names + sections = collect(TensorKit.timer().root.children) + for section in sections + append!(sections, section.children) + end + @test any(s -> s.name == "bookkeeping: cache treebraider", sections) + @test any(s -> s.name == "symmetry: compute treebraider", sections) + summary = TensorKit.timer_summary(nothing) @test summary[:dense].ncalls > 0 @test summary[:symmetry].ncalls > 0 @@ -77,4 +84,8 @@ end TensorKit.disable_timers!() end @test TensorKit.timeit_debug_enabled() === false + TensorKit.reset_timers!() + V = SU2Space(0 => 2, 1 // 2 => 2) + permute(rand(V ⊗ V ← V ⊗ V), ((2, 1), (4, 3))) + @test isempty(TensorKit.timer().root.children) end