matrixmultiply 0.3.10

General matrix multiplication for f32 and f64 matrices. Operates on matrices with general layout (they can use arbitrary row and column stride). Detects and uses AVX or SSE2 on x86 platforms transparently for higher performance. Uses a microkernel strategy, so that the implementation is easy to parallelize and optimize. Supports multithreading.
Documentation
extern crate matrixmultiply;
pub use matrixmultiply::dgemm;
pub use matrixmultiply::sgemm;

#[macro_use]
extern crate bencher;

// Compute GFlop/s
// by flop / s = 2 M N K / time

benchmark_main!(mat_mul_f32, mat_mul_f64, layout_f32_032, layout_f64_032);

macro_rules! mat_mul {
    ($modname:ident, $gemm:ident, $(($name:ident, $m:expr, $n:expr, $k:expr))+) => {
        mod $modname {
            use bencher::{Bencher};
            use crate::$gemm;
            $(
            pub fn $name(bench: &mut Bencher)
            {
                let a = vec![0.; $m * $n];
                let b = vec![0.; $n * $k];
                let mut c = vec![0.; $m * $k];
                bench.iter(|| {
                    unsafe {
                        $gemm(
                            $m, $n, $k,
                            1.,
                            a.as_ptr(), $n, 1,
                            b.as_ptr(), $k, 1,
                            0.,
                            c.as_mut_ptr(), $k, 1,
                            )
                    }
                });
            }
            )+
        }
        benchmark_group!{ $modname, $($modname::$name),+ }
    };
}

mat_mul! {mat_mul_f32, sgemm,
    (m004, 4, 4, 4)
    (m006, 6, 6, 6)
    (m008, 8, 8, 8)
    (m012, 12, 12, 12)
    (m016, 16, 16, 16)
    (m032, 32, 32, 32)
    (m064, 64, 64, 64)
    (m127, 127, 127, 127)
    /*
    (m256, 256, 256, 256)
    (m512, 512, 512, 512)
    (mix16x4, 32, 4, 32)
    (mix32x2, 32, 2, 32)
    (mix97, 97, 97, 125)
    (mix128x10000x128, 128, 10000, 128)
    */
}

mat_mul! {mat_mul_f64, dgemm,
    (m004, 4, 4, 4)
    (m006, 6, 6, 6)
    (m008, 8, 8, 8)
    (m012, 12, 12, 12)
    (m016, 16, 16, 16)
    (m032, 32, 32, 32)
    (m064, 64, 64, 64)
    (m127, 127, 127, 127)
    /*
    (m256, 256, 256, 256)
    (m512, 512, 512, 512)
    (mix16x4, 32, 4, 32)
    (mix32x2, 32, 2, 32)
    (mix97, 97, 97, 125)
    (mix128x10000x128, 128, 10000, 128)
    */
}

/// benchmarks combinations of inputs using various layouts
/// row-major ("c") vs column-major ("f") experiments
///
/// These benchmarks give information about
///
/// 1. Matrix packing sensitivity to input layouts (A, B layouts)
/// 2. Microkernel sensitivity to output layouts (C layout)
///    and performance for beta != 0. vs == 0.
///
/// Clike: all elements spaced at least 2 apart
/// Flike: all elements spaced at least 2 apart

enum Layout {
    C,
    F,
    Clike,
    Flike,
}
use self::Layout::*;

impl Layout {
    fn spread(&self) -> usize {
        match *self {
            C | F => 1,
            Clike | Flike => 2,
        }
    }

    fn strides(&self, rs: isize, cs: isize) -> (isize, isize) {
        let spread = self.spread() as isize;
        match *self {
            C | Clike => (rs * spread, cs * spread),
            F | Flike => (cs * spread, rs * spread),
        }
    }
}

