Caching#

Why there is a cache#

YASTN separates a tensor operation into two stages: computing the metadata — which blocks of the operands meet, how their charges combine, where each block lands in the output buffer — and then moving the data according to that plan. The metadata stage is pure Python over the tensor’s structural descriptors (struct, fusion history hfs, signatures), which are hashable and, crucially, do not depend on the numerical values inside the blocks.

That makes the metadata stage cacheable. Tensor-network algorithms overwhelmingly repeat the same block structure step after step — DMRG sweeps, CTMRG iterations, or the combo loop of Large contractions all contract tensors whose charge sectors are fixed while their entries change. Each metadata builder is therefore wrapped in functools.lru_cache() keyed on structure, so the Python bookkeeping is paid once and every later call with the same structure is a dictionary lookup.

The consequence worth remembering: the cache is what keeps YASTN’s Python overhead from scaling with the number of iterations. If a hot loop shows an unexpectedly low hit rate, something is perturbing the block structure between steps.

Controlling the caches#

yastn.set_cache_maxsize(maxsize=0)[source]#

Rebind every yastn metadata cache with a new lru_cache maxsize, discarding whatever those caches currently hold.

Each builder is created with maxsize=1024. Note the default here is 0: a bare set_cache_maxsize() disables caching, since lru_cache(0) stores nothing. Pass maxsize=None for unbounded caches.

Caches of backend modules that are not loaded yet are covered too – the requested size is remembered and applied when such a backend registers (see yastn._cache_registry).

yastn.clear_cache()[source]#

Drop the contents of every yastn metadata cache, leaving their maxsize unchanged.

Besides freeing host memory, this releases the device-resident index tensors that the torch backend’s fusion cache pins after its first GPU use.

yastn.get_cache_info()[source]#

Return statistics of lru_caches used in yastn.

Backend-side entries (registered via yastn._cache_registry.register_cache()) are present only for backends that are actually loaded, so the key set depends on which backend is in use.

Reading the statistics#

yastn.get_cache_info() returns a dict mapping a short name to that cache’s functools.CacheInfo — a named tuple (hits, misses, maxsize, currsize):

import yastn

yastn.set_cache_maxsize(maxsize=1024)
...  # run a few iterations of your algorithm
info = yastn.get_cache_info()
print(info["fuse_hard"])       # CacheInfo(hits=..., misses=..., maxsize=1024, currsize=...)

# a healthy steady-state loop is dominated by hits
hits, misses = info["fuse_hard"].hits, info["fuse_hard"].misses
print(f"hit rate {hits / (hits + misses):.2%}")

A currsize sitting at maxsize together with a climbing misses count means the cache is thrashing — the working set of distinct structures is larger than the cache, and raising maxsize will help. Persistent misses at low currsize mean the structures themselves keep changing, which no cache size will fix.

The entries, grouped by the area they serve:

Area

Keys

Fusion

fuse_hard, unfuse_hard, intersect_hfs, combine_leg_structure

Contraction

tensordot_f2m, tensordot_fc, tensordot_nf, tensordot_cutensor_cpu, tensordot_cutensor_gpu, broadcast, mask, trace, vdot, swap_gate, swap_gate_charge, ncon

Algebra

addition

Block structure

get_blocks, get_blocks_charges_all, get_trimmed_struct_engine, get_trimmed_struct_engine_gpu

Backend (torch only)

pack_transpose_and_merge_params

The key set depends on the backend#

The tensor-side keys above are always present. Backend-side keys appear only once that backend has actually been imported: a backend module registers its own caches when it is loaded, so yastn itself never has to import a backend — and never forces an optional dependency such as torch on users of another backend. Code that inspects a particular backend key should therefore use .get():

fusion_params = yastn.get_cache_info().get("pack_transpose_and_merge_params")
if fusion_params is not None:      # torch backend is loaded
    ...

yastn.set_cache_maxsize() still covers backends loaded later: the requested size is remembered and applied at registration time, so the ordering of set_cache_maxsize and yastn.make_config() does not matter.

Memory held by the caches#

Most entries are small — tuples of integers and short NumPy index arrays. The exception is the torch backend’s pack_transpose_and_merge_params, which caches the index maps used by the GPU fuse/unfuse path (see GPU execution: hybrid scatter/loop). Those index tensors are moved onto the data’s device on first use and stay there, so the entry pins GPU memory proportional to the size of the fused buffers it has seen.

yastn.clear_cache() is the lever for releasing them — worth calling between phases of a calculation that use very different tensor sizes, and worth knowing about when several worker processes each accumulate their own copy.