memory_utils

GPU memory arena and accounting for the render loop.

ManualMemory is a bump-allocator arena for render-time tensors. Callers snapshot the allocation pointer, allocate freely, and restore it to free everything since the snapshot – deterministic and far cheaper than relying on the caching allocator across a frame batch.

The rest of the module is the accounting that keeps a render inside its budget: available-byte queries, ensure_render_headroom, CUDA peak-tracking scopes, and InsufficientMemoryException / is_cuda_oom for the retry path that shrinks the frame window when a batch does not fit.

AllocationRecorder and ManualMemory.scope() are diagnostics only – they attribute arena usage per stage when you are investigating, and do not participate in batch sizing. That is algan.rendering.memory_model’s job, which measures the arena’s actual high-water mark rather than modelling it.

Classes

AllocationRecorder	Captures every ManualMemory allocation into a scope tree.
ManualMemory	A bump-allocator arena the renderer draws its per-frame tensors from.
ScopeRecord	One recorded ManualMemory.scope() frame.
TempMemoryContext	Restore the arena pointers on exit.

Functions

auto_record_enabled()

Whether calibration recording is armed (see set_auto_record()).

Lets callers skip work that only a calibration run needs, rather than computing it and discarding it on every render.

begin_cuda_peak(device)

Start measuring a region’s peak torch CUDA allocation.

torch.cuda.reset_peak_memory_stats is process-global, so a component that measures its own peak destroys the number the profiler reports for the whole render – which is why the GPU merge’s peak tracking had to default off. Peaks are absolute high-water marks, so the displaced value is not lost, merely remembered: peak_allocated() returns the max of the live counter and everything these regions have reset away. Nesting is safe.

Returns an opaque token for end_cuda_peak(), or None off CUDA. Only torch’s allocator is visible – Taichi’s cuMemAllocAsync pool is not (see is_cuda_oom()) – so callers still need slack.

clear_recorded_arenas()

cuda_peak_scope(device)

Block form of begin_cuda_peak().

Yields a callable returning the region’s peak bytes above entry; it keeps reporting the final value after the block exits.

end_cuda_peak(token)

Finish a begin_cuda_peak() region.

Returns the bytes allocated above its entry point (0 off CUDA).

ensure_render_headroom(device, min_free_fraction=0.15)

Return torch’s reserved-but-free CUDA blocks to the driver when free VRAM is low, so a following Taichi kernel launch has room.

Taichi allocates from its own CUDA pool (cuMemAllocAsync), which cannot draw on torch’s caching allocator. When both share a device and free memory runs low, a Taichi launch (typically the post-process tonemap) OOMs even though torch is holding plenty of reclaimable cached blocks – see is_cuda_oom. The retry loops recover from that, but only after gc + re-rendering the chunk; reclaiming proactively here avoids the round-trip.

Gated on driver-level free memory so the common (plentiful) case pays only a cheap mem_get_info probe: torch.cuda.empty_cache() (~ms, and it forces the next batch to re-acquire blocks from the driver) runs only when free VRAM drops below min_free_fraction of the device total – exactly the regime where a Taichi launch is at risk. no-op off CUDA. Returns True iff it actually reclaimed.

get_num_available_bytes(device=device(type='cuda'))

is_cuda_oom(exc)

True if exc is an out-of-memory failure from either GPU allocator.

PyTorch raises torch.OutOfMemoryError, but Taichi kernel launches allocate from their own CUDA pool (cuMemAllocAsync) and surface exhaustion as a plain RuntimeError wrapping the driver string (CUDA_ERROR_OUT_OF_MEMORY). The render arena bump-allocator never hits the driver mid-render, so a batch that over-committed VRAM fails inside a Taichi launch (typically the post-process tonemap) rather than as a torch OOM – and the retry loops, which only knew the torch type, let it escape. Matching the driver message lets the same release_torch_memory + window-split retry recover it (torch.cuda.empty_cache hands torch’s reserved-but-free blocks back to the driver, which is exactly the memory Taichi needs).

note_nonarena_peak(name, input_bytes, peak_bytes)

