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

cubecl_server/storage/
bytes_cpu.rs
