luma_tensor/device/cpu/
storage.rs

1//! CPU storage: one flat `Vec<T>` per precision, wrapped in an enum so the
2//! concrete precision of a `Float`/`Int` tensor is a runtime choice.
3
4use crate::{
5    Bool, DType, Float, Int, Storage,
6    dtype::{BoolDType, FloatDType, IntDType},
7};
8
9use super::Cpu;
10
11/// Backing buffer for a `Float`-kind tensor on the CPU.
12#[derive(Debug, Clone)]
13pub enum CpuFloatStorage {
14    F32(Vec<f32>),
15    F64(Vec<f64>),
16}
17
18impl Storage<Cpu, Float> for CpuFloatStorage {
19    fn dtype(&self) -> FloatDType {
20        match self {
21            Self::F32(_) => FloatDType::F32,
22            Self::F64(_) => FloatDType::F64,
23        }
24    }
25
26    fn device(&self) -> &Cpu {
27        &Cpu
28    }
29}
30
31impl CpuFloatStorage {
32    pub fn dtype(&self) -> DType {
33        match self {
34            CpuFloatStorage::F32(_) => DType::F32,
35            CpuFloatStorage::F64(_) => DType::F64,
36        }
37    }
38
39    pub fn len(&self) -> usize {
40        match self {
41            CpuFloatStorage::F32(v) => v.len(),
42            CpuFloatStorage::F64(v) => v.len(),
43        }
44    }
45
46    pub fn is_empty(&self) -> bool {
47        self.len() == 0
48    }
49}
50
51/// Backing buffer for an `Int`-kind tensor on the CPU.
52#[derive(Debug, Clone)]
53pub enum CpuIntStorage {
54    I32(Vec<i32>),
55    U32(Vec<u32>),
56    U8(Vec<u8>),
57}
58
59impl Storage<Cpu, Int> for CpuIntStorage {
60    fn dtype(&self) -> IntDType {
61        match self {
62            Self::I32(_) => IntDType::I32,
63            Self::U32(_) => IntDType::U32,
64            Self::U8(_) => IntDType::U8,
65        }
66    }
67
68    fn device(&self) -> &Cpu {
69        &Cpu
70    }
71}
72
73impl CpuIntStorage {
74    pub fn dtype(&self) -> DType {
75        match self {
76            CpuIntStorage::I32(_) => DType::I32,
77            CpuIntStorage::U32(_) => DType::U32,
78            CpuIntStorage::U8(_) => DType::U8,
79        }
80    }
81
82    pub fn len(&self) -> usize {
83        match self {
84            CpuIntStorage::I32(v) => v.len(),
85            CpuIntStorage::U32(v) => v.len(),
86            CpuIntStorage::U8(v) => v.len(),
87        }
88    }
89
90    pub fn is_empty(&self) -> bool {
91        self.len() == 0
92    }
93}
94
95/// Backing buffer for a `Bool`-kind tensor on the CPU.
96#[derive(Debug, Clone)]
97pub struct CpuBoolStorage(pub Vec<bool>);
98
99impl CpuBoolStorage {
100    pub fn len(&self) -> usize {
101        self.0.len()
102    }
103
104    pub fn is_empty(&self) -> bool {
105        self.0.is_empty()
106    }
107}
108
109impl Storage<Cpu, Bool> for CpuBoolStorage {
110    fn dtype(&self) -> BoolDType {
111        BoolDType::Bool
112    }
113
114    fn device(&self) -> &Cpu {
115        &Cpu
116    }
117}
luma_tensor/device/cpu/storage.rs

luma_tensor/device/cpu/
storage.rs
