diskann-linalg 0.53.0

/*
 * Copyright (c) Microsoft Corporation.
 * Licensed under the MIT license.
 */

use std::fmt;

pub mod common;
pub use common::Transpose;

mod faer;
use faer::{random_distance_preserving_matrix_impl, sgemm_impl, svd_into_impl};
use rand::Rng;

/// Matrix identifier for SGEMM operations.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MatrixName {
    A,
    B,
    C,
}

impl fmt::Display for MatrixName {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            MatrixName::A => write!(f, "a (m * k)"),
            MatrixName::B => write!(f, "b (k * n)"),
            MatrixName::C => write!(f, "c (m * n)"),
        }
    }
}

/// Error type for SGEMM operations.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum SgemmError {
    /// Matrix has incorrect dimensions.
    InvalidMatrixDimensions {
        matrix_name: MatrixName,
        expected_rows: usize,
        expected_cols: usize,
        actual_len: usize,
    },
    /// Dimension overflow when computing matrix size.
    DimensionOverflow {
        matrix_name: MatrixName,
        rows: usize,
        cols: usize,
    },
}

impl fmt::Display for SgemmError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            SgemmError::InvalidMatrixDimensions {
                matrix_name,
                expected_rows,
                expected_cols,
                actual_len,
            } => write!(
                f,
                "expected {}x{} matrix {} to have length {}, instead got {}",
                expected_rows,
                expected_cols,
                matrix_name,
                expected_rows * expected_cols,
                actual_len
            ),
            SgemmError::DimensionOverflow {
                matrix_name,
                rows,
                cols,
            } => write!(
                f,
                "dimension overflow in matrix {}: {} * {} would overflow usize",
                matrix_name, rows, cols
            ),
        }
    }
}

impl std::error::Error for SgemmError {}

// Make the reference implementation available for internal testing.
#[cfg(test)]
mod reference;

