cubecl_runtime/storage/
bytes_cpu.rs

1use super::{ComputeStorage, StorageHandle, StorageId, StorageUtilization};
2use alloc::alloc::{Layout, alloc, dealloc};
3use hashbrown::HashMap;
4
5/// The bytes storage maps ids to pointers of bytes in a contiguous layout.
6#[derive(Default)]
7pub struct BytesStorage {
8    memory: HashMap<StorageId, AllocatedBytes>,
9}
10
11impl core::fmt::Debug for BytesStorage {
12    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
13        f.write_str("BytesStorage")
14    }
15}
16
17/// Can send to other threads.
18unsafe impl Send for BytesStorage {}
19unsafe impl Send for BytesResource {}
20
21/// This struct is a pointer to a memory chunk or slice.
22pub struct BytesResource {
23    ptr: *mut u8,
24    utilization: StorageUtilization,
25}
26
27/// This struct refers to a specific (contiguous) layout of bytes.
28struct AllocatedBytes {
29    ptr: *mut u8,
30    layout: Layout,
31}
32
33impl BytesResource {
34    fn get_exact_location_and_length(&self) -> (*mut u8, usize) {
35        unsafe {
36            (
37                self.ptr.add(self.utilization.offset as usize),
38                self.utilization.size as usize,
39            )
40        }
41    }
42
43    /// Returns the resource as a mutable slice of bytes.
44    pub fn write<'a>(&self) -> &'a mut [u8] {
45        let (ptr, len) = self.get_exact_location_and_length();
46
47        unsafe { core::slice::from_raw_parts_mut(ptr, len) }
48    }
49
50    /// Returns the resource as an immutable slice of bytes.
51    pub fn read<'a>(&self) -> &'a [u8] {
52        let (ptr, len) = self.get_exact_location_and_length();
53
54        unsafe { core::slice::from_raw_parts(ptr, len) }
55    }
56}
57
58impl ComputeStorage for BytesStorage {
59    type Resource = BytesResource;
60
61    const ALIGNMENT: u64 = 4;
62
63    fn get(&mut self, handle: &StorageHandle) -> Self::Resource {
64        let allocated_bytes = self.memory.get(&handle.id).unwrap();
65
66        BytesResource {
67            ptr: allocated_bytes.ptr,
68            utilization: handle.utilization.clone(),
69        }
70    }
71
72    fn alloc(&mut self, size: u64) -> StorageHandle {
73        let id = StorageId::new();
74        let handle = StorageHandle {
75            id,
76            utilization: StorageUtilization { offset: 0, size },
77        };
78
79        unsafe {
80            let layout = Layout::array::<u8>(size as usize).unwrap();
81            let ptr = alloc(layout);
82            let memory = AllocatedBytes { ptr, layout };
83
84            self.memory.insert(id, memory);
85        }
86
87        handle
88    }
89
90    fn dealloc(&mut self, id: StorageId) {
91        if let Some(memory) = self.memory.remove(&id) {
92            unsafe {
93                dealloc(memory.ptr, memory.layout);
94            }
95        }
96    }
97}
98
99#[cfg(test)]
100mod tests {
101    use super::*;
102
103    #[test]
104    fn test_can_alloc_and_dealloc() {
105        let mut storage = BytesStorage::default();
106        let handle_1 = storage.alloc(64);
107
108        assert_eq!(handle_1.size(), 64);
109        storage.dealloc(handle_1.id);
110    }
111
112    #[test]
113    fn test_slices() {
114        let mut storage = BytesStorage::default();
115        let handle_1 = storage.alloc(64);
116        let handle_2 = StorageHandle::new(
117            handle_1.id,
118            StorageUtilization {
119                offset: 24,
120                size: 8,
121            },
122        );
123
124        storage
125            .get(&handle_1)
126            .write()
127            .iter_mut()
128            .enumerate()
129            .for_each(|(i, b)| {
130                *b = i as u8;
131            });
132
133        let bytes = storage.get(&handle_2).read().to_vec();
134        storage.dealloc(handle_1.id);
135        assert_eq!(bytes, &[24, 25, 26, 27, 28, 29, 30, 31]);
136    }
137}
cubecl_runtime/storage/bytes_cpu.rs

cubecl_runtime/storage/
bytes_cpu.rs
