Trait cubecl_runtime::memory_management::MemoryManagement

source ·

pub trait MemoryManagement<Storage: ComputeStorage>: Send + Debug {
    type Handle: MemoryHandle<Self::Binding>;
    type Binding: MemoryBinding;

    // Required methods
    fn get(&mut self, binding: Self::Binding) -> Storage::Resource;
    fn reserve<Sync: FnOnce()>(
        &mut self,
        size: usize,
        sync: Sync,
    ) -> Self::Handle;
    fn alloc<Sync: FnOnce()>(&mut self, size: usize, sync: Sync) -> Self::Handle;
    fn dealloc(&mut self, binding: Self::Binding);
    fn storage(&mut self) -> &mut Storage;
}

Expand description

The MemoryManagement trait encapsulates strategies for (de)allocating memory. It is bound to the ComputeStorage trait, which does the actual (de)allocations.

The MemoryManagement can only reserve memory space or get the resource located at a space. Modification of the resource data should be done directly on the resource.

Required Associated Types§

source

type Handle: MemoryHandle<Self::Binding>

The associated type that must implement MemoryHandle.

source

type Binding: MemoryBinding

The associated type that must implement MemoryBinding

Required Methods§

source

fn get(&mut self, binding: Self::Binding) -> Storage::Resource

Returns the resource from the storage at the specified handle

source

fn reserve<Sync: FnOnce()>(&mut self, size: usize, sync: Sync) -> Self::Handle

Finds a spot in memory for a resource with the given size in bytes, and returns a handle to it

source

fn alloc<Sync: FnOnce()>(&mut self, size: usize, sync: Sync) -> Self::Handle

Bypass the memory allocation algorithm to allocate data directly.

§Notes

Can be useful for servers that want specific control over memory.

source

fn dealloc(&mut self, binding: Self::Binding)

Bypass the memory allocation algorithm to deallocate data directly.

§Notes

Can be useful for servers that want specific control over memory.

source

fn storage(&mut self) -> &mut Storage

Fetch the storage used by the memory manager.

§Notes

The storage should probably not be used for allocations since the handles won’t be compatible with the ones provided by the current trait. Prefer using the alloc and dealloc functions.

This is useful if you need to time the deallocations based on async computation, or to change the mode of storage for different reasons.