tenferro-internal-cpu-kernels 0.8.0

//! Scalar-agnostic entry points over the ordinary CPU numerical bodies.
//!
//! The preset scalar types reach the ordinary CPU kernels through the typed
//! pool, which hands out the destination buffer. An external scalar type cannot
//! use that pool, so these entry points take a caller-provided destination and
//! the caller's own arithmetic instead, and run the same elementwise and
//! reduction bodies the preset path uses.
//!
//! Nothing here inspects a dtype tag: the element type and the arithmetic are
//! the caller's. Support for the preset scalars is therefore not a precondition
//! for using the ordinary CPU numerical path, and no set-specific numerical
//! specialization is introduced.

pub mod op;

pub use op::{AddOp, BinaryScalarOp, MulOp, SubOp};

use strided_kernel::{reduce, zip_map2_into, Exec, StridedView, StridedViewMut};
use tenferro_tensor::col_major_strides;
use tenferro_tensor::{DynRank, Host, StrideVec, TypedTensor};

fn strides_for(shape: &[usize]) -> crate::Result<StrideVec> {
    col_major_strides(shape)
}

fn require_same_shape(op: &'static str, lhs: &[usize], rhs: &[usize]) -> crate::Result<()> {
    if lhs == rhs {
        Ok(())
    } else {
        Err(crate::Error::shape_mismatch(op, lhs.to_vec(), rhs.to_vec()))
    }
}

/// Apply a named binary operation elementwise into a caller-owned destination.
///
/// The destination, the two operands, and the operation are all the caller's.
/// The traversal is the same `zip_map2_into` body the preset scalar types use
/// through the typed pool; only the destination's origin differs. The operation
/// is a type rather than a closure so the instantiation depends on the element
/// type and the operation, not on the call site.
///
/// # Examples
///
/// ```rust
/// use strided_kernel::Exec;
/// use tenferro_internal_cpu_kernels::scalar_ops::{scalar_binary_into, AddOp};
/// use tenferro_tensor::{DynRank, Host, TypedTensor};
///
/// let lhs = TypedTensor::<_, DynRank, Host>::from_host_vec_col_major(vec![2], vec![1.0_f64, 2.0])?;
/// let rhs = TypedTensor::<_, DynRank, Host>::from_host_vec_col_major(vec![2], vec![10.0_f64, 20.0])?;
/// let mut out = TypedTensor::<_, DynRank, Host>::from_host_vec_col_major(vec![2], vec![0.0_f64, 0.0])?;
/// scalar_binary_into::<f64, AddOp>("add", &mut out, &lhs, &rhs, Exec::serial())?;
/// assert_eq!(out.as_slice(), &[11.0, 22.0]);
/// # Ok::<(), Box<dyn std::error::Error>>(())
/// ```
///
/// # Errors
///
/// Returns [`Error::Validation`] when the destination shape does not match the operands',
/// or [`Error::BackendSource`] when the underlying strided traversal rejects the views.
pub fn scalar_binary_into<T, Op>(
    op: &'static str,
    destination: &mut TypedTensor<T, DynRank, Host>,
    lhs: &TypedTensor<T, DynRank, Host>,
    rhs: &TypedTensor<T, DynRank, Host>,
    exec: Exec<'_>,
) -> crate::Result<()>
where
    T: Copy + Send + Sync,
    Op: BinaryScalarOp<T>,
{
    require_same_shape(op, destination.shape(), lhs.shape())?;
    require_same_shape(op, lhs.shape(), rhs.shape())?;

    let strides = strides_for(lhs.shape())?;
    // INVARIANT: the first `require_same_shape(op, destination.shape(),
    // lhs.shape())` above proved destination == lhs, so the destination view can
    // borrow the lhs shape instead of copying it into a fresh `Vec` each call.
    let mut destination_view: StridedViewMut<'_, T> =
        StridedViewMut::new(destination.host_data_mut(), lhs.shape(), &strides, 0)
            .map_err(|err| crate::Error::backend_source(op, err))?;
    let lhs_view: StridedView<'_, T> = StridedView::new(lhs.as_slice(), lhs.shape(), &strides, 0)
        .map_err(|err| crate::Error::backend_source(op, err))?;
    let rhs_view: StridedView<'_, T> = StridedView::new(rhs.as_slice(), rhs.shape(), &strides, 0)
        .map_err(|err| crate::Error::backend_source(op, err))?;

    zip_map2_into(
        &mut destination_view,
        &lhs_view,
        &rhs_view,
        |a, b| Op::apply(a, b),
        exec,
    )
    .map_err(|err| crate::Error::backend_source(op, err))
}

/// Fold every element of a caller-owned tensor with a named associative
/// operation, starting from `init`.
///
/// The accumulation order is the backend's and is not part of the contract. The
/// arithmetic is the caller's, so an extended-precision scalar keeps its low
/// components exactly as a preset scalar keeps its own.
///
/// # Examples
///
/// ```rust
/// use strided_kernel::Exec;
/// use tenferro_internal_cpu_kernels::scalar_ops::{scalar_fold, AddOp};
/// use tenferro_tensor::{DynRank, Host, TypedTensor};
///
/// let values = TypedTensor::<_, DynRank, Host>::from_host_vec_col_major(vec![3], vec![1.0_f64, 2.0, 3.0])?;
/// let total = scalar_fold::<f64, AddOp>("sum", &values, 0.0_f64, Exec::serial())?;
/// assert_eq!(total, 6.0);
/// # Ok::<(), Box<dyn std::error::Error>>(())
/// ```
///
/// # Errors
///
/// Returns [`Error::BackendSource`] when the strided reduction rejects the view, which
/// happens when the source's shape has an invalid stride layout.
pub fn scalar_fold<T, Op>(
    op: &'static str,
    source: &TypedTensor<T, DynRank, Host>,
    init: T,
    exec: Exec<'_>,
) -> crate::Result<T>
where
    T: Copy + Send + Sync,
    Op: BinaryScalarOp<T>,
{
    let strides = strides_for(source.shape())?;
    let view: StridedView<'_, T> = StridedView::new(source.as_slice(), source.shape(), &strides, 0)
        .map_err(|err| crate::Error::backend_source(op, err))?;
    reduce(&view, |element| element, |a, b| Op::apply(a, b), init, exec)
        .map_err(|err| crate::Error::backend_source(op, err))
}

#[cfg(test)]
mod tests;