burn_compute/storage/
base.rs

1use crate::storage_id_type;
2
3// This ID is used to map a handle to its actual data.
4storage_id_type!(StorageId);
5
6/// Defines if data uses a full memory chunk or a slice of it.
7#[derive(Clone)]
8pub enum StorageUtilization {
9    /// Full memory chunk of specified size
10    Full(usize),
11    /// Slice of memory chunk with start index and size.
12    Slice(usize, usize),
13}
14
15/// Contains the [storage id](StorageId) of a resource and the way it is used.
16#[derive(new)]
17pub struct StorageHandle {
18    /// Storage id.
19    pub id: StorageId,
20    /// How the storage is used.
21    pub utilization: StorageUtilization,
22}
23
24impl StorageHandle {
25    /// Returns the size the handle is pointing to in memory.
26    pub fn size(&self) -> usize {
27        match self.utilization {
28            StorageUtilization::Full(size) => size,
29            StorageUtilization::Slice(_, size) => size,
30        }
31    }
32}
33
34/// Storage types are responsible for allocating and deallocating memory.
35pub trait ComputeStorage: Send {
36    /// The resource associated type determines the way data is implemented and how
37    /// it can be accessed by kernels.
38    type Resource: Send;
39
40    /// Returns the underlying resource for a specified storage handle
41    fn get(&mut self, handle: &StorageHandle) -> Self::Resource;
42
43    /// Allocates `size` units of memory and returns a handle to it
44    fn alloc(&mut self, size: usize) -> StorageHandle;
45
46    /// Deallocates the memory pointed by the given storage id.
47    fn dealloc(&mut self, id: StorageId);
48}
burn_compute/storage/base.rs

burn_compute/storage/
base.rs
