topos 0.13.0

//! The Accelerate backend: dense GEMM on Apple's matrix units.
//!
//! One `cblas_sgemm`/`cblas_dgemm` call per task, executing on the
//! AMX/SME coprocessor on Apple Silicon (AVX kernels on Intel Macs)
//! with function-call latency — no device, no queue, no state. The
//! module is a pure function from task to product: classification of
//! the task's strides into BLAS transpose flags plus leading
//! dimensions, one foreign call, done. Tasks the mapping cannot
//! express (a stride-0 broadcast, dimensions beyond `i32`) and tasks
//! below the profitability threshold decline to the built-in paths.
//!
//! This is the crate's only `unsafe` code in an accelerate-only
//! build (every other backend feature carries its own); each
//! backend's `kernels` submodule is scope-allowed under the
//! crate-wide `deny(unsafe_code)`, while the descriptor half stays
//! outside the allow.

use crate::backend::operand::{Operand, classify};
use crate::{BatchNormTask, GemmTask, MapOperation, Normalized};

// Row-major CBLAS constants.
const ROW_MAJOR: i32 = 101;
const NO_TRANSPOSE: i32 = 111;
const TRANSPOSE: i32 = 112;

/// Returns the cblas transpose constant for a classified operand.
fn transpose(operand: &Operand) -> i32 {
    if operand.transposed {
        TRANSPOSE
    } else {
        NO_TRANSPOSE
    }
}

/// Below this many floating-point operations (`2 * m * n * k`) the
/// built-in slice path wins on latency alone; the crossover sits
/// around n = 16 square and everything real is far above it.
const FLOP_THRESHOLD: usize = 1 << 13;

/// Below this many elements a vForce call's setup outweighs the
/// scalar loop; the crossover is small and flat, so the constant is
/// conservative rather than tuned.
const MAP_THRESHOLD: usize = 1 << 7;

#[link(name = "Accelerate", kind = "framework")]
unsafe extern "C" {
    fn cblas_sgemm(
        order: i32,
        transpose_a: i32,
        transpose_b: i32,
        m: i32,
        n: i32,
        k: i32,
        alpha: f32,
        a: *const f32,
        leading_a: i32,
        b: *const f32,
        leading_b: i32,
        beta: f32,
        c: *mut f32,
        leading_c: i32,
    );
    fn cblas_dgemm(
        order: i32,
        transpose_a: i32,
        transpose_b: i32,
        m: i32,
        n: i32,
        k: i32,
        alpha: f64,
        a: *const f64,
        leading_a: i32,
        b: *const f64,
        leading_b: i32,
        beta: f64,
        c: *mut f64,
        leading_c: i32,
    );
    // vForce: vectorized transcendentals over whole buffers, the
    // library form of the loops libm calls keep scalar.
    fn vvexpf(mapped: *mut f32, elements: *const f32, count: *const i32);
    fn vvlogf(mapped: *mut f32, elements: *const f32, count: *const i32);
    fn vvsqrtf(mapped: *mut f32, elements: *const f32, count: *const i32);
    fn vvtanhf(mapped: *mut f32, elements: *const f32, count: *const i32);
    fn vvsinf(mapped: *mut f32, elements: *const f32, count: *const i32);
    fn vvcosf(mapped: *mut f32, elements: *const f32, count: *const i32);
    fn vvlog1pf(mapped: *mut f32, elements: *const f32, count: *const i32);
    fn vvexpm1f(mapped: *mut f32, elements: *const f32, count: *const i32);
    fn vvexp(mapped: *mut f64, elements: *const f64, count: *const i32);
    fn vvlog(mapped: *mut f64, elements: *const f64, count: *const i32);
    fn vvsqrt(mapped: *mut f64, elements: *const f64, count: *const i32);
    fn vvtanh(mapped: *mut f64, elements: *const f64, count: *const i32);
    fn vvsin(mapped: *mut f64, elements: *const f64, count: *const i32);
    fn vvcos(mapped: *mut f64, elements: *const f64, count: *const i32);
    fn vvlog1p(mapped: *mut f64, elements: *const f64, count: *const i32);
    fn vvexpm1(mapped: *mut f64, elements: *const f64, count: *const i32);
    // vDSP: the contiguous row passes behind the batch-normalization
    // kernel. Strides are in elements; counts are element counts.
    // Note the argument order of `vsub`: it computes `C = B - A`.
    fn vDSP_vadd(
        a: *const f32,
        a_stride: isize,
        b: *const f32,
        b_stride: isize,
        sum: *mut f32,
        sum_stride: isize,
        count: usize,
    );
    fn vDSP_vsub(
        a: *const f32,
        a_stride: isize,
        b: *const f32,
        b_stride: isize,
        difference: *mut f32,
        difference_stride: isize,
        count: usize,
    );
    fn vDSP_vma(
        a: *const f32,
        a_stride: isize,
        b: *const f32,
        b_stride: isize,
        c: *const f32,
        c_stride: isize,
        result: *mut f32,
        result_stride: isize,
        count: usize,
    );
    fn vDSP_vaddD(
        a: *const f64,
        a_stride: isize,
        b: *const f64,
        b_stride: isize,
        sum: *mut f64,
        sum_stride: isize,
        count: usize,
    );
    fn vDSP_vsubD(
        a: *const f64,
        a_stride: isize,
        b: *const f64,
        b_stride: isize,
        difference: *mut f64,
        difference_stride: isize,
        count: usize,
    );
    fn vDSP_vmaD(
        a: *const f64,
        a_stride: isize,
        b: *const f64,
        b_stride: isize,
        c: *const f64,
        c_stride: isize,
        result: *mut f64,
        result_stride: isize,
        count: usize,
    );
}

