From 467649f3a940634a11156018598b3cf0720f4a5f Mon Sep 17 00:00:00 2001 From: lkdvos Date: Mon, 5 Oct 2026 18:26:47 -0400 Subject: [PATCH 01/10] 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 | 223 ++-------------------- 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(+), 235 deletions(-) diff --git a/Project.toml b/Project.toml index d1e0c834d..49a0b34ed 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.207" EnzymeTestUtils = "0.2.8" FiniteDifferences = "0.12" GPUArrays = "11.4.1" -LRUCache = "1.6" LinearAlgebra = "1" MatrixAlgebraKit = "0.6.9" 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 e24c8fb0c..14373b62a 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 948f9fa75..3019e0372 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 6bdbebfb5..4b4494ef4 100644 --- a/src/TensorKit.jl +++ b/src/TensorKit.jl @@ -119,7 +119,6 @@ const TO = TensorOperations using MatrixAlgebraKit using Dictionaries: Dictionaries, Dictionary, Indices, gettoken, gettokenvalue -using LRUCache using OhMyThreads using ScopedValues using TimerOutputs: TimerOutputs, TimerOutput, @timeit_debug diff --git a/src/auxiliary/caches.jl b/src/auxiliary/caches.jl index bdc07ca08..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,211 +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) - hasparams = true - params = fcall.args[2:end] - fcall = fcall.args[1] - else - hasparams = false - end - if Meta.isexpr(fcall, :(::)) - typed = true - typeex = fcall.args[2] - fcall = fcall.args[1] - else - typed = false - end - Meta.isexpr(fcall, :call) || _cached_error("expected a function signature", fcall) - fname = fcall.args[1] - # qualified names such as `TensorKit.treebraider` add methods to a function of another module - basename = Meta.isexpr(fname, :.) ? fname.args[end].value : fname - basename isa Symbol || - _cached_error("expected a function name or a qualified name such as `Module.f`", fname) - fargs = map(_cached_canonicalize_arg, fcall.args[2:end]) - fcall = Expr(:call, fname, fargs...) - fargnames = map(arg -> arg isa Symbol ? arg : arg.args[1], fargs) - # timer labels for the cache lookup and the miss-path construction - lookuplabel = string("bookkeeping: cache ", basename) - misslabel = string(_cached_category(basename), ": compute ", basename) - _fbody = ex.args[2] - - # actual implenetation, with underscore name - _fname = Symbol(:_, basename) - _fcall = Expr(:call, _fname, fargs...) - if hasparams - _fcall = Expr(:where, _fcall, params...) - end - _fex = Expr(:function, _fcall, _fbody) - - # implementation that chooses the cache style - newfcall = fcall - if hasparams - newfcall = Expr(:where, newfcall, params...) - end - cachestylevar = gensym(:cachestyle) - cachestyleex = Expr( - :(=), cachestylevar, Expr(:call, CacheStyle, fname, fargnames...) - ) - newfbody = Expr( - :block, cachestyleex, Expr(:call, fname, fargnames..., cachestylevar) - ) - newfex = Expr(:function, newfcall, newfbody) - - # nocache implementation - fnocachecall = Expr(:call, fname, fargs..., :(::$NoCache)) - if hasparams - fnocachecall = Expr(:where, fnocachecall, params...) - end - fnocachebody = _timeit_expr(misslabel, Expr(:call, _fname, fargnames...)) - if typed - T = gensym(:T) - fnocachebody = Expr(:block, Expr(:(=), T, typeex), Expr(:(::), fnocachebody, T)) - end - fnocacheex = Expr(:function, fnocachecall, fnocachebody) - - # tasklocal cache implementation - Dvar = gensym(:D) - flocalcachecall = Expr(:call, fname, fargs..., :(::$TaskLocalCache{$Dvar})) - if hasparams - flocalcachecall = Expr(:where, flocalcachecall, params..., Dvar) - else - flocalcachecall = Expr(:where, flocalcachecall, Dvar) - end - globalcachename = Symbol(:GLOBAL_, uppercase(string(basename)), :_CACHE) - localcachename = Symbol(:_tasklocal_, globalcachename) - cachevar = gensym(:cache) - getlocalcacheex = :( - $cachevar::$Dvar = get!(task_local_storage(), $(QuoteNode(localcachename))) do - return $Dvar() - end - ) - valvar = gensym(:val) - if length(fargnames) == 1 - key = fargnames[1] - else - key = Expr(:tuple, fargnames...) - end - missex = _timeit_expr(misslabel, Expr(:call, _fname, fargnames...)) - getvalex = _timeit_expr(lookuplabel, :(get!(() -> $missex, $cachevar, $key))) - if typed - T = gensym(:T) - flocalcachebody = Expr( - :block, - getlocalcacheex, - Expr(:(=), T, typeex), - Expr(:(=), Expr(:(::), valvar, T), getvalex), - Expr(:return, valvar) - ) - else - flocalcachebody = Expr( - :block, - getlocalcacheex, - Expr(:(=), valvar, getvalex), - Expr(:return, valvar) - ) - end - flocalcacheex = Expr(:function, flocalcachecall, flocalcachebody) - - # # global cache implementation - fglobalcachecall = Expr(:call, fname, fargs..., :(::$GlobalLRUCache)) - if hasparams - fglobalcachecall = Expr(:where, fglobalcachecall, params...) - end - getglobalcachex = Expr(:(=), cachevar, globalcachename) - if typed - T = gensym(:T) - fglobalcachebody = Expr( - :block, - getglobalcachex, - Expr(:(=), T, typeex), - Expr(:(=), Expr(:(::), valvar, T), getvalex), - Expr(:return, valvar) - ) - else - fglobalcachebody = Expr( - :block, - getglobalcachex, - Expr(:(=), valvar, getvalex), - Expr(:return, valvar) - ) - end - fglobalcacheex = Expr(:function, fglobalcachecall, fglobalcachebody) - fglobalcachedef = Expr( - :const, - Expr(:(=), globalcachename, :($LRU{Any, Any}(; maxsize = $DEFAULT_GLOBALCACHE_SIZE[]))) - ) - # Other modules register their caches explicitly from `__init__`, since - # mutations of this registry during their precompilation do not survive loading. - fglobalcacheregister = mod === (@__MODULE__) ? Expr( - :call, :push!, GlobalRef(@__MODULE__, :GLOBAL_CACHES), - :($(QuoteNode(globalcachename)) => $globalcachename) - ) : nothing - - # # total expression - return esc( - Expr( - :block, _fex, newfex, fnocacheex, flocalcacheex, - fglobalcachedef, fglobalcacheregister, fglobalcacheex - ) - ) -end diff --git a/src/fusiontrees/braiding_manipulations.jl b/src/fusiontrees/braiding_manipulations.jl index ce08f510c..b975a2194 100644 --- a/src/fusiontrees/braiding_manipulations.jl +++ b/src/fusiontrees/braiding_manipulations.jl @@ -340,7 +340,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 c595afdc2..60bda7ff5 100644 --- a/src/fusiontrees/duality_manipulations.jl +++ b/src/fusiontrees/duality_manipulations.jl @@ -568,9 +568,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 990e50d2d..f9a1dac46 100644 --- a/src/tensors/treetransformers.jl +++ b/src/tensors/treetransformers.jl @@ -210,4 +210,3 @@ end return TreeTransformer(A, fusiontreetransform, p, Vdst, Vsrc, conjsrc) end -# default cachestyle is GlobalLRUCache From ed1557d16af58a0a4db755b2d5bebbc2a4bb5b54 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Mon, 5 Oct 2026 18:26:48 -0400 Subject: [PATCH 02/10] 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 b975a2194..0b4ccb66e 100644 --- a/src/fusiontrees/braiding_manipulations.jl +++ b/src/fusiontrees/braiding_manipulations.jl @@ -284,42 +284,45 @@ function braid(src::Union{FusionTreePair, FusionTreeBlock}, p::Index2Tuple, leve @assert numind(src) == length(p[1]) + length(p[2]) @assert numout(src) == length(levels[1]) && numin(src) == length(levels[2]) @assert TupleTools.isperm((p[1]..., p[2]...)) - return fsbraid((src, p, levels)) + return fsbraid(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) @@ -339,7 +342,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 60bda7ff5..a9c2ae633 100644 --- a/src/fusiontrees/duality_manipulations.jl +++ b/src/fusiontrees/duality_manipulations.jl @@ -524,25 +524,13 @@ 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 -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)) @@ -567,9 +555,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 36ddd21d75f4be1ebb4c1102ff724324220fab93 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Mon, 5 Oct 2026 18:26:59 -0400 Subject: [PATCH 03/10] 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 f9a1dac46..d22f20acf 100644 --- a/src/tensors/treetransformers.jl +++ b/src/tensors/treetransformers.jl @@ -210,3 +210,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 63031375d59c21720f319aa68244729cc64f13a5 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Mon, 5 Oct 2026 18:58:40 -0400 Subject: [PATCH 04/10] 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 7b97223ec..644089579 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 d22f20acf..0ad662854 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, Vsrc, conjsrc::Bool) where {A} t₀ = Base.time() @@ -210,7 +226,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 2e4545d12af6f773aad0b3ad57f3fba2c88569e4 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Mon, 12 Oct 2026 14:51:39 -0400 Subject: [PATCH 05/10] 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 49a0b34ed..235b3e56c 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.9" +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 14373b62a..f94648eb1 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 644089579..fccad1596 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 0ad662854..a83e0d726 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) + @@ -226,10 +226,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 b807074f0eaad7c89e27e05eedc22bc23eaa740c Mon Sep 17 00:00:00 2001 From: lkdvos Date: Mon, 12 Oct 2026 17:11:38 -0400 Subject: [PATCH 06/10] 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 f94648eb1..11d8b07e6 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 0b4ccb66e..eaad39010 100644 --- a/src/fusiontrees/braiding_manipulations.jl +++ b/src/fusiontrees/braiding_manipulations.jl @@ -284,12 +284,12 @@ 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 # 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) @@ -301,10 +301,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 @@ -342,7 +342,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 a9c2ae633..53d207ecd 100644 --- a/src/fusiontrees/duality_manipulations.jl +++ b/src/fusiontrees/duality_manipulations.jl @@ -524,12 +524,24 @@ 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 -@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₂ @@ -555,7 +567,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 a83e0d726..81d75ea02 100644 --- a/src/tensors/treetransformers.jl +++ b/src/tensors/treetransformers.jl @@ -228,7 +228,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 3e5ccc7eaaa11892e6bd9f3e6f8aab3aad2d5139 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Mon, 12 Oct 2026 17:31:58 -0400 Subject: [PATCH 07/10] 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 2a853124a..fa858f922 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 ad93305f0b467672c3fc01be25544a9f26a0cca4 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 13 Oct 2026 09:15:37 -0400 Subject: [PATCH 08/10] Fix GPU transformers for tensors without indices `_dense_strides(())` called `Base.front` on an empty tuple. Co-Authored-By: Claude Opus 5.5 --- ext/TensorKitGPUArraysExt.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/ext/TensorKitGPUArraysExt.jl b/ext/TensorKitGPUArraysExt.jl index 3019e0372..653a3e668 100644 --- a/ext/TensorKitGPUArraysExt.jl +++ b/ext/TensorKitGPUArraysExt.jl @@ -196,6 +196,7 @@ end # strides of a dense array of shape `sz` _dense_strides(size::Dims) = (1, Base.front(cumprod(size))...) +_dense_strides(::Tuple{}) = () # `permute(Vsrc, p) == Vdst` is enforced when the transformer is built, so the permuted # source shape always matches `sz_dst` and the two views share Cartesian inds. From 85df9e7a5cfdc937d2500aba2fb5955055f56a72 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 13 Oct 2026 09:34:31 -0400 Subject: [PATCH 09/10] 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 | 20 ++++++++++---------- src/fusiontrees/duality_manipulations.jl | 4 ++-- 3 files changed, 13 insertions(+), 14 deletions(-) diff --git a/ext/TensorKitEnzymeExt/utility.jl b/ext/TensorKitEnzymeExt/utility.jl index 11d8b07e6..22ea3febc 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.treetransposer), ::Any...) = nothing @inline EnzymeRules.inactive(::typeof(TensorKit.treebraider), ::Any...) = nothing @inline EnzymeRules.inactive(::typeof(TensorKit.insertleftunit), ::HomSpace, ::Any) = nothing diff --git a/src/fusiontrees/braiding_manipulations.jl b/src/fusiontrees/braiding_manipulations.jl index eaad39010..f2a702684 100644 --- a/src/fusiontrees/braiding_manipulations.jl +++ b/src/fusiontrees/braiding_manipulations.jl @@ -305,19 +305,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 53d207ecd..b53267989 100644 --- a/src/fusiontrees/duality_manipulations.jl +++ b/src/fusiontrees/duality_manipulations.jl @@ -542,10 +542,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 177e4ac3829b1348ef9d7817a63d8cd8b056d70d Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 13 Oct 2026 10:51:17 -0400 Subject: [PATCH 10/10] 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 fa858f922..58e06ac27 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 22ea3febc..52a49579b 100644 --- a/ext/TensorKitEnzymeExt/utility.jl +++ b/ext/TensorKitEnzymeExt/utility.jl @@ -147,6 +147,8 @@ end @inline EnzymeRules.inactive(::typeof(TensorKit.fusiontreetranspose), ::Any, ::Any) = nothing @inline EnzymeRules.inactive(::typeof(TensorKit.treetransposer), ::Any...) = nothing @inline EnzymeRules.inactive(::typeof(TensorKit.treebraider), ::Any...) = nothing +@inline EnzymeRules.inactive(::typeof(TensorKit.device_treetransposer), ::Any...) = nothing +@inline EnzymeRules.inactive(::typeof(TensorKit.device_treebraider), ::Any...) = nothing @inline EnzymeRules.inactive(::typeof(TensorKit.insertleftunit), ::HomSpace, ::Any) = nothing @inline EnzymeRules.inactive(::typeof(TensorKit.insertrightunit), ::HomSpace, ::Any) = nothing @inline EnzymeRules.inactive(::typeof(TensorKit.removeunit), ::HomSpace, ::Any) = nothing diff --git a/ext/TensorKitGPUArraysExt.jl b/ext/TensorKitGPUArraysExt.jl index 653a3e668..4ad4d75ca 100644 --- a/ext/TensorKitGPUArraysExt.jl +++ b/ext/TensorKitGPUArraysExt.jl @@ -272,8 +272,16 @@ function DeviceGenericTreeTransformer( ) end +TensorKit.treebraider( + A::Type{<:AnyGPUVector}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, + conjsrc::Bool, levels::IndexTuple +) = TensorKit.device_treebraider(A, Vdst, Vsrc, p, conjsrc, levels) +TensorKit.treetransposer( + A::Type{<:AnyGPUVector}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, conjsrc::Bool +) = TensorKit.device_treetransposer(A, Vdst, Vsrc, p, conjsrc) + # Cache the device transformer directly, constructing the cached host version only on a miss. -@cached function TensorKit.treebraider( +@cached function TensorKit.device_treebraider( A::Type{TA}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, conjsrc::Bool, levels::IndexTuple ) where {TA <: AnyGPUVector} @@ -281,7 +289,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 81d75ea02..27ebe98ae 100644 --- a/src/tensors/treetransformers.jl +++ b/src/tensors/treetransformers.jl @@ -226,9 +226,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 973118cac..d859981d9 100644 --- a/test/cuda/tensors.jl +++ b/test/cuda/tensors.jl @@ -573,3 +573,17 @@ end end end end + +@timedtestset "Device transformers are cached separately" begin + TensorKit.empty_globalcaches!() + V = SU2Space(0 => 2, 1 // 2 => 2) + t = CUDA.rand(Float64, V ⊗ V ← V ⊗ V) + p = ((2, 1), (4, 3)) + @test permute(t, p) ≈ adapt(CuArray, permute(ad(t), p)) + devicecache = only(TensorKit.MemoizationKit.cache_info(TensorKit.device_treebraider)).second + hostcache = only(TensorKit.MemoizationKit.cache_info(TensorKit.treebraider)).second + @test length(devicecache) == 1 + @test all(key -> !(first(key) <: CuArray), keys(hostcache)) + TensorKit.empty_globalcaches!() + @test isempty(devicecache) +end