matrixmultiply_mt 0.2.1

extern crate matrixmultiply_mt;
use matrixmultiply_mt::{dgemm, sgemm};

use std::fmt::{Debug, Display};

trait Float: Copy + Display + Debug + PartialEq {
	fn zero() -> Self;
	fn one() -> Self;
	fn from(int: i64) -> Self;
	fn nan() -> Self;
}

impl Float for f32 {
	fn zero() -> Self {
		0.
	}
	fn one() -> Self {
		1.
	}
	fn from(x: i64) -> Self {
		x as Self
	}
	fn nan() -> Self {
		f32::NAN
	}
}

impl Float for f64 {
	fn zero() -> Self {
		0.
	}
	fn one() -> Self {
		1.
	}
	fn from(x: i64) -> Self {
		x as Self
	}
	fn nan() -> Self {
		f64::NAN
	}
}

#[allow(clippy::too_many_arguments, clippy::many_single_char_names)]
trait Gemm: Sized {
	unsafe fn gemm(
		m: usize,
		k: usize,
		n: usize,
		alpha: Self,
		a: *const Self,
		rsa: isize,
		csa: isize,
		b: *const Self,
		rsb: isize,
		csb: isize,
		beta: Self,
		c: *mut Self,
		rsc: isize,
		csc: isize,
	);
}

#[allow(clippy::too_many_arguments, clippy::many_single_char_names)]
impl Gemm for f32 {
	unsafe fn gemm(
		m: usize,
		k: usize,
		n: usize,
		alpha: Self,
		a: *const Self,
		rsa: isize,
		csa: isize,
		b: *const Self,
		rsb: isize,
		csb: isize,
		beta: Self,
		c: *mut Self,
		rsc: isize,
		csc: isize,
	) {
		sgemm(m, k, n, alpha, a, rsa, csa, b, rsb, csb, beta, c, rsc, csc)
	}
}

#[allow(clippy::too_many_arguments, clippy::many_single_char_names)]
impl Gemm for f64 {
	unsafe fn gemm(
		m: usize,
		k: usize,
		n: usize,
		alpha: Self,
		a: *const Self,
		rsa: isize,
		csa: isize,
		b: *const Self,
		rsb: isize,
		csb: isize,
		beta: Self,
		c: *mut Self,
		rsc: isize,
		csc: isize,
	) {
		dgemm(m, k, n, alpha, a, rsa, csa, b, rsb, csb, beta, c, rsc, csc)
	}
}

#[test]
fn test_sgemm() {
	test_gemm::<f32>();
}
#[test]
fn test_dgemm() {
	test_gemm::<f64>();
}

fn test_gemm<F>()
where
	F: Gemm + Float,
{
	test_mul_with_id::<F>(4, 4, true);
	test_mul_with_id::<F>(8, 8, true);
	test_mul_with_id::<F>(32, 32, false);
	test_mul_with_id::<F>(128, 128, false);
	test_mul_with_id::<F>(17, 128, false);
	for i in 0..16 {
		for j in 0..16 {
			test_mul_with_id::<F>(i, j, true);
		}
	}
	/*
	 */
	test_mul_with_id::<F>(17, 257, false);
	test_mul_with_id::<F>(24, 512, false);
	for i in 0..10 {
		for j in 0..10 {
			test_mul_with_id::<F>(i * 4, j * 4, true);
		}
	}
	test_mul_with_id::<F>(266, 265, false);
	test_mul_id_with::<F>(4, 4, true);
	for i in 0..12 {
		for j in 0..12 {
			test_mul_id_with::<F>(i, j, true);
		}
	}
	test_mul_id_with::<F>(266, 265, false);
	test_scale::<F>(4, 4, 4, true);
	test_scale::<F>(15, 15, 15, true);
	test_scale::<F>(16, 16, 16, true);
	test_scale::<F>(19, 20, 16, true);
	test_scale::<F>(150, 140, 128, false);
}

