tenferro_internal_cpu_kernels/
scalar_ops.rs

1//! Scalar-agnostic entry points over the ordinary CPU numerical bodies.
2//!
3//! The preset scalar types reach the ordinary CPU kernels through the typed
4//! pool, which hands out the destination buffer. An external scalar type cannot
5//! use that pool, so these entry points take a caller-provided destination and
6//! the caller's own arithmetic instead, and run the same elementwise and
7//! reduction bodies the preset path uses.
8//!
9//! Nothing here inspects a dtype tag: the element type and the arithmetic are
10//! the caller's. Support for the preset scalars is therefore not a precondition
11//! for using the ordinary CPU numerical path, and no set-specific numerical
12//! specialization is introduced.
13
14pub mod op;
15
16pub use op::{AddOp, BinaryScalarOp, MulOp, SubOp};
17
18use strided_kernel::{reduce, zip_map2_into, Exec, StridedView, StridedViewMut};
19use tenferro_tensor::col_major_strides;
20use tenferro_tensor::{DynRank, Host, StrideVec, TypedTensor};
21
22fn strides_for(shape: &[usize]) -> crate::Result<StrideVec> {
23    col_major_strides(shape)
24}
25
26fn require_same_shape(op: &'static str, lhs: &[usize], rhs: &[usize]) -> crate::Result<()> {
27    if lhs == rhs {
28        Ok(())
29    } else {
30        Err(crate::Error::shape_mismatch(op, lhs.to_vec(), rhs.to_vec()))
31    }
32}
33
34/// Apply a named binary operation elementwise into a caller-owned destination.
35///
36/// The destination, the two operands, and the operation are all the caller's.
37/// The traversal is the same `zip_map2_into` body the preset scalar types use
38/// through the typed pool; only the destination's origin differs. The operation
39/// is a type rather than a closure so the instantiation depends on the element
40/// type and the operation, not on the call site.
41///
42/// # Examples
43///
44/// ```rust
45/// use strided_kernel::Exec;
46/// use tenferro_internal_cpu_kernels::scalar_ops::{scalar_binary_into, AddOp};
47/// use tenferro_tensor::{DynRank, Host, TypedTensor};
48///
49/// let lhs = TypedTensor::<_, DynRank, Host>::from_host_vec_col_major(vec![2], vec![1.0_f64, 2.0])?;
50/// let rhs = TypedTensor::<_, DynRank, Host>::from_host_vec_col_major(vec![2], vec![10.0_f64, 20.0])?;
51/// let mut out = TypedTensor::<_, DynRank, Host>::from_host_vec_col_major(vec![2], vec![0.0_f64, 0.0])?;
52/// scalar_binary_into::<f64, AddOp>("add", &mut out, &lhs, &rhs, Exec::serial())?;
53/// assert_eq!(out.as_slice(), &[11.0, 22.0]);
54/// # Ok::<(), Box<dyn std::error::Error>>(())
55/// ```
56///
57/// # Errors
58///
59/// Returns [`Error::Validation`] when the destination shape does not match the operands',
60/// or [`Error::BackendSource`] when the underlying strided traversal rejects the views.
61pub fn scalar_binary_into<T, Op>(
62    op: &'static str,
63    destination: &mut TypedTensor<T, DynRank, Host>,
64    lhs: &TypedTensor<T, DynRank, Host>,
65    rhs: &TypedTensor<T, DynRank, Host>,
66    exec: Exec<'_>,
67) -> crate::Result<()>
68where
69    T: Copy + Send + Sync,
70    Op: BinaryScalarOp<T>,
71{
72    require_same_shape(op, destination.shape(), lhs.shape())?;
73    require_same_shape(op, lhs.shape(), rhs.shape())?;
74
75    let strides = strides_for(lhs.shape())?;
76    // INVARIANT: the first `require_same_shape(op, destination.shape(),
77    // lhs.shape())` above proved destination == lhs, so the destination view can
78    // borrow the lhs shape instead of copying it into a fresh `Vec` each call.
79    let mut destination_view: StridedViewMut<'_, T> =
80        StridedViewMut::new(destination.host_data_mut(), lhs.shape(), &strides, 0)
81            .map_err(|err| crate::Error::backend_source(op, err))?;
82    let lhs_view: StridedView<'_, T> = StridedView::new(lhs.as_slice(), lhs.shape(), &strides, 0)
83        .map_err(|err| crate::Error::backend_source(op, err))?;
84    let rhs_view: StridedView<'_, T> = StridedView::new(rhs.as_slice(), rhs.shape(), &strides, 0)
85        .map_err(|err| crate::Error::backend_source(op, err))?;
86
87    zip_map2_into(
88        &mut destination_view,
89        &lhs_view,
90        &rhs_view,
91        |a, b| Op::apply(a, b),
92        exec,
93    )
94    .map_err(|err| crate::Error::backend_source(op, err))
95}
96
97/// Fold every element of a caller-owned tensor with a named associative
98/// operation, starting from `init`.
99///
100/// The accumulation order is the backend's and is not part of the contract. The
101/// arithmetic is the caller's, so an extended-precision scalar keeps its low
102/// components exactly as a preset scalar keeps its own.
103///
104/// # Examples
105///
106/// ```rust
107/// use strided_kernel::Exec;
108/// use tenferro_internal_cpu_kernels::scalar_ops::{scalar_fold, AddOp};
109/// use tenferro_tensor::{DynRank, Host, TypedTensor};
110///
111/// let values = TypedTensor::<_, DynRank, Host>::from_host_vec_col_major(vec![3], vec![1.0_f64, 2.0, 3.0])?;
112/// let total = scalar_fold::<f64, AddOp>("sum", &values, 0.0_f64, Exec::serial())?;
113/// assert_eq!(total, 6.0);
114/// # Ok::<(), Box<dyn std::error::Error>>(())
115/// ```
116///
117/// # Errors
118///
119/// Returns [`Error::BackendSource`] when the strided reduction rejects the view, which
120/// happens when the source's shape has an invalid stride layout.
121pub fn scalar_fold<T, Op>(
122    op: &'static str,
123    source: &TypedTensor<T, DynRank, Host>,
124    init: T,
125    exec: Exec<'_>,
126) -> crate::Result<T>
127where
128    T: Copy + Send + Sync,
129    Op: BinaryScalarOp<T>,
130{
131    let strides = strides_for(source.shape())?;
132    let view: StridedView<'_, T> = StridedView::new(source.as_slice(), source.shape(), &strides, 0)
133        .map_err(|err| crate::Error::backend_source(op, err))?;
134    reduce(&view, |element| element, |a, b| Op::apply(a, b), init, exec)
135        .map_err(|err| crate::Error::backend_source(op, err))
136}
137
138#[cfg(test)]
139mod tests;
tenferro_internal_cpu_kernels/scalar_ops.rs

tenferro_internal_cpu_kernels/
scalar_ops.rs
