Struct CudaAllocator 

pub struct CudaAllocator { /* private fields */ }

A caching GPU memory allocator with block pools, splitting, coalescing, and stream-aware reuse.

Wraps a GpuDevice and maintains two block pools (small and large). Allocation requests are served from cached free blocks when possible; only on cache miss does the allocator call through to the CUDA driver. Freed blocks are returned to the pool and coalesced with neighbors to reduce fragmentation.

CL-323

Implementations

impl CudaAllocator

pub fn new(device: Arc<GpuDevice>) -> Self

Create a new caching allocator for the given device.

pub fn alloc_zeros<T>(&self, count: usize) -> GpuResult<CudaBuffer<T>>

where T: DeviceRepr + ValidAsZeroBits,

Allocate count zero-initialized elements of type T on the device.

The returned CudaBuffer is tracked by this allocator. When you are done with it, pass it to free so the statistics stay accurate. (Dropping the buffer directly still frees GPU memory, but the allocated_bytes counter will be too high.)

Errors

Returns [GpuError::Driver] if the underlying CUDA allocation fails.

pub fn alloc_copy<T>(&self, data: &[T]) -> GpuResult<CudaBuffer<T>>

where T: DeviceRepr,

Copy a host slice to device memory, tracking the allocation.

This is the allocator-aware equivalent of crate::transfer::cpu_to_gpu.

Errors

Returns [GpuError::Driver] if the CUDA memcpy or allocation fails.

pub fn free<T>(&self, buffer: CudaBuffer<T>)

Return a buffer to the allocator, freeing the GPU memory and updating the statistics.

This is preferred over simply dropping the buffer so that memory_allocated stays accurate.

pub fn memory_allocated(&self) -> usize

Bytes currently allocated (live) on the device through this allocator.

pub fn max_memory_allocated(&self) -> usize

Peak bytes ever allocated since creation or the last reset_peak_stats.

pub fn memory_reserved(&self) -> usize

Total bytes reserved from the CUDA driver (cached + in-use).

pub fn reset_peak_stats(&self)

Reset the peak counter to the current allocation level.

pub fn empty_cache(&self)

Release all cached (free) blocks back to the CUDA driver.

After this call, memory_reserved() drops to memory_allocated() (only blocks currently in use remain). This is useful when another component needs GPU memory and the cache is holding onto freed blocks.

CL-323

pub fn device(&self) -> &GpuDevice

The underlying device.

pub fn record_stream_on_block(&self, block_idx: usize, stream: StreamId)

Record that a block was used on stream, preventing reuse until work on that stream completes.

This is the Rust equivalent of PyTorch’s recordStream().

CL-323

pub fn block_count(&self) -> usize

Number of blocks in the arena (for debugging/testing).

pub fn free_block_count(&self) -> usize

Number of free blocks in both pools (for debugging/testing).

pub fn cache_stats(&self) -> (usize, usize)

(hits, misses) cache statistics.

pub fn cached_bytes(&self) -> usize

Total cached (free, reusable) bytes.

pub fn cache_find( &self, size: usize, stream: StreamId, ) -> Option<(usize, usize)>

Try to find a cached block of at least size bytes on stream.

Returns the block index and its actual size if found. The block is marked as allocated and removed from the free pool. If the block is significantly larger than needed, it is split.

CL-323

pub fn cache_insert( &self, requested_size: usize, driver_alloc_size: usize, ptr: usize, stream: StreamId, ) -> (usize, usize)

Register a new block from a fresh driver allocation.

Called when cache_find returns None and the caller has obtained memory from the CUDA driver. The full driver allocation is registered as a block; if it’s larger than the requested size, the remainder is split off and placed in the free pool.

Returns (block_idx, actual_block_size).

CL-323

pub fn cache_free(&self, block_idx: usize)

Return a block to the cache (free it back to a pool).

The block is coalesced with any adjacent free blocks and inserted into the appropriate pool for future reuse.

CL-323

pub fn driver_alloc_size(size: usize) -> usize

Get the driver allocation size for a given request size.

Callers use this to know how many bytes to request from the driver when cache_find misses.

Trait Implementations

impl Debug for CudaAllocator

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.