Skip to content

Device memory plan

Jump to bottom
Berk Geveci edited this page Nov 9, 2026 · 2 revisions

Device memory: accounting and spilling plan

Issue #387, 0.9. Status: planned (2026-11-08).

Tack does not know how much device memory it holds, and when a GPU runs out it fails the allocation. This plan adds two things to CUDA, HIP and Level Zero:

  • Accounting: what Tack holds on each device, its peak, and what its caches keep.
  • Spilling: when memory runs short, Tack copies the arrays it used least recently to host memory, frees their device memory, and brings each one back the next time something uses it.

Metal is out of scope for spilling, since its memory is the host's and macOS already pages it. Metal and the CPU get accounting only.

What happens today

  • Allocation. Every array's storage comes from tack.empty, through _allocate (lang/array.py). On an out-of-memory error, every GPU backend raises DeviceMemoryError. _allocate then tries once more:

    1. collect Python's garbage (a dropped array held in a reference cycle keeps its memory until then);
    2. wait for the device, so deferred frees complete;
    3. empty the backend's cache (HIP's cache of dropped buffers, Metal's pool);
    4. allocate again.

    If that fails, the error reaches the caller, usually from a temporary deep inside a filter.

  • Nothing is counted. There is no record of live bytes and no peak. The only query of the device's free memory is one hipMemGetInfo call in the HIP backend.

  • Derived data stays on the device while its dataset lives: faces, edges, point links, locators, cell geometry.

  • No graceful degradation: no chunking, no host fallback, no out-of-core path. Peak memory has been fixed case by case: the M6 dataset on a 16 GB GPU (#179, #180); AhmedML's face derivation, 6.7 GB of temporaries reduced to 2.0 GB.

Where it attaches

  • Buffers. CUDABuffer, HIPBuffer and L0Buffer are each created in one place, through _allocate, and freed in __del__:

    • CUDA by a stream-ordered free (cuMemFreeAsync, or the deferred list without memory pools);
    • HIP by its cache or hipFreeAsync;
    • Level Zero by a deferred zeMemFreeExt.

    Views of an array (reshapes) share its buffer.

  • Launches read one property. A kernel launch on these three backends reads each array's pointer through buffer.device_ptr, and so do the native reductions and texture uploads. The buffer's own methods (from_numpy, to_numpy, read_range, fill_bytes, device copies) use _device_ptr directly.

  • Everything is in order on one stream or queue. Each device has one stream (CUDA, HIP) or one in-order immediate command list (Level Zero), so a spill (copy out, free) and a restore (allocate, copy in) take their place among the launches.

  • Memory Tack cannot move: arrays wrapping someone else's pointer (wrap_ptr, DLPack imports), memory handed out (DLPack export, export_memory, share(), VTK's copy=False), MPI buffers (allocate_mpi_array), the scalar pack buffers, and texture images.

Accounting (step 1)

Per device, counted where buffers are created and freed:

quantity meaning
in_use bytes of live Tack-owned buffers on the device
peak the most in_use has been since the last reset
cached bytes Tack's caches keep for reuse (HIP's cache, Metal's pool); on CUDA the driver pool's reserve beyond use (CU_MEMPOOL_ATTR_RESERVED_MEM_CURRENT less USED)
spilled bytes of Tack's arrays now held on the host instead (step 2)
device_free, device_total the driver's answer where it gives one (cuMemGetInfo, hipMemGetInfo, Level Zero's device properties or Sysman)

The API:

  • tack.memory_info(device=None) returns these as a small named record.

  • tack.reset_peak_memory(device=None) resets the peak.

  • A scope for the peak of a block of code:

    with tack.memory_peak() as peak:
    td.extract_surface(data)
print(peak.bytes)

The CPU counts its own buffers too; Metal counts its own and reports currentAllocatedSize and recommendedMaxWorkingSetSize.

What this is for:

  • A filter's peak, measured directly. Today it is guessed from the process's RSS, which does not see GPU memory.
  • Tests that hold a filter to its memory.
  • The spilling decisions below.

