refactor: optimize nd-dot by avoiding copy on standard layout, and iterate axis for non-contiguous

Author: EbrahimEldesoky

src/linalg/impl_linalg.rs

@@ -237,24 +237,39 @@ macro_rules! impl_dot_nd_ix2 {
                 let k2 = rhs.shape()[0];
                 let n = rhs.shape()[1];
                 if k != k2 {
-                    dot_shape_error(self.len() / k, k, k2, n);
+                    panic!(
+                        "shapes {:?} and {:?} are not compatible for nd dot \
+                         (last axis of lhs must equal first axis of rhs)",
+                        self.shape(),
+                        rhs.shape()
+                    );
                 }
-                let rows = self.len() / k;
-                let lhs_2d = self
-                    .to_shape((rows, k))
-                    .expect("ndarray: to_shape failed in nd dot");
-                let result_2d = lhs_2d.dot(rhs);
-
-                let mut out_dim = <$dim>::zeros(ndim);
-                for i in 0..ndim - 1 {
-                    out_dim[i] = self.shape()[i];
-                }
-                out_dim[ndim - 1] = n;
 
-                result_2d
-                    .to_shape(out_dim)
-                    .expect("ndarray: to_shape failed reshaping nd dot result")
-                    .into_owned()
+                if self.is_standard_layout() {
+                    // C-contiguous: to_shape returns a *view* (no copy of LHS data).
+                    // Safety: rows * k == self.len() by construction.
+                    let rows = self.len() / k;
+                    let lhs_2d = self.to_shape((rows, k)).unwrap();
+                    let result_2d = lhs_2d.dot(rhs);
+
+                    let mut out_dim = <$dim>::zeros(ndim);
+                    for i in 0..ndim - 1 {
+                        out_dim[i] = self.shape()[i];
+                    }
+                    out_dim[ndim - 1] = n;
+
+                    // result_2d is a fresh C-contiguous owned array;
+                    // into_shape_with_order is free.
+                    result_2d.into_shape_with_order(out_dim).unwrap()
+                } else {
+                    // Non-contiguous: iterate over the first axis so no whole-array
+                    // copy is needed. Each sub-array is (ndim-1)-D; the impl for that
+                    // dimension is already compiled (macro invocations are in order).
+                    let sub_results: Vec<_> =
+                        self.axis_iter(Axis(0)).map(|lane| lane.dot(rhs)).collect();
+                    let views: Vec<_> = sub_results.iter().map(|a| a.view()).collect();
+                    crate::stack(Axis(0), &views).unwrap()
+                }
             }
         }
 
@@ -280,21 +295,35 @@ where A: LinalgScalar
         let k2 = rhs.shape()[0];
         let n = rhs.shape()[1];
         if k != k2 {
-            dot_shape_error(self.len() / k, k, k2, n);
+            panic!(
+                "shapes {:?} and {:?} are not compatible for nd dot \
+                 (last axis of lhs must equal first axis of rhs)",
+                self.shape(),
+                rhs.shape()
+            );
+        }
+
+        if self.is_standard_layout() {
+            // C-contiguous: to_shape returns a *view* (no copy of LHS data).
+            // Safety: rows * k == self.len() by construction.
+            let rows = self.len() / k;
+            let lhs_2d = self.to_shape((rows, k)).unwrap();
+            let result_2d = lhs_2d.dot(rhs);
+
+            let mut out_shape = self.shape().to_vec();
+            *out_shape.last_mut().unwrap() = n;
+
+            // result_2d is a fresh C-contiguous owned array;
+            // into_shape_with_order is free.
+            result_2d.into_shape_with_order(IxDyn(&out_shape)).unwrap()
+        } else {
+            // Non-contiguous: iterate over the first axis so no whole-array
+            // copy is needed. Each sub-array is (ndim-1)-D IxDyn; recursion
+            // eventually reaches contiguous 2-D which terminates the recursion.
+            let sub_results: Vec<_> = self.axis_iter(Axis(0)).map(|lane| lane.dot(rhs)).collect();
+            let views: Vec<_> = sub_results.iter().map(|a| a.view()).collect();
+            crate::stack(Axis(0), &views).unwrap()
         }
-        let rows = self.len() / k;
-        let lhs_2d = self
-            .to_shape((rows, k))
-            .expect("ndarray: to_shape failed in nd dot (IxDyn)");
-        let result_2d = lhs_2d.dot(rhs);
-
-        let mut out_shape = self.shape().to_vec();
-        *out_shape.last_mut().unwrap() = n;
-
-        result_2d
-            .to_shape(IxDyn(&out_shape))
-            .expect("ndarray: to_shape failed reshaping nd dot result (IxDyn)")
-            .into_owned()
     }
 }
 

tests/oper.rs

@@ -939,7 +939,7 @@ fn dot_3d_by_2d_integer()
 }
 
 #[test]
-#[should_panic(expected = "not compatible for matrix multiplication")]
+#[should_panic(expected = "are not compatible for nd dot")]
 fn dot_3d_by_2d_shape_mismatch()
 {
     let lhs: Array3<f64> = Array3::zeros((3, 4, 5));
@@ -1019,7 +1019,23 @@ fn dot_dyn_5d_by_2d()
 }
 
 #[test]
-#[should_panic(expected = "not compatible for matrix multiplication")]
+fn dot_dyn_3d_by_2d_non_contiguous()
+{
+    // Slice with stride 2 to get a non-contiguous layout.
+    let base: Array3<f64> = ArrayBuilder::new((6, 4, 5)).build();
+    let lhs = base.slice(s![..;2, .., ..]).into_dyn();
+    let rhs: Array2<f64> = ArrayBuilder::new((5, 7)).build();
+
+    let result = lhs.dot(&rhs);
+    assert_eq!(result.shape(), &[3, 4, 7]);
+
+    let lhs_fixed: ArrayView3<f64> = lhs.into_dimensionality::<Ix3>().unwrap();
+    let expected = lhs_fixed.dot(&rhs);
+    assert_eq!(result.into_dimensionality::<Ix3>().unwrap(), expected);
+}
+
+#[test]
+#[should_panic(expected = "are not compatible for nd dot")]
 fn dot_dyn_shape_mismatch()
 {
     let lhs: ArrayD<f64> = ArrayD::zeros(IxDyn(&[3, 4, 5]));