mdarray-linalg-blas 0.2.0

use cblas_sys::{CBLAS_LAYOUT, CBLAS_TRANSPOSE};
use mdarray::{Dim, Layout, Shape, Slice};
use mdarray_linalg::utils::into_i32;
use num_complex::ComplexFloat;

use super::scalar::BlasScalar;
use crate::trans_stride;

pub(super) fn gemv<T, D0: Dim, D1: Dim, La, Lx, Ly>(
    alpha: T,
    a: &Slice<T, (D0, D1), La>,
    x: &Slice<T, (D1,), Lx>,
    beta: T,
    y: &mut Slice<T, (D0,), Ly>,
) where
    T: BlasScalar + ComplexFloat,
    La: Layout,
    Lx: Layout,
    Ly: Layout,
{
    let ash = *a.shape();
    let (m, n) = (ash.dim(0), ash.dim(1));

    if a.stride(1) == 1 {
        assert_eq!(x.len(), n, "x length must match number of columns in a");
    } else {
        assert_eq!(x.len(), m, "x length must match number of rows in a");
    }

    assert_eq!(
        y.len(),
        if a.stride(1) == 1 { m } else { n },
        "y length must match the output dimension"
    );

    let row_major = a.stride(1) == 1;
    assert!(
        row_major || a.stride(0) == 1,
        "a must be contiguous in one dimension"
    );

    let (same_order, other_order) = if row_major {
        (CBLAS_TRANSPOSE::CblasNoTrans, CBLAS_TRANSPOSE::CblasTrans)
    } else {
        (CBLAS_TRANSPOSE::CblasTrans, CBLAS_TRANSPOSE::CblasNoTrans)
    };
    let (a_trans, a_stride) = trans_stride(a, same_order, other_order);

    let x_inc = into_i32(x.stride(0));
    let y_inc = into_i32(y.stride(0));

    unsafe {
        T::cblas_gemv(
            if row_major {
                CBLAS_LAYOUT::CblasRowMajor
            } else {
                CBLAS_LAYOUT::CblasColMajor
            },
            a_trans,
            into_i32(m),
            into_i32(n),
            alpha,
            a.as_ptr(),
            a_stride,
            x.as_ptr(),
            x_inc,
            beta,
            y.as_mut_ptr(),
            y_inc,
        )
    }
}

pub(super) fn ger<T, La, Lx, Ly, D0: Dim, D1: Dim>(
    beta: T,
    x: &Slice<T, (D0,), Lx>,
    y: &Slice<T, (D1,), Ly>,
    a: &mut Slice<T, (D0, D1), La>,
) where
    T: BlasScalar + ComplexFloat,
    La: Layout,
    Lx: Layout,
    Ly: Layout,
{
    let ash = *a.shape();
    let (m, n) = (ash.dim(0), ash.dim(1));

    assert_eq!(x.len(), m, "x length must match number of rows in a");
    assert_eq!(y.len(), n, "y length must match number of columns in a");

    let x_inc = into_i32(x.stride(0));
    let y_inc = into_i32(y.stride(0));

    let row_major = a.stride(1) == 1;
    assert!(
        row_major || a.stride(0) == 1,
        "a must be contiguous in one dimension"
    );

    let lda = if row_major {
        into_i32(a.stride(0))
    } else {
        into_i32(a.stride(1))
    };

    unsafe {
        T::cblas_ger(
            if row_major {
                CBLAS_LAYOUT::CblasRowMajor
            } else {
                CBLAS_LAYOUT::CblasColMajor
            },
            into_i32(m),
            into_i32(n),
            beta,
            x.as_ptr(),
            x_inc,
            y.as_ptr(),
            y_inc,
            a.as_mut_ptr(),
            lda,
        )
    }
}

pub(super) fn asum<T, D1: Dim, Lx>(x: &Slice<T, (D1,), Lx>) -> T::Real
where
    T: BlasScalar + ComplexFloat,
    Lx: Layout,
{
    let n = into_i32(x.len());
    let incx = into_i32(x.stride(0));

    unsafe { T::cblas_asum(n, x.as_ptr(), incx) }
}

pub(super) fn axpy<T, D1: Dim, Lx, Ly>(alpha: T, x: &Slice<T, (D1,), Lx>, y: &mut Slice<T, (D1,), Ly>)
where
    T: BlasScalar + ComplexFloat,
    Lx: Layout,
    Ly: Layout,
{
    assert_eq!(x.len(), y.len(), "Vector lengths must match");

    let n = into_i32(x.len());
    let incx = into_i32(x.stride(0));
    let incy = into_i32(y.stride(0));

    unsafe { T::cblas_axpy(n, alpha, x.as_ptr(), incx, y.as_mut_ptr(), incy) }
}

pub(super) fn nrm2<T, D1: Dim, Lx>(x: &Slice<T, (D1,), Lx>) -> T::Real
where
    T: BlasScalar + ComplexFloat,
    Lx: Layout,
{
    let n = into_i32(x.len());
    let incx = into_i32(x.stride(0));

    unsafe { T::cblas_nrm2(n, x.as_ptr(), incx) }
}

pub(super) fn dotu<T: BlasScalar + ComplexFloat, D1: Dim, Lx: Layout, Ly: Layout>(
    x: &Slice<T, (D1,), Lx>,
    y: &Slice<T, (D1,), Ly>,
) -> T {
    assert_eq!(x.len(), y.len());
    let n = into_i32(x.len());
    unsafe {
        T::cblas_dotu_or_dot(
            n,
            x.as_ptr(),
            into_i32(x.stride(0)),
            y.as_ptr(),
            into_i32(y.stride(0)),
        )
    }
}

pub(super) fn dotc<T: BlasScalar + ComplexFloat, D1: Dim, Lx: Layout, Ly: Layout>(
    x: &Slice<T, (D1,), Lx>,
    y: &Slice<T, (D1,), Ly>,
) -> T {
    assert_eq!(x.len(), y.len());
    let n = into_i32(x.len());
    unsafe {
        T::cblas_dotc_or_dot(
            n,
            x.as_ptr(),
            into_i32(x.stride(0)),
            y.as_ptr(),
            into_i32(y.stride(0)),
        )
    }
}

pub(super) fn amax<T, S, L>(x: &Slice<T, S, L>) -> usize
where
    T: BlasScalar + ComplexFloat,
    S: Shape,
    L: Layout,
{
    assert!(!x.is_empty(), "Cannot find amax of empty slice");

    let n = into_i32(x.len());
    let incx = if x.rank() == 1 {
        into_i32(x.stride(0))
    } else {
        1
    };

    (unsafe { T::cblas_amax(n, x.as_ptr(), incx) } as usize)
}