Merge pull request #288 from bluss/zip-parallel

Rayon parallelization for Zip / azip!

Author: bluss

.travis.yml

@@ -31,7 +31,7 @@ script:
       cargo test --release --verbose --no-default-features &&
       cargo build --verbose --features "$FEATURES" &&
       cargo test --verbose --features "$FEATURES" &&
+      CARGO_TARGET_DIR=target/ cargo test --manifest-path=parallel/Cargo.toml --verbose &&
       CARGO_TARGET_DIR=target/ cargo test --manifest-path=serialization-tests/Cargo.toml --verbose &&
       CARGO_TARGET_DIR=target/ cargo test --manifest-path=numeric-tests/Cargo.toml --verbose &&
-      CARGO_TARGET_DIR=target/ cargo test --manifest-path=parallel/Cargo.toml --verbose &&
       ([ "$BENCH" != 1 ] || cargo bench --no-run --verbose --features "$FEATURES")

parallel/Cargo.toml

@@ -13,7 +13,7 @@ keywords = ["data-structure", "multidimensional", "parallel", "concurrent"]
 categories = ["data-structures", "science", "concurrency"]
 
 [dependencies]
-rayon = "0.6"
+rayon = { version = "0.7.0" }
 ndarray = { version = "0.9.0-alpha.1", path = "../" }
 
 [dev-dependencies]

parallel/README.rst

@@ -47,6 +47,12 @@ How to use with cargo::
 Recent Changes (ndarray-parallel)
 ---------------------------------
 
+- 0.3.0-alpha.1
+
+  - ParallelIterator for Zip, including ``.par_apply``.
+  - ``.par_map_inplace`` and ``.par_mav_inplace`` for arrays
+  - Require ndarray 0.9 (when released) and rayon 0.7
+
 - 0.2.0
 
   - Require for ndarray 0.8

parallel/benches/rayon.rs

@@ -12,13 +12,17 @@ extern crate ndarray_parallel;
 use ndarray::prelude::*;
 use ndarray_parallel::prelude::*;
 
+use ndarray::Zip;
+
 const EXP_N: usize = 128;
+const ADDN: usize = 1024;
 
 use std::cmp::max;
 
 fn set_threads() {
     let n = max(1, num_cpus::get() / 2);
-    let cfg = rayon::Configuration::new().set_num_threads(n);
+    //println!("Using {} threads", n);
+    let cfg = rayon::Configuration::new().num_threads(n);
     let _ = rayon::initialize(cfg);
 }
 
@@ -112,3 +116,42 @@ fn rayon_fastexp_by_axis(bench: &mut Bencher)
             .for_each(|mut sheet| sheet.mapv_inplace(fastexp));
     });
 }
+
+#[bench]
+fn rayon_fastexp_zip(bench: &mut Bencher)
+{
+    set_threads();
+    let mut a = Array2::<f64>::zeros((FASTEXP, FASTEXP));
+    bench.iter(|| {
+        Zip::from(&mut a).into_par_iter().for_each(|(elt, )| *elt = fastexp(*elt));
+    });
+}
+
+#[bench]
+fn add(bench: &mut Bencher)
+{
+    let mut a = Array2::<f64>::zeros((ADDN, ADDN));
+    let b = Array2::<f64>::zeros((ADDN, ADDN));
+    let c = Array2::<f64>::zeros((ADDN, ADDN));
+    let d = Array2::<f64>::zeros((ADDN, ADDN));
+    bench.iter(|| {
+        Zip::from(&mut a).and(&b).and(&c).and(&d).apply(|a, &b, &c, &d| {
+            *a += b.exp() + c.exp() + d.exp();
+        })
+    });
+}
+
+#[bench]
+fn rayon_add(bench: &mut Bencher)
+{
+    set_threads();
+    let mut a = Array2::<f64>::zeros((ADDN, ADDN));
+    let b = Array2::<f64>::zeros((ADDN, ADDN));
+    let c = Array2::<f64>::zeros((ADDN, ADDN));
+    let d = Array2::<f64>::zeros((ADDN, ADDN));
+    bench.iter(|| {
+        Zip::from(&mut a).and(&b).and(&c).and(&d).into_par_iter().for_each(|(a, &b, &c, &d)| {
+            *a += b.exp() + c.exp() + d.exp();
+        })
+    });
+}

parallel/src/ext_traits.rs