Later, optionally: tag allocations by origin (a cached derived structure, a field, a filter's temporary), to answer "what is holding my GPU memory".

Spilling (steps 2 and 3)

When it runs:

  • With a budget: tack.set_memory_budget(bytes, device=None). When an allocation would take in_use over the budget, Tack first spills least-recently-used arrays until it fits.
  • Without one: spilling still runs when an allocation fails. In _allocate, after today's garbage collection and cache release, Tack spills least-recently-used arrays until the allocation succeeds or nothing is left to spill; only then does DeviceMemoryError reach the caller.

What is spilled: any Tack-owned buffer except those listed under "Memory Tack cannot move" and the arrays of the launch being prepared. Derived data (faces, edges, point links, locators, cell geometry) is spilled like any other array, not dropped and recomputed.

How a buffer is spilled:

  1. Copy it to ordinary host memory (a NumPy array).
  2. Free its device memory.
  3. Mark it spilled, with the host copy.

On CUDA and HIP, a copy out of the device into pageable memory waits for the queued work before it, so the spill is ordered after every launch that uses the buffer. The free that follows is stream-ordered as today.

Ordinary memory makes each spill and restore a copy the host waits for; pinned memory would let them overlap other work. That is a later choice, if measurements call for it.

Host access uses the spilled copy. A spilled array stays spilled when the host uses it:

  • reads (numpy(), read_range, a[i]) copy from the host copy;
  • host writes (copyto, fill, uploads) write into it.

Reading results back after a large computation, when arrays are most likely to be spilled, then moves nothing to the device only to copy it straight back.

Only the device brings a buffer back. device_ptr, which kernel launches, native reductions and texture uploads read, and the buffer's device-to-device operations make a spilled buffer resident:

  1. allocate (which may spill others);
  2. copy the host copy back;
  3. drop the host copy.

The host copy is not handed out. No host pointer or zero-copy view of spilled data leaves the buffer, because a restore frees the host copy. A holder would then read freed memory, or stale values after a kernel wrote the device copy. Sharing it safely would mean keeping the array spilled while the pointer lives, as share() keeps handed-out device memory in place: a different feature, and not planned.

Before binding a launch's arguments, GPUBackend.execute makes every array argument resident in one pass, and keeps those arrays from being spilled while it allocates the others. A device class's held arrays are among the arguments once its template is expanded.

Least recently used: each buffer keeps a counter its backend advances whenever the buffer is used (each device_ptr read, each host access, on the device or of its spilled copy). The backend keeps its spillable buffers in a weak registry, and spilling takes the lowest counters first.

What this costs a warm launch: one attribute check per array (spilled or not) and the counter's update. That must stay within the dispatch fast path's budget (docs/design/specialization-and-caching.md).

Threads: spilling and restoring run under the backend's lock: CUDA and HIP each launch under their variant's dispatch_lock, Level Zero under its _launch_lock. Two threads must not restore or spill one buffer at once.

Limits:

  • Spilling helps when all live arrays together do not fit but each step's arrays do. One launch's arrays must still fit on the device.
  • A buffer spilled and restored over and over (thrashing) is slow but correct. Accounting counts spills and restores, so it shows.

Considered and not taken

  • CUDA managed memory or HIP's hipMallocManaged oversubscription. The driver would page instead of Tack. But page faults cost per access, not per array, and are hard to predict. Its support varies (Windows WDDM, AMD's XNACK, Level Zero's shared memory). And it would change every allocation, including the memory Tack hands to other libraries.
  • Dropping derived data instead of spilling it. Recomputing is sometimes cheaper than a round trip to the host, but spilling every array the same way is simpler and predictable.
  • Pinned host memory for spilled copies. It is faster and asynchronous, but limited and costly to allocate. Ordinary memory first.

Steps

  1. Accounting on every backend. Counters in each buffer's creation and free, the caches' bytes, the driver's free and total, the API above, docs, tests (counts after allocations, views and frees; the peak scope).

  2. Spilling when an allocation fails on CUDA, HIP and Level Zero:

    • each buffer's states (resident, spilled) and the restore hook on device_ptr and the buffer methods;
    • the weak registry and least-recently-used order;
    • the exclusions;
    • spilling in _allocate's retry;
    • execute making a launch's arrays resident.

    Tests use a small budget (step 3) to force spills, on every GPU backend through the tickets, and checks that a launch never runs on a spilled pointer.

  3. The budget: tack.set_memory_budget, spilling before an allocation that would pass it.

  4. Validation on the GPUs. AhmedML's sweep under a budget well below its peak, for example 8 GB on the RTX 6000 Ada: every step completes, the outputs are the same as without a budget, and the cost is measured. Also the M6 dataset on a 16 GB GPU, and HIP and Level Zero through their machines' tickets.

  5. Later: attribution by origin; pinned host memory if spills dominate.

Related

  • #386: CUDA, HIP and Level Zero clear every new buffer. Independent, but it touches the same allocation code.
  • Array views (a[i:j] over the same memory, on core/array-views). A view with an offset must take its pointer from its base buffer at launch, so that spilling the base is seen through the view. Check when views merge.

Clone this wiki locally