luma_tensor/
dtype.rs

1use crate::{Device, FloatMeta, grad::TensorMeta};
2use std::fmt::Debug;
3
4/// Binds a tensor kind to its per-device storage type and its autograd metadata.
5///
6/// This is the bridge between the compile-time kind marker and the device's
7/// runtime storage: `Tensor<D, K>` stores `K::Storage` and `K::Meta`.
8pub trait DTypeKind<D: Device>: Sized {
9    type Scalar: Send + Sync + Clone + Copy + 'static;
10    type Storage: Storage<D, Self>;
11    type Meta: TensorMeta<D, Self>;
12    type DType: Send + Sync + Clone + Copy + 'static + PartialEq + Eq + Debug + Default;
13
14    /// Runtime discriminant for this kind.
15    const KIND: KindTag;
16}
17
18pub trait Storage<D: Device, K: DTypeKind<D>>: Send + Sync + 'static {
19    fn dtype(&self) -> K::DType;
20    fn device(&self) -> &D;
21}
22
23/// Runtime element type. The tensor *kind* (`Float`/`Int`/`Bool`) is a compile-time
24/// generic; the concrete precision inside a kind is this enum, decided at runtime.
25#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
26pub enum DType {
27    // Float kind
28    F32,
29    F64,
30    // Int kind
31    I32,
32    U8,
33    U32,
34    // Bool kind
35    Bool,
36}
37
38impl DType {
39    pub fn is_float(&self) -> bool {
40        matches!(self, DType::F32 | DType::F64)
41    }
42
43    pub fn is_int(&self) -> bool {
44        matches!(self, DType::I32 | DType::U8 | DType::U32)
45    }
46
47    pub fn is_bool(&self) -> bool {
48        matches!(self, DType::Bool)
49    }
50
51    /// The tensor kind this dtype belongs to.
52    pub fn kind(&self) -> KindTag {
53        match self {
54            DType::F32 | DType::F64 => KindTag::Float,
55            DType::I32 | DType::U8 | DType::U32 => KindTag::Int,
56            DType::Bool => KindTag::Bool,
57        }
58    }
59
60    /// Size of a single element in bytes.
61    pub fn size_in_bytes(&self) -> usize {
62        match self {
63            DType::F32 => 4,
64            DType::F64 => 8,
65            DType::I32 => 4,
66            DType::U8 => 1,
67            DType::U32 => 4,
68            DType::Bool => 1,
69        }
70    }
71
72    /// Convert to `FloatDType` if this is a float dtype, otherwise panic.
73    pub fn as_float(&self) -> FloatDType {
74        match self {
75            DType::F32 => FloatDType::F32,
76            DType::F64 => FloatDType::F64,
77            _ => panic!("DType::{:?} is not a float dtype", self),
78        }
79    }
80
81    /// Convert to `IntDType` if this is an int dtype, otherwise panic.
82    pub fn as_int(&self) -> IntDType {
83        match self {
84            DType::I32 => IntDType::I32,
85            DType::U8 => IntDType::U8,
86            DType::U32 => IntDType::U32,
87            _ => panic!("DType::{:?} is not an int dtype", self),
88        }
89    }
90}
91
92/// Runtime discriminant mirroring the compile-time kind markers.
93#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
94pub enum KindTag {
95    Float,
96    Int,
97    Bool,
98}
99
100/// The three tensor kinds, used only as zero-sized compile-time markers on
101/// `Tensor<D, K>`. They never hold data; behaviour lives in the device ops traits.
102pub struct Float;
103pub struct Int;
104pub struct Bool;
105
106impl<D: Device> DTypeKind<D> for Float {
107    type Scalar = f64;
108    type Storage = D::FloatStorage;
109    type Meta = FloatMeta<D>;
110    type DType = FloatDType;
111    const KIND: KindTag = KindTag::Float;
112}
113
114impl<D: Device> DTypeKind<D> for Int {
115    type Scalar = i64;
116    type Storage = D::IntStorage;
117    type Meta = ();
118    type DType = IntDType;
119    const KIND: KindTag = KindTag::Int;
120}
121
122impl<D: Device> DTypeKind<D> for Bool {
123    type Scalar = bool;
124    type Storage = D::BoolStorage;
125    type Meta = ();
126    type DType = BoolDType;
127    const KIND: KindTag = KindTag::Bool;
128}
129
130#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
131pub enum FloatDType {
132    #[default]
133    F32,
134    F64,
135}
136
137#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
138pub enum IntDType {
139    #[default]
140    I32,
141    U8,
142    U32,
143}
144
145#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
146pub enum BoolDType {
147    #[default]
148    Bool,
149}
150
151impl Into<DType> for FloatDType {
152    fn into(self) -> DType {
153        match self {
154            FloatDType::F32 => DType::F32,
155            FloatDType::F64 => DType::F64,
156        }
157    }
158}
159
160impl Into<DType> for IntDType {
161    fn into(self) -> DType {
162        match self {
163            IntDType::I32 => DType::I32,
164            IntDType::U8 => DType::U8,
165            IntDType::U32 => DType::U32,
166        }
167    }
168}
169
170impl Into<DType> for BoolDType {
171    fn into(self) -> DType {
172        match self {
173            BoolDType::Bool => DType::Bool,
174        }
175    }
176}
luma_tensor/dtype.rs

luma_tensor/
dtype.rs
