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

tenferro_cpu_basic/
lib.rs
