tenferro-cpu 0.8.0

use std::ops::{Add, Mul};
/// The Rust scalar type behind a preset variant name a macro received.
macro_rules! preset_scalar {
    (F32) => {
        f32
    };
    (F64) => {
        f64
    };
    (I32) => {
        i32
    };
    (I64) => {
        i64
    };
    (Bool) => {
        bool
    };
    (C32) => {
        num_complex::Complex32
    };
    (C64) => {
        num_complex::Complex64
    };
}
use num_complex::{Complex32, Complex64};
use tenferro_tensor::{
    ContractionScalar, DType, MemoryKind, TensorRead, TensorView, TensorViewMut, TensorWrite,
};

use crate::buffer_pool::BufferPool;
use threadteam::Exec;

pub(crate) fn validate_cpu_read(op: &'static str, input: &TensorRead<'_>) -> crate::Result<()> {
    if input.backend_family().is_some() || input.placement().memory_kind == MemoryKind::Device {
        Err(crate::cpu_backend_buffer_error(op))
    } else {
        Ok(())
    }
}

pub(crate) fn axpby_read_into_accum(
    _buffers: &mut BufferPool,
    alpha: ContractionScalar,
    x: TensorRead<'_>,
    beta: ContractionScalar,
    y: TensorWrite<'_>,
    exec: Exec<'_>,
) -> crate::Result<()> {
    validate_cpu_read("Engine::axpby_read_into_accum", &x)?;
    validate_cpu_read("Engine::axpby_read_into_accum", &y.as_read())?;
    if !(x.is_col_major_contiguous()? && y.as_read().is_col_major_contiguous()?) {
        // One side is strided: the pinned update primitive reads a strided `x`
        // and writes a strided `y`, so a strided operand is never staged
        // (#2214 item 7). A compact owned destination is exposed as its own
        // view, which copies nothing. The dense path below then covers the
        // both-compact case only.
        return match (x.dtype(), alpha, beta) {
            (DType::F32, ContractionScalar::F32(alpha), ContractionScalar::F32(beta)) => {
                axpby_strided_write_f32(x, y, alpha, beta, exec)
            }
            (DType::F64, ContractionScalar::F64(alpha), ContractionScalar::F64(beta)) => {
                axpby_strided_write_f64(x, y, alpha, beta, exec)
            }
            (DType::C32, ContractionScalar::C32(alpha), ContractionScalar::C32(beta)) => {
                axpby_strided_write_c32(x, y, alpha, beta, exec)
            }
            (DType::C64, ContractionScalar::C64(alpha), ContractionScalar::C64(beta)) => {
                axpby_strided_write_c64(x, y, alpha, beta, exec)
            }
            // INVARIANT: shared validation proves x, y, alpha, and beta have one
            // identical supported dtype before execution dispatch.
            _ => Err(crate::Error::Internal(
                "validated AXPBY dtype changed before execution".into(),
            )),
        };
    }
    match (x.dtype(), alpha, beta) {
        (DType::F32, ContractionScalar::F32(alpha), ContractionScalar::F32(beta)) => {
            let x_data = x.as_slice::<f32>()?;
            axpby_write_f32(x_data, y, alpha, beta, exec)
        }
        (DType::F64, ContractionScalar::F64(alpha), ContractionScalar::F64(beta)) => {
            let x_data = x.as_slice::<f64>()?;
            axpby_write_f64(x_data, y, alpha, beta, exec)
        }
        (DType::C32, ContractionScalar::C32(alpha), ContractionScalar::C32(beta)) => {
            let x_data = x.as_slice::<Complex32>()?;
            axpby_write_c32(x_data, y, alpha, beta, exec)
        }
        (DType::C64, ContractionScalar::C64(alpha), ContractionScalar::C64(beta)) => {
            let x_data = x.as_slice::<Complex64>()?;
            axpby_write_c64(x_data, y, alpha, beta, exec)
        }
        // INVARIANT: shared validation proves x, y, alpha, and beta have one
        // identical supported dtype before execution dispatch.
        _ => Err(crate::Error::Internal(
            "validated AXPBY dtype changed before execution".into(),
        )),
    }
}

fn axpby_typed<T>(
    x_data: &[T],
    y_data: &mut [T],
    alpha: T,
    beta: T,
    exec: Exec<'_>,
) -> crate::Result<()>
where
    T: Copy + Send + Sync + Add<Output = T> + Mul<Output = T>,
{
    strided_kernel::axpby_accum(y_data, x_data, alpha, beta, exec)
        .map_err(|err| crate::Error::backend_source("axpby", err))
}

