fuel_cpu_backend/
byte_storage.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2//! Byte-shaped CPU storage — Phase 7.5 storage-unification target.
3//!
4//! `CpuStorageBytes` is the new CPU storage type that replaces the
5//! legacy [`crate::CpuStorage`] (HostBuffer-based) once kernel
6//! migration completes. Both types coexist during the migration:
7//!
8//! - **Legacy `CpuStorage`** (`dyn_impl::CpuStorage`): wraps
9//!   `HostBuffer` enum, dtype implicit in variant identity. Used by
10//!   every existing op kernel via match-on-variant.
11//! - **`CpuStorageBytes`** (this module): wraps `Arc<AlignedBytes>`,
12//!   dtype lives on the `Storage` wrapper (fuel-storage), kernels
13//!   read typed slices via `bytemuck::cast_slice`.
14//!
15//! Per-op kernels migrate one family at a time during Phase B/C.
16//! When the last kernel migrates, the legacy CpuStorage retires.
17
18mod aligned;
19
20pub use aligned::AlignedBytes;
21
22use std::sync::Arc;
23
24use bytemuck::Pod;
25use fuel_backend_contract::backend::BackendStorage;
26use fuel_ir::{Error, Result};
27
28/// Required alignment for CPU storage allocations. AVX-512-friendly;
29/// also sufficient for AVX2 (32) and NEON (16).
30pub const CPU_ALIGN_BYTES: usize = 64;
31
32/// CPU storage holding bytes addressable by Rust pointer.
33///
34/// Cloning is cheap — bumps the inner `Arc`'s refcount without
35/// copying bytes. Mutating through `bytes_mut` / `as_slice_mut`
36/// performs copy-on-write: if the `Arc` is uniquely held, the bytes
37/// are mutated in place; if shared, the bytes clone before mutation
38/// so other holders see the unchanged data.
39#[derive(Debug, Clone)]
40pub struct CpuStorageBytes {
41    bytes: Arc<AlignedBytes>,
42}
43
44impl CpuStorageBytes {
45    /// Allocate `len_bytes` zero-initialized bytes, 64-byte aligned.
46    pub fn from_zero_bytes(len_bytes: usize) -> Self {
47        Self {
48            bytes: Arc::new(AlignedBytes::new_zeroed(len_bytes, CPU_ALIGN_BYTES)),
49        }
50    }
51
52    /// Phase 7.5 A4 substrate alloc. Allocates `byte_count` zero-
53    /// initialized bytes, 64-byte aligned. Synonym for
54    /// [`Self::from_zero_bytes`]; named `alloc` for symmetry with
55    /// the GPU backend allocators.
56    pub fn alloc(byte_count: usize) -> Self {
57        Self::from_zero_bytes(byte_count)
58    }
59
60    /// Allocate a fresh aligned buffer and copy `src` into it.
61    pub fn from_bytes(src: &[u8]) -> Self {
62        let mut buf = AlignedBytes::new_zeroed(src.len(), CPU_ALIGN_BYTES);
63        if !src.is_empty() {
64            buf.as_slice_mut().copy_from_slice(src);
65        }
66        Self {
67            bytes: Arc::new(buf),
68        }
69    }
70
71    /// Allocate a fresh aligned buffer and copy a typed slice into
72    /// it. `T` must be `Pod` (Copy + Zeroable + 'static + no
73    /// padding); covers every standard ML dtype.
74    pub fn from_slice<T: Pod>(data: &[T]) -> Self {
75        Self::from_bytes(bytemuck::cast_slice(data))
76    }
77
78    /// Total byte count.
79    pub fn len_bytes(&self) -> usize {
80        self.bytes.len()
81    }
82
83    /// Alignment in bytes (always [`CPU_ALIGN_BYTES`]).
84    pub fn align(&self) -> usize {
85        self.bytes.align()
86    }
87
88    /// Borrow the raw bytes immutably.
89    pub fn bytes(&self) -> &[u8] {
90        self.bytes.as_slice()
91    }
92
93    /// Borrow the bytes as a typed slice. Returns `Err` if the byte
94    /// length isn't a multiple of `T`'s size or the buffer alignment
95    /// can't satisfy `T`'s alignment requirement.
96    ///
97    /// Production-correct: no panic. The 64-byte allocation
98    /// alignment is sufficient for every `Pod` type up to and
99    /// including `u64`/`f64`, so the alignment-check failure path
100    /// is unreachable for standard dtypes — but we still surface it
101    /// as `Result` for unusual `Pod` types (e.g., `[u8; 128]`).
102    pub fn as_slice<T: Pod>(&self) -> Result<&[T]> {
103        bytemuck::try_cast_slice(self.bytes()).map_err(|e| {
104            Error::Msg(format!(
105                "CpuStorageBytes::as_slice<{}>: cast failed ({e:?}); \
106                 len_bytes={}, T::size={}, T::align={}",
107                std::any::type_name::<T>(),
108                self.len_bytes(),
109                std::mem::size_of::<T>(),
110                std::mem::align_of::<T>(),
111            ))
112            .bt()
113        })
114    }
115
116    /// Copy-on-write mutable byte access. If this storage is the
117    /// sole holder of its `Arc`, returns a `&mut [u8]` into the
118    /// existing buffer. If the `Arc` is shared, clones the bytes
119    /// into a fresh `Arc` first so other holders are unaffected.
120    ///
121    /// Centralized CoW boundary: every mutating call site goes
122    /// through here, so callers don't need to remember to clone
123    /// manually. Cost when `Arc` is unique: one atomic load (no
124    /// copy). Cost when shared: one allocation + one memcpy.
125    pub fn bytes_mut(&mut self) -> &mut [u8] {
126        Arc::make_mut(&mut self.bytes).as_slice_mut()
127    }
128
129    /// Copy-on-write typed mutable view. Same CoW semantics as
130    /// `bytes_mut`; returns `Err` on a `bytemuck` cast failure.
131    pub fn as_slice_mut<T: Pod>(&mut self) -> Result<&mut [T]> {
132        let len_bytes = self.len_bytes();
133        let size = std::mem::size_of::<T>();
134        let align = std::mem::align_of::<T>();
135        if size != 0 && !len_bytes.is_multiple_of(size) {
136            return Err(Error::Msg(format!(
137                "CpuStorageBytes::as_slice_mut<{}>: byte length {} not a multiple of T size {}",
138                std::any::type_name::<T>(),
139                len_bytes,
140                size,
141            ))
142            .bt());
143        }
144        if align > CPU_ALIGN_BYTES {
145            return Err(Error::Msg(format!(
146                "CpuStorageBytes::as_slice_mut<{}>: T alignment {} exceeds buffer alignment {}",
147                std::any::type_name::<T>(),
148                align,
149                CPU_ALIGN_BYTES,
150            ))
151            .bt());
152        }
153        let bytes = self.bytes_mut();
154        bytemuck::try_cast_slice_mut(bytes).map_err(|e| {
155            Error::Msg(format!(
156                "CpuStorageBytes::as_slice_mut<{}>: cast failed unexpectedly: {e:?}",
157                std::any::type_name::<T>(),
158            ))
159            .bt()
160        })
161    }
162
163    /// Whether this storage's bytes are uniquely owned (no other
164    /// `CpuStorageBytes` shares the same `Arc`). Useful for tests and
165    /// for reasoning about CoW behavior.
166    pub fn is_uniquely_owned(&self) -> bool {
167        Arc::strong_count(&self.bytes) == 1
168    }
169}
170
171impl BackendStorage for CpuStorageBytes {
172    fn len_bytes(&self) -> usize {
173        self.len_bytes()
174    }
175}
176
177#[cfg(test)]
178mod tests {
179    use super::*;
180    use half::{bf16, f16};
181
182    #[test]
183    fn from_zero_bytes_round_trip() {
184        let s = CpuStorageBytes::from_zero_bytes(8);
185        assert_eq!(s.len_bytes(), 8);
186        assert_eq!(s.bytes(), &[0u8; 8]);
187        assert_eq!(s.align(), CPU_ALIGN_BYTES);
188    }
189
190    #[test]
191    fn from_bytes_round_trip() {
192        let s = CpuStorageBytes::from_bytes(&[1, 2, 3, 4]);
193        assert_eq!(s.len_bytes(), 4);
194        assert_eq!(s.bytes(), &[1, 2, 3, 4]);
195    }
196
197    #[test]
198    fn from_slice_typed() {
199        let data = vec![1.0_f32, 2.0, 3.0, 4.0];
200        let s = CpuStorageBytes::from_slice(&data);
201        assert_eq!(s.len_bytes(), 16);
202        let typed: &[f32] = s.as_slice().expect("f32 cast");
203        assert_eq!(typed, &[1.0, 2.0, 3.0, 4.0]);
204    }
205
206    #[test]
207    fn as_slice_supports_bf16_and_f16() {
208        let bf: Vec<bf16> = vec![bf16::from_f32(1.5), bf16::from_f32(-2.5)];
209        let s = CpuStorageBytes::from_slice(&bf);
210        let typed: &[bf16] = s.as_slice().expect("bf16 cast");
211        assert_eq!(typed.len(), 2);
212        assert_eq!(typed[0].to_f32(), 1.5);
213
214        let h: Vec<f16> = vec![f16::from_f32(0.25)];
215        let s2 = CpuStorageBytes::from_slice(&h);
216        let typed2: &[f16] = s2.as_slice().expect("f16 cast");
217        assert_eq!(typed2.len(), 1);
218    }
219
220    #[test]
221    fn as_slice_errors_on_size_mismatch() {
222        let s = CpuStorageBytes::from_bytes(&[0u8; 5]);
223        let result: Result<&[f32]> = s.as_slice();
224        assert!(result.is_err(), "size mismatch should error, not panic");
225    }
226
227    #[test]
228    fn buffer_is_64_byte_aligned() {
229        for size in [1, 4, 16, 64, 1024, 4096] {
230            let s = CpuStorageBytes::from_zero_bytes(size);
231            let addr = s.bytes().as_ptr() as usize;
232            assert_eq!(addr % 64, 0, "size={size} not 64-byte aligned");
233        }
234    }
235
236    #[test]
237    fn clone_is_cheap_and_shares_bytes() {
238        let a = CpuStorageBytes::from_bytes(&[42u8; 1024]);
239        assert!(a.is_uniquely_owned());
240        let b = a.clone();
241        assert!(!a.is_uniquely_owned(), "shared after clone");
242        assert!(!b.is_uniquely_owned());
243        assert_eq!(a.bytes().as_ptr(), b.bytes().as_ptr());
244    }
245
246    #[test]
247    fn cow_isolates_shared_writers() {
248        let a = CpuStorageBytes::from_bytes(&[1u8, 2, 3, 4]);
249        let mut b = a.clone();
250        assert_eq!(a.bytes().as_ptr(), b.bytes().as_ptr());
251
252        b.bytes_mut()[0] = 99;
253
254        assert_eq!(a.bytes(), &[1, 2, 3, 4]);
255        assert_eq!(b.bytes(), &[99, 2, 3, 4]);
256        assert_ne!(a.bytes().as_ptr(), b.bytes().as_ptr());
257    }
258
259    #[test]
260    fn cow_no_clone_when_unique() {
261        let mut a = CpuStorageBytes::from_bytes(&[1u8, 2, 3, 4]);
262        let original_ptr = a.bytes().as_ptr();
263        a.bytes_mut()[0] = 99;
264        assert_eq!(a.bytes().as_ptr(), original_ptr);
265        assert_eq!(a.bytes(), &[99, 2, 3, 4]);
266    }
267
268    #[test]
269    fn typed_mut_view_round_trip() {
270        let mut s = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0]);
271        {
272            let view: &mut [f32] = s.as_slice_mut().expect("f32 mut cast");
273            view[1] = 99.0;
274        }
275        let view: &[f32] = s.as_slice().expect("f32 ref cast");
276        assert_eq!(view, &[1.0, 99.0, 3.0]);
277    }
278
279    #[test]
280    fn typed_mut_view_errors_on_size_mismatch() {
281        let mut s = CpuStorageBytes::from_bytes(&[0u8; 5]);
282        let result: Result<&mut [f32]> = s.as_slice_mut();
283        assert!(result.is_err(), "size mismatch must error, not panic");
284    }
285
286    #[test]
287    fn impl_backend_storage() {
288        // Verify the BackendStorage trait impl works.
289        let s = CpuStorageBytes::from_zero_bytes(32);
290        let bs: &dyn BackendStorage = &s;
291        assert_eq!(bs.len_bytes(), 32);
292    }
293}
fuel_cpu_backend/byte_storage.rs

fuel_cpu_backend/
byte_storage.rs