/// multiply a M x N matrix with an N x N id matrix
#[cfg(test)]
#[allow(clippy::many_single_char_names)]
fn test_mul_with_id<F>(m: usize, n: usize, small: bool)
where
	F: Gemm + Float,
{
	let (m, k, n) = (m, n, n);
	let mut a = vec![F::zero(); m * k];
	let mut b = vec![F::zero(); k * n];
	let mut c = vec![F::zero(); m * n];
	println!("test matrix with id input M={}, N={}", m, n);

	for (i, elt) in a.iter_mut().enumerate() {
		*elt = F::from(i as i64);
	}
	for i in 0..k {
		b[i + i * k] = F::one();
	}

	unsafe {
		F::gemm(
			m,
			k,
			n,
			F::one(),
			a.as_ptr(),
			k as isize,
			1,
			b.as_ptr(),
			n as isize,
			1,
			F::zero(),
			c.as_mut_ptr(),
			n as isize,
			1,
		)
	}
	for (i, (x, y)) in a.iter().zip(&c).enumerate() {
		if x != y {
			if k != 0 && n != 0 && small {
				for row in a.chunks(k) {
					println!("{:?}", row);
				}
				for row in b.chunks(n) {
					println!("{:?}", row);
				}
				for row in c.chunks(n) {
					println!("{:?}", row);
				}
			}
			panic!("mismatch at index={}, x: {}, y: {} (matrix input M={}, N={})", i, x, y, m, n);
		}
	}
	println!("passed matrix with id input M={}, N={}", m, n);
}

/// multiply a K x K id matrix with an K x N matrix
#[cfg(test)]
#[allow(clippy::many_single_char_names)]
fn test_mul_id_with<F>(k: usize, n: usize, small: bool)
where
	F: Gemm + Float,
{
	let (m, k, n) = (k, k, n);
	let mut a = vec![F::zero(); m * k];
	let mut b = vec![F::zero(); k * n];
	let mut c = vec![F::zero(); m * n];

	for i in 0..k {
		a[i + i * k] = F::one();
	}
	for (i, elt) in b.iter_mut().enumerate() {
		*elt = F::from(i as i64);
	}

	unsafe {
		F::gemm(
			m,
			k,
			n,
			F::one(),
			a.as_ptr(),
			k as isize,
			1,
			b.as_ptr(),
			n as isize,
			1,
			F::zero(),
			c.as_mut_ptr(),
			n as isize,
			1,
		)
	}
	for (i, (x, y)) in b.iter().zip(&c).enumerate() {
		if x != y {
			if k != 0 && n != 0 && small {
				for row in a.chunks(k) {
					println!("{:?}", row);
				}
				for row in b.chunks(n) {
					println!("{:?}", row);
				}
				for row in c.chunks(n) {
					println!("{:?}", row);
				}
			}
			panic!("mismatch at index={}, x: {}, y: {} (matrix input M={}, N={})", i, x, y, m, n);
		}
	}
	println!("passed id with matrix input K={}, N={}", k, n);
}

#[cfg(test)]
#[allow(clippy::many_single_char_names)]
fn test_scale<F>(m: usize, k: usize, n: usize, small: bool)
where
	F: Gemm + Float,
{
	let (m, k, n) = (m, k, n);
	let mut a = vec![F::zero(); m * k];
	let mut b = vec![F::zero(); k * n];
	let mut c1 = vec![F::one(); m * n];
	let mut c2 = vec![F::nan(); m * n];
	// init c2 with NaN to test the overwriting behavior when beta = 0.

	for (i, elt) in a.iter_mut().enumerate() {
		*elt = F::from(i as i64);
	}
	for (i, elt) in b.iter_mut().enumerate() {
		*elt = F::from(i as i64);
	}

	unsafe {
		// 3 A B
		F::gemm(
			m,
			k,
			n,
			F::from(3),
			a.as_ptr(),
			k as isize,
			1,
			b.as_ptr(),
			n as isize,
			1,
			F::zero(),
			c1.as_mut_ptr(),
			n as isize,
			1,
		);

		// A B
		F::gemm(
			m,
			k,
			n,
			F::one(),
			a.as_ptr(),
			k as isize,
			1,
			b.as_ptr(),
			n as isize,
			1,
			F::zero(),
			c2.as_mut_ptr(),
			n as isize,
			1,
		);
		// (2 A B) + A B
		F::gemm(
			m,
			k,
			n,
			F::one(),
			a.as_ptr(),
			k as isize,
			1,
			b.as_ptr(),
			n as isize,
			1,
			F::from(2),
			c2.as_mut_ptr(),
			n as isize,
			1,
		);
	}
	for (i, (x, y)) in c1.iter().zip(&c2).enumerate() {
		if x != y {
			if k != 0 && n != 0 && small {
				for row in a.chunks(k) {
					println!("{:?}", row);
				}
				for row in b.chunks(n) {
					println!("{:?}", row);
				}
				for row in c1.chunks(n) {
					println!("{:?}", row);
				}
				for row in c2.chunks(n) {
					println!("{:?}", row);
				}
			}
			panic!("mismatch at index={}, x: {}, y: {} (matrix input M={}, N={})", i, x, y, m, n);
		}
	}
	println!("passed matrix with id input M={}, N={}", m, n);
}