Merge pull request #280 from bluss/axpy-fixup

BUG: Fix blas + axpy by only using it for certain cases

Author: bluss

benches/bench1.rs

@@ -457,26 +457,11 @@ fn iadd_2d_strided(bench: &mut test::Bencher)
     });
 }
 
-const SCALE_ADD_SZ: usize = 64;
-
 #[bench]
 fn scaled_add_2d_f32_regular(bench: &mut test::Bencher)
 {
-    let mut av = Array::<f32, _>::zeros((SCALE_ADD_SZ, SCALE_ADD_SZ));
-    let bv = Array::<f32, _>::zeros((SCALE_ADD_SZ, SCALE_ADD_SZ));
-    let scalar = 3.1415926535;
-    bench.iter(|| {
-        av.scaled_add(scalar, &bv);
-    });
-}
-
-#[bench]
-fn scaled_add_2d_f32_stride(bench: &mut test::Bencher)
-{
-    let mut av = Array::<f32, _>::zeros((SCALE_ADD_SZ, 2 * SCALE_ADD_SZ));
-    let bv = Array::<f32, _>::zeros((SCALE_ADD_SZ, 2 * SCALE_ADD_SZ));
-    let mut av = av.slice_mut(s![.., ..;2]);
-    let bv = bv.slice(s![.., ..;2]);
+    let mut av = Array::<f32, _>::zeros((64, 64));
+    let bv = Array::<f32, _>::zeros((64, 64));
     let scalar = 3.1415926535;
     bench.iter(|| {
         av.scaled_add(scalar, &bv);

src/dimension/dimension_trait.rs

@@ -269,34 +269,20 @@ pub unsafe trait Dimension : Clone + Eq + Debug + Send + Sync + Default +
             return true;
         }
         if dim.ndim() == 1 { return false; }
-
-        match Self::equispaced_stride(dim, strides) {
-            Some(1) => true,
-            _ => false,
-        }
-    }
-
-    /// Return the equispaced stride between all the array elements.
-    ///
-    /// Returns `Some(n)` if the strides in all dimensions are equispaced. Returns `None` if not.
-    #[doc(hidden)]
-    fn equispaced_stride(dim: &Self, strides: &Self) -> Option<isize> {
         let order = strides._fastest_varying_stride_order();
-        let base_stride = strides[order[0]];
+        let strides = strides.slice();
 
         // FIXME: Negative strides
         let dim_slice = dim.slice();
-        let mut next_stride = base_stride;
-        let strides = strides.slice();
+        let mut cstride = 1;
         for &i in order.slice() {
             // a dimension of length 1 can have unequal strides
-            if dim_slice[i] != 1 && strides[i] != next_stride {
-                return None;
+            if dim_slice[i] != 1 && strides[i] != cstride {
+                return false;
             }
-            next_stride *= dim_slice[i];
+            cstride *= dim_slice[i];
         }
-
-       Some(base_stride as isize)
+        true
     }
 
     /// Return the axis ordering corresponding to the fastest variation

src/linalg/impl_linalg.rs

@@ -292,72 +292,9 @@ impl<A, S, D> ArrayBase<S, D>
               S2: Data<Elem=A>,
               A: LinalgScalar,
               E: Dimension,
-    {
-        self.scaled_add_impl(alpha, rhs);
-    }
-
-    fn scaled_add_generic<S2, E>(&mut self, alpha: A, rhs: &ArrayBase<S2, E>)
-        where S: DataMut,
-              S2: Data<Elem=A>,
-              A: LinalgScalar,
-              E: Dimension,
     {
         self.zip_mut_with(rhs, move |y, &x| *y = *y + (alpha * x));
     }
-
-    #[cfg(not(feature = "blas"))]
-    fn scaled_add_impl<S2, E>(&mut self, alpha: A, rhs: &ArrayBase<S2, E>)
-        where S: DataMut,
-              S2: Data<Elem=A>,
-              A: LinalgScalar,
-              E: Dimension,
-    {
-        self.scaled_add_generic(alpha, rhs);
-    }
-
-    #[cfg(feature = "blas")]
-    fn scaled_add_impl<S2, E>(&mut self, alpha: A, rhs: &ArrayBase<S2, E>)
-        where S: DataMut,
-              S2: Data<Elem=A>,
-              A: LinalgScalar,
-              E: Dimension,
-    {
-        debug_assert_eq!(self.len(), rhs.len());
-        assert!(self.len() == rhs.len());
-        {
-            macro_rules! axpy {
-                ($ty:ty, $func:ident) => {{
-            if blas_compat::<$ty, _, _>(self) && blas_compat::<$ty, _, _>(rhs) {
-                let order = Dimension::_fastest_varying_stride_order(&self.strides);
-                let incx = self.strides()[order[0]];
-
-                let order = Dimension::_fastest_varying_stride_order(&rhs.strides);
-                let incy = self.strides()[order[0]];
-
-                unsafe {
-                    let (lhs_ptr, n, incx) = blas_1d_params(self.ptr,
-                                                            self.len(),
-                                                            incx);
-                    let (rhs_ptr, _, incy) = blas_1d_params(rhs.ptr,
-                                                            rhs.len(),
-                                                            incy);
-                    blas_sys::c::$func(
-                        n,
-                        cast_as(&alpha),
-                        rhs_ptr as *const $ty,
-                        incy,
-                        lhs_ptr as *mut $ty,
-                        incx);
-                    return;
-                }
-            }
-                }}
-            }
-            axpy!{f32, cblas_saxpy};
-            axpy!{f64, cblas_daxpy};
-        }
-        self.scaled_add_generic(alpha, rhs);
-    }
 }
 
 // mat_mul_impl uses ArrayView arguments to send all array kinds into
@@ -594,32 +531,6 @@ fn blas_compat_1d<A, S>(a: &ArrayBase<S, Ix1>) -> bool
     true
 }
 
