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::{DType, TensorRank, TensorScalar, 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: TensorScalar>(
84    op: &'static str,
85    tensor: &'a TypedTensor<T>,
86) -> Result<&'a [T]> {
87    if tensor.backend_buffer().is_some() {
88        return Err(cpu_backend_buffer_error(op));
89    }
90    tensor.host_data()
91}
92
93pub(crate) fn typed_view<'a, T: Copy + TensorScalar>(
94    op: &'static str,
95    tensor: &'a TypedTensor<T>,
96) -> Result<StridedView<'a, T>> {
97    if tensor.backend_buffer().is_some() {
98        return Err(cpu_backend_buffer_error(op));
99    }
100    let data = tensor.host_data()?;
101    let strides = kernel_col_major_strides(tensor.shape());
102    StridedView::new(data, tensor.shape(), &strides, 0)
103        .map_err(|err| Error::backend_source(op, err))
104}
105
106pub(crate) fn typed_view_from_view<'a, T: Copy + 'static, R: TensorRank>(
107    op: &'static str,
108    view: &TypedTensorView<'a, T, R>,
109) -> Result<StridedView<'a, T>> {
110    if view.backend_buffer().is_some() {
111        return Err(cpu_backend_buffer_error(op));
112    }
113    StridedView::new(
114        view.host_storage()?,
115        view.shape(),
116        view.strides(),
117        view.offset(),
118    )
119    .map_err(|err| Error::backend_source(op, err))
120}
121
122#[cfg(test)]
123pub(crate) fn materialize_tensor_read(
124    buffers: &mut BufferPool,
125    op: &'static str,
126    input: TensorRead<'_>,
127) -> Result<Tensor> {
128    match input {
129        TensorRead::Tensor(tensor) => clone_host_tensor_read(op, tensor),
130        TensorRead::View(view) => materialize_tensor_view(buffers, op, view),
131    }
132}
133
134#[cfg(test)]
135fn clone_host_tensor_read(op: &'static str, tensor: &Tensor) -> Result<Tensor> {
136    macro_rules! clone_host {
137        ($variant:ident, $tensor:expr) => {{
138            typed_host_data(op, $tensor)?;
139            Ok(Tensor::$variant($tensor.duplicate()?))
140        }};
141    }
142
143    match tensor {
144        Tensor::F32(tensor) => clone_host!(F32, tensor),
145        Tensor::F64(tensor) => clone_host!(F64, tensor),
146        Tensor::I32(tensor) => clone_host!(I32, tensor),
147        Tensor::I64(tensor) => clone_host!(I64, tensor),
148        Tensor::Bool(tensor) => clone_host!(Bool, tensor),
149        Tensor::C32(tensor) => clone_host!(C32, tensor),
150        Tensor::C64(tensor) => clone_host!(C64, tensor),
151    }
152}
153
154#[cfg(test)]
155fn materialize_tensor_view(
156    buffers: &mut BufferPool,
157    op: &'static str,
158    view: TensorView<'_>,
159) -> Result<Tensor> {
160    macro_rules! materialize {
161        ($variant:ident, $view:expr) => {{
162            Ok(Tensor::$variant(typed_materialize_view_for_tests(
163                buffers, &$view, op,
164            )?))
165        }};
166    }
167
168    match view {
169        TensorView::F32(view) => materialize!(F32, view),
170        TensorView::F64(view) => materialize!(F64, view),
171        TensorView::I32(view) => materialize!(I32, view),
172        TensorView::I64(view) => materialize!(I64, view),
173        TensorView::Bool(view) => materialize!(Bool, view),
174        TensorView::C32(view) => materialize!(C32, view),
175        TensorView::C64(view) => materialize!(C64, view),
176    }
177}
178
179#[cfg(test)]
180fn typed_materialize_view_for_tests<T, R>(
181    buffers: &mut BufferPool,
182    view: &TypedTensorView<'_, T, R>,
183    op: &'static str,
184) -> Result<TypedTensor<T, R>>
185where
186    T: Copy + Clone + PoolScalar + 'static,
187    R: TensorRank,
188{
189    if view.backend_buffer().is_some() {
190        return Err(cpu_backend_buffer_error(op));
191    }
192    let src: StridedView<'_, T, Identity> = StridedView::new(
193        view.host_storage()?,
194        view.shape(),
195        view.strides(),
196        view.offset(),
197    )
198    .map_err(|err| Error::backend_source(op, err))?;
199    let mut out = PooledUninitOutput::<T>::new(buffers, view.shape().to_vec())?;
200    map_into(&mut out.as_uninit_view_mut()?, &src, |x| {
201        MaybeUninit::new(x)
202    })
203    .map_err(|err| Error::backend_source(op, err))?;
204    // SAFETY: the successful materialize map replay writes every logical destination element and retains no destination view.
205    let out = unsafe { out.assume_init_as::<R>()? };
206    let shape = R::shape_from_vec(view.shape().to_vec().into())
207        .map_err(|err| Error::backend_source(op, err))?;
208    let mut tensor = TypedTensor::from_vec_col_major(shape, out.into_vec_col_major()?.1)?;
209    tensor.set_placement(view.placement().clone());
210    Ok(tensor)
211}
tenferro_internal_cpu_kernels/lib.rs

tenferro_internal_cpu_kernels/
lib.rs
