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}