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

tenferro_cpu/
lib.rs