@@ -0,0 +1,83 @@
+
+use ndarray::{
+    Dimension,
+    NdProducer,
+    Zip,
+    ArrayBase,
+    DataMut,
+};
+
+use prelude::*;
+
+// Arrays
+
+/// Parallel versions of `map_inplace` and `mapv_inplace`.
+pub trait ParMap {
+    type Item;
+    fn par_map_inplace<F>(&mut self, f: F)
+        where F: Fn(&mut Self::Item) + Sync;
+    fn par_mapv_inplace<F>(&mut self, f: F)
+        where F: Fn(Self::Item) -> Self::Item + Sync,
+              Self::Item: Clone;
+}
+
+impl<A, S, D> ParMap for ArrayBase<S, D>
+    where S: DataMut<Elem=A>,
+          D: Dimension,
+          A: Send + Sync,
+{
+    type Item = A;
+    fn par_map_inplace<F>(&mut self, f: F)
+        where F: Fn(&mut Self::Item) + Sync
+    {
+        self.view_mut().into_par_iter().for_each(f)
+    }
+    fn par_mapv_inplace<F>(&mut self, f: F)
+        where F: Fn(Self::Item) -> Self::Item + Sync,
+              Self::Item: Clone
+    {
+        self.view_mut().into_par_iter()
+            .for_each(move |x| *x = f(x.clone()))
+    }
+}
+
+
+
+
+// Zip
+
+macro_rules! zip_impl {
+    ($([$name:ident $($p:ident)*],)+) => {
+        $(
+        /// The `par_apply` method for `Zip`.
+        ///
+        /// This is a shorthand for using `.into_par_iter().for_each()` on
+        /// `Zip`.
+        pub trait $name<$($p),*> {
+            fn par_apply<F>(self, function: F)
+                where F: Fn($($p),*) + Sync;
+        }
+
+        #[allow(non_snake_case)]
+        impl<Dim: Dimension, $($p: NdProducer<Dim=Dim>),*> $name<$($p::Item),*> for Zip<($($p,)*), Dim>
+            where $($p::Item : Send , )*
+                  $($p : Send , )*
+        {
+            fn par_apply<F>(self, function: F)
+                where F: Fn($($p::Item),*) + Sync
+            {
+                self.into_par_iter().for_each(move |($($p,)*)| function($($p),*))
+            }
+        }
+        )+
+    }
+}
+
+zip_impl!{
+    [ParApply1 P1],
+    [ParApply2 P1 P2],
+    [ParApply3 P1 P2 P3],
+    [ParApply4 P1 P2 P3 P4],
+    [ParApply5 P1 P2 P3 P4 P5],
+    [ParApply6 P1 P2 P3 P4 P5 P6],
+}

parallel/src/into_traits.rs

@@ -1,5 +1,5 @@
 
-use rayon::par_iter::ParallelIterator;
+use rayon::iter::ParallelIterator;
 
 pub trait NdarrayIntoParallelIterator {
     type Iter: ParallelIterator<Item=Self::Item>;

parallel/src/lib.rs

@@ -8,11 +8,17 @@
 //! `.axis_iter()` and `.axis_iter_mut()` also have parallel counterparts,
 //! and their parallel iterators are indexed (and thus ordered) and exact length.
 //!
+//! `Zip` also implements `NdarrayIntoParallelIterator`, and there is an
+//! extension trait so that it can use a method `.par_apply` directly.
+//!
 //! (*) This regime of a custom trait instead of rayon’s own is since we
 //! use this intermediate ndarray-parallel crate.
 //!
 //! # Examples
 //!
+//!
+//! ## Arrays and array views
+//!
 //! Compute the exponential of each element in an array, parallelized.
 //!
 //! ```
@@ -24,10 +30,18 @@
 //!
 //! fn main() {
 //!     let mut a = Array2::<f64>::zeros((128, 128));
+//!
+//!     // Parallel versions of regular array methods (ParMap trait)
+//!     a.par_map_inplace(|x| *x = x.exp());
+//!     a.par_mapv_inplace(f64::exp);
+//!
+//!     // You can also use the parallel iterator directly
 //!     a.par_iter_mut().for_each(|x| *x = x.exp());
 //! }
 //! ```
 //!
+//! ## Axis iterators
+//!
 //! Use the parallel `.axis_iter()` to compute the sum of each row.
 //!
 //! ```
@@ -49,10 +63,39 @@
 //!     assert_eq!(sums, [120., 376., 632., 888.]);
 //! }
 //! ```
+//!
+//! ## Zip
+//!
+//! Use zip for lock step function application across several arrays
+//!
+//! ```
+//! extern crate ndarray;
+//! extern crate ndarray_parallel;
+//!
+//! use ndarray::Array3;
+//! use ndarray::Zip;
+//! use ndarray_parallel::prelude::*;
+//!
+//! type Array3f64 = Array3<f64>;
+//!
+//! fn main() {
+//!     const N: usize = 128;
+//!     let a = Array3f64::from_elem((N, N, N), 1.);
+//!     let b = Array3f64::from_elem(a.dim(), 2.);
+//!     let mut c = Array3f64::zeros(a.dim());
+//!
+//!     Zip::from(&mut c)
+//!         .and(&a)
+//!         .and(&b)
+//!         .par_apply(|c, &a, &b| {
+//!             *c += a - b;
+//!         });
+//! }
+//! ```
 
 
-extern crate ndarray;
-extern crate rayon;
+pub extern crate ndarray;
+pub extern crate rayon;
 
 /// Into- traits for creating parallelized iterators.
 pub mod prelude {
@@ -63,6 +106,16 @@ pub mod prelude {
 
     #[doc(no_inline)]
     pub use rayon::prelude::{ParallelIterator, IndexedParallelIterator, ExactParallelIterator};
+
+    pub use ext_traits::{
+        ParApply1,
+        ParApply2,
+        ParApply3,
+        ParApply4,
+        ParApply5,
+        ParApply6,
+    };
+    pub use ext_traits::ParMap;
 }
 
 pub use par::Parallel;
@@ -73,5 +126,6 @@ pub use into_traits::{
 };
 
 mod par;
+mod ext_traits;
 mod into_traits;
 mod into_impls;