sparse by dense parallel products #298

.github/workflows/ci.yml

@@ -72,7 +72,7 @@ jobs:
         include:
         - build: msrv
           os: ubuntu-18.04
-          rust: 1.49.0
+          rust: 1.52.0
         - build: stable
           os: ubuntu-18.04
           rust: stable

Cargo.toml

@@ -17,7 +17,7 @@ exclude = [
 
 [features]
 default = ["alga", "multi_thread"]
-multi_thread = ["rayon", "num_cpus"]
+multi_thread = ["rayon", "num_cpus", "ndarray/rayon"]
 
 [dependencies]
 num-traits = "0.2.0"

README.md

@@ -84,7 +84,7 @@ See the [changelog](changelog.rst).
 
 ## Minimum Supported Rust Version
 
-The minimum supported Rust version currently is 1.49. Prior to a 1.0 version,
+The minimum supported Rust version currently is 1.52. Prior to a 1.0 version,
 bumping the MSRV will not be considered a breaking change, but breakage will
 be avoided on a best effort basis.
 

src/sparse/csmat.rs

@@ -1901,7 +1901,7 @@ where
 impl<'a, 'b, N, I, Iptr, IpS, IS, DS, DS2> Mul<&'b ArrayBase<DS2, Ix2>>
     for &'a CsMatBase<N, I, IpS, IS, DS, Iptr>
 where
-    N: 'a + crate::MulAcc + num_traits::Zero + Clone,
+    N: 'a + crate::MulAcc + num_traits::Zero + Clone + Send + Sync,
     I: 'a + SpIndex,
     Iptr: 'a + SpIndex,
     IpS: 'a + Deref<Target = [Iptr]>,
@@ -1962,7 +1962,13 @@ where
 impl<'a, 'b, N, I, IpS, IS, DS, DS2> Dot<CsMatBase<N, I, IpS, IS, DS>>
     for ArrayBase<DS2, Ix2>
 where
-    N: 'a + Clone + crate::MulAcc + num_traits::Zero + std::fmt::Debug,
+    N: 'a
+        + Clone
+        + crate::MulAcc
+        + num_traits::Zero
+        + std::fmt::Debug
+        + Send
+        + Sync,
     I: 'a + SpIndex,
     IpS: 'a + Deref<Target = [I]>,
     IS: 'a + Deref<Target = [I]>,
@@ -2013,7 +2019,7 @@ where
 impl<'a, 'b, N, I, Iptr, IpS, IS, DS, DS2> Dot<ArrayBase<DS2, Ix2>>
     for CsMatBase<N, I, IpS, IS, DS, Iptr>
 where
-    N: 'a + Clone + crate::MulAcc + num_traits::Zero,
+    N: 'a + Clone + crate::MulAcc + num_traits::Zero + Send + Sync,
     I: 'a + SpIndex,
     Iptr: 'a + SpIndex,
     IpS: 'a + Deref<Target = [Iptr]>,
@@ -2031,7 +2037,7 @@ where
 impl<'a, 'b, N, I, Iptr, IpS, IS, DS, DS2> Mul<&'b ArrayBase<DS2, Ix1>>
     for &'a CsMatBase<N, I, IpS, IS, DS, Iptr>
 where
-    N: 'a + Clone + crate::MulAcc + num_traits::Zero,
+    N: 'a + Clone + crate::MulAcc + num_traits::Zero + Send + Sync,
     I: 'a + SpIndex,
     Iptr: 'a + SpIndex,
     IpS: 'a + Deref<Target = [Iptr]>,
@@ -2072,7 +2078,7 @@ where
 impl<'a, 'b, N, I, Iptr, IpS, IS, DS, DS2> Dot<ArrayBase<DS2, Ix1>>
     for CsMatBase<N, I, IpS, IS, DS, Iptr>
 where
-    N: 'a + Clone + crate::MulAcc + num_traits::Zero,
+    N: 'a + Clone + crate::MulAcc + num_traits::Zero + Send + Sync,
     I: 'a + SpIndex,
     Iptr: 'a + SpIndex,
     IpS: 'a + Deref<Target = [Iptr]>,

src/sparse/prod.rs

@@ -190,7 +190,9 @@ pub fn csr_mulacc_dense_rowmaj<'a, N, A, B, I, Iptr>(
     rhs: ArrayView<B, Ix2>,
     mut out: ArrayViewMut<'a, N, Ix2>,
 ) where
