cubecl_runtime/storage/
base.rs

1use core::fmt::Debug;
2
3use crate::{server::Binding, storage_id_type};
4
5// This ID is used to map a handle to its actual data.
6storage_id_type!(StorageId);
7
8/// Defines if data uses a full memory chunk or a slice of it.
9#[derive(Clone, Debug)]
10pub struct StorageUtilization {
11    /// The offset in bytes from the chunk start.
12    pub offset: u64,
13    /// The size of the slice in bytes.
14    pub size: u64,
15}
16
17/// Contains the [storage id](StorageId) of a resource and the way it is used.
18#[derive(new, Clone, Debug)]
19pub struct StorageHandle {
20    /// Storage id.
21    pub id: StorageId,
22    /// How the storage is used.
23    pub utilization: StorageUtilization,
24}
25
26impl StorageHandle {
27    /// Returns the size the handle is pointing to in memory.
28    pub fn size(&self) -> u64 {
29        self.utilization.size
30    }
31
32    /// Returns the size the handle is pointing to in memory.
33    pub fn offset(&self) -> u64 {
34        self.utilization.offset
35    }
36
37    /// Increase the current offset with the given value in bytes.
38    pub fn offset_start(&self, offset_bytes: u64) -> Self {
39        let utilization = StorageUtilization {
40            offset: self.offset() + offset_bytes,
41            size: self.size() - offset_bytes,
42        };
43
44        Self {
45            id: self.id,
46            utilization,
47        }
48    }
49
50    /// Reduce the size of the memory handle..
51    pub fn offset_end(&self, offset_bytes: u64) -> Self {
52        let utilization = StorageUtilization {
53            offset: self.offset(),
54            size: self.size() - offset_bytes,
55        };
56
57        Self {
58            id: self.id,
59            utilization,
60        }
61    }
62}
63
64/// Storage types are responsible for allocating and deallocating memory.
65pub trait ComputeStorage: Send {
66    /// The resource associated type determines the way data is implemented and how
67    /// it can be accessed by kernels.
68    type Resource: Send;
69
70    /// The alignment memory is allocated with in this storage.
71    fn alignment(&self) -> usize;
72
73    /// Returns the underlying resource for a specified storage handle
74    fn get(&mut self, handle: &StorageHandle) -> Self::Resource;
75
76    /// Allocates `size` units of memory and returns a handle to it
77    fn alloc(&mut self, size: u64) -> StorageHandle;
78
79    /// Deallocates the memory pointed by the given storage id.
80    ///
81    /// These deallocations might need to be flushed with [`Self::perform_deallocations`].
82    fn dealloc(&mut self, id: StorageId);
83}
84
85/// Access to the underlying resource for a given binding.
86#[derive(new, Debug)]
87pub struct BindingResource<Resource: Send> {
88    // This binding is here just to keep the underlying allocation alive.
89    // If the underlying allocation becomes invalid, someone else might
90    // allocate into this resource which could lead to bad behaviour.
91    #[allow(unused)]
92    binding: Binding,
93    resource: Resource,
94}
95
96impl<Resource: Send> BindingResource<Resource> {
97    /// access the underlying resource. Note: The resource might be bigger
98    /// than just the original allocation for the binding. Only the part
99    /// for the original binding is guaranteed to remain, other parts
100    /// of the resource *will* be re-used.
101    pub fn resource(&self) -> &Resource {
102        &self.resource
103    }
104}
cubecl_runtime/storage/base.rs

cubecl_runtime/storage/
base.rs
