fuel_ir/
cpu_storage.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2//! Host-side tensor storage types.
3//!
4//! The canonical type is [`HostBuffer`] — a dtype-tagged owned `Vec<T>`
5//! representing contiguous tensor data in host-addressable RAM. The
6//! type alias [`CpuStorage`] = [`HostBuffer`] preserves backwards
7//! compatibility so existing code doesn't need to change.
8
9use crate::{DType, Error, Result, WithDType};
10use float8::F8E4M3;
11use half::{bf16, f16};
12
13/// Host-addressable tensor storage holding a typed vector of elements.
14///
15/// This is the universal interchange format for moving tensor data
16/// between backends: every `BackendStorage` can produce one via
17/// `to_host_buffer()`, and every `BackendDevice` can ingest one via
18/// `storage_from_host_buffer()`.
19///
20/// Previously named `CpuStorage` — that name remains as a type alias.
21#[derive(Debug, Clone)]
22pub enum HostBuffer {
23    U8(Vec<u8>),
24    I8(Vec<i8>),
25    U32(Vec<u32>),
26    I16(Vec<i16>),
27    I32(Vec<i32>),
28    I64(Vec<i64>),
29    BF16(Vec<bf16>),
30    F16(Vec<f16>),
31    F32(Vec<f32>),
32    F64(Vec<f64>),
33    F8E4M3(Vec<F8E4M3>),
34    // Dummy types that store raw bytes
35    F6E2M3(Vec<u8>),
36    F6E3M2(Vec<u8>),
37    F4(Vec<u8>),
38    F8E8M0(Vec<u8>),
39    /// Boolean mask storage (GAP-168(c)). One byte per element (`0`/`1`),
40    /// byte-identical to `U8` — but a DISTINCT variant so a `Bool` buffer can
41    /// never be silently consumed as `U8`. The comparison kernels write `0`/`1`
42    /// bytes here; `where_cond`/`masked_fill` read it as a byte mask.
43    Bool(Vec<u8>),
44}
45
46/// Backwards-compatibility alias. New code should use [`HostBuffer`].
47pub type CpuStorage = HostBuffer;
48
49/// A borrowed reference to host tensor storage.
50#[derive(Debug, Clone)]
51pub enum HostBufferRef<'a> {
52    U8(&'a [u8]),
53    I8(&'a [i8]),
54    U32(&'a [u32]),
55    I16(&'a [i16]),
56    I32(&'a [i32]),
57    I64(&'a [i64]),
58    BF16(&'a [bf16]),
59    F16(&'a [f16]),
60    F32(&'a [f32]),
61    F64(&'a [f64]),
62    F8E4M3(&'a [F8E4M3]),
63    // Dummy types that store raw bytes
64    F6E2M3(&'a [u8]),
65    F6E3M2(&'a [u8]),
66    F4(&'a [u8]),
67    F8E8M0(&'a [u8]),
68    /// Boolean mask (GAP-168(c)) — byte-backed like the raw-byte dummies, but a
69    /// distinct variant so `Bool` is never read back as `U8`.
70    Bool(&'a [u8]),
71}
72
73/// Backwards-compatibility alias. New code should use [`HostBufferRef`].
74pub type CpuStorageRef<'a> = HostBufferRef<'a>;
75
76/// A CPU device handle (unit struct — the CPU needs no state).
77#[derive(Debug, Clone)]
78pub struct CpuDevice;
79
80impl HostBuffer {
81    /// Returns the [`DType`] of the elements stored.
82    pub fn dtype(&self) -> DType {
83        match self {
84            Self::U8(_) => DType::U8,
85            Self::I8(_) => DType::I8,
86            Self::U32(_) => DType::U32,
87            Self::I16(_) => DType::I16,
88            Self::I32(_) => DType::I32,
89            Self::I64(_) => DType::I64,
90            Self::BF16(_) => DType::BF16,
91            Self::F16(_) => DType::F16,
92            Self::F32(_) => DType::F32,
93            Self::F64(_) => DType::F64,
94            Self::F8E4M3(_) => DType::F8E4M3,
95            Self::F6E2M3(_) => DType::F6E2M3,
96            Self::F6E3M2(_) => DType::F6E3M2,
97            Self::F4(_) => DType::F4,
98            Self::F8E8M0(_) => DType::F8E8M0,
99            Self::Bool(_) => DType::Bool,
100        }
101    }
102
103    /// Returns a typed slice of the stored data.
104    pub fn as_slice<D: HostDType>(&self) -> crate::Result<&[D]> {
105        D::cpu_storage_as_slice(self)
106    }
107
108    /// Concatenates multiple storages of the same dtype into one.
109    pub fn concat(storages: &[CpuStorage]) -> crate::Result<CpuStorage> {
110        let storage0 = &storages[0];
111        let s = match storage0 {
112            Self::U8(_) => {
113                let storages = storages
114                    .iter()
115                    .map(|s| match s {
116                        Self::U8(s) => Ok(s.as_slice()),
117                        _ => crate::bail!("dtype mismatch"),
118                    })
119                    .collect::<crate::Result<Vec<_>>>()?;
120                Self::U8(storages.concat())
121            }
122            Self::I8(_) => {
123                let storages = storages
124                    .iter()
125                    .map(|s| match s {
126                        Self::I8(s) => Ok(s.as_slice()),
127                        _ => crate::bail!("dtype mismatch"),
128                    })
129                    .collect::<crate::Result<Vec<_>>>()?;
130                Self::I8(storages.concat())
131            }
132            Self::U32(_) => {
133                let storages = storages
134                    .iter()
135                    .map(|s| match s {
136                        Self::U32(s) => Ok(s.as_slice()),
137                        _ => crate::bail!("dtype mismatch"),
138                    })
139                    .collect::<crate::Result<Vec<_>>>()?;
140                Self::U32(storages.concat())
141            }
142            Self::I16(_) => {
143                let storages = storages
144                    .iter()
145                    .map(|s| match s {
146                        Self::I16(s) => Ok(s.as_slice()),
147                        _ => crate::bail!("dtype mismatch"),
148                    })
149                    .collect::<crate::Result<Vec<_>>>()?;
150                Self::I16(storages.concat())
151            }
152            Self::I32(_) => {
153                let storages = storages
154                    .iter()
155                    .map(|s| match s {
156                        Self::I32(s) => Ok(s.as_slice()),
157                        _ => crate::bail!("dtype mismatch"),
158                    })
159                    .collect::<crate::Result<Vec<_>>>()?;
160                Self::I32(storages.concat())
161            }
162            Self::I64(_) => {
163                let storages = storages
164                    .iter()
165                    .map(|s| match s {
166                        Self::I64(s) => Ok(s.as_slice()),
167                        _ => crate::bail!("dtype mismatch"),
168                    })
169                    .collect::<crate::Result<Vec<_>>>()?;
170                Self::I64(storages.concat())
171            }
172            Self::BF16(_) => {
173                let storages = storages
174                    .iter()
175                    .map(|s| match s {
176                        Self::BF16(s) => Ok(s.as_slice()),
177                        _ => crate::bail!("dtype mismatch"),
178                    })
179                    .collect::<crate::Result<Vec<_>>>()?;
180                Self::BF16(storages.concat())
181            }
182            Self::F16(_) => {
183                let storages = storages
184                    .iter()
185                    .map(|s| match s {
186                        Self::F16(s) => Ok(s.as_slice()),
187                        _ => crate::bail!("dtype mismatch"),
188                    })
189                    .collect::<crate::Result<Vec<_>>>()?;
190                Self::F16(storages.concat())
191            }
192            Self::F32(_) => {
193                let storages = storages
194                    .iter()
195                    .map(|s| match s {
196                        Self::F32(s) => Ok(s.as_slice()),
197                        _ => crate::bail!("dtype mismatch"),
198                    })
199                    .collect::<crate::Result<Vec<_>>>()?;
200                Self::F32(storages.concat())
201            }
202            Self::F64(_) => {
203                let storages = storages
204                    .iter()
205                    .map(|s| match s {
206                        Self::F64(s) => Ok(s.as_slice()),
207                        _ => crate::bail!("dtype mismatch"),
208                    })
209                    .collect::<crate::Result<Vec<_>>>()?;
210                Self::F64(storages.concat())
211            }
212            Self::F8E4M3(_) => {
213                let storages = storages
214                    .iter()
215                    .map(|s| match s {
216                        Self::F8E4M3(s) => Ok(s.as_slice()),
217                        _ => crate::bail!("dtype mismatch"),
218                    })
219                    .collect::<crate::Result<Vec<_>>>()?;
220                Self::F8E4M3(storages.concat())
221            }
222            Self::F6E2M3(_) => {
223                let storages = storages
224                    .iter()
225                    .map(|s| match s {
226                        Self::F6E2M3(s) => Ok(s.as_slice()),
227                        _ => crate::bail!("dtype mismatch"),
228                    })
229                    .collect::<crate::Result<Vec<_>>>()?;
230                Self::F6E2M3(storages.concat())
231            }
232            Self::F6E3M2(_) => {
233                let storages = storages
234                    .iter()
235                    .map(|s| match s {
236                        Self::F6E3M2(s) => Ok(s.as_slice()),
237                        _ => crate::bail!("dtype mismatch"),
238                    })
239                    .collect::<crate::Result<Vec<_>>>()?;
240                Self::F6E3M2(storages.concat())
241            }
242            Self::F4(_) => {
243                let storages = storages
244                    .iter()
245                    .map(|s| match s {
246                        Self::F4(s) => Ok(s.as_slice()),
247                        _ => crate::bail!("dtype mismatch"),
248                    })
249                    .collect::<crate::Result<Vec<_>>>()?;
250                Self::F4(storages.concat())
251            }
252            Self::F8E8M0(_) => {
253                let storages = storages
254                    .iter()
255                    .map(|s| match s {
256                        Self::F8E8M0(s) => Ok(s.as_slice()),
257                        _ => crate::bail!("dtype mismatch"),
258                    })
259                    .collect::<crate::Result<Vec<_>>>()?;
260                Self::F8E8M0(storages.concat())
261            }
262            Self::Bool(_) => {
263                let storages = storages
264                    .iter()
265                    .map(|s| match s {
266                        Self::Bool(s) => Ok(s.as_slice()),
267                        _ => crate::bail!("dtype mismatch"),
268                    })
269                    .collect::<crate::Result<Vec<_>>>()?;
270                Self::Bool(storages.concat())
271            }
272        };
273        Ok(s)
274    }
275
276    /// Borrow this buffer as a [`HostBufferRef`] (zero-copy).
277    ///
278    /// The ref form is the common shape for `HostStorage::as_host_buffer_ref`;
279    /// owned-buffer implementors delegate through here.
280    pub fn as_ref(&self) -> HostBufferRef<'_> {
281        match self {
282            Self::U8(v) => HostBufferRef::U8(v),
283            Self::I8(v) => HostBufferRef::I8(v),
284            Self::U32(v) => HostBufferRef::U32(v),
285            Self::I16(v) => HostBufferRef::I16(v),
286            Self::I32(v) => HostBufferRef::I32(v),
287            Self::I64(v) => HostBufferRef::I64(v),
288            Self::BF16(v) => HostBufferRef::BF16(v),
289            Self::F16(v) => HostBufferRef::F16(v),
290            Self::F32(v) => HostBufferRef::F32(v),
291            Self::F64(v) => HostBufferRef::F64(v),
292            Self::F8E4M3(v) => HostBufferRef::F8E4M3(v),
293            Self::F6E2M3(v) => HostBufferRef::F6E2M3(v),
294            Self::F6E3M2(v) => HostBufferRef::F6E3M2(v),
295            Self::F4(v) => HostBufferRef::F4(v),
296            Self::F8E8M0(v) => HostBufferRef::F8E8M0(v),
297            Self::Bool(v) => HostBufferRef::Bool(v),
298        }
299    }
300}
301
302impl<'a> HostBufferRef<'a> {
303    /// The [`DType`] of the borrowed data.
304    pub fn dtype(&self) -> DType {
305        match self {
306            Self::U8(_) => DType::U8,
307            Self::I8(_) => DType::I8,
308            Self::U32(_) => DType::U32,
309            Self::I16(_) => DType::I16,
310            Self::I32(_) => DType::I32,
311            Self::I64(_) => DType::I64,
312            Self::BF16(_) => DType::BF16,
313            Self::F16(_) => DType::F16,
314            Self::F32(_) => DType::F32,
315            Self::F64(_) => DType::F64,
316            Self::F8E4M3(_) => DType::F8E4M3,
317            Self::F6E2M3(_) => DType::F6E2M3,
318            Self::F6E3M2(_) => DType::F6E3M2,
319            Self::F4(_) => DType::F4,
320            Self::F8E8M0(_) => DType::F8E8M0,
321            Self::Bool(_) => DType::Bool,
322        }
323    }
324
325    /// Number of elements in the borrowed slice.
326    pub fn len(&self) -> usize {
327        match self {
328            Self::U8(v) => v.len(),
329            Self::I8(v) => v.len(),
330            Self::U32(v) => v.len(),
331            Self::I16(v) => v.len(),
332            Self::I32(v) => v.len(),
333            Self::I64(v) => v.len(),
334            Self::BF16(v) => v.len(),
335            Self::F16(v) => v.len(),
336            Self::F32(v) => v.len(),
337            Self::F64(v) => v.len(),
338            Self::F8E4M3(v) => v.len(),
339            Self::F6E2M3(v) => v.len(),
340            Self::F6E3M2(v) => v.len(),
341            Self::F4(v) => v.len(),
342            Self::F8E8M0(v) => v.len(),
343            Self::Bool(v) => v.len(),
344        }
345    }
346
347    /// `true` if the borrowed slice is empty.
348    pub fn is_empty(&self) -> bool {
349        self.len() == 0
350    }
351
352    /// Materialize into an owned [`HostBuffer`] (copies the slice).
353    ///
354    /// Use when a caller needs ownership and the storage backend is not
355    /// itself buffer-owning (mmap-backed, pinned, remote, …). For
356    /// `CpuBackendStorage` the cheaper path is `into_host_buffer`, which
357    /// transfers the existing `Vec<T>` without allocating.
358    pub fn to_owned(&self) -> HostBuffer {
359        match self {
360            Self::U8(v) => HostBuffer::U8(v.to_vec()),
361            Self::I8(v) => HostBuffer::I8(v.to_vec()),
362            Self::U32(v) => HostBuffer::U32(v.to_vec()),
363            Self::I16(v) => HostBuffer::I16(v.to_vec()),
364            Self::I32(v) => HostBuffer::I32(v.to_vec()),
365            Self::I64(v) => HostBuffer::I64(v.to_vec()),
366            Self::BF16(v) => HostBuffer::BF16(v.to_vec()),
367            Self::F16(v) => HostBuffer::F16(v.to_vec()),
368            Self::F32(v) => HostBuffer::F32(v.to_vec()),
369            Self::F64(v) => HostBuffer::F64(v.to_vec()),
370            Self::F8E4M3(v) => HostBuffer::F8E4M3(v.to_vec()),
371            Self::F6E2M3(v) => HostBuffer::F6E2M3(v.to_vec()),
372            Self::F6E3M2(v) => HostBuffer::F6E3M2(v.to_vec()),
373            Self::F4(v) => HostBuffer::F4(v.to_vec()),
374            Self::F8E8M0(v) => HostBuffer::F8E8M0(v.to_vec()),
375            Self::Bool(v) => HostBuffer::Bool(v.to_vec()),
376        }
377    }
378}
379
380/// Host-buffer conversion contract for a tensor element type.
381///
382/// Split off `WithDType` in B0.4 (the weld break): `WithDType` stays pure
383/// vocabulary, while these `HostBuffer`/`HostBufferRef` conversions live here,
384/// next to the buffer types they produce. Generic code that needs to materialize
385/// a `&[T]` / `Vec<T>` into host storage bounds on `T: HostDType` (which implies
386/// `WithDType`). This breaks the old `dtype <-> cpu_storage` dependency cycle.
387pub trait HostDType: WithDType {
388    fn cpu_storage_ref(data: &[Self]) -> HostBufferRef<'_>;
389    fn to_cpu_storage_owned(data: Vec<Self>) -> HostBuffer;
390
391    fn to_cpu_storage(data: &[Self]) -> HostBuffer {
392        Self::to_cpu_storage_owned(data.to_vec())
393    }
394
395    fn cpu_storage_as_slice(s: &HostBuffer) -> Result<&[Self]>;
396    fn cpu_storage_data(s: HostBuffer) -> Result<Vec<Self>>;
397}
398
399macro_rules! host_dtype {
400    ($ty:ty, $dtype:ident) => {
401        impl HostDType for $ty {
402            fn cpu_storage_ref(data: &[Self]) -> HostBufferRef<'_> {
403                HostBufferRef::$dtype(data)
404            }
405
406            fn to_cpu_storage_owned(data: Vec<Self>) -> HostBuffer {
407                HostBuffer::$dtype(data)
408            }
409
410            fn cpu_storage_data(s: HostBuffer) -> Result<Vec<Self>> {
411                match s {
412                    HostBuffer::$dtype(data) => Ok(data),
413                    _ => Err(Error::UnexpectedDType {
414                        expected: DType::$dtype,
415                        got: s.dtype(),
416                        msg: "unexpected dtype",
417                    }
418                    .bt()),
419                }
420            }
421
422            fn cpu_storage_as_slice(s: &HostBuffer) -> Result<&[Self]> {
423                match s {
424                    HostBuffer::$dtype(data) => Ok(data),
425                    _ => Err(Error::UnexpectedDType {
426                        expected: DType::$dtype,
427                        got: s.dtype(),
428                        msg: "unexpected dtype",
429                    }
430                    .bt()),
431                }
432            }
433        }
434    };
435}
436
437host_dtype!(u8, U8);
438host_dtype!(i8, I8);
439host_dtype!(u32, U32);
440host_dtype!(i16, I16);
441host_dtype!(i32, I32);
442host_dtype!(i64, I64);
443host_dtype!(f16, F16);
444host_dtype!(bf16, BF16);
445host_dtype!(f32, F32);
446host_dtype!(f64, F64);
447host_dtype!(F8E4M3, F8E4M3);
fuel_ir/cpu_storage.rs

fuel_ir/
cpu_storage.rs