/// Matrix-matrix multiplication for implicit row-major matrices `a` and `b` using the
/// implicit row-major matrix `c` as the destination.
///
/// Performs one of the following operations:
/// ```ignore
/// 1. c = [beta * c] + alpha * a * b
/// 2. c = [beta * c] + alpha * a' * b
/// 3. c = [beta * c] + alpha * a * b'
/// 3. c = [beta * c] + alpha * a' * b'
/// ```
/// Where `x'` indicates the ordinary transpose of `x`.
///
/// If `beta` is `None`, the destination `c` is completely over-written.
///
/// * `atranspose`: Whether `a` should be interpreted as an in-place transpose.
/// * `btranspose`: Whether `b` should be interpreted as an in-place transpose.
/// * `m`: The number of rows in `c`. Additionally:
///     - If `!atranspose.is_transpose()`, this is the number of rows in `a`.
///     - If `atranspose.is_transpose()`, this is the number of rows in `a`.
/// * `n`: The number of columns `c`. Additionally:
///     - If `!btranspose.is_transpose()`, this is the number of columns in `b`.
///     - If `btranspose.is_transpose()`, this is the number of columns in `b`.
/// * `k`: The number of columns in matrix `a` and the number of rows in matrix `b`.
/// * `k`: Refer to the following:
///     - If `!atranspose.is_transpose()`, this is the number of columns in `a`.
///       Otherwise, this is the number of rows in `a`.
///     - If `!btranspose.is_transpose()`, this is the number of rows in `b`.
///       Otherwise, this is the number of columns in `b`.
/// * `alpha`: Scaling parameter for the operation `a * b`.
/// * `a`: The matrix `a` with dimension `m x k` (potentially after transposing).
/// * `b`: The matrix `b` with dimension `k x n` (potentially after transposing).
/// * `beta`: Optional scaling parameter for the matrix `c`. If `None`, then `c` will be
///   overwritten entirely.
/// * `c`: The output matrix with dimension `m x n`.
///
/// # Note
///
/// This inteface is a simplified version of the full cblas `sgemm` interface, namely that it
///
/// 1. Does not support column-major layouts
/// 2. Does not allow for arbitrary strides in the leading dimension of the matrices.
///
/// This is to support the common-case in DiskANN that uses a Row-Major layout and always
/// uses dense matrices.
///
/// If the more esoteric features of the cblas `sgemm` API are needed, we can provide
/// that as an interface extension.
///
/// # Errors
///
/// Returns an error if:
/// * `m * k` would overflow `usize`
/// * `k * n` would overflow `usize`
/// * `m * n` would overflow `usize`
/// * `a.len() != m * k`
/// * `b.len() != k * n`
/// * `c.len() != m * n`
#[allow(clippy::too_many_arguments)]
pub fn sgemm(
    atranspose: Transpose,
    btranspose: Transpose,
    m: usize,
    n: usize,
    k: usize,
    alpha: f32,
    a: &[f32],
    b: &[f32],
    beta: Option<f32>,
    c: &mut [f32],
) -> Result<(), SgemmError> {
    // Check size requirements with overflow protection.
    let expected_a_len = m.checked_mul(k).ok_or(SgemmError::DimensionOverflow {
        matrix_name: MatrixName::A,
        rows: m,
        cols: k,
    })?;

    if a.len() != expected_a_len {
        return Err(SgemmError::InvalidMatrixDimensions {
            matrix_name: MatrixName::A,
            expected_rows: m,
            expected_cols: k,
            actual_len: a.len(),
        });
    }

    let expected_b_len = k.checked_mul(n).ok_or(SgemmError::DimensionOverflow {
        matrix_name: MatrixName::B,
        rows: k,
        cols: n,
    })?;

    if b.len() != expected_b_len {
        return Err(SgemmError::InvalidMatrixDimensions {
            matrix_name: MatrixName::B,
            expected_rows: k,
            expected_cols: n,
            actual_len: b.len(),
        });
    }

    let expected_c_len = m.checked_mul(n).ok_or(SgemmError::DimensionOverflow {
        matrix_name: MatrixName::C,
        rows: m,
        cols: n,
    })?;

    if c.len() != expected_c_len {
        return Err(SgemmError::InvalidMatrixDimensions {
            matrix_name: MatrixName::C,
            expected_rows: m,
            expected_cols: n,
            actual_len: c.len(),
        });
    }

    // Invoke the actual implementation.
    sgemm_impl(atranspose, btranspose, m, n, k, alpha, a, b, beta, c);
    Ok(())
}

/// Compute the SVD of the provided matrix implicit row-major matrix `data`.
///
/// * `m`: The number of rows in `a`.
/// * `n`: The number of columns in `a`.
/// * `a`: The data matrix to decompose with dimensiuon `m x n` stored in Row-Major order [Note 1].
/// * `singular_values`: Contains the singular values of `a` sorted so that
///   `singular_values[i] ≥ singular_values[i+1]`.
/// * `u`: Contains the `m x m` unitary matrix in Row-Major order.
/// * `vt`: Contains the `n x n` unitary matrix in Column-Major order [Note 2].
///
/// # Notes
///
/// 1. Due to the contract offered by `lapacke`, callers of this function must assume that
///    the contents of `a` are left in an undefined state after this function.
///
///    See: https://netlib.org/lapack/explore-html//df/d22/group__gesdd_gab9ffdde22b38f0cc442e44cbea23818f.html
///
/// 2. Similar to #1, the restriction that `vt` is transposed is a lapack byproduct.
///
/// # Panics
///
/// Panics if
///
/// * `a.len() != m * n`
/// * `singular_values.len() != min(m, n)`
/// * `u.len() != m * m`.
/// * `vt.len() != n * n`.
///
/// Additionally, if MKL is used, panics if any either `m` or `n` is not representable
/// as a signed 32-bit integer due to `cblas` limitations.
pub fn svd_into(
    m: usize,
    n: usize,
    a: &mut [f32],
    singular_values: &mut [f32],
    u: &mut [f32],
    vt: &mut [f32],
) -> Result<(), impl std::error::Error + 'static> {
    // Check size requirements.
    assert_eq!(a.len(), m * n);
    assert_eq!(singular_values.len(), m.min(n));
    assert_eq!(u.len(), m * m);
    assert_eq!(vt.len(), n * n);

    // Invoke the actual implementation.
    svd_into_impl(m, n, a, singular_values, u, vt)
}