/// Below this many elements the per-feature vDSP call overhead
/// outweighs the composed path; the constant is provisional until
/// the measurement plan runs.
const BATCH_NORM_THRESHOLD: usize = 1 << 12;

/// It runs a whole training-mode batch normalization through vDSP,
/// in row-major passes so every call is contiguous: accumulate the
/// per-feature sums row by row, accumulate the centered squares for
/// the biased variance, then apply the affine as one fused
/// multiply-add per row — or declines with `None` below the
/// threshold. A per-feature strided formulation measured flat
/// against the composed path (one cache line per element); the row
/// passes are what pay.
pub(crate) fn batch_norm_f32(task: &BatchNormTask<'_, f32>) -> Option<Normalized<f32>> {
    let (batch, features) = (task.batch(), task.features());
    if batch * features < BATCH_NORM_THRESHOLD {
        return None;
    }
    let input = task.input();
    let scale = task.scale();
    let shift = task.shift();
    let epsilon = *task.epsilon();
    let mut output = vec![0.0_f32; batch * features];
    let mut mean = vec![0.0_f32; features];
    let mut variance = vec![0.0_f32; features];
    let mut centered = vec![0.0_f32; features];
    // SAFETY: every pointer addresses a live buffer of at least
    // `features` elements — rows of the validated
    // `batch * features` input, or the `features`-sized scratch and
    // output vectors — with unit strides, so each call touches
    // exactly `features` contiguous elements; the accumulating
    // calls write in place through the documented vDSP in-place
    // contract (equal strides, output aliasing one input); all
    // other outputs are exclusively borrowed.
    unsafe {
        for row in 0..batch {
            let elements = input[row * features..].as_ptr();
            vDSP_vadd(
                elements,
                1,
                mean.as_ptr(),
                1,
                mean.as_mut_ptr(),
                1,
                features,
            );
        }
        let inverse_batch = 1.0 / batch as f32;
        for entry in &mut mean {
            *entry *= inverse_batch;
        }
        for row in 0..batch {
            let elements = input[row * features..].as_ptr();
            vDSP_vsub(
                mean.as_ptr(),
                1,
                elements,
                1,
                centered.as_mut_ptr(),
                1,
                features,
            );
            vDSP_vma(
                centered.as_ptr(),
                1,
                centered.as_ptr(),
                1,
                variance.as_ptr(),
                1,
                variance.as_mut_ptr(),
                1,
                features,
            );
        }
        // The per-feature affine: `output = input * a + b` with
        // `a = scale / sqrt(variance + epsilon)` and
        // `b = shift - mean * a`.
        let mut multiplier = vec![0.0_f32; features];
        let mut addend = vec![0.0_f32; features];
        for feature in 0..features {
            variance[feature] *= inverse_batch;
            multiplier[feature] = scale[feature] / (variance[feature] + epsilon).sqrt();
            addend[feature] = shift[feature] - mean[feature] * multiplier[feature];
        }
        for row in 0..batch {
            let elements = input[row * features..].as_ptr();
            vDSP_vma(
                elements,
                1,
                multiplier.as_ptr(),
                1,
                addend.as_ptr(),
                1,
                output[row * features..].as_mut_ptr(),
                1,
                features,
            );
        }
    }
    Some(Normalized {
        output,
        mean,
        variance,
    })
}

