tenferro_cpu_basic/
lib.rs

1#![doc(hidden)]
2
3//! Shared host CPU resources and low-level adapters.
4
5/// Internal result alias for shared CPU adapters.
6///
7/// # Examples
8///
9/// ```rust
10/// use tenferro_cpu_basic::Result;
11/// let result: Result<()> = Ok(());
12/// assert!(result.is_ok());
13/// ```
14pub type Result<T> = tenferro_tensor::Result<T>;
15use tenferro_tensor::CompareDir;
16pub use tenferro_tensor::{
17    CacheStats, DType, Error, ErrorKind, TensorRank, TensorScalar, TypedTensor, TypedTensorView,
18};
19
20pub mod buffer_pool;
21pub use buffer_pool::{BufferPool, PoolScalar};
22mod pooled_uninit_output;
23pub use pooled_uninit_output::PooledUninitOutput;
24
25use num_complex::{Complex32, Complex64};
26use std::mem::MaybeUninit;
27use std::ptr::NonNull;
28use strided_basic::{
29    col_major_strides as kernel_col_major_strides, ErasedRawStridedMut, ErasedRawStridedPtr,
30    ErasedRawStridedRef, ErasedRawStridedUninitMut, KernelDType, StridedView,
31};
32
33/// Error returned when a CPU-only operation receives backend storage.
34#[doc(hidden)]
35pub fn cpu_backend_buffer_error(op: &'static str) -> Error {
36    Error::runtime_state(
37        op,
38        "CPU backend received backend buffer; download to host before CPU execution",
39    )
40}
41
42#[derive(Debug, thiserror::Error)]
43#[doc(hidden)]
44pub enum CpuNumericalError {
45    #[error("{op} detected division by zero for dtype {dtype:?}")]
46    DivisionByZero { op: &'static str, dtype: DType },
47}
48
49#[doc(hidden)]
50pub fn cpu_division_by_zero(op: &'static str, dtype: DType) -> Error {
51    Error::extension(
52        op,
53        "cpu",
54        ErrorKind::NumericalFailure,
55        CpuNumericalError::DivisionByZero { op, dtype },
56    )
57}
58
59#[doc(hidden)]
60pub trait ConjElem {
61    /// Apply complex conjugation, or leave a real scalar unchanged.
62    ///
63    /// # Examples
64    ///
65    /// ```rust
66    /// use tenferro_cpu_basic::ConjElem;
67    /// assert_eq!(ConjElem::conj_elem(2.0_f64), 2.0);
68    /// ```
69    fn conj_elem(self) -> Self;
70}
71
72impl ConjElem for f32 {
73    fn conj_elem(self) -> Self {
74        self
75    }
76}
77impl ConjElem for f64 {
78    fn conj_elem(self) -> Self {
79        self
80    }
81}
82impl ConjElem for Complex32 {
83    fn conj_elem(self) -> Self {
84        self.conj()
85    }
86}
87impl ConjElem for Complex64 {
88    fn conj_elem(self) -> Self {
89        self.conj()
90    }
91}
92
93#[doc(hidden)]
94pub fn is_complex_dtype(dtype: DType) -> bool {
95    matches!(dtype, DType::C32 | DType::C64)
96}
97
98#[doc(hidden)]
99pub fn ordered_complex_error(op: &'static str) -> Error {
100    Error::unsupported(
101        op,
102        "complex tensors do not have a total order; compute abs/norm explicitly before ordered operations",
103    )
104}
105
106#[doc(hidden)]
107pub fn reject_complex_ordered_dtypes(op: &'static str, dtypes: &[DType]) -> Result<()> {
108    if dtypes.iter().copied().any(is_complex_dtype) {
109        return Err(ordered_complex_error(op));
110    }
111    Ok(())
112}
113
114#[doc(hidden)]
115pub fn reject_complex_unsupported_compare_dtypes(dir: &CompareDir, dtypes: &[DType]) -> Result<()> {
116    if *dir != CompareDir::Eq {
117        reject_complex_ordered_dtypes("compare", dtypes)?;
118    }
119    Ok(())
120}
121
122#[doc(hidden)]
123pub fn typed_host_data<'a, T: TensorScalar>(
124    op: &'static str,
125    tensor: &'a TypedTensor<T>,
126) -> Result<&'a [T]> {
127    if tensor.backend_buffer().is_some() {
128        return Err(cpu_backend_buffer_error(op));
129    }
130    tensor.host_data()
131}
132
133#[doc(hidden)]
134pub fn typed_view<'a, T: Copy + TensorScalar>(
135    op: &'static str,
136    tensor: &'a TypedTensor<T>,
137) -> Result<StridedView<'a, T>> {
138    if tensor.backend_buffer().is_some() {
139        return Err(cpu_backend_buffer_error(op));
140    }
141    let data = tensor.host_data()?;
142    let strides = kernel_col_major_strides(tensor.shape());
143    StridedView::new(data, tensor.shape(), &strides, 0)
144        .map_err(|err| Error::backend_source(op, err))
145}
146
147#[doc(hidden)]
148pub fn typed_view_from_view<'a, T: Copy + 'static, R: TensorRank>(
149    op: &'static str,
150    view: &TypedTensorView<'a, T, R>,
151) -> Result<StridedView<'a, T>> {
152    if view.backend_buffer().is_some() {
153        return Err(cpu_backend_buffer_error(op));
154    }
155    StridedView::new(
156        view.host_storage()?,
157        view.shape(),
158        view.strides(),
159        view.offset(),
160    )
161    .map_err(|err| Error::backend_source(op, err))
162}
163
164/// Construct an erased read-only strided view over initialized typed storage.
165///
166/// # Safety
167/// `dtype` must match the aligned initialized storage. All reachable offsets
168/// must be in bounds and the storage/metadata must remain valid for `'a`.
169#[doc(hidden)]
170pub unsafe fn erased_raw_strided_ref<'a>(
171    dtype: KernelDType,
172    data: &'a [u8],
173    dims: &'a [usize],
174    strides: &'a [isize],
175    offset: isize,
176) -> strided_basic::Result<ErasedRawStridedRef<'a>> {
177    let data_ptr = NonNull::new(data.as_ptr().cast_mut()).unwrap_or_else(NonNull::dangling);
178    // SAFETY: the caller supplies the documented dtype, bounds, alignment and lifetime invariants.
179    unsafe {
180        ErasedRawStridedRef::from_raw_parts(dtype, data_ptr, data.len(), dims, strides, offset)
181    }
182}
183
184/// Construct an erased pointer view over initialized storage.
185///
186/// # Safety
187/// The caller must uphold the dtype, alignment, bounds, lifetime and
188/// initialization requirements of the returned descriptor.
189#[doc(hidden)]
190pub unsafe fn erased_raw_strided_ptr<'a>(
191    dtype: KernelDType,
192    data: &'a [u8],
193    dims: &'a [usize],
194    strides: &'a [isize],
195    offset: isize,
196) -> strided_basic::Result<ErasedRawStridedPtr<'a>> {
197    let data_ptr = NonNull::new(data.as_ptr().cast_mut()).unwrap_or_else(NonNull::dangling);
198    // SAFETY: the caller supplies the documented dtype, bounds, alignment and lifetime invariants.
199    unsafe {
200        ErasedRawStridedPtr::from_raw_parts(dtype, data_ptr, data.len(), dims, strides, offset)
201    }
202}
203
204/// Construct an erased writable view over exclusively owned initialized storage.
205///
206/// # Safety
207/// The caller must uphold the dtype, alignment, bounds, lifetime and exclusive
208/// access requirements of the returned descriptor.
209#[doc(hidden)]
210pub unsafe fn erased_raw_strided_mut<'a>(
211    dtype: KernelDType,
212    data: &'a mut [u8],
213    dims: &'a [usize],
214    strides: &'a [isize],
215    offset: isize,
216) -> strided_basic::Result<ErasedRawStridedMut<'a>> {
217    let data_ptr = NonNull::new(data.as_mut_ptr()).unwrap_or_else(NonNull::dangling);
218    // SAFETY: the caller supplies the documented dtype, bounds, alignment and exclusive-lifetime invariants.
219    unsafe {
220        ErasedRawStridedMut::from_raw_parts(dtype, data_ptr, data.len(), dims, strides, offset)
221    }
222}
223
224/// Construct an erased writable strided view over exclusively owned uninitialized storage.
225///
226/// # Safety
227/// `dtype` must match the aligned storage; reachable offsets must be in bounds;
228/// the borrow must remain exclusive; and storage must stay uninitialized until
229/// every reachable element has been written.
230#[doc(hidden)]
231pub unsafe fn erased_raw_strided_uninit_mut<'a>(
232    dtype: KernelDType,
233    data: &'a mut [MaybeUninit<u8>],
234    dims: &'a [usize],
235    strides: &'a [isize],
236    offset: isize,
237) -> strided_basic::Result<ErasedRawStridedUninitMut<'a>> {
238    let data_ptr = NonNull::new(data.as_mut_ptr().cast::<u8>()).unwrap_or_else(NonNull::dangling);
239    // SAFETY: the caller supplies the documented dtype, bounds, alignment and exclusive-lifetime invariants.
240    unsafe {
241        ErasedRawStridedUninitMut::from_raw_parts(
242            dtype,
243            data_ptr,
244            data.len(),
245            dims,
246            strides,
247            offset,
248        )
249    }
250}
tenferro_cpu_basic/lib.rs

tenferro_cpu_basic/
lib.rs
