Add complex GEMM and simple test

Author: emmatyping

src/linalg/impl_linalg.rs

@@ -29,6 +29,9 @@ use cblas_sys as blas_sys;
 #[cfg(feature = "blas")]
 use cblas_sys::{CblasNoTrans, CblasRowMajor, CblasTrans, CBLAS_LAYOUT};
 
+#[cfg(feature = "blas")]
+use num_complex::{Complex32 as c32, Complex64 as c64};
+
 /// len of vector before we use blas
 #[cfg(feature = "blas")]
 const DOT_BLAS_CUTOFF: usize = 32;
@@ -374,6 +377,7 @@ fn mat_mul_impl<A>(
 ) where
     A: LinalgScalar,
 {
+
     // size cutoff for using BLAS
     let cut = GEMM_BLAS_CUTOFF;
     let ((mut m, a), (_, mut n)) = (lhs.dim(), rhs.dim());
@@ -455,6 +459,8 @@ fn mat_mul_impl<A>(
         }
         gemm!(f32, cblas_sgemm);
         gemm!(f64, cblas_dgemm);
+        gemm!(c32, cblas_cgemm);
+        gemm!(c64, cblas_zgemm);
     }
     mat_mul_general(alpha, lhs, rhs, beta, c)
 }

xtest-blas/Cargo.toml

@@ -11,6 +11,7 @@ test = false
 approx = "0.4"
 defmac = "0.2"
 num-traits = "0.2"
+num-complex = { version = "0.4", default-features = false }
 
 [dependencies]
 ndarray = { path = "../", features = ["approx", "blas"] }

xtest-blas/tests/oper.rs

@@ -3,6 +3,7 @@ extern crate defmac;
 extern crate ndarray;
 extern crate num_traits;
 extern crate blas_src;
+extern crate num_complex;
 
 use ndarray::prelude::*;
 
@@ -12,6 +13,8 @@ use ndarray::{Data, Ix, LinalgScalar};
 
 use approx::assert_relative_eq;
 use defmac::defmac;
+use num_complex::Complex32;
+use num_complex::Complex64;
 
 #[test]
 fn mat_vec_product_1d() {
@@ -52,6 +55,20 @@ fn range_mat64(m: Ix, n: Ix) -> Array2<f64> {
         .unwrap()
 }
 
+fn range_mat_complex(m: Ix, n: Ix) -> Array2<Complex32> {
+    Array::linspace(0., (m * n) as f32 - 1., m * n)
+        .into_shape((m, n))
+        .unwrap()
+        .map_mut(|&mut f| Complex32::new(f, 0.))
+}
+
+fn range_mat_complex64(m: Ix, n: Ix) -> Array2<Complex64> {
+    Array::linspace(0., (m * n) as f64 - 1., m * n)
+        .into_shape((m, n))
+        .unwrap()
+        .map_mut(|&mut f| Complex64::new(f, 0.))
+}
+
 fn range1_mat64(m: Ix) -> Array1<f64> {
     Array::linspace(0., m as f64 - 1., m)
 }
@@ -250,6 +267,30 @@ fn gemm_64_1_f() {
     assert_relative_eq!(y, answer, epsilon = 1e-12, max_relative = 1e-7);
 }
 
+#[test]
+fn gemm_c64_1_f() {
+    let a = range_mat_complex64(64, 64).reversed_axes();
+    let (m, n) = a.dim();
+    // m x n  times n x 1  == m x 1
+    let x = range_mat_complex64(n, 1);
+    let mut y = range_mat_complex64(m, 1);
+    let answer = reference_mat_mul(&a, &x) + &y;
+    general_mat_mul(Complex64::new(1.0, 0.), &a, &x, Complex64::new(1.0, 0.), &mut y);
+    assert_relative_eq!(y.mapv(|i| i.norm_sqr()), answer.mapv(|i| i.norm_sqr()), epsilon = 1e-12, max_relative = 1e-7);
+}
+
+#[test]
+fn gemm_c32_1_f() {
+    let a = range_mat_complex(64, 64).reversed_axes();
+    let (m, n) = a.dim();
+    // m x n  times n x 1  == m x 1
+    let x = range_mat_complex(n, 1);
+    let mut y = range_mat_complex(m, 1);
+    let answer = reference_mat_mul(&a, &x) + &y;
+    general_mat_mul(Complex32::new(1.0, 0.), &a, &x, Complex32::new(1.0, 0.), &mut y);
+    assert_relative_eq!(y.mapv(|i| i.norm_sqr()), answer.mapv(|i| i.norm_sqr()), epsilon = 1e-12, max_relative = 1e-7);
+}
+
 #[test]
 fn gen_mat_vec_mul() {
     let alpha = -2.3;