Record a transient peak taken outside the arena (calibration only).

The GPU merge and the projection build out of place in pool headroom, so the arena recorder cannot see them and their size has to come from torch’s allocator counters instead. Ignored unless set_auto_record() is on, so a normal render pays one boolean test.

peak_allocated(device=None)

Process-wide peak torch CUDA allocation, including peaks that an intervening begin_cuda_peak() region reset off the live counter.

recorded_arenas()

Managed arenas armed by set_auto_record(), oldest first.

recorded_nonarena_peaks()

(name, input_bytes, peak_bytes) samples, oldest first.

release_torch_memory(force_gc=True)

Reclaim freed memory back to the allocators.

gc.collect() walks the entire Python object graph and dominates this call (~0.2s each on a large scene; it was costing ~40% of a small render when called several times per frame batch). It is only needed to break reference cycles – reference counting already frees the (explicitly nulled) geometry tensors immediately – so it is skipped unless the GPU or host is actually under memory pressure (where reclaiming cyclic garbage matters for avoiding OOM) or force_gc is set. A render additionally freezes the authored scene out of collection entirely (scene_excluded_from_gc()), which is what makes the surviving collections cheap.

torch.cuda.empty_cache() is not cheap: it drains the device and hands every cached block back to the driver, which on Windows/WDDM costs tens of milliseconds per call whether or not there is anything to hand back (measured at ~79 ms a call, 33 s of a four-minute render, once the gc above stopped dominating it). It is therefore gated on the same memory pressure as the collection, plus a worthwhile amount actually being reclaimable – or on the caller forcing it, which every failure/retry path does. This is self-regulating: a cache left unreclaimed shows up as driver-level used memory, so it raises the pressure that triggers the next reclaim.

Host/native reclamation is more destructive and therefore has its own real pressure gate. Under host/cgroup pressure Linux first calls malloc_trim(0) to return glibc-retained free pages without losing JIT state. Pressure is measured again; only if it remains does the default Quadrants backend reset its live runtime, and never while a render is active. A successful reset is followed by one more Linux trim because the freed LLVM/JIT allocations can otherwise remain in glibc’s arenas. A reset discards the in-process compiled specializations, which the next render reloads from the offline cache, so this is an OOM-avoidance last resort rather than routine cleanup. force_gc does not force either native step.

Sizing decisions are unaffected either way: get_num_available_bytes() reclaims unconditionally before it measures, so every batch and chunk still sees the same free-byte figure.

reset_peak_floor()

Forget displaced peaks, alongside a reset_peak_memory_stats call.

Callers that reset the process counter to start a fresh measurement window (the profiler, between runs) must clear the floor too, or the previous window’s peak leaks into this one’s.

scene_excluded_from_gc()

Keep the authored scene out of every collection made inside the block.

release_torch_memory runs gc.collect() several times per frame batch to break the reference cycles a batch leaves behind before the device runs out of memory, and a collection walks every tracked object in the process. An authored scene is millions of them – one per Mob, per recorded edit, per retained tensor – all live from the first frame to the last, so each collection re-walked the whole scene to find the handful of cycles the batch actually produced. Measured on this project’s reference scene: 0.24 s per call, ~100 s of a four-minute render, every second of it holding the GIL against the batch-prep worker that is supposed to be running concurrently.

gc.freeze() moves everything that already exists into a permanent generation collections skip, so the render’s collections walk only what the render itself allocated. Nothing leaks: the frozen objects are the scene, which is live throughout the render, and gc.unfreeze() returns them to the ordinary generations on the way out (including on an error, so a failed render does not leave the process with collection disabled for its scene). The one collection before freezing keeps pre-existing garbage collectable.

set_auto_record(enabled)

Record every managed arena created from now on (calibration only).

The render arena is built deep inside get_frames and dropped when the job ends, so a driver that wants a whole render’s allocation stream cannot reach in and arm it. Arenas armed this way are retained in a registry that outlives the render; call clear_recorded_arenas() between runs.

Exceptions

InsufficientMemoryException