Skip to content

Replace TensorKit cache machinery with MemoizationKit.jl - #568

Draft
lkdvos wants to merge 10 commits into
mainfrom
cached-integration
Draft

lkdvos wants to merge 10 commits into
mainfrom
cached-integration

Conversation

@lkdvos

@lkdvos lkdvos commented Oct 5, 2026 •

Copy link
Copy Markdown
Member

TensorKit currently maintains its own cache machinery and key wrappers. This change delegates those responsibilities to MemoizationKit.jl (registered, v0.1).

  • Replace the internal cache macro and containers with MemoizationKit, using its native container preferences and module-wide inspection and clearing.
  • Use MemoizationKit.Hashed and cachekey for shared sector structures; simplify fusion-tree cache signatures while preserving anyonic braiding inference.
  • Rename the cached fusion-tree functions fsbraid/fstranspose to fusiontreebraid/fusiontreetranspose, since they are now user-facing through the MemoizationKit APIs and dashboard. Their return types are annotated, so that cache hits stay inferred.
  • Integrate MemoizationKit lookup and computation timers with TensorKit profiling, and remove duplicate machinery and upstream-covered tests.
  • Supply efficient MemoizationKit.cachesize methods for TensorKit's cached values, measuring live buffers without traversing inline coefficients. Shared buffers are counted conservatively across entries; keys and allocator/container overhead are excluded.

Following #556, treebraider and treetransposer forward device storage types to device_treebraider and device_treetransposer, whose @cached methods are added by the GPU extension. These are owned by TensorKit, so they are inspected, emptied and timed together with TensorKit's other caches, and the separate GLOBAL_CACHES registry is removed. Being separate functions, the device transformers keep their own cache, which can be sized independently of the host transformers. They currently use the global default (10k entries) instead of the smaller limit (1000 entries) of #556: a package-provided default that still respects user preferences needs QuantumKitHub/MemoizationKit.jl#19.
This also fixes the GPU transformers for tensors without indices (_dense_strides(())).

The documentation now refers to the MemoizationKit.jl docs for the dashboard, disk caching and the full configuration options.


Main point is to separate this out into its own package, but also to tackle the configurability and visualization, as MemoizationKit.jl supplies a dashboard for working with these caches:

using TensorKit, MemoizationKit, Tachikoma
MemoizationKit.cache_dashboard()

Example RAM view rendered from repeated SU(2) tensor permutations and transposes, with treebraider configured for a 1 MiB cached-value budget. The numbers describe this small example workload, not a benchmark:

[RAM] │  Disk
 6 caches · 42 entries · 13KiB in byte-measured caches · recent hit rate 97.0%
 Name ▲                        Hit rate      Size                   Kind       Activity
▌TensorKit.degeneracystructure █████97%████▌ ▏░░░░░░░12/10k░░░░░░░░ Clock      ▅▆▇███▆▇
 TensorKit.fusiontreebraid     ████100%█████ ▏░░░░░░░8/10k░░░░░░░░░ Clock      ▇█
 TensorKit.fusiontreetranspose ████100%█████ ▏░░░░░░░8/10k░░░░░░░░░ Clock      ▇█
 TensorKit.sectorstructure     ████100%█████ ▏░░░░░░░2/10k░░░░░░░░░ Clock      ▇█▂▂▂▂▁▂
 TensorKit.treebraider         █████96%████▌ ▎░░░░13KiB/1.0MiB░░░░░ Clock      ▄▅▇███▆▇
 TensorKit.treetransposer      █████96%████▌ ▏░░░░░░░6/10k░░░░░░░░░ Clock      ▄▅▇███▆▇



╭─  TensorKit.degeneracystructure  ────────────────────────────────────────────────────────────────────────────────────╮
│Clock · 12 of 10000 entries · 276 hits · 12 misses                                                                    │
│hit rate: recent 96.6% · lifetime 95.8% · 104738.6 hits/s                                                             │
│ ▅▆▅▆▇███▆▇                                                                                                           │
╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯
 ⇥ tab  ↑↓ select  ⏎ resize  e empty  s sort  r reverse  / filter  g refresh  q quit

For example, apply a 1 GiB cached-value budget immediately with:

MemoizationKit.set_cache_size!(TensorKit.treebraider, 2^30; by = MemoizationKit.cachesize)

The dashboard shows live hit rates, cache occupancy and activity; select a row and press Enter to resize it or e to empty it. Cache limits can count entries or cached-value bytes, and package/per-function defaults can be saved through MemoizationKit.set_cache_preferences!.

Closes #382 (better control over cache sizes).
Closes #385 (memory limits for caches).
Closes #386 (cache inspection and monitoring).

@lkdvos
lkdvos force-pushed the cached-integration branch from b4ccaab to 00afd35 Compare October 5, 2026 23:00
@lkdvos

lkdvos commented Oct 5, 2026

Copy link
Copy Markdown
Member Author

Rebased onto remote main (6f135471) and added a separate commit for efficient cached-value byte accounting. The GPU cache implementation remains unchanged apart from its Cached LRU import.

Small container microbenchmarks on this branch: median nanoseconds per get! lookup (lower is better), with the same keys, access order, precomputed transformer value and Any, Any containers used by TensorKit. This isolates cache access; it does not measure tensor contraction speed.

Workload ClockCache LRU LRU time / Clock time
Warm integer keys, 10,000 entries 42.5 ns 56.1 ns 1.32×
Warm TensorKit transformer keys, 10,000 entries 357.3 ns 370.5 ns 1.04×
Mixed integer hits/misses, 1,024-entry cache over 2,048 keys 92.2 ns 91.0 ns 0.99×

Both warm-hit workloads allocate zero bytes. Clock is faster on these warm hits, while the eviction workload is effectively tied and slightly favors LRU. The small transformer-key difference should not be read as a guaranteed application-level improvement.

