Skip to content

Use axpy for scaled_add #278

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Mar 22, 2017
Merged

Use axpy for scaled_add #278

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
5e6a6b4
Check if all array elements are equidistant
SuperFluffy Mar 18, 2017
35a7dc6
Use axpy for scaled_add
SuperFluffy Mar 18, 2017
38cdcc3
Move axpy into a generic scaled_add interface
SuperFluffy Mar 19, 2017
d1f50b0
Add bench for scaled add
SuperFluffy Mar 21, 2017
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 17 additions & 2 deletions benches/bench1.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -457,11 +457,26 @@ 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((64, 64));
let bv = Array::<f32, _>::zeros((64, 64));
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 scalar = 3.1415926535;
bench.iter(|| {
av.scaled_add(scalar, &bv);
Expand Down
26 changes: 20 additions & 6 deletions src/dimension/dimension_trait.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -269,20 +269,34 @@ 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 strides = strides.slice();
let base_stride = strides[order[0]];

// FIXME: Negative strides
let dim_slice = dim.slice();
let mut cstride = 1;
let mut next_stride = base_stride;
let strides = strides.slice();
for &i in order.slice() {
// a dimension of length 1 can have unequal strides
if dim_slice[i] != 1 && strides[i] != cstride {
return false;
if dim_slice[i] != 1 && strides[i] != next_stride {
return None;
}
cstride *= dim_slice[i];
next_stride *= dim_slice[i];
}
true

Some(base_stride as isize)
}

/// Return the axis ordering corresponding to the fastest variation
Expand Down
89 changes: 89 additions & 0 deletions src/linalg/impl_linalg.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -292,9 +292,72 @@ 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());
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Lol, I missed a lot. This is not what we want to do.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We don't? :( I basically took this from dot_impl:

if self.len() >= DOT_BLAS_CUTOFF {
    debug_assert_eq!(self.len(), rhs.len());
    assert!(self.len() == rhs.len());

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

rhs is broadcast to self's dim if they mismatch. Yeah, dot doesn't broadcast.

Copy link
Member

@bluss bluss Mar 22, 2017 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sorry, the mess is on my side. I've added more rigorous tests so that it has something to aspire to.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Right, the assertion there doesn't make any sense. With broadcasting, one stride would be 0 so that equispaced_stride returns None (as per the argument above), which would skip axpy and fall through to the default implementation. The assertion would simply prevent the broadcasting with the generic scaled_add to trigger.

{
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
Expand Down Expand Up @@ -531,6 +594,32 @@ 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) {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is it possible that this stride is 0 here? I don't think blas supports that.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Depends, can the stride ever be 0 in any dimension and for any kind of vector? A n-dimensional vector with one dimension set to 0 doesn't really make sense, so we should make sure that such a state can never be reached.

A special case is a 0-dimensional array. It looks like its stride is set to 0.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes it's very common, that's how broadcasting works (same in numpy)

>>> import numpy as np
>>> np.broadcast_to(np.arange(9.), (9, 9)).strides
(0, 8)

Copy link
Contributor Author

@SuperFluffy SuperFluffy Mar 22, 2017 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ok, so the way the equispaced_stride method is implemented right now is that it grabs the smallest stride and calculates the rest based on that. That means that if we indeed have one stride set to 0, the demanded/calculated stride will always be 0 irrespective of the dimension (because it predicts the necessary stride s_i+1 as s_i * d_i). It follows that the demanded stride will never match the rest of the stride info in the struct. Thus the function will return None and we won't be able to use blas.

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,
Expand Down
12 changes: 12 additions & 0 deletions tests/dimension.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,18 @@ 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]);
Expand Down