cubecl_runtime/storage/
base.rs

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

cubecl_runtime/storage/
base.rs
