Replace TensorKit cache machinery with MemoizationKit.jl

Project.toml

@@ -9,9 +9,9 @@ projects = ["test", "docs"]
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 Dictionaries = "85a47980-9c8c-11e8-2b9f-f7ca1fa99fb4"
-LRUCache = "8ac3fa9e-de4c-5943-b1dc-09c6b5f20637"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MatrixAlgebraKit = "6c742aac-3347-4629-af66-fc926824e5e4"
+MemoizationKit = "1b238080-9255-4fe9-b224-89eb24efe93b"
 OhMyThreads = "67456a42-1dca-4109-a031-0a68de7e3ad5"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Preferences = "21216c6a-2e73-6563-6e65-726566657250"
@@ -55,9 +55,9 @@ Enzyme = "0.13.207"
 EnzymeTestUtils = "0.2.8"
 FiniteDifferences = "0.12"
 GPUArrays = "11.4.1"
-LRUCache = "1.6"
 LinearAlgebra = "1"
 MatrixAlgebraKit = "0.6.9"
+MemoizationKit = "0.1"
 Mooncake = "0.5.27"
 OhMyThreads = "0.8.0"
 PrecompileTools = "1"

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

docs/src/man/fusiontrees.md

@@ -260,7 +260,29 @@ This is simply given by
 
 The `braid` and `permute` routines for double fusion trees will be the main access point for corresponding manipulations on tensors.
 As a consequence, results from these routines are memoized via a `@cached` mechanism.
