Merge pull request #965 from SparrowLii/to_dim

Improve performance of "no-broadcasting-needed" scenario in &array +&array operation

Author: bluss

benches/bench1.rs

@@ -11,7 +11,7 @@ extern crate test;
 
 use std::mem::MaybeUninit;
 
-use ndarray::ShapeBuilder;
+use ndarray::{ShapeBuilder, Array3, Array4};
 use ndarray::{arr0, arr1, arr2, azip, s};
 use ndarray::{Array, Array1, Array2, Axis, Ix, Zip};
 use ndarray::{Ix1, Ix2, Ix3, Ix5, IxDyn};
@@ -998,3 +998,21 @@ fn into_dyn_dyn(bench: &mut test::Bencher) {
     let a = a.view();
     bench.iter(|| a.clone().into_dyn());
 }
+
+#[bench]
+fn broadcast_same_dim(bench: &mut test::Bencher) {
+    let s = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12];
+    let s = Array4::from_shape_vec((2, 2, 3, 2), s.to_vec()).unwrap();
+    let a = s.slice(s![.., ..;-1, ..;2, ..]);
+    let b = s.slice(s![.., .., ..;2, ..]);
+    bench.iter(|| &a + &b);
+}
+
+#[bench]
+fn broadcast_one_side(bench: &mut test::Bencher) {
+    let s = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12];
+    let s2 = [1 ,2 ,3 ,4 ,5 ,6];
+    let a = Array4::from_shape_vec((4, 1, 3, 2), s.to_vec()).unwrap();
+    let b = Array3::from_shape_vec((1, 3, 2), s2.to_vec()).unwrap();
+    bench.iter(|| &a + &b);
+}

src/impl_methods.rs

@@ -1804,12 +1804,21 @@ where
         E: Dimension,
     {
         let shape = co_broadcast::<D, E, <D as DimMax<E>>::Output>(&self.dim, &other.dim)?;
-        if let Some(view1) = self.broadcast(shape.clone()) {
-            if let Some(view2) = other.broadcast(shape) {
-                return Ok((view1, view2));
-            }
-        }
-        Err(from_kind(ErrorKind::IncompatibleShape))
+        let view1 = if shape.slice() == self.dim.slice() {
+            self.view().into_dimensionality::<<D as DimMax<E>>::Output>().unwrap()
+        } else if let Some(view1) = self.broadcast(shape.clone()) {
+            view1
+        } else {
+            return Err(from_kind(ErrorKind::IncompatibleShape))
+        };
+        let view2 = if shape.slice() == other.dim.slice() {
+            other.view().into_dimensionality::<<D as DimMax<E>>::Output>().unwrap()
+        } else if let Some(view2) = other.broadcast(shape) {
+            view2
+        } else {
+            return Err(from_kind(ErrorKind::IncompatibleShape))
+        };
+        Ok((view1, view2))
     }
 
     /// Swap axes `ax` and `bx`.

src/impl_ops.rs

@@ -179,8 +179,14 @@ where
 {
     type Output = Array<A, <D as DimMax<E>>::Output>;
     fn $mth(self, rhs: &'a ArrayBase<S2, E>) -> Self::Output {
-        let (lhs, rhs) = self.broadcast_with(rhs).unwrap();
-        Zip::from(&lhs).and(&rhs).map_collect(clone_opf(A::$mth))
+        let (lhs, rhs) = if self.ndim() == rhs.ndim() && self.shape() == rhs.shape() {
+            let lhs = self.view().into_dimensionality::<<D as DimMax<E>>::Output>().unwrap();
+            let rhs = rhs.view().into_dimensionality::<<D as DimMax<E>>::Output>().unwrap();
+            (lhs, rhs)
+        } else {
+            self.broadcast_with(rhs).unwrap()
+        };
+        Zip::from(lhs).and(rhs).map_collect(clone_opf(A::$mth))
     }
 }
 

tests/array.rs

@@ -1627,6 +1627,12 @@ fn arithmetic_broadcast() {
         sa2 + &sb2 + sc2.into_owned(),
         arr3(&[[[15, 19], [32, 36]], [[14, 18], [31, 35]]])
     );
+
+    // Same shape
+    let a = s.slice(s![..;-1, ..;2, ..]);
+    let b = s.slice(s![.., ..;2, ..]);
+    assert_eq!(a.shape(), b.shape());
+    assert_eq!(&a + &b, arr3(&[[[3, 7], [19, 23]], [[3, 7], [19, 23]]]));
 }
 
 #[test]