tenferro_cpu/
lib.rs

1//! CPU backend, kernels, provider selection, and CPU resource pools.
2//!
3//! # Examples
4//!
5//! ```rust
6//! use tenferro_cpu::{add, CpuBackend};
7//! use tenferro_tensor::{Tensor, TensorBackend, TensorElementwise};
8//!
9//! let mut backend = CpuBackend::new();
10//! let a = Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0])?;
11//! let b = Tensor::from_vec_col_major(vec![2], vec![3.0_f64, 4.0])?;
12//! let c = backend.add(&a, &b)?;
13//! assert_eq!(c.as_slice::<f64>().unwrap(), &[4.0, 6.0]);
14//! let direct = add(&a, &b)?;
15//! assert_eq!(direct.shape(), &[2]);
16//! # Ok::<(), tenferro_tensor::Error>(())
17//! ```
18
19#[cfg(not(any(feature = "cpu-faer", feature = "cpu-blas")))]
20compile_error!("enable at least one CPU backend: cpu-faer or cpu-blas");
21
22#[cfg(all(feature = "provider-inject", not(feature = "cpu-blas")))]
23compile_error!("provider-inject requires cpu-blas");
24
25#[cfg(any(
26    all(feature = "blas-openblas", feature = "blas-accelerate"),
27    all(feature = "blas-openblas", feature = "blas-mkl"),
28    all(feature = "blas-accelerate", feature = "blas-mkl"),
29))]
30compile_error!(
31    "enable at most one explicit BLAS provider feature: blas-openblas, blas-accelerate, or blas-mkl"
32);
33
34#[cfg(all(
35    feature = "provider-inject",
36    any(
37        feature = "blas-openblas",
38        feature = "blas-accelerate",
39        feature = "blas-mkl"
40    )
41))]
42compile_error!("provider-inject cannot be combined with explicit BLAS provider features");
43
44pub mod affinity;
45mod analytic;
46pub mod backend;
47mod buffer_pool;
48mod capability;
49pub mod context;
50mod elementwise;
51mod exec_session;
52mod gemm;
53mod indexing;
54mod indexing_alloc;
55#[cfg(feature = "provider-inject")]
56pub mod inject;
57mod reduction;
58mod structural;
59
60use strided_kernel::{col_major_strides as kernel_col_major_strides, StridedArray, StridedView};
61
62use crate::buffer_pool::{BufferPool, PoolScalar};
63pub(crate) use tenferro_tensor::*;
64
65#[cfg(feature = "provider-src")]
66extern crate blas_src as _;
67#[cfg(feature = "provider-inject")]
68extern crate cblas_inject as _;
69#[cfg(feature = "provider-src")]
70extern crate cblas_src as _;
71#[cfg(feature = "provider-inject")]
72extern crate lapack_inject as _;
73#[cfg(feature = "provider-src")]
74extern crate lapack_src as _;
75
76pub use affinity::{available_parallelism, process_cpu_affinity_count};
77pub use analytic::pow;
78pub use backend::{CpuBackend, CpuBackendKind};
79pub use buffer_pool::BufferPoolStats;
80pub use capability::cpu_capabilities;
81pub use context::CpuContext;
82pub use elementwise::{
83    abs, add, clamp, compare, conj, div, maximum, minimum, mul, neg, rem, select, sign, sub,
84};
85pub use indexing::{dynamic_slice, dynamic_update_slice, gather, pad, scatter};
86pub use reduction::{reduce_max, reduce_min, reduce_prod, reduce_sum};
87pub use structural::{
88    broadcast_in_dim, convert, embed_diagonal, extract_diagonal, reshape, transpose, tril, triu,
89};
90
91/// Owner-scoped CPU scratch-pool API for operation-family crates.
92///
93/// This module is not an application-facing tensor API. It exists so
94/// operation crates that implement CPU kernels can share `CpuBackend`'s
95/// allocation pool without exposing the pool as a general public contract.
96#[doc(hidden)]
97pub mod linalg_interop {
98    pub use crate::buffer_pool::{BufferPool, PoolScalar};
99}
100
101pub(crate) fn cpu_backend_buffer_error(op: &'static str) -> crate::Error {
102    crate::Error::backend_failure(
103        op,
104        "CPU backend received backend buffer; download to host before CPU execution",
105    )
106}
107
108pub(crate) trait ConjElem {
109    fn conj_elem(self) -> Self;
110}
111
112impl ConjElem for f32 {
113    fn conj_elem(self) -> Self {
114        self
115    }
116}
117
118impl ConjElem for f64 {
119    fn conj_elem(self) -> Self {
120        self
121    }
122}
123
124impl ConjElem for num_complex::Complex32 {
125    fn conj_elem(self) -> Self {
126        self.conj()
127    }
128}
129
130impl ConjElem for num_complex::Complex64 {
131    fn conj_elem(self) -> Self {
132        self.conj()
133    }
134}
135
136pub(crate) fn typed_host_data<'a, T>(
137    op: &'static str,
138    tensor: &'a TypedTensor<T>,
139) -> crate::Result<&'a [T]> {
140    match tensor.buffer() {
141        Buffer::Host(data) => Ok(data.as_slice()),
142        Buffer::Backend(_) => Err(cpu_backend_buffer_error(op)),
143    }
144}
145
146pub(crate) fn typed_view<'a, T: Copy>(
147    op: &'static str,
148    tensor: &'a TypedTensor<T>,
149) -> crate::Result<StridedView<'a, T>> {
150    match tensor.buffer() {
151        Buffer::Host(data) => {
152            let strides = kernel_col_major_strides(tensor.shape());
153            StridedView::new(data.as_slice(), tensor.shape(), &strides, 0)
154                .map_err(|err| crate::Error::backend_failure(op, err))
155        }
156        Buffer::Backend(_) => Err(cpu_backend_buffer_error(op)),
157    }
158}
159
160pub(crate) fn typed_view_from_view<'a, T: Copy + 'static, R: TensorRank>(
161    op: &'static str,
162    view: &TypedTensorView<'a, T, R>,
163) -> crate::Result<StridedView<'a, T>> {
164    if view.backend_buffer().is_some() {
165        return Err(cpu_backend_buffer_error(op));
166    }
167    StridedView::new(
168        view.host_storage()?,
169        view.shape(),
170        view.strides(),
171        view.offset(),
172    )
173    .map_err(|err| crate::Error::backend_failure(op, err))
174}
175
176pub(crate) fn materialize_tensor_read(
177    op: &'static str,
178    input: TensorRead<'_>,
179) -> crate::Result<Tensor> {
180    match input {
181        TensorRead::Tensor(tensor) => clone_host_tensor_read(op, tensor),
182        TensorRead::View(view) => materialize_tensor_view(op, view),
183    }
184}
185
186fn clone_host_tensor_read(op: &'static str, tensor: &Tensor) -> crate::Result<Tensor> {
187    macro_rules! clone_host {
188        ($variant:ident, $tensor:expr) => {{
189            typed_host_data(op, $tensor)?;
190            Ok(Tensor::$variant($tensor.clone()))
191        }};
192    }
193
194    match tensor {
195        Tensor::F32(tensor) => clone_host!(F32, tensor),
196        Tensor::F64(tensor) => clone_host!(F64, tensor),
197        Tensor::I32(tensor) => clone_host!(I32, tensor),
198        Tensor::I64(tensor) => clone_host!(I64, tensor),
199        Tensor::Bool(tensor) => clone_host!(Bool, tensor),
200        Tensor::C32(tensor) => clone_host!(C32, tensor),
201        Tensor::C64(tensor) => clone_host!(C64, tensor),
202    }
203}
204
205fn materialize_tensor_view(op: &'static str, view: TensorView<'_>) -> crate::Result<Tensor> {
206    macro_rules! materialize {
207        ($variant:ident, $view:expr) => {{
208            if $view.backend_buffer().is_some() {
209                return Err(cpu_backend_buffer_error(op));
210            }
211            Ok(Tensor::$variant($view.to_contiguous()?))
212        }};
213    }
214
215    match view {
216        TensorView::F32(view) => materialize!(F32, view),
217        TensorView::F64(view) => materialize!(F64, view),
218        TensorView::I32(view) => materialize!(I32, view),
219        TensorView::I64(view) => materialize!(I64, view),
220        TensorView::Bool(view) => materialize!(Bool, view),
221        TensorView::C32(view) => materialize!(C32, view),
222        TensorView::C64(view) => materialize!(C64, view),
223    }
224}
225
226/// Create an output array WITHOUT initializing element values.
227///
228/// # Safety
229/// Caller must write every element before reading. The returned array
230/// contains uninitialized data.
231#[allow(clippy::uninit_vec)]
232#[cfg(test)]
233pub(crate) unsafe fn typed_array_uninit<T>(shape: &[usize]) -> StridedArray<T> {
234    let total: usize = shape.iter().product();
235    let strides = kernel_col_major_strides(shape);
236    let mut data = Vec::with_capacity(total);
237    // SAFETY: test-only helper is used for outputs whose elements are fully overwritten.
238    unsafe { data.set_len(total) };
239    // Invariant: `kernel_col_major_strides(shape)` and `total` describe the
240    // compact column-major array for this validated test output shape.
241    StridedArray::from_parts(data, shape, &strides, 0).expect("column-major output array")
242}
243
244/// Create an output array from the CPU buffer pool WITHOUT initializing values.
245///
246/// # Safety
247/// Caller must write every element before reading. The returned array contains
248/// uninitialized or stale data acquired from `buffers`.
249pub(crate) unsafe fn typed_array_uninit_from_pool<T>(
250    buffers: &mut BufferPool,
251    shape: &[usize],
252) -> crate::Result<StridedArray<T>>
253where
254    T: PoolScalar,
255{
256    let total = tenferro_tensor::validate::checked_shape_product(
257        "typed_array_uninit_from_pool",
258        "shape",
259        shape,
260    )?;
261    let strides = kernel_col_major_strides(shape);
262    // SAFETY: callers use this only for operation outputs that fully overwrite every element.
263    let data = unsafe { T::pool_acquire(buffers, total) };
264    // Invariant: callers pass validated tensor-derived or prechecked output
265    // shapes, and `strides` is their compact column-major layout.
266    StridedArray::from_parts(data, shape, &strides, 0)
267        .map_err(|err| crate::Error::backend_failure("typed_array_uninit_from_pool", err))
268}
269
270pub(crate) fn tensor_from_array<T: Clone>(array: StridedArray<T>) -> TypedTensor<T> {
271    // Invariant: `StridedArray` owns data whose length matches its validated dimensions.
272    TypedTensor::from_vec_col_major(array.dims().to_vec(), array.into_data())
273        .expect("strided array dimensions match owned data length")
274}
275
276pub(crate) fn default_placement() -> Placement {
277    Placement {
278        memory_kind: MemoryKind::UnpinnedHost,
279        device: None,
280    }
281}
282
283pub(crate) fn flat_to_multi(mut flat: usize, shape: &[usize], out: &mut [usize]) {
284    assert_eq!(shape.len(), out.len());
285    for (axis, &dim) in shape.iter().enumerate() {
286        if dim == 0 {
287            out[axis] = 0;
288        } else {
289            out[axis] = flat % dim;
290            flat /= dim;
291        }
292    }
293}
294
295#[cfg(test)]
296mod tests;
tenferro_cpu/lib.rs

tenferro_cpu/
lib.rs
