From f7b180fc77d20f61a3d435703109265c57c54b0f Mon Sep 17 00:00:00 2001 From: lkdvos Date: Mon, 5 Oct 2026 18:26:47 -0400 Subject: [PATCH 01/11] Replace TensorKit cache machinery with MemoizationKit.jl Track MemoizationKit.jl main through Pkg sources and delegate cache storage, container selection, and module-wide inspection and clearing to MemoizationKit. Preserve the registry for the existing GPU transformer cache and use MemoizationKit.LRU for its container. Co-Authored-By: Claude Opus 5.5 --- Project.toml | 7 +- docs/src/man/fusiontrees.md | 2 +- ext/TensorKitEnzymeExt/utility.jl | 2 +- ext/TensorKitGPUArraysExt.jl | 18 -- src/TensorKit.jl | 1 - src/auxiliary/caches.jl | 221 ++-------------------- src/fusiontrees/braiding_manipulations.jl | 2 +- src/fusiontrees/duality_manipulations.jl | 4 +- src/spaces/structure.jl | 4 - src/tensors/treetransformers.jl | 1 - 10 files changed, 29 insertions(+), 233 deletions(-) diff --git a/Project.toml b/Project.toml index e8d17e10f..3f5289b96 100644 --- a/Project.toml +++ b/Project.toml @@ -8,8 +8,8 @@ projects = ["test", "docs"] [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" +Cached = "1b238080-9255-4fe9-b224-89eb24efe93b" Dictionaries = "85a47980-9c8c-11e8-2b9f-f7ca1fa99fb4" -LRUCache = "8ac3fa9e-de4c-5943-b1dc-09c6b5f20637" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" MatrixAlgebraKit = "6c742aac-3347-4629-af66-fc926824e5e4" OhMyThreads = "67456a42-1dca-4109-a031-0a68de7e3ad5" @@ -45,9 +45,13 @@ TensorKitFiniteDifferencesExt = "FiniteDifferences" TensorKitGPUArraysExt = "GPUArrays" TensorKitMooncakeExt = "Mooncake" +[sources] +Cached = {url = "https://github.com/lkdvos/Cached.jl.git", rev = "main"} + [compat] AMDGPU = "2" Adapt = "4" +Cached = "0.1" CUDA = "6" ChainRulesCore = "1" Dictionaries = "0.4" @@ -55,7 +59,6 @@ Enzyme = "0.13.216" EnzymeTestUtils = "0.2.8" FiniteDifferences = "0.12" GPUArrays = "11.4.1" -LRUCache = "1.6" LinearAlgebra = "1" MatrixAlgebraKit = "0.6.13" Mooncake = "0.5.27" diff --git a/docs/src/man/fusiontrees.md b/docs/src/man/fusiontrees.md index 7797bc8bd..1e6a97ada 100644 --- a/docs/src/man/fusiontrees.md +++ b/docs/src/man/fusiontrees.md @@ -260,7 +260,7 @@ 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 `Cached.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 Cached.jl's configured container (ClockCache by default), as computing the transformation matrix can be expensive and results are reused frequently. 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/ext/TensorKitEnzymeExt/utility.jl b/ext/TensorKitEnzymeExt/utility.jl index bb9a21db0..a686d989f 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.Cached.AbstractCache}) = true function EnzymeRules.augmented_primal( config::EnzymeRules.RevConfigWidth{1}, diff --git a/ext/TensorKitGPUArraysExt.jl b/ext/TensorKitGPUArraysExt.jl index 45f6df640..28c5157ae 100644 --- a/ext/TensorKitGPUArraysExt.jl +++ b/ext/TensorKitGPUArraysExt.jl @@ -287,24 +287,6 @@ end 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 2d1d806db..97031a665 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 d6753ead9..290c96c59 100644 --- a/src/auxiliary/caches.jl +++ b/src/auxiliary/caches.jl @@ -1,16 +1,25 @@ +# Caching of the fusion tree manipulations and tensor structures is provided by Cached.jl. +# `CacheStyle` is imported (not just used) so that the `CacheStyle(::typeof(f), ...)` +# specializations in TensorKit, and those of users, extend `Cached.CacheStyle`. +import Cached: CacheStyle +using Cached: Cached, @cached, NoCache, GlobalCache, LRU + """ const GLOBAL_CACHES -Registry of the global caches TensorKit maintains, as `name => cache` pairs. - -See also [`empty_globalcaches!`](@ref) and [`global_cache_info`](@ref). +Registry of additional global caches, which are not managed by `@cached` (e.g. by +extensions), as `name => cache` pairs. They are emptied by [`empty_globalcaches!`](@ref) +and listed by [`global_cache_info`](@ref), together with the caches of TensorKit's cached +functions. """ const GLOBAL_CACHES = Pair{Symbol, Any}[] +const DEFAULT_GLOBALCACHE_SIZE = Ref(10^4) + """ 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,6 +28,7 @@ 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!() + Cached.empty_caches!(TensorKit) foreach(empty! ∘ last, GLOBAL_CACHES) return nothing end @@ -26,209 +36,16 @@ 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 (f, cache) in Cached.cache_info(TensorKit) + println(io, nameof(f), ":\t", cache) + end for (name, cache) in GLOBAL_CACHES - println(io, name, ":\t", LRUCache.cache_info(cache)) + println(io, name, ":\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) - params = fcall.args[2:end] - fcall = fcall.args[1] - else - params = nothing - end - if Meta.isexpr(fcall, :(::)) - typeex = fcall.args[2] - fcall = fcall.args[1] - else - typeex = nothing - 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 !isnothing(params) - _fcall = Expr(:where, _fcall, params...) - end - _fex = Expr(:function, _fcall, _fbody) - - # implementation that chooses the cache style - newfcall = fcall - if !isnothing(params) - 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 !isnothing(params) - fnocachecall = Expr(:where, fnocachecall, params...) - end - fnocachebody = _timeit_expr(misslabel, Expr(:call, _fname, fargnames...)) - if !isnothing(typeex) - 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 !isnothing(params) - 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 !isnothing(typeex) - 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 !isnothing(params) - fglobalcachecall = Expr(:where, fglobalcachecall, params...) - end - getglobalcachex = Expr(:(=), cachevar, globalcachename) - if !isnothing(typeex) - 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/fusiontrees/braiding_manipulations.jl b/src/fusiontrees/braiding_manipulations.jl index 0ad6353ed..abdb16853 100644 --- a/src/fusiontrees/braiding_manipulations.jl +++ b/src/fusiontrees/braiding_manipulations.jl @@ -347,7 +347,7 @@ end end CacheStyle(::typeof(fsbraid), k::Union{FSPBraidKey{I}, FSBBraidKey{I}}) where {I} = - FusionStyle(I) isa UniqueFusion ? NoCache() : GlobalLRUCache() + 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 cae44aca3..0b179985b 100644 --- a/src/fusiontrees/duality_manipulations.jl +++ b/src/fusiontrees/duality_manipulations.jl @@ -575,9 +575,9 @@ end end CacheStyle(::typeof(fstranspose), k::FSPTransposeKey{I}) where {I} = - FusionStyle(I) isa UniqueFusion ? NoCache() : GlobalLRUCache() + FusionStyle(I) isa UniqueFusion ? NoCache() : GlobalCache() CacheStyle(::typeof(fstranspose), k::FSBTransposeKey{I}) where {I} = - FusionStyle(I) isa UniqueFusion ? NoCache() : GlobalLRUCache() + FusionStyle(I) isa UniqueFusion ? NoCache() : GlobalCache() # COMPOSITE DUALITY MANIPULATIONS PART 2: Planar traces #------------------------------------------------------------------- diff --git a/src/spaces/structure.jl b/src/spaces/structure.jl index 9a4ddf5c0..a26517e7c 100644 --- a/src/spaces/structure.jl +++ b/src/spaces/structure.jl @@ -71,8 +71,6 @@ sectorstructure(W::HomSpace) = sectorstructure(Hashed(W, sectorhash, sectorequal return SectorStructure{I, F}(Indices(bs), Indices(trees)) end -CacheStyle(::typeof(sectorstructure), ::Hashed{<:HomSpace}) = GlobalLRUCache() - # DegeneracyStructure: degeneracy-dependent characterization of HomSpaces # ----------------------------------------------------------------------- """ @@ -194,5 +192,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 b274c6f8c..9eea333dd 100644 --- a/src/tensors/treetransformers.jl +++ b/src/tensors/treetransformers.jl @@ -202,4 +202,3 @@ end return TreeTransformer(A, fusiontreetransform, p, Vdst, Vsrc, conjsrc) end -# default cachestyle is GlobalLRUCache From 3eb08063be001a8eb093f1e468458adc00ac62a0 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Mon, 5 Oct 2026 18:26:48 -0400 Subject: [PATCH 02/11] Use MemoizationKit keys and remove duplicated cache helpers Use MemoizationKit.Hashed and cachekey to share sector structures across degeneracy dimensions without changing the function signature. Pass fusion-tree arguments directly to MemoizationKit, retaining the braiding result annotation needed for anyonic inference. Remove the upstream-covered Hashed tests, add TensorKit integration checks, and document native cache configuration and optional extensions. --- docs/src/man/fusiontrees.md | 17 ++++++++ src/auxiliary/caches.jl | 2 +- src/auxiliary/dicts.jl | 22 ---------- src/fusiontrees/braiding_manipulations.jl | 49 ++++++++++++----------- src/fusiontrees/duality_manipulations.jl | 26 +++--------- src/spaces/structure.jl | 11 +++-- test/other/caches.jl | 46 +++++++++++++++++++++ test/other/hashed.jl | 48 ---------------------- 8 files changed, 101 insertions(+), 120 deletions(-) create mode 100644 test/other/caches.jl delete mode 100644 test/other/hashed.jl diff --git a/docs/src/man/fusiontrees.md b/docs/src/man/fusiontrees.md index 1e6a97ada..ac46b1f6c 100644 --- a/docs/src/man/fusiontrees.md +++ b/docs/src/man/fusiontrees.md @@ -262,6 +262,23 @@ The `braid` and `permute` routines for double fusion trees will be the main acce As a consequence, results from these routines are memoized via a `@cached` mechanism. The caching strategy is controlled by `Cached.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 Cached.jl's configured container (ClockCache by default), as computing the transformation matrix can be expensive and results are reused frequently. +TensorKit also uses Cached.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 `Cached.cachekey`. +The Cached APIs can inspect, resize and configure these caches: + +```julia +using TensorKit, Cached +Cached.cache_info(TensorKit) +Cached.set_cache_size!(TensorKit.treebraider, 50_000) +Cached.set_cache_preferences!(TensorKit; maxsize = 50_000, measure = "bytes") +Cached.set_cache_preferences!(; container = "LRU") +``` + +Preferences are saved in `LocalPreferences.toml` and apply after restarting Julia; a runtime size change applies immediately. +`Cached.uncached(f, args...)` computes a result without consulting or filling the cache. +Optional Cached features, such as the Tachikoma dashboard and SQLite disk caching, remain available through Cached's extensions. +Disk caching requires loading SQLite and defining `Cached.DiskCacheStyle` for the desired function; it is disabled by default, and `empty_globalcaches!()` clears only in-memory caches. + 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. ## Inspecting fusion trees as tensors diff --git a/src/auxiliary/caches.jl b/src/auxiliary/caches.jl index 290c96c59..441f6468d 100644 --- a/src/auxiliary/caches.jl +++ b/src/auxiliary/caches.jl @@ -2,7 +2,7 @@ # `CacheStyle` is imported (not just used) so that the `CacheStyle(::typeof(f), ...)` # specializations in TensorKit, and those of users, extend `Cached.CacheStyle`. import Cached: CacheStyle -using Cached: Cached, @cached, NoCache, GlobalCache, LRU +using Cached: Cached, @cached, NoCache, GlobalCache, Hashed, LRU """ const GLOBAL_CACHES 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/fusiontrees/braiding_manipulations.jl b/src/fusiontrees/braiding_manipulations.jl index abdb16853..50b759f49 100644 --- a/src/fusiontrees/braiding_manipulations.jl +++ b/src/fusiontrees/braiding_manipulations.jl @@ -284,7 +284,7 @@ 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 fsbraid(src, p, levels) end function braid(src::ProductSectorBlock, p::Index2Tuple, levels::Index2Tuple) isempty(src) && return fsbraid((src, p, levels)) @@ -294,39 +294,42 @@ function braid(src::ProductSectorBlock, p::Index2Tuple, levels::Index2Tuple) end::_fsdicttype(typeof((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 fsbraidtype( + 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₂)) +@cached function fsbraid( + src::Union{FusionTreePair{I}, FusionTreeBlock{I}}, + p::Index2Tuple{N₁, N₂}, levels::Index2Tuple + )::fsbraidtype(src, p) where {I, N₁, N₂} + p1, p2 = p + l1, l2 = levels + if src isa FusionTreePair + f₁, f₂ = src + 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′, coeff2 = braid(f, p′, levels) (f₁′, f₂′), coeff3 = repartition((f′, f0), N₁) return (f₁′, f₂′) => coeff1 * coeff2 * coeff3 else - src, (p1, p2), (l1, l2) = key - - p = linearizepermutation(p1, p2, numout(src), numin(src)) + p′ = linearizepermutation(p1, p2, numout(src), numin(src)) levels = (l1..., reverse(l2)...) 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) @@ -346,7 +349,7 @@ end end end -CacheStyle(::typeof(fsbraid), k::Union{FSPBraidKey{I}, FSBBraidKey{I}}) where {I} = +CacheStyle(::typeof(fsbraid), src::Union{FusionTreePair{I}, FusionTreeBlock{I}}, args...) where {I} = FusionStyle(I) isa UniqueFusion ? NoCache() : GlobalCache() """ diff --git a/src/fusiontrees/duality_manipulations.jl b/src/fusiontrees/duality_manipulations.jl index 0b179985b..73b7b46c9 100644 --- a/src/fusiontrees/duality_manipulations.jl +++ b/src/fusiontrees/duality_manipulations.jl @@ -524,7 +524,7 @@ 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 fstranspose(src, p) end function Base.transpose(src::ProductSectorBlock, p::Index2Tuple) isempty(src) && return fstranspose((src, p)) @@ -534,22 +534,10 @@ function Base.transpose(src::ProductSectorBlock, p::Index2Tuple) end::_fsdicttype(typeof((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}} -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 fstranspose( + src::Union{FusionTreePair{I}, FusionTreeBlock{I}}, p::Index2Tuple{N₁, N₂} + ) where {I, N₁, N₂} + p1, p2 = p N = N₁ + N₂ p = linearizepermutation(p1, p2, numout(src), numin(src)) @@ -574,9 +562,7 @@ end return dst => U end -CacheStyle(::typeof(fstranspose), k::FSPTransposeKey{I}) where {I} = - FusionStyle(I) isa UniqueFusion ? NoCache() : GlobalCache() -CacheStyle(::typeof(fstranspose), k::FSBTransposeKey{I}) where {I} = +CacheStyle(::typeof(fstranspose), 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/spaces/structure.jl b/src/spaces/structure.jl index a26517e7c..fc3849afb 100644 --- a/src/spaces/structure.jl +++ b/src/spaces/structure.jl @@ -21,7 +21,7 @@ struct SectorStructure{I <: Sector, F <: FusionTreePair{I}} fusiontrees::Indices{F} end -Base.@assume_effects :foldable function sectorstructuretype(key::Hashed{S}) where {S <: HomSpace} +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 +35,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,6 +68,8 @@ sectorstructure(W::HomSpace) = sectorstructure(Hashed(W, sectorhash, sectorequal return SectorStructure{I, F}(Indices(bs), Indices(trees)) end +Cached.cachekey(::typeof(sectorstructure), W::HomSpace) = (Hashed(W, sectorhash, sectorequal),) + # DegeneracyStructure: degeneracy-dependent characterization of HomSpaces # ----------------------------------------------------------------------- """ diff --git a/test/other/caches.jl b/test/other/caches.jl new file mode 100644 index 000000000..09a71657b --- /dev/null +++ b/test/other/caches.jl @@ -0,0 +1,46 @@ +using Test, TestExtras +using TensorKit +using TensorKit: Cached, sectorstructure, degeneracystructure + +@testset "Cached 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 = Cached.uncached(sectorstructure, W1) + @test s1.blocksectors == fresh.blocksectors + @test s1.fusiontrees == fresh.fusiontrees + @test parent(only(Cached.cachekey(sectorstructure, W1))) === W1 + @test length(only(Cached.cache_info(sectorstructure)).second) == 1 + @test (@constinferred degeneracystructure(W1)).totaldim != + (@constinferred degeneracystructure(W2)).totaldim + @test length(only(Cached.cache_info(degeneracystructure)).second) == 2 +end + +# A cache owned by another module must survive TensorKit's cache clearing. +Cached.@cached unrelated_cache(x) = [x] + +@testset "TensorKit cache management" begin + retained = unrelated_cache(1) + extra = Dict(:key => :value) + push!(TensorKit.GLOBAL_CACHES, :extension_cache => extra) + 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 occursin("extension_cache", info) + @test empty_globalcaches!() === nothing + @test all(isempty(cache) for (_, cache) in Cached.cache_info(TensorKit)) + @test isempty(extra) + @test unrelated_cache(1) === retained + finally + pop!(TensorKit.GLOBAL_CACHES) + Cached.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 From 9cb3a0d5d8f73d6f9d816dc80b1e4910f1aca440 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Mon, 5 Oct 2026 18:26:59 -0400 Subject: [PATCH 03/11] Integrate MemoizationKit timers with TensorKit profiling Enable and disable MemoizationKit instrumentation with TensorKit timers, using the shared timer and existing profiling categories. Verify lookup and computation sections are recorded and that disabling timers stops recording. --- docs/src/man/profiling.md | 1 + src/auxiliary/timers.jl | 2 ++ src/tensors/treetransformers.jl | 8 ++++++++ test/other/timers.jl | 11 +++++++++++ 4 files changed, 22 insertions(+) diff --git a/docs/src/man/profiling.md b/docs/src/man/profiling.md index 81b36abfc..2b26d34ef 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 Cached.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/src/auxiliary/timers.jl b/src/auxiliary/timers.jl index 2c3977458..c8bef53d5 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) + Cached.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) + Cached.disable_cache_timers!(TensorKit) return nothing end diff --git a/src/tensors/treetransformers.jl b/src/tensors/treetransformers.jl index 9eea333dd..77646eedb 100644 --- a/src/tensors/treetransformers.jl +++ b/src/tensors/treetransformers.jl @@ -202,3 +202,11 @@ end return TreeTransformer(A, fusiontreetransform, p, Vdst, Vsrc, conjsrc) end + +# Keep TensorKit's profiling categories for the phases instrumented by Cached. +const TensorKitCachedFunction = Union{ + typeof(fsbraid), typeof(fstranspose), typeof(treebraider), typeof(treetransposer), + typeof(sectorstructure), typeof(degeneracystructure), +} +Cached.instrument_label(f::TensorKitCachedFunction, ::Val{:lookup}) = "bookkeeping: cache $(nameof(f))" +Cached.instrument_label(f::TensorKitCachedFunction, ::Val{:compute}) = "symmetry: compute $(nameof(f))" 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 From 350cf8e491cb2a7c69627b0d4681d080bca0b673 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Mon, 5 Oct 2026 18:58:40 -0400 Subject: [PATCH 04/11] Measure TensorKit cached values by buffer sizes --- docs/src/man/fusiontrees.md | 6 +++++ src/auxiliary/caches.jl | 10 +++++++ src/fusiontrees/fusiontrees.jl | 9 +++++++ src/spaces/structure.jl | 5 ++++ src/tensors/treetransformers.jl | 17 +++++++++++- test/other/caches.jl | 48 +++++++++++++++++++++++++++++++++ 6 files changed, 94 insertions(+), 1 deletion(-) diff --git a/docs/src/man/fusiontrees.md b/docs/src/man/fusiontrees.md index ac46b1f6c..8ebe183fb 100644 --- a/docs/src/man/fusiontrees.md +++ b/docs/src/man/fusiontrees.md @@ -270,12 +270,18 @@ The Cached APIs can inspect, resize and configure these caches: using TensorKit, Cached Cached.cache_info(TensorKit) Cached.set_cache_size!(TensorKit.treebraider, 50_000) +Cached.set_cache_size!(TensorKit.treebraider, 2^30; by = Cached.cachesize) # 1 GiB of cached values Cached.set_cache_preferences!(TensorKit; maxsize = 50_000, measure = "bytes") Cached.set_cache_preferences!(; container = "LRU") ``` Preferences are saved in `LocalPreferences.toml` and apply after restarting Julia; a runtime size change applies immediately. `Cached.uncached(f, args...)` computes a result without consulting or filling the cache. +TensorKit supplies efficient `Cached.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. Optional Cached features, such as the Tachikoma dashboard and SQLite disk caching, remain available through Cached's extensions. Disk caching requires loading SQLite and defining `Cached.DiskCacheStyle` for the desired function; it is disabled by default, and `empty_globalcaches!()` clears only in-memory caches. diff --git a/src/auxiliary/caches.jl b/src/auxiliary/caches.jl index 441f6468d..b91c332f0 100644 --- a/src/auxiliary/caches.jl +++ b/src/auxiliary/caches.jl @@ -16,6 +16,16 @@ const GLOBAL_CACHES = Pair{Symbol, Any}[] const DEFAULT_GLOBALCACHE_SIZE = Ref(10^4) +# 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) + +# 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!() diff --git a/src/fusiontrees/fusiontrees.jl b/src/fusiontrees/fusiontrees.jl index ea8334462..34cd67640 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 +Cached.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. +Cached.cachesize(p::Pair{<:FusionTreePair, <:Number}) = + isbitstype(typeof(p)) ? sizeof(p) : Base.summarysize(p) +Cached.cachesize(p::Pair{<:FusionTreeBlock, <:Matrix}) = + sizeof(p) + Cached.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 fc3849afb..415b0228e 100644 --- a/src/spaces/structure.jl +++ b/src/spaces/structure.jl @@ -21,6 +21,9 @@ struct SectorStructure{I <: Sector, F <: FusionTreePair{I}} fusiontrees::Indices{F} end +Cached.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)) @@ -94,6 +97,8 @@ struct DegeneracyStructure{N} subblockstructure::Vector{StridedStructure{N}} end +Cached.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} diff --git a/src/tensors/treetransformers.jl b/src/tensors/treetransformers.jl index 77646eedb..2ed9b9099 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)) + +Cached.cachesize(t::UniqueTreeTransformer) = sizeof(t) + _cache_payload_size(t.data) + _cache_structure_size(t) + +function Cached.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::HomSpace, Vsrc::HomSpace, conjsrc::Bool) where {A} t₀ = Base.time() @@ -202,7 +218,6 @@ end return TreeTransformer(A, fusiontreetransform, p, Vdst, Vsrc, conjsrc) end - # Keep TensorKit's profiling categories for the phases instrumented by Cached. const TensorKitCachedFunction = Union{ typeof(fsbraid), typeof(fstranspose), typeof(treebraider), typeof(treetransposer), diff --git a/test/other/caches.jl b/test/other/caches.jl index 09a71657b..ca7e9c2e7 100644 --- a/test/other/caches.jl +++ b/test/other/caches.jl @@ -19,6 +19,54 @@ using TensorKit: Cached, sectorstructure, degeneracystructure @test length(only(Cached.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(W, W, p, false, levels) + for value in (sectorstructure(W), degeneracystructure(W), transformer) + bytes = @constinferred Cached.cachesize(value) + @test 0 < bytes <= Base.summarysize(value) + end + for src in TensorKit.fusionblocks(W) + result = TensorKit.fsbraid(src, p, ((1, 2), (3, 4))) + @test 0 < (@constinferred Cached.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. + Cached.set_cache_size!(TensorKit.treebraider, Cached.cachesize(transformer); by = Cached.cachesize) + try + TensorKit.treebraider(W, W, p, false, levels) + TensorKit.treebraider(W, W, ((1, 2), (3, 4)), false, levels) + Cached.empty_caches!(degeneracystructure) + cache = only(Cached.cache_info(TensorKit.treebraider)).second + stats = Cached.cache_stats(cache) + @test stats.by === Cached.cachesize + @test stats.currentsize <= stats.maxsize + @test stats.length <= 1 + @test stats.currentsize == sum(Cached.cachesize, values(cache); init = 0) + + Cached.set_cache_size!(TensorKit.treebraider, 0; by = Cached.cachesize) + @test isempty(cache) + @test TensorKit.treebraider(W, W, p, false, levels).data == transformer.data + @test isempty(cache) + finally + Cached.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 Cached.cachesize(empty_transformer) > 0 + big_transformer = TensorKit.UniqueTreeTransformer([(BigFloat(1), 1, 1)], strides, strides) + @test Cached.cachesize(big_transformer) >= Base.summarysize(big_transformer.data) +end + # A cache owned by another module must survive TensorKit's cache clearing. Cached.@cached unrelated_cache(x) = [x] From cfb5c7f05035e2cca453574345709357c4f3492e Mon Sep 17 00:00:00 2001 From: lkdvos Date: Mon, 12 Oct 2026 14:51:39 -0400 Subject: [PATCH 05/11] Switch to registered MemoizationKit.jl Cached.jl was registered as MemoizationKit.jl under QuantumKitHub; depend on the registered release and drop the `[sources]` entry. Co-Authored-By: Claude Opus 5.5 --- Project.toml | 7 ++---- docs/src/man/fusiontrees.md | 28 ++++++++++----------- docs/src/man/profiling.md | 2 +- ext/TensorKitEnzymeExt/utility.jl | 2 +- src/auxiliary/caches.jl | 12 ++++----- src/auxiliary/timers.jl | 4 +-- src/fusiontrees/fusiontrees.jl | 8 +++--- src/spaces/structure.jl | 6 ++--- src/tensors/treetransformers.jl | 10 ++++---- test/other/caches.jl | 42 +++++++++++++++---------------- 10 files changed, 59 insertions(+), 62 deletions(-) diff --git a/Project.toml b/Project.toml index 3f5289b96..5548a0e0d 100644 --- a/Project.toml +++ b/Project.toml @@ -8,10 +8,10 @@ projects = ["test", "docs"] [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" -Cached = "1b238080-9255-4fe9-b224-89eb24efe93b" Dictionaries = "85a47980-9c8c-11e8-2b9f-f7ca1fa99fb4" 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" @@ -45,13 +45,9 @@ TensorKitFiniteDifferencesExt = "FiniteDifferences" TensorKitGPUArraysExt = "GPUArrays" TensorKitMooncakeExt = "Mooncake" -[sources] -Cached = {url = "https://github.com/lkdvos/Cached.jl.git", rev = "main"} - [compat] AMDGPU = "2" Adapt = "4" -Cached = "0.1" CUDA = "6" ChainRulesCore = "1" Dictionaries = "0.4" @@ -61,6 +57,7 @@ FiniteDifferences = "0.12" GPUArrays = "11.4.1" LinearAlgebra = "1" MatrixAlgebraKit = "0.6.13" +MemoizationKit = "0.1" Mooncake = "0.5.27" OhMyThreads = "0.8.0" PrecompileTools = "1" diff --git a/docs/src/man/fusiontrees.md b/docs/src/man/fusiontrees.md index 8ebe183fb..eb13194b6 100644 --- a/docs/src/man/fusiontrees.md +++ b/docs/src/man/fusiontrees.md @@ -260,30 +260,30 @@ 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 `Cached.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 Cached.jl's configured container (ClockCache by default), 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 Cached.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 `Cached.cachekey`. -The Cached APIs can inspect, resize and configure these caches: +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, Cached -Cached.cache_info(TensorKit) -Cached.set_cache_size!(TensorKit.treebraider, 50_000) -Cached.set_cache_size!(TensorKit.treebraider, 2^30; by = Cached.cachesize) # 1 GiB of cached values -Cached.set_cache_preferences!(TensorKit; maxsize = 50_000, measure = "bytes") -Cached.set_cache_preferences!(; container = "LRU") +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. -`Cached.uncached(f, args...)` computes a result without consulting or filling the cache. -TensorKit supplies efficient `Cached.cachesize` methods for its structure, fusion-tree and transformer results. +`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. -Optional Cached features, such as the Tachikoma dashboard and SQLite disk caching, remain available through Cached's extensions. -Disk caching requires loading SQLite and defining `Cached.DiskCacheStyle` for the desired function; it is disabled by default, and `empty_globalcaches!()` clears only in-memory caches. +Optional MemoizationKit features, such as the Tachikoma dashboard and SQLite disk caching, remain available through MemoizationKit's extensions. +Disk caching requires loading SQLite and defining `MemoizationKit.DiskCacheStyle` for the desired function; it is disabled by default, and `empty_globalcaches!()` clears only in-memory caches. 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 2b26d34ef..236ffe968 100644 --- a/docs/src/man/profiling.md +++ b/docs/src/man/profiling.md @@ -1,7 +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 Cached.jl's cache lookup and computation timers, recorded in the same timer tree and categories. +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 a686d989f..120a9dab3 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.Cached.AbstractCache}) = true +@inline EnzymeRules.inactive_type(::Type{<:TensorKit.MemoizationKit.AbstractCache}) = true function EnzymeRules.augmented_primal( config::EnzymeRules.RevConfigWidth{1}, diff --git a/src/auxiliary/caches.jl b/src/auxiliary/caches.jl index b91c332f0..d66582842 100644 --- a/src/auxiliary/caches.jl +++ b/src/auxiliary/caches.jl @@ -1,8 +1,8 @@ -# Caching of the fusion tree manipulations and tensor structures is provided by Cached.jl. +# 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 `Cached.CacheStyle`. -import Cached: CacheStyle -using Cached: Cached, @cached, NoCache, GlobalCache, Hashed, LRU +# specializations in TensorKit, and those of users, extend `MemoizationKit.CacheStyle`. +import MemoizationKit: CacheStyle +using MemoizationKit: MemoizationKit, @cached, NoCache, GlobalCache, Hashed, LRU """ const GLOBAL_CACHES @@ -38,7 +38,7 @@ 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!() - Cached.empty_caches!(TensorKit) + MemoizationKit.empty_caches!(TensorKit) foreach(empty! ∘ last, GLOBAL_CACHES) return nothing end @@ -51,7 +51,7 @@ Print the hit/miss statistics and current size of every global cache of TensorKi See also [`empty_globalcaches!`](@ref). """ function global_cache_info(io::IO = stdout) - for (f, cache) in Cached.cache_info(TensorKit) + for (f, cache) in MemoizationKit.cache_info(TensorKit) println(io, nameof(f), ":\t", cache) end for (name, cache) in GLOBAL_CACHES diff --git a/src/auxiliary/timers.jl b/src/auxiliary/timers.jl index c8bef53d5..77bde1c60 100644 --- a/src/auxiliary/timers.jl +++ b/src/auxiliary/timers.jl @@ -61,7 +61,7 @@ host-side dispatch of asynchronous kernels unless the workload is explicitly syn """ function enable_timers!() TimerOutputs.enable_debug_timings(TensorKit) - Cached.enable_cache_timers!(TensorKit, GLOBAL_TIMER) + MemoizationKit.enable_cache_timers!(TensorKit, GLOBAL_TIMER) return nothing end @@ -73,7 +73,7 @@ Also triggers recompilation of the instrumented methods on first use. """ function disable_timers!() TimerOutputs.disable_debug_timings(TensorKit) - Cached.disable_cache_timers!(TensorKit) + MemoizationKit.disable_cache_timers!(TensorKit) return nothing end diff --git a/src/fusiontrees/fusiontrees.jl b/src/fusiontrees/fusiontrees.jl index 34cd67640..7cb2b3a0c 100644 --- a/src/fusiontrees/fusiontrees.jl +++ b/src/fusiontrees/fusiontrees.jl @@ -122,14 +122,14 @@ struct FusionTreeBlock{I, N₁, N₂, F <: FusionTreePair{I, N₁, N₂}} trees::Vector{F} end -Cached.cachesize(b::FusionTreeBlock) = sizeof(b) + _cache_payload_size(b.trees) +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. -Cached.cachesize(p::Pair{<:FusionTreePair, <:Number}) = +MemoizationKit.cachesize(p::Pair{<:FusionTreePair, <:Number}) = isbitstype(typeof(p)) ? sizeof(p) : Base.summarysize(p) -Cached.cachesize(p::Pair{<:FusionTreeBlock, <:Matrix}) = - sizeof(p) + Cached.cachesize(first(p)) + _cache_payload_size(last(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}}, diff --git a/src/spaces/structure.jl b/src/spaces/structure.jl index 415b0228e..1485e4691 100644 --- a/src/spaces/structure.jl +++ b/src/spaces/structure.jl @@ -21,7 +21,7 @@ struct SectorStructure{I <: Sector, F <: FusionTreePair{I}} fusiontrees::Indices{F} end -Cached.cachesize(s::SectorStructure) = +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} @@ -71,7 +71,7 @@ See also [`degeneracystructure`](@ref), [`fusiontrees`](@ref), [`blocksectors`]( return SectorStructure{I, F}(Indices(bs), Indices(trees)) end -Cached.cachekey(::typeof(sectorstructure), W::HomSpace) = (Hashed(W, sectorhash, sectorequal),) +MemoizationKit.cachekey(::typeof(sectorstructure), W::HomSpace) = (Hashed(W, sectorhash, sectorequal),) # DegeneracyStructure: degeneracy-dependent characterization of HomSpaces # ----------------------------------------------------------------------- @@ -97,7 +97,7 @@ struct DegeneracyStructure{N} subblockstructure::Vector{StridedStructure{N}} end -Cached.cachesize(s::DegeneracyStructure) = sizeof(s) + sizeof(s.blockstructure) + sizeof(s.subblockstructure) +MemoizationKit.cachesize(s::DegeneracyStructure) = sizeof(s) + sizeof(s.blockstructure) + sizeof(s.subblockstructure) function degeneracystructuretype(W::HomSpace) N = length(codomain(W)) + length(domain(W)) diff --git a/src/tensors/treetransformers.jl b/src/tensors/treetransformers.jl index 2ed9b9099..d3bef41ab 100644 --- a/src/tensors/treetransformers.jl +++ b/src/tensors/treetransformers.jl @@ -73,9 +73,9 @@ end _cache_structure_size(t::TreeTransformer) = sizeof(t.structure_src) + (t.structure_src === t.structure_dst ? 0 : sizeof(t.structure_dst)) -Cached.cachesize(t::UniqueTreeTransformer) = sizeof(t) + _cache_payload_size(t.data) + _cache_structure_size(t) +MemoizationKit.cachesize(t::UniqueTreeTransformer) = sizeof(t) + _cache_payload_size(t.data) + _cache_structure_size(t) -function Cached.cachesize(t::GenericTreeTransformer) +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) + @@ -218,10 +218,10 @@ end return TreeTransformer(A, fusiontreetransform, p, Vdst, Vsrc, conjsrc) end -# Keep TensorKit's profiling categories for the phases instrumented by Cached. +# Keep TensorKit's profiling categories for the phases instrumented by MemoizationKit. const TensorKitCachedFunction = Union{ typeof(fsbraid), typeof(fstranspose), typeof(treebraider), typeof(treetransposer), typeof(sectorstructure), typeof(degeneracystructure), } -Cached.instrument_label(f::TensorKitCachedFunction, ::Val{:lookup}) = "bookkeeping: cache $(nameof(f))" -Cached.instrument_label(f::TensorKitCachedFunction, ::Val{:compute}) = "symmetry: compute $(nameof(f))" +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/other/caches.jl b/test/other/caches.jl index ca7e9c2e7..f9db88c3d 100644 --- a/test/other/caches.jl +++ b/test/other/caches.jl @@ -1,22 +1,22 @@ using Test, TestExtras using TensorKit -using TensorKit: Cached, sectorstructure, degeneracystructure +using TensorKit: MemoizationKit, sectorstructure, degeneracystructure -@testset "Cached structure keys" begin +@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 = Cached.uncached(sectorstructure, W1) + fresh = MemoizationKit.uncached(sectorstructure, W1) @test s1.blocksectors == fresh.blocksectors @test s1.fusiontrees == fresh.fusiontrees - @test parent(only(Cached.cachekey(sectorstructure, W1))) === W1 - @test length(only(Cached.cache_info(sectorstructure)).second) == 1 + @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(Cached.cache_info(degeneracystructure)).second) == 2 + @test length(only(MemoizationKit.cache_info(degeneracystructure)).second) == 2 end @testset "TensorKit cached-value sizes" begin @@ -27,34 +27,34 @@ end p, levels = ((2, 1), (4, 3)), (1, 2, 3, 4) transformer = TensorKit.treebraider(W, W, p, false, levels) for value in (sectorstructure(W), degeneracystructure(W), transformer) - bytes = @constinferred Cached.cachesize(value) + bytes = @constinferred MemoizationKit.cachesize(value) @test 0 < bytes <= Base.summarysize(value) end for src in TensorKit.fusionblocks(W) result = TensorKit.fsbraid(src, p, ((1, 2), (3, 4))) - @test 0 < (@constinferred Cached.cachesize(result)) <= Base.summarysize(result) + @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. - Cached.set_cache_size!(TensorKit.treebraider, Cached.cachesize(transformer); by = Cached.cachesize) + MemoizationKit.set_cache_size!(TensorKit.treebraider, MemoizationKit.cachesize(transformer); by = MemoizationKit.cachesize) try TensorKit.treebraider(W, W, p, false, levels) TensorKit.treebraider(W, W, ((1, 2), (3, 4)), false, levels) - Cached.empty_caches!(degeneracystructure) - cache = only(Cached.cache_info(TensorKit.treebraider)).second - stats = Cached.cache_stats(cache) - @test stats.by === Cached.cachesize + 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(Cached.cachesize, values(cache); init = 0) + @test stats.currentsize == sum(MemoizationKit.cachesize, values(cache); init = 0) - Cached.set_cache_size!(TensorKit.treebraider, 0; by = Cached.cachesize) + MemoizationKit.set_cache_size!(TensorKit.treebraider, 0; by = MemoizationKit.cachesize) @test isempty(cache) @test TensorKit.treebraider(W, W, p, false, levels).data == transformer.data @test isempty(cache) finally - Cached.set_cache_size!(TensorKit.treebraider, 10_000) + MemoizationKit.set_cache_size!(TensorKit.treebraider, 10_000) end end @@ -62,13 +62,13 @@ end # 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 Cached.cachesize(empty_transformer) > 0 + @test MemoizationKit.cachesize(empty_transformer) > 0 big_transformer = TensorKit.UniqueTreeTransformer([(BigFloat(1), 1, 1)], strides, strides) - @test Cached.cachesize(big_transformer) >= Base.summarysize(big_transformer.data) + @test MemoizationKit.cachesize(big_transformer) >= Base.summarysize(big_transformer.data) end # A cache owned by another module must survive TensorKit's cache clearing. -Cached.@cached unrelated_cache(x) = [x] +MemoizationKit.@cached unrelated_cache(x) = [x] @testset "TensorKit cache management" begin retained = unrelated_cache(1) @@ -84,11 +84,11 @@ Cached.@cached unrelated_cache(x) = [x] @test occursin("treebraider", info) @test occursin("extension_cache", info) @test empty_globalcaches!() === nothing - @test all(isempty(cache) for (_, cache) in Cached.cache_info(TensorKit)) + @test all(isempty(cache) for (_, cache) in MemoizationKit.cache_info(TensorKit)) @test isempty(extra) @test unrelated_cache(1) === retained finally pop!(TensorKit.GLOBAL_CACHES) - Cached.empty_caches!(unrelated_cache) + MemoizationKit.empty_caches!(unrelated_cache) end end From a3b793ae1ccb9818c4918b032cea39a5d6209b90 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Mon, 12 Oct 2026 17:11:38 -0400 Subject: [PATCH 06/11] Address review: descriptive cache names and transpose inference - Rename `fsbraid`/`fstranspose` to `fusiontreebraid`/`fusiontreetranspose`, as these are now user-facing through the MemoizationKit APIs and dashboard. - Restore the return type annotation on `fusiontreetranspose`, without which the `FusionTreePair` path infers `Any` and the cache stores untyped values. - Update the Enzyme inactivity rule for the new three-argument `fusiontreebraid`. - Refer to the MemoizationKit documentation for further features. Co-Authored-By: Claude Opus 5.5 --- docs/src/man/fusiontrees.md | 3 +-- ext/TensorKitEnzymeExt/utility.jl | 3 +-- src/fusiontrees/braiding_manipulations.jl | 10 +++++----- src/fusiontrees/duality_manipulations.jl | 20 ++++++++++++++++---- src/tensors/treetransformers.jl | 2 +- test/other/caches.jl | 2 +- 6 files changed, 25 insertions(+), 15 deletions(-) diff --git a/docs/src/man/fusiontrees.md b/docs/src/man/fusiontrees.md index eb13194b6..b6e9dd0f5 100644 --- a/docs/src/man/fusiontrees.md +++ b/docs/src/man/fusiontrees.md @@ -282,8 +282,7 @@ These estimate bytes from the fixed-size fields and live array buffers; matrix c 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. -Optional MemoizationKit features, such as the Tachikoma dashboard and SQLite disk caching, remain available through MemoizationKit's extensions. -Disk caching requires loading SQLite and defining `MemoizationKit.DiskCacheStyle` for the desired function; it is disabled by default, and `empty_globalcaches!()` clears only in-memory caches. +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/ext/TensorKitEnzymeExt/utility.jl b/ext/TensorKitEnzymeExt/utility.jl index 120a9dab3..94340eacf 100644 --- a/ext/TensorKitEnzymeExt/utility.jl +++ b/ext/TensorKitEnzymeExt/utility.jl @@ -142,8 +142,7 @@ 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 diff --git a/src/fusiontrees/braiding_manipulations.jl b/src/fusiontrees/braiding_manipulations.jl index 50b759f49..bd2a5c27c 100644 --- a/src/fusiontrees/braiding_manipulations.jl +++ b/src/fusiontrees/braiding_manipulations.jl @@ -284,7 +284,7 @@ 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 function braid(src::ProductSectorBlock, p::Index2Tuple, levels::Index2Tuple) isempty(src) && return fsbraid((src, p, levels)) @@ -296,7 +296,7 @@ end # 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 fsbraidtype( +Base.@assume_effects :foldable function fusiontreebraidtype( src::Union{FusionTreePair{I}, FusionTreeBlock{I}}, ::Index2Tuple{N₁, N₂} ) where {I, N₁, N₂} T = sectorscalartype(I) @@ -308,10 +308,10 @@ Base.@assume_effects :foldable function fsbraidtype( end end -@cached function fsbraid( +@cached function fusiontreebraid( src::Union{FusionTreePair{I}, FusionTreeBlock{I}}, p::Index2Tuple{N₁, N₂}, levels::Index2Tuple - )::fsbraidtype(src, p) where {I, N₁, N₂} + )::fusiontreebraidtype(src, p) where {I, N₁, N₂} p1, p2 = p l1, l2 = levels if src isa FusionTreePair @@ -349,7 +349,7 @@ end end end -CacheStyle(::typeof(fsbraid), src::Union{FusionTreePair{I}, FusionTreeBlock{I}}, args...) where {I} = +CacheStyle(::typeof(fusiontreebraid), src::Union{FusionTreePair{I}, FusionTreeBlock{I}}, args...) where {I} = FusionStyle(I) isa UniqueFusion ? NoCache() : GlobalCache() """ diff --git a/src/fusiontrees/duality_manipulations.jl b/src/fusiontrees/duality_manipulations.jl index 73b7b46c9..10f74715f 100644 --- a/src/fusiontrees/duality_manipulations.jl +++ b/src/fusiontrees/duality_manipulations.jl @@ -524,7 +524,7 @@ 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 function Base.transpose(src::ProductSectorBlock, p::Index2Tuple) isempty(src) && return fstranspose((src, p)) @@ -534,9 +534,21 @@ function Base.transpose(src::ProductSectorBlock, p::Index2Tuple) end::_fsdicttype(typeof((src, p))) end -@cached function fstranspose( - src::Union{FusionTreePair{I}, FusionTreeBlock{I}}, p::Index2Tuple{N₁, N₂} +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 fusiontreetranspose( + src::Union{FusionTreePair{I}, FusionTreeBlock{I}}, p::Index2Tuple{N₁, N₂} + )::fusiontreetransposetype(src, p) where {I, N₁, N₂} p1, p2 = p N = N₁ + N₂ @@ -562,7 +574,7 @@ end return dst => U end -CacheStyle(::typeof(fstranspose), src::Union{FusionTreePair{I}, FusionTreeBlock{I}}, args...) where {I} = +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/tensors/treetransformers.jl b/src/tensors/treetransformers.jl index d3bef41ab..655ff108b 100644 --- a/src/tensors/treetransformers.jl +++ b/src/tensors/treetransformers.jl @@ -220,7 +220,7 @@ end # Keep TensorKit's profiling categories for the phases instrumented by MemoizationKit. const TensorKitCachedFunction = Union{ - typeof(fsbraid), typeof(fstranspose), typeof(treebraider), typeof(treetransposer), + typeof(fusiontreebraid), typeof(fusiontreetranspose), typeof(treebraider), typeof(treetransposer), typeof(sectorstructure), typeof(degeneracystructure), } MemoizationKit.instrument_label(f::TensorKitCachedFunction, ::Val{:lookup}) = "bookkeeping: cache $(nameof(f))" diff --git a/test/other/caches.jl b/test/other/caches.jl index f9db88c3d..d527efb27 100644 --- a/test/other/caches.jl +++ b/test/other/caches.jl @@ -31,7 +31,7 @@ end @test 0 < bytes <= Base.summarysize(value) end for src in TensorKit.fusionblocks(W) - result = TensorKit.fsbraid(src, p, ((1, 2), (3, 4))) + result = TensorKit.fusiontreebraid(src, p, ((1, 2), (3, 4))) @test 0 < (@constinferred MemoizationKit.cachesize(result)) <= Base.summarysize(result) end From c1437f39acbbbf9a51b52197b08ecc52bced75c6 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Mon, 12 Oct 2026 17:31:58 -0400 Subject: [PATCH 07/11] Drop the extra cache registry now that extension caches are owned by TensorKit The GPU transformers are cached through `@cached` methods of `treebraider` and `treetransposer`, which share TensorKit's caches. These are inspected, emptied and timed through MemoizationKit, so the separate `GLOBAL_CACHES` registry is no longer needed. Co-Authored-By: Claude Opus 5.5 --- docs/src/Changelog.md | 3 ++- src/auxiliary/caches.jl | 16 ---------------- test/other/caches.jl | 13 ++++--------- 3 files changed, 6 insertions(+), 26 deletions(-) diff --git a/docs/src/Changelog.md b/docs/src/Changelog.md index a6ab14c3b..4c91f261d 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 `treebraider` and `treetransposer`, sharing their cache (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/src/auxiliary/caches.jl b/src/auxiliary/caches.jl index d66582842..c327f34ba 100644 --- a/src/auxiliary/caches.jl +++ b/src/auxiliary/caches.jl @@ -4,18 +4,6 @@ import MemoizationKit: CacheStyle using MemoizationKit: MemoizationKit, @cached, NoCache, GlobalCache, Hashed, LRU -""" - const GLOBAL_CACHES - -Registry of additional global caches, which are not managed by `@cached` (e.g. by -extensions), as `name => cache` pairs. They are emptied by [`empty_globalcaches!`](@ref) -and listed by [`global_cache_info`](@ref), together with the caches of TensorKit's cached -functions. -""" -const GLOBAL_CACHES = Pair{Symbol, Any}[] - -const DEFAULT_GLOBALCACHE_SIZE = Ref(10^4) - # 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) @@ -39,7 +27,6 @@ See also [`global_cache_info`](@ref) to display the status. """ function empty_globalcaches!() MemoizationKit.empty_caches!(TensorKit) - foreach(empty! ∘ last, GLOBAL_CACHES) return nothing end @@ -54,8 +41,5 @@ function global_cache_info(io::IO = stdout) for (f, cache) in MemoizationKit.cache_info(TensorKit) println(io, nameof(f), ":\t", cache) end - for (name, cache) in GLOBAL_CACHES - println(io, name, ":\t", cache) - end return end diff --git a/test/other/caches.jl b/test/other/caches.jl index d527efb27..ec750d48e 100644 --- a/test/other/caches.jl +++ b/test/other/caches.jl @@ -25,7 +25,7 @@ end 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(W, W, p, false, levels) + 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) @@ -39,8 +39,8 @@ end # 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(W, W, p, false, levels) - TensorKit.treebraider(W, W, ((1, 2), (3, 4)), false, levels) + 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) @@ -51,7 +51,7 @@ end MemoizationKit.set_cache_size!(TensorKit.treebraider, 0; by = MemoizationKit.cachesize) @test isempty(cache) - @test TensorKit.treebraider(W, W, p, false, levels).data == transformer.data + @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) @@ -72,8 +72,6 @@ MemoizationKit.@cached unrelated_cache(x) = [x] @testset "TensorKit cache management" begin retained = unrelated_cache(1) - extra = Dict(:key => :value) - push!(TensorKit.GLOBAL_CACHES, :extension_cache => extra) try V = SU2Space(0 => 2, 1 // 2 => 2) t = rand(V ⊗ V ← V ⊗ V) @@ -82,13 +80,10 @@ MemoizationKit.@cached unrelated_cache(x) = [x] @test (@constinferred permute(permute(t, p), p)) ≈ t info = sprint(TensorKit.global_cache_info) @test occursin("treebraider", info) - @test occursin("extension_cache", info) @test empty_globalcaches!() === nothing @test all(isempty(cache) for (_, cache) in MemoizationKit.cache_info(TensorKit)) - @test isempty(extra) @test unrelated_cache(1) === retained finally - pop!(TensorKit.GLOBAL_CACHES) MemoizationKit.empty_caches!(unrelated_cache) end end From 3e5e5790fe3625d5ff095469b84f08a3d27413d3 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 13 Oct 2026 09:34:31 -0400 Subject: [PATCH 08/11] Use renamed fusiontreetranspose in Enzyme rules, unicode subscripts Co-Authored-By: Claude Opus 5.5 --- ext/TensorKitEnzymeExt/utility.jl | 3 +-- src/fusiontrees/braiding_manipulations.jl | 24 +++++++++++------------ src/fusiontrees/duality_manipulations.jl | 8 ++++---- 3 files changed, 17 insertions(+), 18 deletions(-) diff --git a/ext/TensorKitEnzymeExt/utility.jl b/ext/TensorKitEnzymeExt/utility.jl index 94340eacf..57b75c771 100644 --- a/ext/TensorKitEnzymeExt/utility.jl +++ b/ext/TensorKitEnzymeExt/utility.jl @@ -144,8 +144,7 @@ end @inline EnzymeRules.inactive(::typeof(TensorKit.fusiontrees), ::Any, ::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.product_transform), ::Any...) = nothing @inline EnzymeRules.inactive(::typeof(TensorKit.treetransposer), ::Any...) = nothing @inline EnzymeRules.inactive(::typeof(TensorKit.treebraider), ::Any...) = nothing diff --git a/src/fusiontrees/braiding_manipulations.jl b/src/fusiontrees/braiding_manipulations.jl index bd2a5c27c..4974ea1f3 100644 --- a/src/fusiontrees/braiding_manipulations.jl +++ b/src/fusiontrees/braiding_manipulations.jl @@ -287,11 +287,11 @@ function braid(src::Union{FusionTreePair, FusionTreeBlock}, p::Index2Tuple, leve return fusiontreebraid(src, p, levels) end function braid(src::ProductSectorBlock, p::Index2Tuple, levels::Index2Tuple) - isempty(src) && return fsbraid((src, p, levels)) + isempty(src) && return fusiontreebraid(src, p, levels) # the return type is asserted since inference gives up on the recursion on older Julia versions return product_transform(sectorscalartype(sectortype(src)), src) do x return braid(x, p, levels) - end::_fsdicttype(typeof((src, p, levels))) + end::fusiontreebraidtype(src, p) end # Braiding can promote real fusion coefficients to complex values. Keep the result @@ -312,19 +312,19 @@ end src::Union{FusionTreePair{I}, FusionTreeBlock{I}}, p::Index2Tuple{N₁, N₂}, levels::Index2Tuple )::fusiontreebraidtype(src, p) where {I, N₁, N₂} - p1, p2 = p - l1, l2 = levels + p₁, p₂ = p + l₁, l₂ = levels if src isa FusionTreePair f₁, f₂ = src - 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 + 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 - 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)) diff --git a/src/fusiontrees/duality_manipulations.jl b/src/fusiontrees/duality_manipulations.jl index 10f74715f..e3ed06c9f 100644 --- a/src/fusiontrees/duality_manipulations.jl +++ b/src/fusiontrees/duality_manipulations.jl @@ -527,11 +527,11 @@ function Base.transpose(src::Union{FusionTreePair, FusionTreeBlock}, p::Index2Tu return fusiontreetranspose(src, p) end function Base.transpose(src::ProductSectorBlock, p::Index2Tuple) - isempty(src) && return fstranspose((src, p)) + isempty(src) && return fusiontreetranspose(src, p) # the return type is asserted since inference gives up on the recursion on older Julia versions return product_transform(fusionscalartype(sectortype(src)), src) do x return transpose(x, p) - end::_fsdicttype(typeof((src, p))) + end::fusiontreetransposetype(src, p) end Base.@assume_effects :foldable function fusiontreetransposetype( @@ -549,10 +549,10 @@ end @cached function fusiontreetranspose( src::Union{FusionTreePair{I}, FusionTreeBlock{I}}, p::Index2Tuple{N₁, N₂} )::fusiontreetransposetype(src, p) where {I, N₁, N₂} - p1, p2 = p + 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 From ba0a49edb06b2aa5aa26afcedb07386e7269a3ca Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 13 Oct 2026 10:51:17 -0400 Subject: [PATCH 09/11] Cache device transformers in separate functions `treebraider` and `treetransposer` forward device storage types to `device_treebraider` and `device_treetransposer`, whose methods are added by the GPU extension. These have their own MemoizationKit cache, so their size can be configured separately from the host transformers. Co-Authored-By: Claude Opus 5.5 --- docs/src/Changelog.md | 2 +- ext/TensorKitEnzymeExt/utility.jl | 2 ++ ext/TensorKitGPUArraysExt.jl | 12 ++++++++++-- src/tensors/treetransformers.jl | 13 +++++++++++++ test/cuda/tensors.jl | 14 ++++++++++++++ 5 files changed, 40 insertions(+), 3 deletions(-) diff --git a/docs/src/Changelog.md b/docs/src/Changelog.md index 4c91f261d..efd4207fe 100644 --- a/docs/src/Changelog.md +++ b/docs/src/Changelog.md @@ -29,7 +29,7 @@ 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 - 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 `treebraider` and `treetransposer`, sharing their cache (internal) +- 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/ext/TensorKitEnzymeExt/utility.jl b/ext/TensorKitEnzymeExt/utility.jl index 57b75c771..301d7e2c1 100644 --- a/ext/TensorKitEnzymeExt/utility.jl +++ b/ext/TensorKitEnzymeExt/utility.jl @@ -148,6 +148,8 @@ end @inline EnzymeRules.inactive(::typeof(TensorKit.product_transform), ::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 28c5157ae..8e9f81157 100644 --- a/ext/TensorKitGPUArraysExt.jl +++ b/ext/TensorKitGPUArraysExt.jl @@ -271,8 +271,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,7 +288,7 @@ 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) diff --git a/src/tensors/treetransformers.jl b/src/tensors/treetransformers.jl index 655ff108b..06f2f0612 100644 --- a/src/tensors/treetransformers.jl +++ b/src/tensors/treetransformers.jl @@ -218,9 +218,22 @@ end return TreeTransformer(A, fusiontreetransform, p, Vdst, Vsrc, conjsrc) end +""" + 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))" diff --git a/test/cuda/tensors.jl b/test/cuda/tensors.jl index 59196d52c..f8e9164bd 100644 --- a/test/cuda/tensors.jl +++ b/test/cuda/tensors.jl @@ -576,3 +576,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 From daaf1b2b9d60e353ef7548bc3d334be03935e72e Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 30 Oct 2026 11:44:48 -0400 Subject: [PATCH 10/11] Require MemoizationKit v0.1.1 Co-Authored-By: Claude Opus 5.5 --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 5548a0e0d..2809e07cf 100644 --- a/Project.toml +++ b/Project.toml @@ -57,7 +57,7 @@ FiniteDifferences = "0.12" GPUArrays = "11.4.1" LinearAlgebra = "1" MatrixAlgebraKit = "0.6.13" -MemoizationKit = "0.1" +MemoizationKit = "0.1.1" Mooncake = "0.5.27" OhMyThreads = "0.8.0" PrecompileTools = "1" From 03838971ede0420c86393b086be5eb6a96b84bfb Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 30 Oct 2026 15:42:14 -0400 Subject: [PATCH 11/11] Ignore JET's uncorrelated union split in the tree transformer bodies Co-Authored-By: Claude Opus 5.5 --- test/other/jet.jl | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/test/other/jet.jl b/test/other/jet.jl index d611625e5..3b719d87b 100644 --- a/test/other/jet.jl +++ b/test/other/jet.jl @@ -10,6 +10,14 @@ const IGNORED = ( # `AnyFrameMethod` rather than `LastFrameMethod`: the `@debug` sits inside a `taskforeach` # closure, so the innermost frame is a gensym'd closure name rather than the constructor. JET.AnyFrameMethod(:GenericTreeTransformer), + # `no matching method found convert(::Type{UniqueTreeTransformer}, ::GenericTreeTransformer)` + # (and vice versa) in the uncached bodies of `treebraider` and `treetransposer`: for abstract + # spaces, JET cannot correlate the `FusionStyle` branch of the return type annotation, which + # `@cached` keeps on the body, with that of `TreeTransformer`. + ( + JET.LastFrameMethod(only(methods(TensorKit.MemoizationKit.implementation, Tuple{typeof(f), Vararg{Any}}))) + for f in (TensorKit.treebraider, TensorKit.treetransposer) + )..., ) # also gated in runtests.jl; JET < 0.12 does not define `JET_AVAILABLE`