macro_rules! axpby_write {
    ($name:ident, $variant:ident, $ty:ty) => {
        fn $name(
            x_data: &[$ty],
            y: TensorWrite<'_>,
            alpha: $ty,
            beta: $ty,
            exec: Exec<'_>,
        ) -> crate::Result<()> {
            match y {
                TensorWrite::Tensor(tensor)
                    if tensor.dtype()
                        == <preset_scalar!($variant) as tenferro_tensor::TensorScalar>::dtype() =>
                {
                    let tensor = tensor
                        .as_typed_mut::<preset_scalar!($variant)>()
                        .expect("the dtype guard selects this arm");
                    axpby_typed(x_data, tensor.host_data_mut()?, alpha, beta, exec)
                }
                TensorWrite::View(TensorViewMut::$variant(mut view)) => {
                    let len = x_data.len();
                    if len == 0 {
                        return axpby_typed(x_data, &mut [], alpha, beta, exec);
                    }
                    // INVARIANT: TensorViewMut construction proves reachable
                    // bounds and injectivity; shared AXPBY validation proves a
                    // compact layout with the same shape as x. A non-empty
                    // compact view therefore has a nonnegative offset and one
                    // contiguous in-bounds region of exactly `len` elements.
                    let offset = view.offset() as usize;
                    let storage = view.host_storage_mut()?;
                    // SAFETY: the invariant above proves `offset..offset+len`
                    // is within `storage` and uniquely writable.
                    let y_data = unsafe {
                        std::slice::from_raw_parts_mut(storage.as_mut_ptr().add(offset), len)
                    };
                    axpby_typed(x_data, y_data, alpha, beta, exec)
                }
                // INVARIANT: shared validation proves y has the dtype selected
                // by the outer dispatch before this helper is called.
                _ => Err(crate::Error::Internal(
                    "validated AXPBY output dtype changed before execution".into(),
                )),
            }
        }
    };
}

axpby_write!(axpby_write_f32, F32, f32);
axpby_write!(axpby_write_f64, F64, f64);
axpby_write!(axpby_write_c32, C32, Complex32);
axpby_write!(axpby_write_c64, C64, Complex64);

/// Apply `y <- alpha * x + beta * y` into a strided/offset destination view.
///
/// `strided_kernel::zip_update2_into` owns the traversal. The destination view
/// is write-injective and disjoint from `x` (shared validation plus
/// `TensorViewMut` construction), and the source may itself be strided, so no
/// compact staging buffer is needed on this path.
macro_rules! axpby_strided_write {
    ($name:ident, $variant:ident, $ty:ty) => {
        fn $name(
            x: TensorRead<'_>,
            y: TensorWrite<'_>,
            alpha: $ty,
            beta: $ty,
            exec: Exec<'_>,
        ) -> crate::Result<()> {
            let source = match x.tensor_view() {
                TensorView::$variant(view) => view,
                // INVARIANT: shared validation proves x has the dtype selected
                // by the outer dispatch before this helper is called.
                _ => {
                    return Err(crate::Error::Internal(
                        "validated AXPBY input dtype changed before execution".into(),
                    ))
                }
            };
            let source = strided_kernel::StridedView::<_, strided_kernel::Identity>::new(
                source.host_storage()?,
                source.shape(),
                source.strides(),
                source.offset(),
            )
            .map_err(|err| crate::Error::backend_source("axpby", err))?;
            let run = |destination: &mut strided_kernel::StridedViewMut<'_, $ty>| {
                strided_kernel::zip_update2_into::<
                    _,
                    _,
                    strided_kernel::Identity,
                    strided_kernel::Identity,
                >(destination, &source, |y, x| alpha * x + beta * y, exec)
                .map_err(|err| crate::Error::backend_source("axpby", err))
            };
            match y {
                TensorWrite::View(TensorViewMut::$variant(mut view)) => {
                    let shape = view.shape().to_vec();
                    let strides = view.strides().to_vec();
                    let offset = view.offset();
                    let mut destination = strided_kernel::StridedViewMut::new(
                        view.host_storage_mut()?,
                        &shape,
                        &strides,
                        offset,
                    )
                    .map_err(|err| crate::Error::backend_source("axpby", err))?;
                    run(&mut destination)
                }
                TensorWrite::Tensor(tensor) => {
                    // A compact owned destination is compact column-major, which is
                    // a valid strided destination; wrapping it copies nothing, so a
                    // strided `x` needs no staging buffer either (invariant I1,
                    // #2214 item 7).
                    let typed = tensor.as_typed_mut::<$ty>().ok_or_else(|| {
                        crate::Error::Internal(
                            "validated AXPBY destination dtype changed before execution".into(),
                        )
                    })?;
                    let shape = typed.shape().to_vec();
                    // An owned tensor is compact column-major, so its strides are
                    // the running product of the shape (`INVARIANT`); no staging.
                    let mut strides = Vec::with_capacity(shape.len());
                    let mut stride = 1isize;
                    for &dim in &shape {
                        strides.push(stride);
                        stride *= dim as isize;
                    }
                    let mut destination = strided_kernel::StridedViewMut::new(
                        typed.host_data_mut()?,
                        &shape,
                        &strides,
                        0,
                    )
                    .map_err(|err| crate::Error::backend_source("axpby", err))?;
                    run(&mut destination)
                }
                _ => Err(crate::Error::Internal(
                    "validated AXPBY destination dtype changed before execution".into(),
                )),
            }
        }
    };
}

axpby_strided_write!(axpby_strided_write_f32, F32, f32);
axpby_strided_write!(axpby_strided_write_f64, F64, f64);
axpby_strided_write!(axpby_strided_write_c32, C32, Complex32);
axpby_strided_write!(axpby_strided_write_c64, C64, Complex64);