Experimental divide & conquer in matrix multiply using rayon

Author: bluss

Cargo.toml

@@ -35,6 +35,8 @@ version = "0.3.16"
 optional = true
 
 [dependencies]
+rayon = "0.3.1"
+
 # Use via the `blas` crate feature!
 blas-sys = { version = "0.6", optional = true, default-features = false }
 openblas-provider = { version = "0.4", optional = true, default-features = false }

src/lib.rs

@@ -73,6 +73,7 @@ extern crate rustc_serialize as serialize;
 extern crate blas_sys;
 
 extern crate matrixmultiply;
+extern crate rayon;
 
 extern crate itertools;
 extern crate num as libnum;

src/linalg/impl_linalg.rs

@@ -6,6 +6,7 @@
 // option. This file may not be copied, modified, or distributed
 // except according to those terms.
 
+use rayon;
 use libnum::Zero;
 use itertools::free::enumerate;
 
@@ -417,6 +418,8 @@ fn mat_mul_impl<A>(alpha: A,
     mat_mul_general(alpha, lhs, rhs, beta, c)
 }
 
+const SPLIT: usize = 64;
+
 /// C ← α A B + β C
 fn mat_mul_general<A>(alpha: A,
                       lhs: &ArrayView<A, (Ix, Ix)>,
@@ -425,7 +428,27 @@ fn mat_mul_general<A>(alpha: A,
                       c: &mut ArrayViewMut<A, (Ix, Ix)>)
     where A: LinalgScalar,
 {
-    let ((m, k), (_, n)) = (lhs.dim, rhs.dim);
+    let ((m, k), (k2, n)) = (lhs.dim, rhs.dim);
+
+    debug_assert_eq!(k, k2);
+    if m > SPLIT {
+        // [ A0 ] B = [ C0 ]
+        // [ A1 ]     [ C1 ]
+        let mid = m / 2;
+        let (a0, a1) = lhs.split_at(Axis(0), mid);
+        let (mut c0, mut c1) = c.view_mut().split_at(Axis(0), mid);
+        rayon::join(move || mat_mul_general(alpha, &a0, rhs, beta, &mut c0),
+                    move || mat_mul_general(alpha, &a1, rhs, beta, &mut c1));
+        return;
+    } else if n > SPLIT {
+        // A [ B0 B1 ] = [ C0 C1 ]
+        let mid = n / 2;
+        let (b0, b1) = rhs.split_at(Axis(1), mid);
+        let (mut c0, mut c1) = c.view_mut().split_at(Axis(1), mid);
+        rayon::join(move || mat_mul_general(alpha, lhs, &b0, beta, &mut c0),
+                    move || mat_mul_general(alpha, lhs, &b1, beta, &mut c1));
+        return;
+    }
 
     // common parameters for gemm
     let ap = lhs.as_ptr();