-#[cfg(feature="blas")]
-fn blas_compat<A, S, D>(a: &ArrayBase<S, D>) -> bool
-    where S: Data,
-          A: 'static,
-          S::Elem: 'static,
-          D: Dimension,
-{
-    if !same_type::<A, S::Elem>() {
-        return false;
-    }
-
-    match D::equispaced_stride(&a.raw_dim(), &a.strides) {
-        Some(stride) => {
-            if a.len() as isize * stride > blas_index::max_value() as isize ||
-                stride < blas_index::min_value() as isize {
-                return false;
-            }
-        },
-        None => {
-            return false;
-        }
-    }
-
-    true
-}
-
 #[cfg(feature="blas")]
 fn blas_row_major_2d<A, S>(a: &ArrayBase<S, Ix2>) -> bool
     where S: Data,

tests/dimension.rs

@@ -42,18 +42,6 @@ fn dyn_dimension()
     assert_eq!(z.shape(), &dim[..]);
 }
 
-#[test]
-fn equidistance_strides() {
-    let strides = Dim([4,2,1]);
-    assert_eq!(Dimension::equispaced_stride(&Dim([2,2,2]), &strides), Some(1));
-
-    let strides = Dim([8,4,2]);
-    assert_eq!(Dimension::equispaced_stride(&Dim([2,2,2]), &strides), Some(2));
-
-    let strides = Dim([16,4,1]);
-    assert_eq!(Dimension::equispaced_stride(&Dim([2,2,2]), &strides), None);
-}
-
 #[test]
 fn fastest_varying_order() {
     let strides = Dim([2, 8, 4, 1]);

tests/oper.rs

@@ -5,6 +5,7 @@ use ndarray::prelude::*;
 use ndarray::{rcarr1, rcarr2};
 use ndarray::{LinalgScalar, Data};
 use ndarray::linalg::general_mat_mul;
+use ndarray::Si;
 
 use std::fmt;
 use num_traits::Float;
@@ -494,6 +495,52 @@ fn scaled_add_2() {
     }
 }
 
+#[test]
+fn scaled_add_3() {
+    let beta = -2.3;
+    let sizes = vec![(4, 4, 1, 4),
+                     (8, 8, 1, 8),
+                     (17, 15, 17, 15),
+                     (4, 17, 4, 17),
+                     (17, 3, 1, 3),
+                     (19, 18, 19, 18),
+                     (16, 17, 16, 17),
+                     (15, 16, 15, 16),
+                     (67, 63, 1, 63),
+        ];
+    // test different strides
+    for &s1 in &[1, 2, -1, -2] {
+        for &s2 in &[1, 2, -1, -2] {
+            for &(m, k, n, q) in &sizes {
+                let mut a = range_mat64(m, k);
+                let mut answer = a.clone();
+                let cdim = if n == 1 {
+                    vec![q]
+                } else {
+                    vec![n, q]
+                };
+                let cslice = if n == 1 {
+                    vec![Si(0, None, s2)]
+                } else {
+                    vec![Si(0, None, s1), Si(0, None, s2)]
+                };
+
+                let c = range_mat64(n, q).into_shape(cdim).unwrap();
+
+                {
+                    let mut av = a.slice_mut(s![..;s1, ..;s2]);
+                    let c = c.slice(&cslice);
+
+                    let mut answerv = answer.slice_mut(s![..;s1, ..;s2]);
+                    answerv += &(beta * &c);
+                    av.scaled_add(beta, &c);
+                }
+                assert_close(a.view(), answer.view());
+            }
+        }
+    }
+}
+
 
 #[test]
 fn gen_mat_mul() {