/// The `f64` twin of [`batch_norm_f32`].
pub(crate) fn batch_norm_f64(task: &BatchNormTask<'_, f64>) -> Option<Normalized<f64>> {
    let (batch, features) = (task.batch(), task.features());
    if batch * features < BATCH_NORM_THRESHOLD {
        return None;
    }
    let input = task.input();
    let scale = task.scale();
    let shift = task.shift();
    let epsilon = *task.epsilon();
    let mut output = vec![0.0_f64; batch * features];
    let mut mean = vec![0.0_f64; features];
    let mut variance = vec![0.0_f64; features];
    let mut centered = vec![0.0_f64; features];
    // SAFETY: identical to `batch_norm_f32` — contiguous
    // `features`-length rows of validated buffers, the documented
    // in-place accumulation, exclusive outputs.
    unsafe {
        for row in 0..batch {
            let elements = input[row * features..].as_ptr();
            vDSP_vaddD(
                elements,
                1,
                mean.as_ptr(),
                1,
                mean.as_mut_ptr(),
                1,
                features,
            );
        }
        let inverse_batch = 1.0 / batch as f64;
        for entry in &mut mean {
            *entry *= inverse_batch;
        }
        for row in 0..batch {
            let elements = input[row * features..].as_ptr();
            vDSP_vsubD(
                mean.as_ptr(),
                1,
                elements,
                1,
                centered.as_mut_ptr(),
                1,
                features,
            );
            vDSP_vmaD(
                centered.as_ptr(),
                1,
                centered.as_ptr(),
                1,
                variance.as_ptr(),
                1,
                variance.as_mut_ptr(),
                1,
                features,
            );
        }
        let mut multiplier = vec![0.0_f64; features];
        let mut addend = vec![0.0_f64; features];
        for feature in 0..features {
            variance[feature] *= inverse_batch;
            multiplier[feature] = scale[feature] / (variance[feature] + epsilon).sqrt();
            addend[feature] = shift[feature] - mean[feature] * multiplier[feature];
        }
        for row in 0..batch {
            let elements = input[row * features..].as_ptr();
            vDSP_vmaD(
                elements,
                1,
                multiplier.as_ptr(),
                1,
                addend.as_ptr(),
                1,
                output[row * features..].as_mut_ptr(),
                1,
                features,
            );
        }
    }
    Some(Normalized {
        output,
        mean,
        variance,
    })
}

/// It runs a `f32` task through `cblas_sgemm`, or declines with
/// `None` when the task is below the threshold or outside the
/// mapping.
pub(crate) fn gemm_f32(task: &GemmTask<'_, f32>) -> Option<Vec<f32>> {
    if flops(task.m(), task.n(), task.k()) < FLOP_THRESHOLD {
        return None;
    }
    executed_f32(task)
}

/// It runs a `f64` task through `cblas_dgemm`, with the same decline
/// rules as the `f32` twin.
pub(crate) fn gemm_f64(task: &GemmTask<'_, f64>) -> Option<Vec<f64>> {
    if flops(task.m(), task.n(), task.k()) < FLOP_THRESHOLD {
        return None;
    }
    executed_f64(task)
}

/// Returns the task's floating-point operation count, saturating —
/// a saturated count is enormous and therefore above any threshold.
fn flops(m: usize, n: usize, k: usize) -> usize {
    2usize.saturating_mul(m).saturating_mul(n).saturating_mul(k)
}

/// The `f32` call without the threshold gate, so tests can drive the
/// mapping over shapes of every size.
pub(super) fn executed_f32(task: &GemmTask<'_, f32>) -> Option<Vec<f32>> {
    let a = classify(task.a_strides(), task.m(), task.k())?;
    let b = classify(task.b_strides(), task.k(), task.n())?;
    let m = i32::try_from(task.m()).ok()?;
    let n = i32::try_from(task.n()).ok()?;
    let k = i32::try_from(task.k()).ok()?;
    let mut product = vec![0.0_f32; task.m() * task.n()];
    // SAFETY: the operand pointers come from live slices whose spans
    // the `GemmTask` constructor validated against the dimensions
    // and strides; `classify` guarantees the leading dimensions
    // satisfy the cblas access-pattern contract, so every read is in
    // bounds; `product` is exclusively borrowed and sized `m * n`
    // with `leading_c = n`; with `beta = 0` cblas only writes `c`
    // and only reads `a` and `b`.
    unsafe {
        cblas_sgemm(
            ROW_MAJOR,
            transpose(&a),
            transpose(&b),
            m,
            n,
            k,
            1.0,
            task.a().as_ptr(),
            a.leading,
            task.b().as_ptr(),
            b.leading,
            0.0,
            product.as_mut_ptr(),
            n,
        );
    }
    Some(product)
}

