Skip to main content

tenferro_cpu/
lib.rs

1//! CPU backend, kernels, provider selection, and CPU resource pools.
2//!
3//! # Examples
4//!
5//! ```rust
6//! use tenferro_cpu::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//! # Ok::<(), tenferro_tensor::Error>(())
15//! ```
16
17// `provider-inject` unit tests deliberately omit the broad default-backend
18// suite below because no fixture has registered its FFI symbols. That makes
19// private helpers referenced only by the broad suite appear unused in this one
20// test build; call-through coverage lives in the registered integration test.
21#![cfg_attr(
22    all(test, feature = "provider-inject"),
23    allow(dead_code, unused_imports)
24)]
25
26#[cfg(not(any(feature = "cpu-faer", feature = "cpu-blas")))]
27compile_error!("enable at least one CPU backend: cpu-faer or cpu-blas");
28
29#[cfg(all(feature = "provider-inject", not(feature = "cpu-blas")))]
30compile_error!("provider-inject requires cpu-blas");
31
32#[cfg(any(
33    all(feature = "blas-openblas", feature = "blas-accelerate"),
34    all(feature = "blas-openblas", feature = "blas-mkl"),
35    all(feature = "blas-accelerate", feature = "blas-mkl"),
36))]
37compile_error!(
38    "enable at most one explicit BLAS provider feature: blas-openblas, blas-accelerate, or blas-mkl"
39);
40
41#[cfg(all(
42    feature = "provider-inject",
43    any(
44        feature = "blas-openblas",
45        feature = "blas-accelerate",
46        feature = "blas-mkl"
47    )
48))]
49compile_error!("provider-inject cannot be combined with explicit BLAS provider features");
50
51pub mod affinity;
52mod affinity_policy;
53mod analytic;
54mod arbiter;
55pub mod backend;
56mod blas1;
57pub(crate) mod buffer_pool {
58    pub use tenferro_cpu_basic::buffer_pool::*;
59}
60mod capability;
61pub mod context;
62// INVARIANT: Task 2 stages crate-private stack adapters here before Task 3 wires
63// them into CpuContext.
64#[allow(dead_code)]
65mod domain_executor;
66#[allow(dead_code)]
67mod dot_runtime;
68pub(crate) use tenferro_cpu_basic::PooledUninitOutput;
69pub(crate) use tenferro_cpu_basic::{erased_raw_strided_ref, erased_raw_strided_uninit_mut};
70pub(crate) use tenferro_internal_cpu_kernels::elementwise;
71mod engine;
72mod exec_session;
73mod gemm;
74mod indexed_plan_cache;
75mod indexing;
76#[cfg(feature = "provider-inject")]
77pub mod inject;
78mod placement;
79pub mod provider;
80mod provider_capability;
81mod reduction;
82mod resource_domain;
83mod runtime_adapter;
84mod structural;
85mod topology;
86
87use std::ptr::NonNull;
88#[cfg(test)]
89use strided_kernel::col_major_strides as kernel_col_major_strides;
90#[cfg(test)]
91use strided_kernel::StridedArray;
92
93use crate::buffer_pool::BufferPool;
94pub(crate) use tenferro_tensor::*;
95
96pub(crate) fn cpu_contraction_unsupported_dtype_message(dtype: DType) -> String {
97    let remedy = matches!(dtype, DType::I32 | DType::I64)
98        .then_some(format!("; convert {dtype:?} to F64 before contraction"));
99    format!(
100        "CPU contraction providers support F32/F64/C32/C64{}",
101        remedy.unwrap_or_default()
102    )
103}
104
105pub(crate) fn erased_raw_strided_mut<'a>(
106    dtype: strided_kernel::KernelDType,
107    data: &'a mut [u8],
108    dims: &'a [usize],
109    strides: &'a [isize],
110    offset: isize,
111) -> strided_kernel::Result<strided_kernel::ErasedRawStridedMut<'a>> {
112    let data_ptr = NonNull::new(data.as_mut_ptr()).unwrap_or_else(NonNull::dangling);
113    // SAFETY: callers derive `data` from a uniquely borrowed initialized host
114    // destination and retain that borrow for the returned descriptor lifetime.
115    unsafe {
116        strided_kernel::ErasedRawStridedMut::from_raw_parts(
117            dtype,
118            data_ptr,
119            data.len(),
120            dims,
121            strides,
122            offset,
123        )
124    }
125}
126
127#[cfg(feature = "provider-src")]
128extern crate blas_src as _;
129#[cfg(feature = "provider-inject")]
130extern crate cblas_inject as _;
131#[cfg(feature = "provider-src")]
132extern crate cblas_src as _;
133#[cfg(feature = "provider-inject")]
134extern crate lapack_inject as _;
135#[cfg(feature = "provider-src")]
136extern crate lapack_src as _;
137
138pub use affinity::{
139    available_parallelism, process_cpu_affinity, process_cpu_affinity_count, CpuAffinityError,
140};
141pub use affinity_policy::{
142    resolve_cpu_affinity, resolve_cpu_affinity_with_override, CpuAffinityInput,
143    CpuAffinityInputError, CpuAffinityPolicy, CpuAffinityResolutionError, CpuAffinitySelection,
144    CpuAffinitySelectionReason,
145};
146pub use backend::{
147    CpuBackend, CpuBackendError, CpuBackendKind, CpuExecutionInfo, CpuExecutionMode,
148    CpuRuntimeIdentity, ExternalCpuDomainRegistryError,
149};
150pub use buffer_pool::BufferPoolStats;
151pub use capability::cpu_capabilities;
152pub use context::{CpuContext, CpuContextError};
153pub use domain_executor::{
154    CpuDomainExecutor, CpuDomainExecutorCapabilities, CpuDomainExecutorError, CpuExecutorAffinity,
155    CpuExecutorReentrancy, CpuExecutorShutdown, CpuInnerParallelism, RayonCpuDomainExecutor,
156    ScopedCpuJob, ScopedCpuJobs,
157};
158pub use dot_runtime::{
159    CpuProviderBundle, CpuProviderBundleBuildError, CpuProviderBundleBuilder,
160    CpuProviderBundleInstallError, CpuProviderSlot, GeneralContractionPolicy,
161};
162#[doc(hidden)]
163pub use exec_session::CpuExecSession;
164pub use indexed_plan_cache::IndexedPlanCacheLimits;
165pub use placement::{
166    CpuEngineConstructionError, CpuPlacement, CpuPlacementError, CpuPlacementGuarantee,
167    ResolvedCpuPlacement,
168};
169pub use provider::{CpuExecutionContext, ParallelMode};
170pub use provider_capability::{
171    CpuPlacementControl, CpuProviderDomainError, CpuProviderExecutionCapabilities,
172    CpuThreadCountControl,
173};
174pub use resource_domain::{
175    CpuAdmissionMode, CpuDomainOwnership, ExternalCpuDomain, ExternalCpuDomainError,
176};
177pub use runtime_adapter::{
178    runtime_engine_id, runtime_engine_registration, runtime_engine_registration_with_id,
179    runtime_hardware_class,
180};
181pub use topology::{
182    discover_cpu_topology, CpuId, CpuNode, CpuSet, CpuSetError, CpuTopology, CpuTopologyError,
183    NumaNodeId,
184};
185
186/// Visit a CPU execution session carried by a type-erased backend session.
187///
188/// This is a backend-leaf capability bridge. The exact session marker is checked
189/// before the erased pointer is reconstructed, and the callback cannot return a
190/// borrow of the session, so the borrowed resource lease remains scoped to the
191/// caller's session closure.
192#[doc(hidden)]
193pub fn with_cpu_exec_session<B, R>(
194    session: &mut B,
195    f: impl for<'a> FnOnce(&'a mut CpuExecSession<'a>) -> R,
196) -> Option<R>
197where
198    B: tenferro_tensor::BackendSession + ?Sized,
199{
200    if session.session_type_id() != std::any::TypeId::of::<exec_session::CpuExecSessionMarker>() {
201        return None;
202    }
203    let data = unsafe { session.session_data_mut() };
204    // SAFETY: the exact marker is supplied by CpuExecSession's explicit
205    // `BackendSession` implementation that produced `session_data_mut`, and
206    // the equality above proves that the erased value is `CpuExecSession`.
207    // The callback is higher-ranked and returns no session borrow, so the
208    // reconstructed reference cannot escape the original session borrow.
209    Some(unsafe { f(&mut *(data.cast::<CpuExecSession<'static>>())) })
210}
211
212/// Invoke a direct faer operation with the parallelism selected by a CPU session.
213///
214/// The `faer::Par` value is scoped to the callback and is derived from the
215/// session's managed thread budget and nesting policy. A non-CPU session, or a
216/// CPU session built without `cpu-faer`, returns a typed unsupported error.
217/// `Par::Seq` remains the portable choice for direct calls outside a session.
218///
219/// # Examples
220///
221/// ```rust
222/// # #[cfg(feature = "cpu-faer")]
223/// # fn example() -> tenferro_tensor::Result<()> {
224/// use tenferro_cpu::{CpuBackend, FaerParallelismExt};
225/// use tenferro_tensor::BackendSessionHost;
226///
227/// let mut backend = CpuBackend::with_threads(2)?;
228/// backend.with_backend_session(|session| {
229///     session.with_faer_parallelism(|parallel| {
230///         let _ = parallel;
231///         Ok(())
232///     })
233/// })?;
234/// # Ok(())
235/// # }
236/// # fn main() {}
237/// ```
238#[cfg(feature = "cpu-faer")]
239#[cfg_attr(docsrs, doc(cfg(feature = "cpu-faer")))]
240pub trait FaerParallelismExt {
241    /// Run a scoped callback with this session's faer parallelism policy.
242    ///
243    /// # Errors
244    ///
245    /// Returns [`tenferro_tensor::Error::Unsupported`] when the session is not
246    /// a CPU/faer execution session, or the callback's own typed error.
247    ///
248    /// # Examples
249    ///
250    /// ```rust
251    /// # #[cfg(feature = "cpu-faer")]
252    /// # fn example(session: &mut dyn tenferro_tensor::BackendSession) -> tenferro_tensor::Result<()> {
253    /// use tenferro_cpu::FaerParallelismExt;
254    /// session.with_faer_parallelism(|parallel| {
255    ///     let _ = parallel;
256    ///     Ok::<_, tenferro_tensor::Error>(())
257    /// })?;
258    /// # Ok(())
259    /// # }
260    /// ```
261    fn with_faer_parallelism(
262        &mut self,
263        callback: impl FnOnce(faer::Par) -> tenferro_tensor::Result<()> + Send,
264    ) -> tenferro_tensor::Result<()>;
265}
266
267#[cfg(feature = "cpu-faer")]
268impl<S> FaerParallelismExt for S
269where
270    S: tenferro_tensor::BackendSession + ?Sized,
271{
272    fn with_faer_parallelism(
273        &mut self,
274        callback: impl FnOnce(faer::Par) -> tenferro_tensor::Result<()> + Send,
275    ) -> tenferro_tensor::Result<()> {
276        with_cpu_exec_session(self, |session| session.with_faer_parallelism(callback))
277            .unwrap_or_else(|| {
278                Err(tenferro_tensor::Error::unsupported(
279                    "with_faer_parallelism",
280                    "selected session is not a CPU/faer execution session",
281                ))
282            })
283    }
284}
285
286// Unit tests exercise the pool-aware kernels through the former convenience
287// names without restoring those names to the production crate surface.
288#[cfg(test)]
289pub(crate) use analytic::pow;
290#[cfg(test)]
291macro_rules! test_elementwise_wrapper {
292    ($name:ident($($arg:ident: $ty:ty),*) => $with_pool:ident) => {
293        pub(crate) fn $name($($arg: $ty),*) -> crate::Result<Tensor> {
294            let mut buffers = BufferPool::new();
295            elementwise::$with_pool(&mut buffers, $($arg),*)
296        }
297    };
298}
299#[cfg(test)]
300test_elementwise_wrapper!(abs(input: &Tensor) => abs_with_pool);
301#[cfg(test)]
302test_elementwise_wrapper!(add(lhs: &Tensor, rhs: &Tensor) => add_with_pool);
303#[cfg(test)]
304test_elementwise_wrapper!(clamp(input: &Tensor, lower: &Tensor, upper: &Tensor) => clamp_with_pool);
305#[cfg(test)]
306test_elementwise_wrapper!(compare(lhs: &Tensor, rhs: &Tensor, dir: &CompareDir) => compare_with_pool);
307#[cfg(test)]
308test_elementwise_wrapper!(conj(input: &Tensor) => conj_with_pool);
309#[cfg(test)]
310test_elementwise_wrapper!(div(lhs: &Tensor, rhs: &Tensor) => div_with_pool);
311#[cfg(test)]
312test_elementwise_wrapper!(maximum(lhs: &Tensor, rhs: &Tensor) => maximum_with_pool);
313#[cfg(test)]
314test_elementwise_wrapper!(minimum(lhs: &Tensor, rhs: &Tensor) => minimum_with_pool);
315#[cfg(test)]
316test_elementwise_wrapper!(mul(lhs: &Tensor, rhs: &Tensor) => mul_with_pool);
317#[cfg(test)]
318test_elementwise_wrapper!(neg(input: &Tensor) => neg_with_pool);
319#[cfg(test)]
320test_elementwise_wrapper!(rem(lhs: &Tensor, rhs: &Tensor) => rem_with_pool);
321#[cfg(test)]
322test_elementwise_wrapper!(select(pred: &Tensor, on_true: &Tensor, on_false: &Tensor) => select_with_pool);
323#[cfg(test)]
324test_elementwise_wrapper!(sign(input: &Tensor) => sign_with_pool);
325#[cfg(test)]
326test_elementwise_wrapper!(sub(lhs: &Tensor, rhs: &Tensor) => sub_with_pool);
327#[cfg(test)]
328pub(crate) use indexing::{dynamic_slice, dynamic_update_slice, gather, pad, scatter};
329#[cfg(test)]
330pub(crate) use reduction::{reduce_max, reduce_min, reduce_prod, reduce_sum, reduce_sum_squares};
331#[cfg(test)]
332pub(crate) use structural::{
333    broadcast_in_dim, embed_diagonal, extract_diagonal, reshape, transpose, tril, triu,
334};
335
336/// Owner-scoped CPU scratch-pool API for operation-family crates.
337///
338/// This module is not an application-facing tensor API. It exists so
339/// operation crates that implement CPU kernels can share `CpuBackend`'s
340/// allocation pool without exposing the pool as a general public contract.
341#[doc(hidden)]
342pub mod linalg_interop {
343    pub use crate::buffer_pool::{BufferPool, PoolScalar};
344    pub use tenferro_cpu_basic::PooledUninitOutput;
345}
346
347#[derive(Debug, thiserror::Error)]
348pub(crate) enum CpuNumericalError {
349    #[error("{op} received a negative integer exponent for dtype {dtype:?}")]
350    NegativeIntegerExponent { op: &'static str, dtype: DType },
351}
352
353pub(crate) fn cpu_negative_integer_exponent(op: &'static str, dtype: DType) -> crate::Error {
354    crate::Error::extension(
355        op,
356        "cpu",
357        ErrorKind::NumericalFailure,
358        CpuNumericalError::NegativeIntegerExponent { op, dtype },
359    )
360}
361
362pub(crate) use tenferro_cpu_basic::{
363    cpu_backend_buffer_error, typed_host_data, typed_view, typed_view_from_view, ConjElem,
364};
365pub(crate) fn materialize_tensor_read(
366    buffers: &mut BufferPool,
367    op: &'static str,
368    input: TensorRead<'_>,
369) -> crate::Result<Tensor> {
370    match input {
371        TensorRead::Tensor(tensor) => clone_host_tensor_read(op, tensor),
372        TensorRead::View(view) => materialize_tensor_view(buffers, op, view),
373    }
374}
375
376pub(crate) fn copy_tensor_read_into(
377    op: &'static str,
378    src: TensorRead<'_>,
379    dst: TensorWrite<'_>,
380) -> crate::Result<()> {
381    let src_dtype = src.dtype();
382    let dst_dtype = dst.dtype();
383    macro_rules! copy_source {
384        ($variant:ident, $src:expr) => {{
385            let src = $src;
386            match dst {
387                TensorWrite::Tensor(Tensor::$variant(dst)) => {
388                    let mut dst = dst.as_view_mut();
389                    structural::typed_copy_view_into(&src, &mut dst, op)
390                }
391                TensorWrite::View(TensorViewMut::$variant(mut dst)) => {
392                    structural::typed_copy_view_into(&src, &mut dst, op)
393                }
394                _ => Err(crate::Error::dtype_mismatch(op, src_dtype, dst_dtype)),
395            }
396        }};
397    }
398
399    match src {
400        TensorRead::Tensor(Tensor::F32(src)) => copy_source!(F32, src.as_view()),
401        TensorRead::Tensor(Tensor::F64(src)) => copy_source!(F64, src.as_view()),
402        TensorRead::Tensor(Tensor::I32(src)) => copy_source!(I32, src.as_view()),
403        TensorRead::Tensor(Tensor::I64(src)) => copy_source!(I64, src.as_view()),
404        TensorRead::Tensor(Tensor::Bool(src)) => copy_source!(Bool, src.as_view()),
405        TensorRead::Tensor(Tensor::C32(src)) => copy_source!(C32, src.as_view()),
406        TensorRead::Tensor(Tensor::C64(src)) => copy_source!(C64, src.as_view()),
407        TensorRead::View(TensorView::F32(src)) => copy_source!(F32, src),
408        TensorRead::View(TensorView::F64(src)) => copy_source!(F64, src),
409        TensorRead::View(TensorView::I32(src)) => copy_source!(I32, src),
410        TensorRead::View(TensorView::I64(src)) => copy_source!(I64, src),
411        TensorRead::View(TensorView::Bool(src)) => copy_source!(Bool, src),
412        TensorRead::View(TensorView::C32(src)) => copy_source!(C32, src),
413        TensorRead::View(TensorView::C64(src)) => copy_source!(C64, src),
414    }
415}
416
417fn clone_host_tensor_read(op: &'static str, tensor: &Tensor) -> crate::Result<Tensor> {
418    macro_rules! clone_host {
419        ($variant:ident, $tensor:expr) => {{
420            structural::validate_cpu_host_placement(op, "source", $tensor.placement())?;
421            typed_host_data(op, $tensor)?;
422            $tensor.duplicate().map(Tensor::$variant)
423        }};
424    }
425
426    match tensor {
427        Tensor::F32(tensor) => clone_host!(F32, tensor),
428        Tensor::F64(tensor) => clone_host!(F64, tensor),
429        Tensor::I32(tensor) => clone_host!(I32, tensor),
430        Tensor::I64(tensor) => clone_host!(I64, tensor),
431        Tensor::Bool(tensor) => clone_host!(Bool, tensor),
432        Tensor::C32(tensor) => clone_host!(C32, tensor),
433        Tensor::C64(tensor) => clone_host!(C64, tensor),
434    }
435}
436
437fn materialize_tensor_view(
438    buffers: &mut BufferPool,
439    op: &'static str,
440    view: TensorView<'_>,
441) -> crate::Result<Tensor> {
442    macro_rules! materialize {
443        ($variant:ident, $view:expr) => {{
444            Ok(Tensor::$variant(
445                structural::typed_materialize_view_with_pool(buffers, &$view, op)?,
446            ))
447        }};
448    }
449
450    match view {
451        TensorView::F32(view) => materialize!(F32, view),
452        TensorView::F64(view) => materialize!(F64, view),
453        TensorView::I32(view) => materialize!(I32, view),
454        TensorView::I64(view) => materialize!(I64, view),
455        TensorView::Bool(view) => materialize!(Bool, view),
456        TensorView::C32(view) => materialize!(C32, view),
457        TensorView::C64(view) => materialize!(C64, view),
458    }
459}
460
461/// Create an output array WITHOUT initializing element values.
462///
463/// # Safety
464/// Caller must write every element before reading. The returned array
465/// contains uninitialized data.
466#[allow(clippy::uninit_vec)]
467#[cfg(test)]
468pub(crate) unsafe fn typed_array_uninit<T>(shape: &[usize]) -> StridedArray<T> {
469    let total: usize = shape.iter().product();
470    let strides = kernel_col_major_strides(shape);
471    let mut data = Vec::with_capacity(total);
472    // SAFETY: test-only helper is used for outputs whose elements are fully overwritten.
473    unsafe { data.set_len(total) };
474    // Invariant: `kernel_col_major_strides(shape)` and `total` describe the
475    // compact column-major array for this validated test output shape.
476    StridedArray::from_parts(data, shape, &strides, 0).expect("column-major output array")
477}
478
479#[cfg(test)]
480pub(crate) fn tensor_from_array<T: Clone + tenferro_tensor::TensorScalar>(
481    array: StridedArray<T>,
482) -> TypedTensor<T> {
483    // Invariant: `StridedArray` owns data whose length matches its validated dimensions.
484    TypedTensor::from_vec_col_major(array.dims().to_vec(), array.into_data())
485        .expect("strided array dimensions match owned data length")
486}
487
488pub(crate) fn flat_to_multi(mut flat: usize, shape: &[usize], out: &mut [usize]) {
489    assert_eq!(shape.len(), out.len());
490    for (axis, &dim) in shape.iter().enumerate() {
491        if dim == 0 {
492            out[axis] = 0;
493        } else {
494            out[axis] = flat % dim;
495            flat /= dim;
496        }
497    }
498}
499
500// `provider-inject` owns call-through coverage in the serialized integration
501// fixture, which registers every BLAS symbol before the first operation.  The
502// broad unit suite selects the compiled default backend and therefore must not
503// call an intentionally unregistered injected symbol.
504#[cfg(all(test, not(feature = "provider-inject")))]
505mod tests;