/// Construct a random `dim x dim` distance preserving matrix.
///
/// Practically speaking, the returned matrix should be orthogonal with a determinant of
/// either +1 or -1.
pub fn random_distance_preserving_matrix<T: Rng + ?Sized>(dim: usize, rng: &mut T) -> Vec<f32> {
    random_distance_preserving_matrix_impl(dim, rng)
}

#[cfg(test)]
mod tests {
    use approx::{assert_abs_diff_eq, assert_relative_eq};
    use rand::{distr::Distribution, rngs::StdRng, SeedableRng};
    use rand_distr::StandardNormal;
    use serde::Deserialize;

    use super::*;
    use crate::reference;

    ////////////////////////
    // Simple SGEMM tests //
    ////////////////////////

    #[test]
    fn test_reference_implementation() {
        let problems = reference::test_sgemm_problems();
        for (i, problem) in problems.iter().enumerate() {
            let result = problem.check(sgemm);
            if let Err(err) = result {
                panic!("{} on iteration {}. Problem: {:?}", err, i, problem);
            }
        }
    }

    #[test]
    fn test_sgemm_invalid_matrix_a_dimensions() {
        let mut c = [0.0f32; 6];
        let err = sgemm(
            Transpose::None,
            Transpose::None,
            2,
            3,
            4,
            1.0,
            &[0.0; 5], // Should be 2 * 4 = 8
            &[0.0; 12],
            None,
            &mut c,
        )
        .unwrap_err();

        assert_eq!(
            err.to_string(),
            "expected 2x4 matrix a (m * k) to have length 8, instead got 5"
        );
    }

    #[test]
    fn test_sgemm_invalid_matrix_b_dimensions() {
        let mut c = [0.0f32; 6];
        let err = sgemm(
            Transpose::None,
            Transpose::None,
            2,
            3,
            4,
            1.0,
            &[0.0; 8],
            &[0.0; 10], // Should be 4 * 3 = 12
            None,
            &mut c,
        )
        .unwrap_err();

        assert_eq!(
            err.to_string(),
            "expected 4x3 matrix b (k * n) to have length 12, instead got 10"
        );
    }

    #[test]
    fn test_sgemm_invalid_matrix_c_dimensions() {
        let mut c = [0.0f32; 5]; // Should be 2 * 3 = 6
        let err = sgemm(
            Transpose::None,
            Transpose::None,
            2,
            3,
            4,
            1.0,
            &[0.0; 8],
            &[0.0; 12],
            None,
            &mut c,
        )
        .unwrap_err();

        assert_eq!(
            err.to_string(),
            "expected 2x3 matrix c (m * n) to have length 6, instead got 5"
        );
    }

    #[test]
    fn test_sgemm_m_times_k_overflow() {
        let mut c = [0.0f32];
        let err = sgemm(
            Transpose::None,
            Transpose::None,
            usize::MAX,
            1,
            2,
            1.0,
            &[],
            &[0.0],
            None,
            &mut c,
        )
        .unwrap_err();

        assert_eq!(
            err.to_string(),
            format!(
                "dimension overflow in matrix a (m * k): {} * 2 would overflow usize",
                usize::MAX
            )
        );
    }

    #[test]
    fn test_sgemm_k_times_n_overflow() {
        let mut c = vec![0.0f32; 10];
        let err = sgemm(
            Transpose::None,
            Transpose::None,
            1,
            usize::MAX,
            10,
            1.0,
            &[0.0f32; 10],
            &[],
            None,
            &mut c,
        )
        .unwrap_err();

        assert_eq!(
            err.to_string(),
            format!(
                "dimension overflow in matrix b (k * n): 10 * {} would overflow usize",
                usize::MAX
            )
        );
    }

