cubecl_server/storage/
bytes_cpu.rs

1use crate::server::IoError;
2
3use super::{ComputeStorage, StorageHandle, StorageId, StorageUtilization};
4use alloc::alloc::{Layout, alloc_zeroed, dealloc};
5use cubecl_environment::backtrace::BackTrace;
6use cubecl_environment::collections::HashMap;
7
8/// The bytes storage maps ids to pointers of bytes in a contiguous layout.
9#[derive(Default)]
10pub struct BytesStorage {
11    memory: HashMap<StorageId, AllocatedBytes>,
12}
13
14impl BytesStorage {
15    /// Base alignment for host buffers, including 512-bit SIMD accesses.
16    pub const ALIGNMENT: usize = 64;
17}
18
19impl core::fmt::Debug for BytesStorage {
20    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
21        f.write_str("BytesStorage")
22    }
23}
24
25/// Can send to other threads.
26unsafe impl Send for BytesStorage {}
27unsafe impl Send for BytesResource {}
28
29/// This struct is a pointer to a memory chunk or slice.
30#[derive(Debug)]
31pub struct BytesResource {
32    ptr: *mut u8,
33    utilization: StorageUtilization,
34}
35
36/// This struct refers to a specific (contiguous) layout of bytes.
37struct AllocatedBytes {
38    ptr: *mut u8,
39    layout: Layout,
40}
41
42impl BytesResource {
43    /// Returns a mutable pointer to the start of the resource and its length.
44    pub fn get_write_ptr_and_length(&self) -> (*mut u8, usize) {
45        (
46            // SAFETY:
47            // - The offset is created to be within the bounds of the allocation.
48            unsafe { self.ptr.add(self.utilization.offset as usize) },
49            self.utilization.size as usize,
50        )
51    }
52
53    /// Returns the resource as a mutable slice of bytes.
54    ///
55    /// The lifetime `'a` is the lifetime of the underlying `BytesStorage` allocation,
56    /// not of `self`. The `&mut self` ensures only one mutable slice is created per
57    /// resource. Multiple resources may point to non-overlapping regions of the same
58    /// allocation (like `split_at_mut`); each resource owns its region exclusively.
59    pub fn write<'a>(&mut self) -> &'a mut [u8] {
60        let (ptr, len) = self.get_write_ptr_and_length();
61
62        // SAFETY:
63        // - ptr is non-null and aligned (from BytesStorage::alloc).
64        // - The region [ptr..ptr+len) is within a single allocation.
65        // - Memory is initialized (BytesStorage uses alloc_zeroed).
66        // - `&mut self` ensures exclusive access to this resource's region.
67        // - `StorageHandle` assigns non-overlapping regions per resource.
68        // - Systems must make sure this is the only `BytesResource` with an outstanding mutable borrow.
69        unsafe { core::slice::from_raw_parts_mut(ptr, len) }
70    }
71
72    /// Returns the resource as an immutable slice of bytes.
73    ///
74    /// See [`write`](Self::write) for lifetime and safety notes.
75    pub fn read<'a>(&self) -> &'a [u8] {
76        let (ptr, len) = self.get_write_ptr_and_length();
77
78        // SAFETY:
79        // - ptr is non-null and aligned (from BytesStorage::alloc).
80        // - The region [ptr..ptr+len) is within a single allocation.
81        // - Memory is initialized (BytesStorage uses alloc_zeroed).
82        unsafe { core::slice::from_raw_parts(ptr, len) }
83    }
84}
85
86impl ComputeStorage for BytesStorage {
87    type Resource = BytesResource;
88
89    fn alignment(&self) -> usize {
90        Self::ALIGNMENT
91    }
92
93    fn get(&mut self, handle: &StorageHandle) -> Result<Self::Resource, IoError> {
94        let allocated_bytes =
95            self.memory
96                .get(&handle.id)
97                .ok_or_else(|| IoError::StorageHandleNotFound {
98                    reason: alloc::format!("{} in the bytes storage", handle.id).into(),
99                    backtrace: BackTrace::capture(),
100                })?;
101
102        Ok(BytesResource {
103            ptr: allocated_bytes.ptr,
104            utilization: handle.utilization.clone(),
105        })
106    }
107
108    #[cfg_attr(
109        feature = "tracing",
110        tracing::instrument(level = "trace", skip(self, size))
111    )]
112    fn alloc(&mut self, size: u64) -> Result<StorageHandle, IoError> {
113        let id = StorageId::new();
114        let handle = StorageHandle {
115            id,
116            utilization: StorageUtilization { offset: 0, size },
117        };
118
119        if size == 0 {
120            // Zero-size allocations are valid handles but don't need real memory.
121            let memory = AllocatedBytes {
122                ptr: core::ptr::without_provenance_mut(Self::ALIGNMENT),
123                layout: Layout::from_size_align(0, Self::ALIGNMENT).unwrap(),
124            };
125            self.memory.insert(id, memory);
126        } else {
127            unsafe {
128                let layout = Layout::from_size_align(size as usize, Self::ALIGNMENT).unwrap();
129
130                // We allocate zeroed memory since we expose it as &[u8] / &mut [u8]
131                // which requires initialization.
132                let ptr = alloc_zeroed(layout);
133                if ptr.is_null() {
134                    return Err(IoError::BufferTooBig {
135                        size,
136                        backtrace: BackTrace::capture(),
137                    });
138                }
139                let memory = AllocatedBytes { ptr, layout };
140                self.memory.insert(id, memory);
141            }
142        }
143
144        Ok(handle)
145    }
146
147    #[cfg_attr(feature = "tracing", tracing::instrument(level = "trace", skip(self)))]
148    fn dealloc(&mut self, id: StorageId) {
149        if let Some(memory) = self.memory.remove(&id)
150            && memory.layout.size() > 0
151        {
152            unsafe {
153                dealloc(memory.ptr, memory.layout);
154            }
155        }
156    }
157
158    fn flush(&mut self) {
159        // We don't wait for dealloc.
160    }
161}
162
163#[cfg(test)]
164mod tests {
165    use super::*;
166
167    #[test_log::test]
168    fn test_can_alloc_and_dealloc() {
169        let mut storage = BytesStorage::default();
170        let handle_1 = storage.alloc(64).unwrap();
171
172        assert_eq!(handle_1.size(), 64);
173        storage.dealloc(handle_1.id);
174    }
175
176    #[test_log::test]
177    fn test_slices() {
178        let mut storage = BytesStorage::default();
179        let handle_1 = storage.alloc(64).unwrap();
180        let handle_2 = StorageHandle::new(
181            handle_1.id,
182            StorageUtilization {
183                offset: 24,
184                size: 8,
185            },
186        );
187
188        storage
189            .get(&handle_1)
190            .unwrap()
191            .write()
192            .iter_mut()
193            .enumerate()
194            .for_each(|(i, b)| {
195                *b = i as u8;
196            });
197
198        let bytes = storage.get(&handle_2).unwrap().read().to_vec();
199
200        storage.dealloc(handle_1.id);
201        assert_eq!(bytes, &[24, 25, 26, 27, 28, 29, 30, 31]);
202    }
203
204    /// Miri catches: "reading memory, but memory is uninitialized"
205    #[test_log::test]
206    fn test_read_after_alloc_without_write() {
207        let mut storage = BytesStorage::default();
208        let handle = storage.alloc(16).unwrap();
209        let resource = storage.get(&handle).unwrap();
210        assert!(resource.read().iter().all(|&b| b == 0));
211        storage.dealloc(handle.id);
212    }
213
214    /// Miri catches: "creating allocation with size 0"
215    #[test_log::test]
216    fn test_zero_size_alloc_and_dealloc() {
217        let mut storage = BytesStorage::default();
218        let handle = storage.alloc(0).unwrap();
219        assert_eq!(handle.size(), 0);
220        storage.dealloc(handle.id);
221    }
222
223    #[test_log::test]
224    fn test_alloc_dealloc_realloc() {
225        let mut storage = BytesStorage::default();
226        let h1 = storage.alloc(32).unwrap();
227        storage.get(&h1).unwrap().write()[0] = 0xAA;
228        storage.dealloc(h1.id);
229        let h2 = storage.alloc(32).unwrap();
230        storage.dealloc(h2.id);
231    }
232
233    #[test_log::test]
234    fn test_multiple_non_overlapping_regions() {
235        let mut storage = BytesStorage::default();
236        let base = storage.alloc(64).unwrap();
237
238        let regions: alloc::vec::Vec<_> = (0..4)
239            .map(|i| {
240                StorageHandle::new(
241                    base.id,
242                    StorageUtilization {
243                        offset: i * 16,
244                        size: 16,
245                    },
246                )
247            })
248            .collect();
249
250        for (i, region) in regions.iter().enumerate() {
251            storage.get(region).unwrap().write().fill(i as u8);
252        }
253        for (i, region) in regions.iter().enumerate() {
254            assert!(
255                storage
256                    .get(region)
257                    .unwrap()
258                    .read()
259                    .iter()
260                    .all(|&b| b == i as u8)
261            );
262        }
263        storage.dealloc(base.id);
264    }
265}
cubecl_server/storage/bytes_cpu.rs

cubecl_server/storage/
bytes_cpu.rs
