Struct rten::TensorPool

pub struct TensorPool { /* private fields */ }

A pool which enables reuse of data buffers from tensors and other containers.

Reusing buffers for operator outputs, as opposed to allocating a fresh buffer from the global allocator and freeing it when no longer needed, can provide a significant performance improvement.

TensorPool implements the Alloc trait, enabling tensors to be allocated from the pool using the various Tensor::*_in methods, eg. [Tensor::zeros_in]. Allocation requests will be satisfied from the pool if there is a suitable buffer available, or it will fall back to the global allocator otherwise.

When a tensor is no longer needed, it’s buffer can be added to the pool using pool.add(tensor.extract_buffer()), making it available for future allocations. A more convenient method is to wrap the tensor in a [PoolRef] smart pointer which will auto-return the tensor to the pool when dropped. A tensor can be wrapped using tensor.auto_return(pool). The [PoolRef] smart pointer can also be used with other container types, by implementing the [ExtractBuffer] trait for them.

Implementations

impl TensorPool

pub fn new() -> TensorPool

Return a new, empty pool.

This is a cheap operation that does not allocate, so it can be used to create a temporary pool to pass to a function that requires one, if the caller does not have a pool otherwise available.

pub fn alloc<T>(&self, capacity: usize) -> Vec<T>

Allocate an empty vec with a given capacity from the pool.

pub fn add<T>(&self, vec: Vec<T>)

Add a data buffer to the pool.

The buffer will be cleared using Vec::clear and then made available to fulfill future allocation requests.

pub fn alloc_count(&self) -> usize

Return the total number of allocation requests.

pub fn hit_count(&self) -> usize

Return the number of allocation requests that were fulfilled using items in the pool.

pub fn len(&self) -> usize

Return the number of buffers currently in the pool.

pub fn is_empty(&self) -> bool

Return true if the pool is empty.

Trait Implementations

impl Alloc for TensorPool

fn alloc<T>(&self, capacity: usize) -> Vec<T>

Allocate storage for an owned tensor.

impl Default for TensorPool

fn default() -> Self

Returns the “default value” for a type.