Merge pull request #809 from rust-ndarray/zip-strided-for-c-and-f

Zip: Handle preferred memory layout of inhomogenous inputs better

benches/iter.rs

@@ -381,12 +381,12 @@ pub fn zip_mut_with(data: &Array3<f32>, out: &mut Array3<f32>) {
 fn zip_mut_with_cc(b: &mut Bencher) {
     let data: Array3<f32> = Array3::zeros((ISZ, ISZ, ISZ));
     let mut out = Array3::zeros(data.dim());
-    b.iter(|| black_box(zip_mut_with(&data, &mut out)));
+    b.iter(|| zip_mut_with(&data, &mut out));
 }
 
 #[bench]
 fn zip_mut_with_ff(b: &mut Bencher) {
     let data: Array3<f32> = Array3::zeros((ISZ, ISZ, ISZ).f());
     let mut out = Array3::zeros(data.dim().f());
-    b.iter(|| black_box(zip_mut_with(&data, &mut out)));
+    b.iter(|| zip_mut_with(&data, &mut out));
 }

benches/zip.rs

@@ -0,0 +1,120 @@
+#![feature(test)]
+extern crate test;
+use test::{black_box, Bencher};
+use ndarray::{Array3, ShapeBuilder, Zip};
+use ndarray::s;
+use ndarray::IntoNdProducer;
+
+pub fn zip_copy<'a, A, P, Q>(data: P, out: Q)
+    where P: IntoNdProducer<Item = &'a A>,
+          Q: IntoNdProducer<Item = &'a mut A, Dim = P::Dim>,
+          A: Copy + 'a
+{
+    Zip::from(data).and(out).apply(|&i, o| {
+        *o = i;
+    });
+}
+
+pub fn zip_copy_split<'a, A, P, Q>(data: P, out: Q)
+    where P: IntoNdProducer<Item = &'a A>,
+          Q: IntoNdProducer<Item = &'a mut A, Dim = P::Dim>,
+          A: Copy + 'a
+{
+    let z = Zip::from(data).and(out);
+    let (z1, z2) = z.split();
+    let (z11, z12) = z1.split();
+    let (z21, z22) = z2.split();
+    let f = |&i: &A, o: &mut A| *o = i;
+    z11.apply(f);
+    z12.apply(f);
+    z21.apply(f);
+    z22.apply(f);
+}
+
+pub fn zip_indexed(data: &Array3<f32>, out: &mut Array3<f32>) {
+    Zip::indexed(data).and(out).apply(|idx, &i, o| {
+        let _ = black_box(idx);
+        *o = i;
+    });
+}
+
+// array size in benchmarks
+const SZ3: (usize, usize, usize) = (100, 110, 100);
+
+#[bench]
+fn zip_cc(b: &mut Bencher) {
+    let data: Array3<f32> = Array3::zeros(SZ3);
+    let mut out = Array3::zeros(data.dim());
+    b.iter(|| zip_copy(&data, &mut out));
+}
+
+#[bench]
+fn zip_cf(b: &mut Bencher) {
+    let data: Array3<f32> = Array3::zeros(SZ3);
+    let mut out = Array3::zeros(data.dim().f());
+    b.iter(|| zip_copy(&data, &mut out));
+}
+
+#[bench]
+fn zip_fc(b: &mut Bencher) {
+    let data: Array3<f32> = Array3::zeros(SZ3.f());
+    let mut out = Array3::zeros(data.dim());
+    b.iter(|| zip_copy(&data, &mut out));
+}
+
+#[bench]
+fn zip_ff(b: &mut Bencher) {
+    let data: Array3<f32> = Array3::zeros(SZ3.f());
+    let mut out = Array3::zeros(data.dim().f());
+    b.iter(|| zip_copy(&data, &mut out));
+}
+
+#[bench]
+fn zip_indexed_cc(b: &mut Bencher) {
+    let data: Array3<f32> = Array3::zeros(SZ3);
+    let mut out = Array3::zeros(data.dim());
+    b.iter(|| zip_indexed(&data, &mut out));
+}
+
+#[bench]
+fn zip_indexed_ff(b: &mut Bencher) {
+    let data: Array3<f32> = Array3::zeros(SZ3.f());
+    let mut out = Array3::zeros(data.dim().f());
+    b.iter(|| zip_indexed(&data, &mut out));
+}
+
+#[bench]
+fn slice_zip_cc(b: &mut Bencher) {
+    let data: Array3<f32> = Array3::zeros(SZ3);
+    let mut out = Array3::zeros(data.dim());
+    let data = data.slice(s![1.., 1.., 1..]);
+    let mut out = out.slice_mut(s![1.., 1.., 1..]);
+    b.iter(|| zip_copy(&data, &mut out));
+}
+
+#[bench]
+fn slice_zip_ff(b: &mut Bencher) {
+    let data: Array3<f32> = Array3::zeros(SZ3.f());
+    let mut out = Array3::zeros(data.dim().f());
+    let data = data.slice(s![1.., 1.., 1..]);
+    let mut out = out.slice_mut(s![1.., 1.., 1..]);
+    b.iter(|| zip_copy(&data, &mut out));
+}
+
+#[bench]
+fn slice_split_zip_cc(b: &mut Bencher) {
+    let data: Array3<f32> = Array3::zeros(SZ3);
+    let mut out = Array3::zeros(data.dim());
+    let data = data.slice(s![1.., 1.., 1..]);
+    let mut out = out.slice_mut(s![1.., 1.., 1..]);
+    b.iter(|| zip_copy_split(&data, &mut out));
+}
+
+#[bench]
+fn slice_split_zip_ff(b: &mut Bencher) {
+    let data: Array3<f32> = Array3::zeros(SZ3.f());
+    let mut out = Array3::zeros(data.dim().f());
+    let data = data.slice(s![1.., 1.., 1..]);
+    let mut out = out.slice_mut(s![1.., 1.., 1..]);
+    b.iter(|| zip_copy_split(&data, &mut out));
+}

src/layout/layoutfmt.rs

@@ -8,7 +8,7 @@
 
 use super::Layout;
 
-const LAYOUT_NAMES: &[&str] = &["C", "F"];
+const LAYOUT_NAMES: &[&str] = &["C", "F", "c", "f"];
 
 use std::fmt;
 

src/layout/mod.rs

@@ -1,54 +1,145 @@
 mod layoutfmt;
 
-// public struct but users don't interact with it
+// Layout it a bitset used for internal layout description of
+// arrays, producers and sets of producers.
+// The type is public but users don't interact with it.
 #[doc(hidden)]
 /// Memory layout description
 #[derive(Copy, Clone)]
 pub struct Layout(u32);
 
 impl Layout {
     #[inline(always)]
-    pub(crate) fn new(x: u32) -> Self {
-        Layout(x)
+    pub(crate) fn is(self, flag: u32) -> bool {
+        self.0 & flag != 0
     }
 
+    /// Return layout common to both inputs
     #[inline(always)]
-    pub(crate) fn is(self, flag: u32) -> bool {
-        self.0 & flag != 0
+    pub(crate) fn intersect(self, other: Layout) -> Layout {
+        Layout(self.0 & other.0)
     }
+
+    /// Return a layout that simultaneously "is" what both of the inputs are
     #[inline(always)]
-    pub(crate) fn and(self, flag: Layout) -> Layout {
-        Layout(self.0 & flag.0)
+    pub(crate) fn also(self, other: Layout) -> Layout {
+        Layout(self.0 | other.0)
     }
 
     #[inline(always)]
-    pub(crate) fn flag(self) -> u32 {
-        self.0
+    pub(crate) fn one_dimensional() -> Layout {
+        Layout::c().also(Layout::f())
     }
-}
 
-impl Layout {
-    #[doc(hidden)]
     #[inline(always)]
-    pub fn one_dimensional() -> Layout {
-        Layout(CORDER | FORDER)
+    pub(crate) fn c() -> Layout {
+        Layout(CORDER | CPREFER)
     }
-    #[doc(hidden)]
+
     #[inline(always)]
-    pub fn c() -> Layout {
-        Layout(CORDER)
+    pub(crate) fn f() -> Layout {
+        Layout(FORDER | FPREFER)
     }
-    #[doc(hidden)]
+
     #[inline(always)]
-    pub fn f() -> Layout {
-        Layout(FORDER)
+    pub(crate) fn cpref() -> Layout {
+        Layout(CPREFER)
     }
+
+    #[inline(always)]
+    pub(crate) fn fpref() -> Layout {
+        Layout(FPREFER)
+    }
+
     #[inline(always)]
-    #[doc(hidden)]
-    pub fn none() -> Layout {
+    pub(crate) fn none() -> Layout {
         Layout(0)
     }
+
+    /// A simple "score" method which scores positive for preferring C-order, negative for F-order
+    /// Subject to change when we can describe other layouts
+    pub(crate) fn tendency(self) -> i32 {
+        (self.is(CORDER) as i32 - self.is(FORDER) as i32) +
+        (self.is(CPREFER) as i32 - self.is(FPREFER) as i32)
+
+    }
 }
 
 pub const CORDER: u32 = 0b01;
 pub const FORDER: u32 = 0b10;
+pub const CPREFER: u32 = 0b0100;
+pub const FPREFER: u32 = 0b1000;
+
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::imp_prelude::*;
+    use crate::NdProducer;
+
+    type M = Array2<f32>;
+
+    #[test]
+    fn contig_layouts() {
+        let a = M::zeros((5, 5));
+        let b = M::zeros((5, 5).f());
+        let ac = a.view().layout();
+        let af = b.view().layout();
+        assert!(ac.is(CORDER) && ac.is(CPREFER));
+        assert!(!ac.is(FORDER) && !ac.is(FPREFER));
+        assert!(!af.is(CORDER) && !af.is(CPREFER));
+        assert!(af.is(FORDER) && af.is(FPREFER));
+    }
+
+    #[test]
+    fn stride_layouts() {
+        let a = M::zeros((5, 5));
+
+        {
+            let v1 = a.slice(s![1.., ..]).layout();
+            let v2 = a.slice(s![.., 1..]).layout();
+
+            assert!(v1.is(CORDER) && v1.is(CPREFER));
+            assert!(!v1.is(FORDER) && !v1.is(FPREFER));
+            assert!(!v2.is(CORDER) && v2.is(CPREFER));
+            assert!(!v2.is(FORDER) && !v2.is(FPREFER));
+        }
+
+        let b = M::zeros((5, 5).f());
+
+        {
+            let v1 = b.slice(s![1.., ..]).layout();
+            let v2 = b.slice(s![.., 1..]).layout();
+
+            assert!(!v1.is(CORDER) && !v1.is(CPREFER));
+            assert!(!v1.is(FORDER) && v1.is(FPREFER));
+            assert!(!v2.is(CORDER) && !v2.is(CPREFER));
+            assert!(v2.is(FORDER) && v2.is(FPREFER));
+        }
+    }
+
+    #[test]
+    fn skip_layouts() {
+        let a = M::zeros((5, 5));
+        {
+            let v1 = a.slice(s![..;2, ..]).layout();
+            let v2 = a.slice(s![.., ..;2]).layout();
+
+            assert!(!v1.is(CORDER) && v1.is(CPREFER));
+            assert!(!v1.is(FORDER) && !v1.is(FPREFER));
+            assert!(!v2.is(CORDER) && !v2.is(CPREFER));
+            assert!(!v2.is(FORDER) && !v2.is(FPREFER));
+        }
+
+        let b = M::zeros((5, 5).f());
+        {
+            let v1 = b.slice(s![..;2, ..]).layout();
+            let v2 = b.slice(s![.., ..;2]).layout();
+
+            assert!(!v1.is(CORDER) && !v1.is(CPREFER));
+            assert!(!v1.is(FORDER) && !v1.is(FPREFER));
+            assert!(!v2.is(CORDER) && !v2.is(CPREFER));
+            assert!(!v2.is(FORDER) && v2.is(FPREFER));
+        }
+    }
+}