/// It maps one transcendental over an `f32` buffer through vForce,
/// declining buffers too small to pay the call or too long for the
/// interface's `i32` count.
pub(crate) fn map_f32(operation: MapOperation, elements: &[f32]) -> Option<Vec<f32>> {
    // vForce has no erf family; the reference path computes it.
    if matches!(operation, MapOperation::Erf | MapOperation::ErfDerivative) {
        return None;
    }
    if elements.len() < MAP_THRESHOLD {
        return None;
    }
    let count = i32::try_from(elements.len()).ok()?;
    let mut mapped = vec![0.0_f32; elements.len()];
    // SAFETY: both pointers address live buffers of exactly `count`
    // elements — `mapped` exclusively — and vForce reads the input
    // and count while writing only the output.
    unsafe {
        match operation {
            MapOperation::Exp => vvexpf(mapped.as_mut_ptr(), elements.as_ptr(), &count),
            MapOperation::Ln => vvlogf(mapped.as_mut_ptr(), elements.as_ptr(), &count),
            MapOperation::Sqrt => vvsqrtf(mapped.as_mut_ptr(), elements.as_ptr(), &count),
            MapOperation::Tanh => vvtanhf(mapped.as_mut_ptr(), elements.as_ptr(), &count),
            MapOperation::Sin => vvsinf(mapped.as_mut_ptr(), elements.as_ptr(), &count),
            MapOperation::Cos => vvcosf(mapped.as_mut_ptr(), elements.as_ptr(), &count),
            MapOperation::Log1p => vvlog1pf(mapped.as_mut_ptr(), elements.as_ptr(), &count),
            MapOperation::Expm1 => vvexpm1f(mapped.as_mut_ptr(), elements.as_ptr(), &count),
            // The entry declines the pair before any dispatch.
            MapOperation::Erf | MapOperation::ErfDerivative => {
                unreachable!("the erf pair has no vForce kernel")
            }
        }
    }
    Some(mapped)
}

/// The `f64` twin of [`map_f32`].
pub(crate) fn map_f64(operation: MapOperation, elements: &[f64]) -> Option<Vec<f64>> {
    if matches!(operation, MapOperation::Erf | MapOperation::ErfDerivative) {
        return None;
    }
    if elements.len() < MAP_THRESHOLD {
        return None;
    }
    let count = i32::try_from(elements.len()).ok()?;
    let mut mapped = vec![0.0_f64; elements.len()];
    // SAFETY: identical to `map_f32` — live buffers of `count`
    // elements, exclusive output, read-only input.
    unsafe {
        match operation {
            MapOperation::Exp => vvexp(mapped.as_mut_ptr(), elements.as_ptr(), &count),
            MapOperation::Ln => vvlog(mapped.as_mut_ptr(), elements.as_ptr(), &count),
            MapOperation::Sqrt => vvsqrt(mapped.as_mut_ptr(), elements.as_ptr(), &count),
            MapOperation::Tanh => vvtanh(mapped.as_mut_ptr(), elements.as_ptr(), &count),
            MapOperation::Sin => vvsin(mapped.as_mut_ptr(), elements.as_ptr(), &count),
            MapOperation::Cos => vvcos(mapped.as_mut_ptr(), elements.as_ptr(), &count),
            MapOperation::Log1p => vvlog1p(mapped.as_mut_ptr(), elements.as_ptr(), &count),
            MapOperation::Expm1 => vvexpm1(mapped.as_mut_ptr(), elements.as_ptr(), &count),
            // The entry declines the pair before any dispatch.
            MapOperation::Erf | MapOperation::ErfDerivative => {
                unreachable!("the erf pair has no vForce kernel")
            }
        }
    }
    Some(mapped)
}

/// The `f64` call without the threshold gate; see [`executed_f32`].
pub(super) fn executed_f64(task: &GemmTask<'_, f64>) -> Option<Vec<f64>> {
    let a = classify(task.a_strides(), task.m(), task.k())?;
    let b = classify(task.b_strides(), task.k(), task.n())?;
    let m = i32::try_from(task.m()).ok()?;
    let n = i32::try_from(task.n()).ok()?;
    let k = i32::try_from(task.k()).ok()?;
    let mut product = vec![0.0_f64; task.m() * task.n()];
    // SAFETY: identical to `executed_f32` — validated spans, checked
    // leading dimensions, an exclusive `m * n` output, and cblas's
    // read/write contract under `beta = 0`.
    unsafe {
        cblas_dgemm(
            ROW_MAJOR,
            transpose(&a),
            transpose(&b),
            m,
            n,
            k,
            1.0,
            task.a().as_ptr(),
            a.leading,
            task.b().as_ptr(),
            b.leading,
            0.0,
            product.as_mut_ptr(),
            n,
        );
    }
    Some(product)
}

#[cfg(test)]
#[path = "tests/accelerate_tests.rs"]
mod tests;