tenferro_internal_cpu_kernels/
lib.rs

1#![doc(hidden)]
2
3//! Internal CPU kernel and scratch-pool implementation crate.
4
5#[cfg(test)]
6use std::mem::MaybeUninit;
7
8pub type Result<T> = tenferro_tensor::Result<T>;
9pub use tenferro_tensor::{CacheStats, Error, ErrorKind};
10
11pub mod buffer_pool;
12mod pooled_uninit_output;
13/// Canonical pooled full-overwrite owner shared with the sibling `tenferro-cpu`
14/// crate. This is the minimum cross-crate surface required for one ownership
15/// contract; `pub(crate)` cannot cross that crate boundary. Construction only
16/// exposes `MaybeUninit` storage, with initialization completed by an explicit
17/// unsafe handoff.
18pub use pooled_uninit_output::PooledUninitOutput;
19pub mod elementwise;
20
21use num_complex::{Complex32, Complex64};
22use strided_kernel::{col_major_strides as kernel_col_major_strides, StridedView};
23#[cfg(test)]
24use strided_kernel::{map_into, Identity};
25use tenferro_tensor::{Buffer, DType, TensorRank, TypedTensor, TypedTensorView};
26#[cfg(test)]
27use tenferro_tensor::{Tensor, TensorRead, TensorView};
28
29#[cfg(test)]
30use crate::buffer_pool::{BufferPool, PoolScalar};
31
32pub(crate) fn cpu_backend_buffer_error(op: &'static str) -> Error {
33    Error::runtime_state(
34        op,
35        "CPU backend received backend buffer; download to host before CPU execution",
36    )
37}
38
39#[derive(Debug, thiserror::Error)]
40pub(crate) enum CpuNumericalError {
41    #[error("{op} detected division by zero for dtype {dtype:?}")]
42    DivisionByZero { op: &'static str, dtype: DType },
43}
44
45pub(crate) fn cpu_division_by_zero(op: &'static str, dtype: DType) -> Error {
46    Error::extension(
47        op,
48        "cpu",
49        ErrorKind::NumericalFailure,
50        CpuNumericalError::DivisionByZero { op, dtype },
51    )
52}
53
54#[doc(hidden)]
55pub trait ConjElem {
56    fn conj_elem(self) -> Self;
57}
58
59impl ConjElem for f32 {
60    fn conj_elem(self) -> Self {
61        self
62    }
63}
64
65impl ConjElem for f64 {
66    fn conj_elem(self) -> Self {
67        self
68    }
69}
70
71impl ConjElem for Complex32 {
72    fn conj_elem(self) -> Self {
73        self.conj()
74    }
75}
76
77impl ConjElem for Complex64 {
78    fn conj_elem(self) -> Self {
79        self.conj()
80    }
81}
82
83pub(crate) fn typed_host_data<'a, T>(
84    op: &'static str,
85    tensor: &'a TypedTensor<T>,
86) -> Result<&'a [T]> {
87    match tensor.buffer() {
88        Buffer::Host(data) => Ok(data.as_slice()),
89        Buffer::Backend(_) => Err(cpu_backend_buffer_error(op)),
90    }
91}
92
93pub(crate) fn typed_view<'a, T: Copy>(
94    op: &'static str,
95    tensor: &'a TypedTensor<T>,
96) -> Result<StridedView<'a, T>> {
97    match tensor.buffer() {
98        Buffer::Host(data) => {
99            let strides = kernel_col_major_strides(tensor.shape());
100            StridedView::new(data.as_slice(), tensor.shape(), &strides, 0)
101                .map_err(|err| Error::backend_source(op, err))
102        }
103        Buffer::Backend(_) => Err(cpu_backend_buffer_error(op)),
104    }
105}
106
107pub(crate) fn typed_view_from_view<'a, T: Copy + 'static, R: TensorRank>(
108    op: &'static str,
109    view: &TypedTensorView<'a, T, R>,
110) -> Result<StridedView<'a, T>> {
111    if view.backend_buffer().is_some() {
112        return Err(cpu_backend_buffer_error(op));
113    }
114    StridedView::new(
115        view.host_storage()?,
116        view.shape(),
117        view.strides(),
118        view.offset(),
119    )
120    .map_err(|err| Error::backend_source(op, err))
121}
122
123#[cfg(test)]
124pub(crate) fn materialize_tensor_read(
125    buffers: &mut BufferPool,
126    op: &'static str,
127    input: TensorRead<'_>,
128) -> Result<Tensor> {
129    match input {
130        TensorRead::Tensor(tensor) => clone_host_tensor_read(op, tensor),
131        TensorRead::View(view) => materialize_tensor_view(buffers, op, view),
132    }
133}
134
135#[cfg(test)]
136fn clone_host_tensor_read(op: &'static str, tensor: &Tensor) -> Result<Tensor> {
137    macro_rules! clone_host {
138        ($variant:ident, $tensor:expr) => {{
139            typed_host_data(op, $tensor)?;
140            Ok(Tensor::$variant($tensor.clone()))
141        }};
142    }
143
144    match tensor {
145        Tensor::F32(tensor) => clone_host!(F32, tensor),
146        Tensor::F64(tensor) => clone_host!(F64, tensor),
147        Tensor::I32(tensor) => clone_host!(I32, tensor),
148        Tensor::I64(tensor) => clone_host!(I64, tensor),
149        Tensor::Bool(tensor) => clone_host!(Bool, tensor),
150        Tensor::C32(tensor) => clone_host!(C32, tensor),
151        Tensor::C64(tensor) => clone_host!(C64, tensor),
152    }
153}
154
155#[cfg(test)]
156fn materialize_tensor_view(
157    buffers: &mut BufferPool,
158    op: &'static str,
159    view: TensorView<'_>,
160) -> Result<Tensor> {
161    macro_rules! materialize {
162        ($variant:ident, $view:expr) => {{
163            Ok(Tensor::$variant(typed_materialize_view_for_tests(
164                buffers, &$view, op,
165            )?))
166        }};
167    }
168
169    match view {
170        TensorView::F32(view) => materialize!(F32, view),
171        TensorView::F64(view) => materialize!(F64, view),
172        TensorView::I32(view) => materialize!(I32, view),
173        TensorView::I64(view) => materialize!(I64, view),
174        TensorView::Bool(view) => materialize!(Bool, view),
175        TensorView::C32(view) => materialize!(C32, view),
176        TensorView::C64(view) => materialize!(C64, view),
177    }
178}
179
180#[cfg(test)]
181fn typed_materialize_view_for_tests<T, R>(
182    buffers: &mut BufferPool,
183    view: &TypedTensorView<'_, T, R>,
184    op: &'static str,
185) -> Result<TypedTensor<T, R>>
186where
187    T: Copy + Clone + PoolScalar + 'static,
188    R: TensorRank,
189{
190    if view.backend_buffer().is_some() {
191        return Err(cpu_backend_buffer_error(op));
192    }
193    let src: StridedView<'_, T, Identity> = StridedView::new(
194        view.host_storage()?,
195        view.shape(),
196        view.strides(),
197        view.offset(),
198    )
199    .map_err(|err| Error::backend_source(op, err))?;
200    let mut out = PooledUninitOutput::<T>::new(buffers, view.shape().to_vec())?;
201    map_into(&mut out.as_uninit_view_mut()?, &src, |x| {
202        MaybeUninit::new(x)
203    })
204    .map_err(|err| Error::backend_source(op, err))?;
205    // SAFETY: the successful materialize map replay writes every logical destination element and retains no destination view.
206    let out = unsafe { out.assume_init_as::<R>()? };
207    let shape = R::shape_from_vec(view.shape().to_vec().into())
208        .map_err(|err| Error::backend_source(op, err))?;
209    TypedTensor::from_buffer_col_major(
210        shape,
211        Buffer::Host(out.into_vec_col_major()?.1),
212        view.placement().clone(),
213    )
214}
tenferro_internal_cpu_kernels/lib.rs

tenferro_internal_cpu_kernels/
lib.rs