    #[test]
    fn test_sgemm_m_times_n_overflow() {
        let mut c = [];
        let err = sgemm(
            Transpose::None,
            Transpose::None,
            2,
            usize::MAX,
            0,
            1.0,
            &[],
            &[],
            None,
            &mut c,
        )
        .unwrap_err();

        assert_eq!(
            err.to_string(),
            format!(
                "dimension overflow in matrix c (m * n): 2 * {} would overflow usize",
                usize::MAX
            )
        );
    }

    /// This test ensures the Result type doesn't grow unexpectedly.
    /// A large Result size increases stack usage even for the Ok(()) case.
    #[test]
    fn test_sgemm_result_size() {
        let mut c = [0.0f32; 6];
        let result = sgemm(
            Transpose::None,
            Transpose::None,
            2,
            3,
            4,
            1.0,
            &[0.0; 5],
            &[0.0; 12],
            None,
            &mut c,
        );

        let result_size = std::mem::size_of_val(&result);
        const EXPECTED_RESULT_SIZE: usize = 32;
        assert_eq!(
            result_size, EXPECTED_RESULT_SIZE,
            "Result size is {} bytes, does not match the expected size of {} bytes.",
            result_size, EXPECTED_RESULT_SIZE
        );
    }

    ///////////////
    // SVD Tests //
    ///////////////

    fn test_file_path(name: &str) -> String {
        format!("{}/test_data/{}", env!("CARGO_MANIFEST_DIR"), name)
    }

    /// The generate set of reference SVD input files.
    const SVD_INPUT_FILE: &str = "reference_svd_inputs.json";

    #[derive(Deserialize, Debug)]
    struct SVDTestCase {
        m: usize,
        n: usize,
        matrix: Vec<f32>,
        singular_values: Vec<f32>,
    }

    impl SVDTestCase {
        fn summary(&self) -> String {
            format!("svd test case with dimension {}x{}", self.m, self.n)
        }
    }

    struct SVDTolerance {
        absolute: f32,
        relative: f32,
    }

    impl SVDTolerance {
        fn check(&self, absolute: f32, relative: f32) -> bool {
            absolute <= self.absolute || relative <= self.relative
        }
    }

    fn materialize_singular_values(singular_values: &[f32], m: usize, n: usize) -> Vec<f32> {
        assert_eq!(singular_values.len(), m.min(n));
        let mut output = vec![0.0; m * n];

        for (i, &s) in singular_values.iter().enumerate() {
            output[n * i + i] = s;
        }
        output
    }

    fn test_svd(
        case: &SVDTestCase,
        singular_value_tolerance: &SVDTolerance,
        reconstructed_tolerance: &SVDTolerance,
        context: &dyn std::fmt::Display,
    ) {
        // Create the output matrices.
        let mut singular_values = vec![0.0; case.m.min(case.n)];
        let mut u = vec![0.0; case.m * case.m];
        let mut vt = vec![0.0; case.n * case.n];

        svd_into(
            case.m,
            case.n,
            &mut case.matrix.clone(),
            &mut singular_values,
            &mut u,
            &mut vt,
        )
        .unwrap();

        // Check the resulting singular values.
        for (i, (&got, &expected)) in
            std::iter::zip(singular_values.iter(), case.singular_values.iter()).enumerate()
        {
            let diff = (got - expected).abs();
            let relative = diff / expected;
            assert!(
                singular_value_tolerance.check(diff, relative),
                "got {} but expected {} (diff: {}, relative: {}) at position {}: {}",
                got,
                expected,
                diff,
                relative,
                i,
                context
            );
        }

        // Test the reconstruction.
        let full_singular_values = materialize_singular_values(&singular_values, case.m, case.n);
        let mut temp = vec![0.0; case.m * case.n];

        // Multiply `u * singular_values`.
        sgemm(
            Transpose::None,
            Transpose::None,
            case.m,
            case.n,
            case.m,
            1.0,
            &u,
            &full_singular_values,
            None,
            &mut temp,
        )
        .unwrap();

        let mut output = vec![0.0; case.m * case.n];
        sgemm(
            Transpose::None,
            Transpose::None,
            case.m,
            case.n,
            case.n,
            1.0,
            &temp,
            &vt,
            None,
            &mut output,
        )
        .unwrap();

        for row in 0..case.m {
            for col in 0..case.n {
                let got = output[case.n * row + col];
                let expected = case.matrix[case.n * row + col];
                let diff = (got - expected).abs();
                let relative = diff / expected;
                assert!(
                    reconstructed_tolerance.check(diff, relative),
                    "mismatch in reconstructed matrix at (row, col) = ({}, {}). \
                     Got {}, expected {} (diff: {}, relative: {}). {}",
                    row,
                    col,
                    got,
                    expected,
                    diff,
                    relative,
                    context
                );
            }
        }
    }

