luma_tensor/device/cpu/
mod.rs

1pub mod kernels;
2mod ops;
3mod storage;
4pub use storage::*;
5
6use crate::{DType, Device};
7
8/// The CPU device: a zero-sized type tag. All ops are associated functions.
9#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
10pub struct Cpu;
11
12impl Device for Cpu {
13    type FloatStorage = CpuFloatStorage;
14    type IntStorage = CpuIntStorage;
15    type BoolStorage = CpuBoolStorage;
16
17    fn name(&self) -> String {
18        "cpu".into()
19    }
20}
21
22// ============================================================================
23// Dispatch macros: map a storage enum to its concrete element type, run the
24// generic kernel, and re-wrap the result.
25// ============================================================================
26
27/// Run `$body` with `$data` bound to the inner `&Vec<t>` of a `CpuFloatStorage`,
28/// then wrap the returned `Vec<t>` back into a `CpuFloatStorage`.
29#[macro_export]
30macro_rules! dispatch_float {
31    ($storage:expr, |$data:ident| $body:expr) => {
32        match $storage {
33            $crate::CpuFloatStorage::F32($data) => $crate::CpuFloatStorage::F32($body),
34            $crate::CpuFloatStorage::F64($data) => $crate::CpuFloatStorage::F64($body),
35        }
36    };
37}
38
39/// Like [`dispatch_float`] but `$body` yields a non-storage value (e.g. Vec<bool>).
40#[macro_export]
41macro_rules! dispatch_float_raw {
42    ($storage:expr, |$data:ident| $body:expr) => {
43        match $storage {
44            $crate::CpuFloatStorage::F32($data) => $body,
45            $crate::CpuFloatStorage::F64($data) => $body,
46        }
47    };
48}
49
50/// Dispatch two float storages of the SAME variant; errors on mismatch.
51#[macro_export]
52macro_rules! dispatch_float2 {
53    ($lhs:expr, $rhs:expr, $op:literal, |$a:ident, $b:ident| $body:expr) => {
54        match ($lhs, $rhs) {
55            ($crate::CpuFloatStorage::F32($a), $crate::CpuFloatStorage::F32($b)) => Ok($crate::CpuFloatStorage::F32($body)),
56            ($crate::CpuFloatStorage::F64($a), $crate::CpuFloatStorage::F64($b)) => Ok($crate::CpuFloatStorage::F64($body)),
57            (l, r) => Err($crate::Error::DTypeMismatch { lhs: l.dtype(), rhs: r.dtype(), op: $op }),
58        }
59    };
60}
61
62/// Dispatch two float storages of the SAME variant, `$body` yields a raw value.
63#[macro_export]
64macro_rules! dispatch_float2_raw {
65    ($lhs:expr, $rhs:expr, $op:literal, |$a:ident, $b:ident| $body:expr) => {
66        match ($lhs, $rhs) {
67            ($crate::CpuFloatStorage::F32($a), $crate::CpuFloatStorage::F32($b)) => Ok($body),
68            ($crate::CpuFloatStorage::F64($a), $crate::CpuFloatStorage::F64($b)) => Ok($body),
69            (l, r) => Err($crate::Error::DTypeMismatch { lhs: l.dtype(), rhs: r.dtype(), op: $op }),
70        }
71    };
72}
73
74#[macro_export]
75macro_rules! dispatch_int {
76    ($storage:expr, |$data:ident| $body:expr) => {
77        match $storage {
78            $crate::CpuIntStorage::I32($data) => $crate::CpuIntStorage::I32($body),
79            $crate::CpuIntStorage::U32($data) => $crate::CpuIntStorage::U32($body),
80            $crate::CpuIntStorage::U8($data) => $crate::CpuIntStorage::U8($body),
81        }
82    };
83}
84
85#[macro_export]
86macro_rules! dispatch_int_raw {
87    ($storage:expr, |$data:ident| $body:expr) => {
88        match $storage {
89            $crate::CpuIntStorage::I32($data) => $body,
90            $crate::CpuIntStorage::U32($data) => $body,
91            $crate::CpuIntStorage::U8($data) => $body,
92        }
93    };
94}
95
96#[macro_export]
97macro_rules! dispatch_int2 {
98    ($lhs:expr, $rhs:expr, $op:literal, |$a:ident, $b:ident| $body:expr) => {
99        match ($lhs, $rhs) {
100            ($crate::CpuIntStorage::I32($a), $crate::CpuIntStorage::I32($b)) => Ok($crate::CpuIntStorage::I32($body)),
101            ($crate::CpuIntStorage::U32($a), $crate::CpuIntStorage::U32($b)) => Ok($crate::CpuIntStorage::U32($body)),
102            ($crate::CpuIntStorage::U8($a), $crate::CpuIntStorage::U8($b)) => Ok($crate::CpuIntStorage::U8($body)),
103            (l, r) => Err($crate::Error::DTypeMismatch { lhs: l.dtype(), rhs: r.dtype(), op: $op }),
104        }
105    };
106}
107
108#[macro_export]
109macro_rules! dispatch_int2_raw {
110    ($lhs:expr, $rhs:expr, $op:literal, |$a:ident, $b:ident| $body:expr) => {
111        match ($lhs, $rhs) {
112            ($crate::CpuIntStorage::I32($a), $crate::CpuIntStorage::I32($b)) => Ok($body),
113            ($crate::CpuIntStorage::U32($a), $crate::CpuIntStorage::U32($b)) => Ok($body),
114            ($crate::CpuIntStorage::U8($a), $crate::CpuIntStorage::U8($b)) => Ok($body),
115            (l, r) => Err($crate::Error::DTypeMismatch { lhs: l.dtype(), rhs: r.dtype(), op: $op }),
116        }
117    };
118}
119
120/// Read an int storage's elements (in `layout` order) as `usize`, mapping the
121/// dtype's MAX sentinel to `kernels::indexing::PAD`. Used by indexing kernels.
122pub(crate) fn int_ids_as_usize(storage: &CpuIntStorage, layout: &crate::Layout) -> Vec<usize> {
123    use kernels::element::{CpuInt, CpuNum};
124    use kernels::indexing::PAD;
125    macro_rules! collect {
126        ($data:expr) => {
127            layout
128                .storage_indices()
129                .map(|i| {
130                    let v = $data[i];
131                    if v == CpuInt::MAX { PAD } else { v.to_usize() }
132                })
133                .collect()
134        };
135    }
136    match storage {
137        CpuIntStorage::I32(d) => collect!(d),
138        CpuIntStorage::U32(d) => collect!(d),
139        CpuIntStorage::U8(d) => collect!(d),
140    }
141}
142
143/// Build an int storage of the given dtype from `usize` indices.
144pub(crate) fn usize_to_int_storage(data: &[usize], dtype: DType) -> CpuIntStorage {
145    match dtype {
146        DType::I32 => CpuIntStorage::I32(data.iter().map(|&v| v as i32).collect()),
147        DType::U32 => CpuIntStorage::U32(data.iter().map(|&v| v as u32).collect()),
148        DType::U8 => CpuIntStorage::U8(data.iter().map(|&v| v as u8).collect()),
149        _ => CpuIntStorage::U32(data.iter().map(|&v| v as u32).collect()),
150    }
151}
luma_tensor/device/cpu/mod.rs

luma_tensor/device/cpu/
mod.rs