-    N: 'a + crate::MulAcc<A, B>,
+    A: Send + Sync,
+    B: Send + Sync,
+    N: 'a + crate::MulAcc<A, B> + Send + Sync,
     I: 'a + SpIndex,
     Iptr: 'a + SpIndex,
 {
@@ -200,6 +202,21 @@ pub fn csr_mulacc_dense_rowmaj<'a, N, A, B, I, Iptr>(
     assert!(lhs.is_csr(), "Storage mismatch");
 
     let axis0 = Axis(0);
+    #[cfg(feature = "multi_thread")]
+    for (line, mut oline) in lhs.outer_iterator().zip(out.axis_iter_mut(axis0))
+    {
+        for (col_ind, lval) in line.iter() {
+            let rline = rhs.row(col_ind);
+            // TODO: call an axpy primitive to benefit from vectorisation?
+            ndarray::Zip::from(&mut oline).and(rline).par_for_each(
+                |oval, rval| {
+                    oval.mul_acc(lval, rval);
+                },
+            );
+        }
+    }
+
+    #[cfg(not(feature = "multi_thread"))]
     for (line, mut oline) in lhs.outer_iterator().zip(out.axis_iter_mut(axis0))
     {
         for (col_ind, lval) in line.iter() {
@@ -220,7 +237,9 @@ pub fn csc_mulacc_dense_rowmaj<'a, N, A, B, I, Iptr>(
     rhs: ArrayView<B, Ix2>,
     mut out: ArrayViewMut<'a, N, Ix2>,
 ) where
-    N: 'a + crate::MulAcc<A, B>,
+    A: Send + Sync,
+    B: Send + Sync,
+    N: 'a + crate::MulAcc<A, B> + Send + Sync,
     I: 'a + SpIndex,
     Iptr: 'a + SpIndex,
 {
@@ -229,6 +248,19 @@ pub fn csc_mulacc_dense_rowmaj<'a, N, A, B, I, Iptr>(
     assert_eq!(rhs.shape()[1], out.shape()[1], "Dimension mismatch");
     assert!(lhs.is_csc(), "Storage mismatch");
 
+    #[cfg(feature = "multi_thread")]
+    for (lcol, rline) in lhs.outer_iterator().zip(rhs.outer_iter()) {
+        for (orow, lval) in lcol.iter() {
+            let oline = out.row_mut(orow);
+            ndarray::Zip::from(oline)
+                .and(rline)
+                .par_for_each(|oval, rval| {
+                    oval.mul_acc(lval, rval);
+                });
+        }
+    }
+
+    #[cfg(not(feature = "multi_thread"))]
     for (lcol, rline) in lhs.outer_iterator().zip(rhs.outer_iter()) {
         for (orow, lval) in lcol.iter() {
             let mut oline = out.row_mut(orow);
@@ -247,7 +279,9 @@ pub fn csc_mulacc_dense_colmaj<'a, N, A, B, I, Iptr>(
     rhs: ArrayView<B, Ix2>,
     mut out: ArrayViewMut<'a, N, Ix2>,
 ) where
-    N: 'a + crate::MulAcc<A, B>,
+    A: Send + Sync,
+    B: Send + Sync,
+    N: 'a + crate::MulAcc<A, B> + Send + Sync,
     I: 'a + SpIndex,
     Iptr: 'a + SpIndex,
 {
@@ -257,6 +291,20 @@ pub fn csc_mulacc_dense_colmaj<'a, N, A, B, I, Iptr>(
     assert!(lhs.is_csc(), "Storage mismatch");
 
     let axis1 = Axis(1);
+    // NOTE: See csr_mulacc_dense_colmaj, same issue
+    #[cfg(feature = "multi_thread")]
+    ndarray::Zip::from(out.axis_iter_mut(axis1))
+        .and(rhs.axis_iter(axis1))
+        .par_for_each(|mut ocol, rcol| {
+            for (rrow, lcol) in lhs.outer_iterator().enumerate() {
+                let rval = &rcol[[rrow]];
+                for (orow, lval) in lcol.iter() {
+                    ocol[[orow]].mul_acc(lval, rval);
+                }
+            }
+        });
+
+    #[cfg(not(feature = "multi_thread"))]
     for (mut ocol, rcol) in out.axis_iter_mut(axis1).zip(rhs.axis_iter(axis1)) {
         for (rrow, lcol) in lhs.outer_iterator().enumerate() {
             let rval = &rcol[[rrow]];
@@ -275,7 +323,9 @@ pub fn csr_mulacc_dense_colmaj<'a, N, A, B, I, Iptr>(
     rhs: ArrayView<B, Ix2>,
     mut out: ArrayViewMut<'a, N, Ix2>,
 ) where
-    N: 'a + crate::MulAcc<A, B>,
+    A: Send + Sync,
+    B: Send + Sync,
+    N: 'a + crate::MulAcc<A, B> + Send + Sync,
     I: 'a + SpIndex,
     Iptr: 'a + SpIndex,
 {
@@ -285,6 +335,23 @@ pub fn csr_mulacc_dense_colmaj<'a, N, A, B, I, Iptr>(
     assert!(lhs.is_csr(), "Storage mismatch");
 
     let axis1 = Axis(1);
+    // NOTE: This is parallel over the columns of the output and rhs
+    // which isn't ideal. This is still sequential for dense vector product.
+    // Ideally CsMat.outer_iterator() should get a par_iter rayon impl
+    #[cfg(feature = "multi_thread")]
+    ndarray::Zip::from(out.axis_iter_mut(axis1))
+        .and(rhs.axis_iter(axis1))
+        .par_for_each(|mut ocol, rcol| {
+            for (orow, lrow) in lhs.outer_iterator().enumerate() {
+                let oval = &mut ocol[[orow]];
+                for (rrow, lval) in lrow.iter() {
+                    let rval = &rcol[[rrow]];
+                    oval.mul_acc(lval, rval);
+                }
+            }
+        });
+
+    #[cfg(not(feature = "multi_thread"))]
     for (mut ocol, rcol) in out.axis_iter_mut(axis1).zip(rhs.axis_iter(axis1)) {
         for (orow, lrow) in lhs.outer_iterator().enumerate() {
             let oval = &mut ocol[[orow]];