Use axpy for scaled_add

Author: SuperFluffy

benches/bench1.rs

@@ -1,6 +1,7 @@
 #![feature(test)]
 #![allow(unused_imports)]
 
+extern crate blas_sys;
 extern crate test;
 #[macro_use(s)]
 extern crate ndarray;
@@ -468,6 +469,18 @@ fn scaled_add_2d_f32_regular(bench: &mut test::Bencher)
     });
 }
 
+#[bench]
+fn scaled_add_2d_f32_axpy(bench: &mut test::Bencher)
+{
+    let mut av = Array::<f32, _>::zeros((64, 64));
+    let bv = Array::<f32, _>::zeros((64, 64));
+    let scalar = 3.1415926535;
+    bench.iter(|| {
+        av.scaled_add_axpy(scalar, &bv);
+    });
+}
+
+
 #[bench]
 fn assign_scalar_2d_corder(bench: &mut test::Bencher)
 {

src/dimension/dimension_trait.rs

@@ -280,7 +280,7 @@ pub unsafe trait Dimension : Clone + Eq + Debug + Send + Sync + Default +
     ///
     /// 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<usize> {
+    fn equispaced_stride(dim: &Self, strides: &Self) -> Option<isize> {
         let order = strides._fastest_varying_stride_order();
         let base_stride = strides[order[0]];
 
@@ -295,7 +295,10 @@ pub unsafe trait Dimension : Clone + Eq + Debug + Send + Sync + Default +
             }
             next_stride *= dim_slice[i];
         }
-        Some(base_stride)
+
+        unsafe {
+           Some(::std::ptr::read(&base_stride as *const _ as *const isize))
+        }
     }
 
     /// Return the axis ordering corresponding to the fastest variation

src/impl_methods.rs

@@ -73,6 +73,11 @@ impl<A, S, D> ArrayBase<S, D> where S: Data<Elem=A>, D: Dimension
         self.dim.clone()
     }
 
+    /// Return the strides of the array as they are stored in the array.
+    pub fn raw_strides(&self) -> D {
+        self.strides.clone()
+    }
+
     /// Return the shape of the array as a slice.
     pub fn shape(&self) -> &[Ix] {
         self.dim.slice()

src/linalg/impl_linalg.rs

@@ -295,6 +295,62 @@ impl<A, S, D> ArrayBase<S, D>
     {
         self.zip_mut_with(rhs, move |y, &x| *y = *y + (alpha * x));
     }
+
+    /// Perform the operation `self += alpha * rhs` efficiently, where
+    /// `alpha` is a scalar and `rhs` is another array. This operation is
+    /// also known as `axpy` in BLAS.
+    ///
+    /// If their shapes disagree, `rhs` is broadcast to the shape of `self`.
+    ///
+    /// **Panics** if broadcasting isn’t possible.
+    #[cfg(feature="blas")]
+    pub fn scaled_add_axpy<S2, E>(&mut self, alpha: A, rhs: &ArrayBase<S2, E>)
+        where S: DataMut,
+              S2: Data<Elem=A>,
+              A: LinalgScalar + ::std::fmt::Debug,
+              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 strides = self.raw_strides();
+                let order = Dimension::_fastest_varying_stride_order(&strides);
+                let incx = self.strides()[order[0]];
+
+                let strides = rhs.raw_strides();
+                let order = Dimension::_fastest_varying_stride_order(&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(alpha, rhs);
+    }
 }
 
 // mat_mul_impl uses ArrayView arguments to send all array kinds into
@@ -531,6 +587,35 @@ 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;
+    }
+
+    if a.len() > blas_index::max_value() as usize {
+        return false;
+    }
+
+    match D::equispaced_stride(&a.raw_dim(), &a.raw_strides()) {
+        Some(stride) => {
+            if 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/oper.rs

@@ -459,6 +459,22 @@ fn scaled_add() {
 
 }
 
+#[cfg(blas)]
+#[test]
+fn scaled_add_axpy() {
+    let a = range_mat(16, 15);
+    let mut b = range_mat(16, 15);
+    b.mapv_inplace(f32::exp);
+
+    let alpha = 0.2_f32;
+    let mut c = a.clone();
+    c.scaled_add(alpha, &b);
+
+    let mut d = a.clone();
+    d.scaled_add_axpy(alpha, &b);
+    assert_eq!(c, d);
+}
+
 #[test]
 fn gen_mat_mul() {
     let alpha = -2.3;