Skip to main content

oxmera_tensor/
storage.rs

1//! Storage: the owned buffer behind one or more tensors.
2
3use oxmera_core::{DType, Device};
4
5/// An owned, reference-counted buffer of elements on one device.
6///
7/// Storage is untyped at the buffer level — a flat byte buffer plus the
8/// `DType` that says how to read it. Typed access is the business of the
9/// backend that allocated it; nothing else may reinterpret the bytes.
10/// Multiple tensors (views) may share one storage; storage never knows how
11/// many.
12#[derive(Debug)]
13pub struct Storage {
14    data: StorageData,
15    dtype: DType,
16    device: Device,
17}
18
19/// Where the bytes actually live.
20///
21/// One variant per device family. GPU variants hold opaque backend handles
22/// once those backends exist; they are deliberately absent until then so
23/// this enum never carries a stub it cannot honor.
24#[derive(Debug)]
25enum StorageData {
26    /// Host memory for the CPU reference backend.
27    // The expect below is a tripwire: solving exercise A1 makes the
28    // constructors build this variant, the expectation stops holding, and
29    // clippy -D warnings forces its removal.
30    #[expect(dead_code, reason = "constructed once exercise A1 is solved")]
31    Cpu(Vec<u8>),
32}
33
34impl Storage {
35    /// Allocate zero-initialized CPU storage for `numel` elements of
36    /// `dtype`.
37    pub fn cpu_zeros(numel: usize, dtype: DType) -> Self {
38        let _ = (numel, dtype);
39        todo!("exercise A1: shape and strides")
40    }
41
42    /// Wrap raw host bytes as CPU storage. The byte length must be a
43    /// multiple of the dtype size; the caller asserts the encoding.
44    pub fn cpu_from_bytes(bytes: Vec<u8>, dtype: DType) -> Self {
45        let _ = (bytes, dtype);
46        todo!("exercise A1: shape and strides")
47    }
48
49    /// The element type of this buffer.
50    pub fn dtype(&self) -> DType {
51        self.dtype
52    }
53
54    /// The device this buffer lives on.
55    pub fn device(&self) -> Device {
56        self.device
57    }
58
59    /// The raw bytes, when the storage is on the CPU.
60    ///
61    /// Backends other than the CPU reference return `None`; they expose
62    /// their own typed access instead.
63    pub fn cpu_bytes(&self) -> Option<&[u8]> {
64        match &self.data {
65            StorageData::Cpu(bytes) => Some(bytes),
66        }
67    }
68}