    #[test]
    fn test_svd_implementation() {
        let path = test_file_path(SVD_INPUT_FILE);
        let file = std::fs::File::open(path.clone())
            .unwrap_or_else(|_| panic!("failed to open file {path}"));

        let reader = std::io::BufReader::new(file);
        let cases: Vec<SVDTestCase> = serde_json::from_reader(reader).unwrap();

        let singular_values_tolerance = SVDTolerance {
            absolute: 2.0e-6,
            relative: 3.0e-6,
        };

        let reconstructed_tolerance = SVDTolerance {
            absolute: 5.0e-5,
            relative: 0.0,
        };

        for (i, case) in cases.iter().enumerate() {
            let context = format!(
                "while processing case {} of {}: {}",
                i + 1,
                cases.len(),
                case.summary()
            );
            test_svd(
                case,
                &singular_values_tolerance,
                &reconstructed_tolerance,
                &context,
            );
        }
    }

    ///////////////////////////
    // Rotation Matrix Tests //
    ///////////////////////////

    const EPSILON: f32 = 1e-5;

    fn test_distance_preserving_matrix_impl(dim: usize, rng: &mut StdRng) {
        // Construct the distance preserving matrix.
        let q = random_distance_preserving_matrix(dim, rng);

        // Check that `q * q'` is close to the identity matrix.
        let qm = ::faer::mat::MatRef::from_row_major_slice(&q, dim, dim);
        let m = qm * qm.transpose();

        for j in 0..dim {
            for i in 0..dim {
                if i == j {
                    assert_abs_diff_eq!(m[(i, j)], 1.0, epsilon = EPSILON);
                } else {
                    assert_abs_diff_eq!(m[(i, j)], 0.0, epsilon = EPSILON);
                }
            }
        }

        // Instead of explicitly checking the determinant, we sample using 100 randomly
        // generated vectors, verifying that the norms are unchanged.
        const RANDOM_TRIALS: usize = 100;
        let mut v = vec![0.0f32; dim];
        for _ in 0..RANDOM_TRIALS {
            v.iter_mut()
                .for_each(|i| *i = StandardNormal {}.sample(rng));
            let vm = ::faer::mat::MatRef::from_row_major_slice(&v, dim, 1);
            let v_norm = vm.squared_norm_l2();
            let t = qm * vm;
            let t_norm = t.squared_norm_l2();

            assert_relative_eq!(v_norm, t_norm, epsilon = EPSILON, max_relative = EPSILON);
            assert_ne!(vm, t);
        }
    }

    #[test]
    fn test_rotation_matrix() {
        let mut rng = StdRng::seed_from_u64(0xc0ff33);
        let num_trials = 5;
        for dim in [2, 100, 256] {
            for _ in 0..num_trials {
                test_distance_preserving_matrix_impl(dim, &mut rng);
            }
        }
    }
}