cubecl_runtime/storage/
base.rs

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

cubecl_runtime/storage/
base.rs