macro_rules! gemm_layout {
    ($modname:ident, $gemm:ident, $(($name:ident, $m:expr))+) => {
        mod $modname {
            use bencher::{Bencher};
            use super::Layout::{self, *};
            use crate::$gemm;
            $(

            fn base(bench: &mut Bencher, al: Layout, bl: Layout, cl: Layout, use_beta: bool)
            {
                let a = vec![0.; $m * $m * al.spread()];
                let b = vec![0.; $m * $m * bl.spread()];
                let mut c = vec![0.; $m * $m * cl.spread()];
                let beta = if use_beta { 0.1 } else { 0. };

                let (rsa, csa) = al.strides($m, 1);
                let (rsb, csb) = bl.strides($m, 1);
                let (rsc, csc) = cl.strides($m, 1);

                let max_stride_a = (rsa as usize) * ($m - 1) + (csa as usize) * ($m - 1);
                let max_stride_b = (rsb as usize) * ($m - 1) + (csb as usize) * ($m - 1);
                let max_stride_c = (rsc as usize) * ($m - 1) + (csc as usize) * ($m - 1);

                debug_assert!(max_stride_a < a.len());
                debug_assert!(max_stride_b < b.len());
                debug_assert!(max_stride_c < c.len());
                bench.iter(|| {
                    unsafe {
                        $gemm(
                            $m, $m, $m,
                            1.,
                            a.as_ptr(), rsa, csa,
                            b.as_ptr(), rsb, csb,
                            beta,
                            c.as_mut_ptr(), rsc, csc,
                            )
                    }
                });
            }

            pub fn nobeta_ccc(bench: &mut Bencher) { base(bench, C, C, C, false); }
            pub fn nobeta_ccf(bench: &mut Bencher) { base(bench, C, C, F, false); }
            pub fn nobeta_fcc(bench: &mut Bencher) { base(bench, F, C, C, false); }
            pub fn nobeta_cfc(bench: &mut Bencher) { base(bench, C, F, C, false); }
            pub fn nobeta_ffc(bench: &mut Bencher) { base(bench, F, F, C, false); }
            pub fn nobeta_cff(bench: &mut Bencher) { base(bench, C, F, F, false); }
            pub fn nobeta_fcf(bench: &mut Bencher) { base(bench, F, C, F, false); }
            pub fn nobeta_fff(bench: &mut Bencher) { base(bench, F, F, F, false); }

            pub fn nobeta_cfc_spread_yyn(bench: &mut Bencher) { base(bench, Clike, Flike, C, false); }
            pub fn nobeta_fcc_spread_yyn(bench: &mut Bencher) { base(bench, Flike, Clike, C, false); }
            pub fn nobeta_fcc_spread_nny(bench: &mut Bencher) { base(bench, C, F, Clike, false); }
            pub fn nobeta_fcf_spread_nny(bench: &mut Bencher) { base(bench, C, F, Flike, false); }

            pub fn beta_ccc(bench: &mut Bencher) { base(bench, C, C, C, true); }
            pub fn beta_ccf(bench: &mut Bencher) { base(bench, C, C, F, true); }
            pub fn beta_fcc(bench: &mut Bencher) { base(bench, F, C, C, true); }
            pub fn beta_cfc(bench: &mut Bencher) { base(bench, C, F, C, true); }
            pub fn beta_ffc(bench: &mut Bencher) { base(bench, F, F, C, true); }
            pub fn beta_cff(bench: &mut Bencher) { base(bench, C, F, F, true); }
            pub fn beta_fcf(bench: &mut Bencher) { base(bench, F, C, F, true); }
            pub fn beta_fff(bench: &mut Bencher) { base(bench, F, F, F, true); }

            pub fn beta_fcc_spread_nny(bench: &mut Bencher) { base(bench, C, F, Clike, true); }
            pub fn beta_fcf_spread_nny(bench: &mut Bencher) { base(bench, C, F, Flike, true); }
            )+
        }
        benchmark_group!{ $modname,
            $modname::nobeta_ccc,
            $modname::nobeta_ccf,
            $modname::nobeta_fcc,
            $modname::nobeta_cfc,
            $modname::nobeta_ffc,
            $modname::nobeta_cff,
            $modname::nobeta_fcf,
            $modname::nobeta_fff,

            $modname::nobeta_cfc_spread_yyn,
            $modname::nobeta_fcc_spread_yyn,
            $modname::nobeta_fcc_spread_nny,
            $modname::nobeta_fcf_spread_nny,

            $modname::beta_ccc,
            $modname::beta_ccf,
            $modname::beta_fcc,
            $modname::beta_cfc,
            $modname::beta_ffc,
            $modname::beta_cff,
            $modname::beta_fcf,
            $modname::beta_fff,

            $modname::beta_fcc_spread_nny,
            $modname::beta_fcf_spread_nny
        }
    };
}

gemm_layout! {layout_f32_032, sgemm,
    (m032, 32)
}

gemm_layout! {layout_f64_032, dgemm,
    (m032, 32)
}

use std::ops::{Add, Mul};

trait Z {
    fn zero() -> Self;
}
impl Z for f32 {
    fn zero() -> Self {
        0.
    }
}
impl Z for f64 {
    fn zero() -> Self {
        0.
    }
}

// simple, slow, correct (hopefully) mat mul (Row Major)
#[inline(never)]
fn reference_mat_mul<A>(m: usize, k: usize, n: usize, a: &[A], b: &[A], c: &mut [A])
where
    A: Z + Add<Output = A> + Mul<Output = A> + Copy,
{
    assert!(a.len() >= m * k);
    assert!(b.len() >= k * n);
    assert!(c.len() >= m * n);

    for i in 0..m {
        for j in 0..n {
            unsafe {
                let celt = c.get_unchecked_mut(i * m + j);
                *celt = (0..k).fold(A::zero(), move |s, x| {
                    s + *a.get_unchecked(i * k + x) * *b.get_unchecked(x * n + j)
                });
            }
        }
    }
}

macro_rules! ref_mat_mul {
    ($modname:ident, $ty:ty, $(($name:ident, $m:expr, $n:expr, $k:expr))+) => {
        mod $modname {
            use bencher::{Bencher};
            use super::reference_mat_mul;
            $(
            pub fn $name(bench: &mut Bencher)
            {
                let a = vec![0. as $ty; $m * $n];
                let b = vec![0.; $n * $k];
                let mut c = vec![0.; $m * $k];
                bench.iter(|| {
                    reference_mat_mul($m, $n, $k, &a, &b, &mut c);
                    c[0]
                });
            }
            )+
        }
        benchmark_group!{ $modname, $($modname::$name),+ }
    };
}
ref_mat_mul! {ref_mat_mul_f32, f32,
    (m004, 4, 4, 4)
    (m005, 5, 5, 5)
    (m006, 6, 6, 6)
    (m007, 7, 7, 7)
    (m008, 8, 8, 8)
    (m009, 9, 9, 9)
    (m012, 12, 12, 12)
    (m016, 16, 16, 16)
    (m032, 32, 32, 32)
    (m064, 64, 64, 64)
}