-The caching strategy is controlled by `CacheStyle`: for `FusionStyle(I) isa UniqueFusion` no caching is used (`NoCache()`), since the result is a single cheap scalar coefficient; for `MultipleFusion` or `GenericFusion` sectors, a global LRU cache (`GlobalLRUCache()`) is used, as computing the transformation matrix can be expensive and results are reused frequently.
+The caching strategy is controlled by `MemoizationKit.CacheStyle`: for `FusionStyle(I) isa UniqueFusion` no caching is used (`NoCache()`), since the result is a single cheap scalar coefficient; for `MultipleFusion` or `GenericFusion` sectors, `GlobalCache()` uses MemoizationKit.jl's configured container (ClockCache by default), as computing the transformation matrix can be expensive and results are reused frequently.
+
+TensorKit also uses MemoizationKit.jl for sector structures, degeneracy structures and fusion tree transformers.
+Sector structures share cache entries across spaces with the same sectors, even if their degeneracy dimensions differ, through `MemoizationKit.cachekey`.
+The MemoizationKit APIs can inspect, resize and configure these caches:
+
+```julia
+using TensorKit, MemoizationKit
+MemoizationKit.cache_info(TensorKit)
+MemoizationKit.set_cache_size!(TensorKit.treebraider, 50_000)
+MemoizationKit.set_cache_size!(TensorKit.treebraider, 2^30; by = MemoizationKit.cachesize) # 1 GiB of cached values
+MemoizationKit.set_cache_preferences!(TensorKit; maxsize = 50_000, measure = "bytes")
+MemoizationKit.set_cache_preferences!(; container = "LRU")
+```
+
+Preferences are saved in `LocalPreferences.toml` and apply after restarting Julia; a runtime size change applies immediately.
+`MemoizationKit.uncached(f, args...)` computes a result without consulting or filling the cache.
+TensorKit supplies efficient `MemoizationKit.cachesize` methods for its structure, fusion-tree and transformer results.
+These estimate bytes from the fixed-size fields and live array buffers; matrix coefficients and inline sector values are not traversed individually.
+Reference-containing scalar or sector arrays fall back to `Base.summarysize`.
+Buffers shared between cache entries count toward each entry, while identical source and destination structure buffers within a transformer count once.
+The budget excludes cache keys, container and allocation overhead, and spare array capacity; it is a cached-value estimate rather than a limit on total process memory.
+Further features, such as the live cache dashboard, disk caching and the full configuration options, are described in the [MemoizationKit.jl documentation](https://quantumkithub.github.io/MemoizationKit.jl/stable/).
 
 This caching implies that potential inefficiencies in the fusion tree manipulations (which we nonetheless try to avoid) will not seriously affect performance of tensor manipulations.
 

docs/src/man/profiling.md

@@ -1,6 +1,7 @@
 # [Profiling and timers](@id s_profiling)
 
 TensorKit's index manipulations, tensor contractions and factorizations are instrumented with [TimerOutputs.jl](https://github.com/KristofferC/TimerOutputs.jl) sections that are compiled away by default, so they incur no runtime cost.
+Enabling TensorKit's timers also enables MemoizationKit.jl's cache lookup and computation timers, recorded in the same timer tree and categories.
 They can be enabled to obtain a detailed breakdown of where time is spent inside these operations, in particular the split between the different kinds of work involved in manipulating symmetric tensors:
 
 | category | contents |

ext/TensorKitEnzymeExt/utility.jl

@@ -26,7 +26,7 @@ pullback_dC!(ΔC, β::Number) = scale!(ΔC, conj(β))
 @inline EnzymeRules.inactive_type(::Type{<:TensorKit.GenericTreeTransformer}) = true
 @inline EnzymeRules.inactive_type(::Type{<:TensorKit.UniqueTreeTransformer}) = true
 @inline EnzymeRules.inactive_type(::Type{<:TensorKit.VectorSpace}) = true
-@inline EnzymeRules.inactive_type(::Type{<:TensorKit.LRU}) = true
+@inline EnzymeRules.inactive_type(::Type{<:TensorKit.MemoizationKit.AbstractCache}) = true
 
 function EnzymeRules.augmented_primal(
         config::EnzymeRules.RevConfigWidth{1},
@@ -142,11 +142,9 @@ end
 @inline EnzymeRules.inactive(::typeof(TensorKit.fusiontrees), ::Any) = nothing
 @inline EnzymeRules.inactive(::typeof(TensorKit.fusiontrees), ::Any, ::Any) = nothing
 @inline EnzymeRules.inactive(::typeof(TensorKit.fusiontrees), ::Any, ::Any, ::Any) = nothing
-@inline EnzymeRules.inactive(::typeof(TensorKit.fsbraid), ::Any) = nothing
-@inline EnzymeRules.inactive(::typeof(TensorKit.fsbraid), ::Any, ::Any) = nothing
+@inline EnzymeRules.inactive(::typeof(TensorKit.fusiontreebraid), ::Any, ::Any, ::Any) = nothing
 @inline EnzymeRules.inactive(::typeof(TensorKit.artin_braid), ::Any, ::Any) = nothing
-@inline EnzymeRules.inactive(::typeof(TensorKit.fstranspose), ::Any) = nothing
-@inline EnzymeRules.inactive(::typeof(TensorKit.fstranspose), ::Any, ::Any) = nothing
+@inline EnzymeRules.inactive(::typeof(TensorKit.fusiontreetranspose), ::Any, ::Any) = nothing
 @inline EnzymeRules.inactive(::typeof(TensorKit.treetransposer), ::Any...) = nothing
 @inline EnzymeRules.inactive(::typeof(TensorKit.treebraider), ::Any...) = nothing
 @inline EnzymeRules.inactive(::typeof(TensorKit.insertleftunit), ::HomSpace, ::Any) = nothing

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.
@@ -287,24 +288,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)

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

src/auxiliary/caches.jl

@@ -1,16 +1,23 @@
-"""
-    const GLOBAL_CACHES
+# Caching of the fusion tree manipulations and tensor structures is provided by MemoizationKit.jl.
+# `CacheStyle` is imported (not just used) so that the `CacheStyle(::typeof(f), ...)`
+# specializations in TensorKit, and those of users, extend `MemoizationKit.CacheStyle`.
+import MemoizationKit: CacheStyle
+using MemoizationKit: MemoizationKit, @cached, NoCache, GlobalCache, Hashed, LRU
 
-Registry of the global caches TensorKit maintains, as `name => cache` pairs.
+# Buffer payloads can be measured without visiting their elements when stored inline.
+# Reference-containing sector/scalar types need the recursive fallback to include their data.
+_cache_payload_size(a::Array{T}) where {T} = isbitstype(T) ? sizeof(a) : Base.summarysize(a)
 
-See also [`empty_globalcaches!`](@ref) and [`global_cache_info`](@ref).
-"""
-const GLOBAL_CACHES = Pair{Symbol, Any}[]
+# Include the index tables as well as the values (Dictionaries 0.4's Indices layout).
+function _cache_payload_size(inds::Indices)
+    return sizeof(getfield(inds, :slots)) + sizeof(getfield(inds, :hashes)) +
+        _cache_payload_size(getfield(inds, :values))
+end
 
 """
     empty_globalcaches!()
 
-Empty every global cache in [`GLOBAL_CACHES`](@ref).
+Empty the global caches of TensorKit.
 These mostly contain bookkeeping for various different index manipulations and tensor structures,
 so clearing this out can free up some memory whenever you have a workflow that involves a large variety of structures.
 For example, you may want to clear the cache when working with different symmetries, or in algorithms that dynamically alter the sizes of tensors.
@@ -19,218 +26,20 @@ Since everything is recomputed on demand, this is purely a memory measure.
 See also [`global_cache_info`](@ref) to display the status.
 """
 function empty_globalcaches!()
-    foreach(empty! ∘ last, GLOBAL_CACHES)
+    MemoizationKit.empty_caches!(TensorKit)
     return nothing
 end
 
 """
     global_cache_info([io::IO = stdout])
 
-Print the hit/miss statistics and current size of every global cache in [`GLOBAL_CACHES`](@ref).
+Print the hit/miss statistics and current size of every global cache of TensorKit.
 
 See also [`empty_globalcaches!`](@ref).
 """
 function global_cache_info(io::IO = stdout)
-    for (name, cache) in GLOBAL_CACHES
-        println(io, name, ":\t", LRUCache.cache_info(cache))
+    for (f, cache) in MemoizationKit.cache_info(TensorKit)
+        println(io, nameof(f), ":\t", cache)
     end
     return
 end
-
-abstract type CacheStyle end
-struct NoCache <: CacheStyle end
-struct TaskLocalCache{D <: AbstractDict} <: CacheStyle end
-struct GlobalLRUCache <: CacheStyle end
-
-const DEFAULT_GLOBALCACHE_SIZE = Ref(10^4)
-
-function CacheStyle(args...)
-    return GlobalLRUCache()
-end
-
-# category of the miss-path (construction) timer section of an `@cached` function
-function _cached_category(fname::Symbol)
-    return fname in (:fsbraid, :fstranspose, :treebraider, :treetransposer) ?
-        "symmetry" : "bookkeeping"
-end
-
-# timer sections are controlled by the `timeit_debug_enabled` switch of the module using `@cached`
-_timeit_expr(label, ex) =
-    Expr(:macrocall, GlobalRef(TimerOutputs, Symbol("@timeit_debug")), LineNumberNode(@__LINE__, @__FILE__), GlobalRef(@__MODULE__, :GLOBAL_TIMER), label, ex)
-
-macro cached(ex)
-    return _cached(__module__, ex)
-end
-
-function _cached_error(msg, ex)
-    ex isa Expr && (ex = Base.remove_linenums!(copy(ex)))
-    Meta.isexpr(ex, :kw) && (ex = Expr(:(=), ex.args...)) # prints as `x = 1`
-    throw(ArgumentError("@cached: $msg, got `$ex`"))
-end
-
-# the arguments form the cache key, so they must be positional; unnamed ones get a name
-function _cached_canonicalize_arg(arg)
-    Meta.isexpr(arg, :parameters) && _cached_error("keyword arguments are not supported", arg)
-    Meta.isexpr(arg, :kw) && _cached_error("default values are not supported", arg)
-    Meta.isexpr(arg, :...) && _cached_error("varargs are not supported", arg)
-    arg isa Symbol && return arg
-    Meta.isexpr(arg, :(::)) || _cached_error("unsupported argument", arg)
-    length(arg.args) == 1 && return Expr(:(::), gensym(:arg), arg.args[1])
-    return arg
-end
-
-function _cached(mod::Module, ex)
-    if !Meta.isexpr(ex, :function)
-        Meta.isexpr(ex, :(=)) && Meta.isexpr(ex.args[1], (:call, :where, :(::))) &&
-            _cached_error("short-form function definitions are not supported, use `function ... end`", ex.args[1])
-        _cached_error("expected a function definition", ex)
-    end
-    fcall = ex.args[1]
-    if Meta.isexpr(fcall, :where)
-        hasparams = true
-        params = fcall.args[2:end]
-        fcall = fcall.args[1]
-    else
-        hasparams = false
-    end
-    if Meta.isexpr(fcall, :(::))
-        typed = true
-        typeex = fcall.args[2]
-        fcall = fcall.args[1]
-    else
-        typed = false
-    end
-    Meta.isexpr(fcall, :call) || _cached_error("expected a function signature", fcall)
-    fname = fcall.args[1]
-    # qualified names such as `TensorKit.treebraider` add methods to a function of another module
-    basename = Meta.isexpr(fname, :.) ? fname.args[end].value : fname
-    basename isa Symbol ||
-        _cached_error("expected a function name or a qualified name such as `Module.f`", fname)
-    fargs = map(_cached_canonicalize_arg, fcall.args[2:end])
-    fcall = Expr(:call, fname, fargs...)
-    fargnames = map(arg -> arg isa Symbol ? arg : arg.args[1], fargs)
-    # timer labels for the cache lookup and the miss-path construction
-    lookuplabel = string("bookkeeping: cache ", basename)
-    misslabel = string(_cached_category(basename), ": compute ", basename)
-    _fbody = ex.args[2]
-
-    # actual implenetation, with underscore name
-    _fname = Symbol(:_, basename)
-    _fcall = Expr(:call, _fname, fargs...)
-    if hasparams
-        _fcall = Expr(:where, _fcall, params...)
-    end
-    _fex = Expr(:function, _fcall, _fbody)
-
-    # implementation that chooses the cache style
-    newfcall = fcall
-    if hasparams
-        newfcall = Expr(:where, newfcall, params...)
-    end
-    cachestylevar = gensym(:cachestyle)
-    cachestyleex = Expr(
-        :(=), cachestylevar, Expr(:call, CacheStyle, fname, fargnames...)
-    )
-    newfbody = Expr(
-        :block, cachestyleex, Expr(:call, fname, fargnames..., cachestylevar)
-    )
-    newfex = Expr(:function, newfcall, newfbody)
-
-    # nocache implementation
-    fnocachecall = Expr(:call, fname, fargs..., :(::$NoCache))
-    if hasparams
-        fnocachecall = Expr(:where, fnocachecall, params...)
-    end
-    fnocachebody = _timeit_expr(misslabel, Expr(:call, _fname, fargnames...))
-    if typed
-        T = gensym(:T)
-        fnocachebody = Expr(:block, Expr(:(=), T, typeex), Expr(:(::), fnocachebody, T))
-    end
-    fnocacheex = Expr(:function, fnocachecall, fnocachebody)
-
-    # tasklocal cache implementation
-    Dvar = gensym(:D)
-    flocalcachecall = Expr(:call, fname, fargs..., :(::$TaskLocalCache{$Dvar}))
-    if hasparams
-        flocalcachecall = Expr(:where, flocalcachecall, params..., Dvar)
-    else
-        flocalcachecall = Expr(:where, flocalcachecall, Dvar)
-    end
-    globalcachename = Symbol(:GLOBAL_, uppercase(string(basename)), :_CACHE)
-    localcachename = Symbol(:_tasklocal_, globalcachename)
-    cachevar = gensym(:cache)
-    getlocalcacheex = :(
-        $cachevar::$Dvar = get!(task_local_storage(), $(QuoteNode(localcachename))) do
-            return $Dvar()
-        end
-    )
-    valvar = gensym(:val)
-    if length(fargnames) == 1
-        key = fargnames[1]
-    else
-        key = Expr(:tuple, fargnames...)
-    end
-    missex = _timeit_expr(misslabel, Expr(:call, _fname, fargnames...))
-    getvalex = _timeit_expr(lookuplabel, :(get!(() -> $missex, $cachevar, $key)))
-    if typed
-        T = gensym(:T)
-        flocalcachebody = Expr(
-            :block,
-            getlocalcacheex,
-            Expr(:(=), T, typeex),
-            Expr(:(=), Expr(:(::), valvar, T), getvalex),
-            Expr(:return, valvar)
-        )
-    else
-        flocalcachebody = Expr(
-            :block,
-            getlocalcacheex,
-            Expr(:(=), valvar, getvalex),
-            Expr(:return, valvar)
-        )
-    end
-    flocalcacheex = Expr(:function, flocalcachecall, flocalcachebody)
-
-    # # global cache implementation
-    fglobalcachecall = Expr(:call, fname, fargs..., :(::$GlobalLRUCache))
-    if hasparams
-        fglobalcachecall = Expr(:where, fglobalcachecall, params...)
-    end
-    getglobalcachex = Expr(:(=), cachevar, globalcachename)
-    if typed
-        T = gensym(:T)
-        fglobalcachebody = Expr(
-            :block,
-            getglobalcachex,
-            Expr(:(=), T, typeex),
-            Expr(:(=), Expr(:(::), valvar, T), getvalex),
-            Expr(:return, valvar)
-        )
-    else
-        fglobalcachebody = Expr(
-            :block,
-            getglobalcachex,
-            Expr(:(=), valvar, getvalex),
-            Expr(:return, valvar)
-        )
-    end
-    fglobalcacheex = Expr(:function, fglobalcachecall, fglobalcachebody)
-    fglobalcachedef = Expr(
-        :const,
-        Expr(:(=), globalcachename, :($LRU{Any, Any}(; maxsize = $DEFAULT_GLOBALCACHE_SIZE[])))
-    )
-    # Other modules register their caches explicitly from `__init__`, since
-    # mutations of this registry during their precompilation do not survive loading.
-    fglobalcacheregister = mod === (@__MODULE__) ? Expr(
-            :call, :push!, GlobalRef(@__MODULE__, :GLOBAL_CACHES),
-            :($(QuoteNode(globalcachename)) => $globalcachename)
-        ) : nothing
-
-    # # total expression
-    return esc(
-        Expr(
-            :block, _fex, newfex, fnocacheex, flocalcacheex,
-            fglobalcachedef, fglobalcacheregister, fglobalcacheex
-        )
-    )
-end