The new byte measurement also helps avoid recursive traversal: for a synthetic transformer containing one 256×256 Float64 matrix, Cached.cachesize takes 4.4 ns, zero allocations, versus 7.61 μs, 57 allocations (2,608 bytes) for Base.summarysize. The former estimates live payload bytes and excludes allocator/container overhead, so these are different measurements.

Environment: Julia 1.13.1, BenchmarkTools 1.8.0, Cached main c8a74b4, Intel Xeon Gold 6244 @ 3.60 GHz, one Julia thread. Each trial uses up to 1,000 samples / 2 seconds, 4,096 accesses per sample, evals=1, seed 568. Mixed-workload caches are reset before each sample.

Reproduction script (requires BenchmarkTools)
using TensorKit, Cached, BenchmarkTools, Random

function lookup_batch(cache::C, keys, order, value::V) where {C, V}
    checksum = 0
    for i in order
        cached = get!(() -> value, cache, keys[i])::V
        checksum += length(cached.data)
    end
    return checksum
end

function filled_cache(C, keys, value; limit = length(keys))
    cache = C{Any, Any}(; maxsize = limit)
    for key in Iterators.take(keys, limit)
        cache[key] = value
    end
    return cache
end

rng = MersenneTwister(568)
V = SU2Space(0 => 2, 1 // 2 => 2)
W = V ⊗ V ← V ⊗ V
p, levels = ((2, 1), (4, 3)), (1, 2, 3, 4)
value = TensorKit.treebraider(W, W, p, false, levels)
tensor_keys = [(let V = SU2Space(0 => i, 1 // 2 => 2); V ⊗ V ← V ⊗ V end,
    W, p, false, levels) for i in 1:10_000]
integer_keys = collect(1:10_000)
order = rand(rng, 1:10_000, 4096)
results = NamedTuple[]
for (workload, keys) in (("Integer keys, warm hits (10k entries)", integer_keys), ("TensorKit transformer keys, warm hits (10k entries)", tensor_keys))
    for (kind, C) in (("ClockCache", Cached.ClockCache), ("LRU", Cached.LRU))
        cache = filled_cache(C, keys, value)
        trial = @benchmark lookup_batch($cache, $keys, $order, $value) seconds = 2 samples = 1000 evals = 1
        estimate = median(trial)
        push!(results, (; workload, kind, ns = estimate.time / length(order),
            bytes = estimate.memory, allocations = estimate.allocs))
    end
end
churn_keys = collect(1:2048)
churn_order = rand(rng, 1:2048, 4096)
for (kind, C) in (("ClockCache", Cached.ClockCache), ("LRU", Cached.LRU))
    trial = @benchmark lookup_batch(cache, $churn_keys, $churn_order, $value) setup = (cache = filled_cache($C, $churn_keys, $value; limit = 1024)) seconds = 2 samples = 1000 evals = 1
    estimate = median(trial)
    push!(results, (; workload = "Integer keys, mixed hits/misses and eviction", kind,
        ns = estimate.time / length(churn_order), bytes = estimate.memory, allocations = estimate.allocs))
end

# Use a large recoupling matrix to compare payload accounting against a full traversal.
synthetic = TensorKit.GenericTreeTransformer(
    [(ones(256, 256), collect(1:10_000), collect(1:10_000))],
    TensorKit.StridedStructure{2}[], TensorKit.StridedStructure{2}[],
)
size_trials = (
    payload = @benchmark(Cached.cachesize($synthetic), seconds = 2, samples = 1000),
    recursive = @benchmark(Base.summarysize($synthetic), seconds = 2, samples = 1000),
)
report = IOBuffer()
println(report, "Julia: ", VERSION, "; BenchmarkTools: ", pkgversion(BenchmarkTools), "; Cached: ", pkgversion(Cached))
println(report, "CPU: ", Sys.cpu_info()[1].model, "; Julia threads: ", Threads.nthreads())
println(report, "Median ns/lookup; 4096 lookups/sample; seed 568; 1000 samples max; 2 s/trial; evals=1")
for r in results
    println(report, r)
end
for (label, trial) in pairs(size_trials)
    estimate = median(trial)
    println(report, label, " size: ", estimate.time, " ns; ", estimate.memory, " bytes; ", estimate.allocs, " allocations")
end
report_text = String(take!(report))
println(report_text)

Validation after rebasing: 143 cache/key/size assertions, 27 timer assertions and 43 remote-source/tensor-operation smoke assertions passed. The broader fusion-tree checks passed before the rebase; the full suite was not run.

@Jutho

Jutho commented Oct 9, 2026

Copy link
Copy Markdown
Member

Does this PR secretly require a full Cached.jl review?

@lkdvos
lkdvos force-pushed the cached-integration branch from 00afd35 to 49e7b01 Compare October 12, 2026 18:51
@lkdvos lkdvos changed the title Replace TensorKit cache machinery with Cached.jl Replace TensorKit cache machinery with MemoizationKit.jl Oct 12, 2026
@lkdvos
lkdvos force-pushed the cached-integration branch from 49e7b01 to ae793a9 Compare October 12, 2026 18:55
lkdvos and others added 8 commits October 13, 2026 09:06
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 <noreply@anthropic.com>
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.
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.
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 <noreply@anthropic.com>
- 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 <noreply@anthropic.com>
…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 <noreply@anthropic.com>
`_dense_strides(())` called `Base.front` on an empty tuple.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@lkdvos
lkdvos force-pushed the cached-integration branch from ae793a9 to ad93305 Compare October 13, 2026 13:26
lkdvos and others added 2 commits October 13, 2026 09:46
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
`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 <noreply@anthropic.com>

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

2 participants