Skip to content

Add methods to collect Zip into an array #797

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.

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 7 commits into from
Apr 18, 2020
Merged
Changes from all commits
Commits
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
18 changes: 15 additions & 3 deletions benches/bench1.rs
Original file line number Diff line number Diff line change
@@ -258,23 +258,35 @@ fn add_2d_zip(bench: &mut test::Bencher) {
}

#[bench]
fn add_2d_alloc(bench: &mut test::Bencher) {
fn add_2d_alloc_plus(bench: &mut test::Bencher) {
let a = Array::<i32, _>::zeros((ADD2DSZ, ADD2DSZ));
let b = Array::<i32, _>::zeros((ADD2DSZ, ADD2DSZ));
bench.iter(|| &a + &b);
}

#[bench]
fn add_2d_zip_alloc(bench: &mut test::Bencher) {
fn add_2d_alloc_zip_uninit(bench: &mut test::Bencher) {
let a = Array::<i32, _>::zeros((ADD2DSZ, ADD2DSZ));
let b = Array::<i32, _>::zeros((ADD2DSZ, ADD2DSZ));
bench.iter(|| unsafe {
let mut c = Array::uninitialized(a.dim());
azip!((&a in &a, &b in &b, c in &mut c) *c = a + b);
azip!((&a in &a, &b in &b, c in c.raw_view_mut())
std::ptr::write(c, a + b)
);
c
});
}

#[bench]
fn add_2d_alloc_zip_collect(bench: &mut test::Bencher) {
let a = Array::<i32, _>::zeros((ADD2DSZ, ADD2DSZ));
let b = Array::<i32, _>::zeros((ADD2DSZ, ADD2DSZ));
bench.iter(|| {
Zip::from(&a).and(&b).apply_collect(|&x, &y| x + y)
});
}


#[bench]
fn add_2d_assign_ops(bench: &mut test::Bencher) {
let mut a = Array::<i32, _>::zeros((ADD2DSZ, ADD2DSZ));
31 changes: 31 additions & 0 deletions src/argument_traits.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
use std::cell::Cell;
use std::mem::MaybeUninit;


/// A producer element that can be assigned to once
pub trait AssignElem<T> {
/// Assign the value `input` to the element that self represents.
fn assign_elem(self, input: T);
}

/// Assignable element, simply `*self = input`.
impl<'a, T> AssignElem<T> for &'a mut T {
fn assign_elem(self, input: T) {
*self = input;
}
}

/// Assignable element, simply `self.set(input)`.
impl<'a, T> AssignElem<T> for &'a Cell<T> {
fn assign_elem(self, input: T) {
self.set(input);
}
}

/// Assignable element, the item in the MaybeUninit is overwritten (prior value, if any, is not
/// read or dropped).
impl<'a, T> AssignElem<T> for &'a mut MaybeUninit<T> {
fn assign_elem(self, input: T) {
*self = MaybeUninit::new(input);
}
}
20 changes: 20 additions & 0 deletions src/impl_constructors.rs
Original file line number Diff line number Diff line change
@@ -13,6 +13,7 @@
#![allow(clippy::match_wild_err_arm)]

use num_traits::{Float, One, Zero};
use std::mem::MaybeUninit;

use crate::dimension;
use crate::error::{self, ShapeError};
@@ -517,3 +518,22 @@ where
Self::from_shape_vec_unchecked(shape, v)
}
}

impl<S, A, D> ArrayBase<S, D>
where
S: DataOwned<Elem = MaybeUninit<A>>,
D: Dimension,
{
pub(crate) fn maybe_uninit<Sh>(shape: Sh) -> Self
where
Sh: ShapeBuilder<Dim = D>,
{
unsafe {
let shape = shape.into_shape();
let size = size_of_shape_checked_unwrap!(&shape.dim);
let mut v = Vec::with_capacity(size);
v.set_len(size);
Self::from_shape_vec_unchecked(shape, v)
}
}
}
34 changes: 34 additions & 0 deletions src/impl_owned_array.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,8 @@
use std::mem::MaybeUninit;
use std::mem::transmute;

use crate::imp_prelude::*;
use crate::OwnedRepr;

/// Methods specific to `Array0`.
///
@@ -57,3 +61,33 @@ where
self.data.0
}
}

/// Methods specific to `Array` of `MaybeUninit`.
///
/// ***See also all methods for [`ArrayBase`]***
///
/// [`ArrayBase`]: struct.ArrayBase.html
impl<A, D> Array<MaybeUninit<A>, D>
where
D: Dimension,
{
/// Assert that the array's storage's elements are all fully initialized, and conver
/// the array from element type `MaybeUninit<A>` to `A`.
pub(crate) unsafe fn assume_init(self) -> Array<A, D> {
// NOTE: Fully initialized includes elements not reachable in current slicing/view.
//
// Should this method be generalized to all array types?
// (Will need a way to map the RawData<Elem=X> to RawData<Elem=Y> of same kind)

let Array { data, ptr, dim, strides } = self;
let data = transmute::<OwnedRepr<MaybeUninit<A>>, OwnedRepr<A>>(data);
let ptr = ptr.cast::<A>();

Array {
data,
ptr,
dim,
strides,
}
}
}
1 change: 0 additions & 1 deletion src/layout/layoutfmt.rs
Original file line number Diff line number Diff line change
@@ -7,7 +7,6 @@
// except according to those terms.

use super::Layout;
use super::LayoutPriv;

const LAYOUT_NAMES: &[&str] = &["C", "F"];

19 changes: 6 additions & 13 deletions src/layout/mod.rs
Original file line number Diff line number Diff line change
@@ -1,35 +1,28 @@
mod layoutfmt;

// public but users don't interact with it
// public struct but users don't interact with it
#[doc(hidden)]
/// Memory layout description
#[derive(Copy, Clone)]
pub struct Layout(u32);

pub trait LayoutPriv: Sized {
fn new(x: u32) -> Self;
fn and(self, flag: Self) -> Self;
fn is(self, flag: u32) -> bool;
fn flag(self) -> u32;
}

impl LayoutPriv for Layout {
impl Layout {
#[inline(always)]
fn new(x: u32) -> Self {
pub(crate) fn new(x: u32) -> Self {
Layout(x)
}

#[inline(always)]
fn is(self, flag: u32) -> bool {
pub(crate) fn is(self, flag: u32) -> bool {
self.0 & flag != 0
}
#[inline(always)]
fn and(self, flag: Layout) -> Layout {
pub(crate) fn and(self, flag: Layout) -> Layout {
Layout(self.0 & flag.0)
}

#[inline(always)]
fn flag(self) -> u32 {
pub(crate) fn flag(self) -> u32 {
self.0
}
}
2 changes: 2 additions & 0 deletions src/lib.rs
Original file line number Diff line number Diff line change
@@ -149,6 +149,8 @@ mod array_approx;
mod array_serde;
mod arrayformat;
mod arraytraits;
mod argument_traits;
pub use crate::argument_traits::AssignElem;
mod data_traits;

pub use crate::aliases::*;
53 changes: 45 additions & 8 deletions src/parallel/impl_par_methods.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use crate::{ArrayBase, DataMut, Dimension, NdProducer, Zip};
use crate::{Array, ArrayBase, DataMut, Dimension, IntoNdProducer, NdProducer, Zip};
use crate::AssignElem;

use crate::parallel::prelude::*;

@@ -43,7 +44,7 @@ where
// Zip

macro_rules! zip_impl {
($([$($p:ident)*],)+) => {
($([$notlast:ident $($p:ident)*],)+) => {
$(
#[allow(non_snake_case)]
impl<D, $($p),*> Zip<($($p,)*), D>
@@ -63,16 +64,52 @@ macro_rules! zip_impl {
{
self.into_par_iter().for_each(move |($($p,)*)| function($($p),*))
}

expand_if!(@bool [$notlast]

/// Apply and collect the results into a new array, which has the same size as the
/// inputs.
///
/// If all inputs are c- or f-order respectively, that is preserved in the output.
///
/// Restricted to functions that produce copyable results for technical reasons; other
/// cases are not yet implemented.
pub fn par_apply_collect<R>(self, f: impl Fn($($p::Item,)* ) -> R + Sync + Send) -> Array<R, D>
where R: Copy + Send
{
let mut output = self.uninitalized_for_current_layout::<R>();
self.par_apply_assign_into(&mut output, f);
unsafe {
output.assume_init()
}
}

/// Apply and assign the results into the producer `into`, which should have the same
/// size as the other inputs.
///
/// The producer should have assignable items as dictated by the `AssignElem` trait,
/// for example `&mut R`.
pub fn par_apply_assign_into<R, Q>(self, into: Q, f: impl Fn($($p::Item,)* ) -> R + Sync + Send)
where Q: IntoNdProducer<Dim=D>,
Q::Item: AssignElem<R> + Send,
Q::Output: Send,
{
self.and(into)
.par_apply(move |$($p, )* output_| {
output_.assign_elem(f($($p ),*));
});
}
);
}
)+
}
}

zip_impl! {
[P1],
[P1 P2],
[P1 P2 P3],
[P1 P2 P3 P4],
[P1 P2 P3 P4 P5],
[P1 P2 P3 P4 P5 P6],
[true P1],
[true P1 P2],
[true P1 P2 P3],
[true P1 P2 P3 P4],
[true P1 P2 P3 P4 P5],
[false P1 P2 P3 P4 P5 P6],
}
47 changes: 46 additions & 1 deletion src/zip/mod.rs
Original file line number Diff line number Diff line change
@@ -9,13 +9,15 @@
#[macro_use]
mod zipmacro;

use std::mem::MaybeUninit;

use crate::imp_prelude::*;
use crate::AssignElem;
use crate::IntoDimension;
use crate::Layout;
use crate::NdIndex;

use crate::indexes::{indices, Indices};
use crate::layout::LayoutPriv;
use crate::layout::{CORDER, FORDER};

/// Return if the expression is a break value.
@@ -579,6 +581,7 @@ pub struct Zip<Parts, D> {
layout: Layout,
}


impl<P, D> Zip<(P,), D>
where
D: Dimension,
@@ -735,6 +738,12 @@ where
self.dimension[unroll_axis] = inner_len;
FoldWhile::Continue(acc)
}

pub(crate) fn uninitalized_for_current_layout<T>(&self) -> Array<MaybeUninit<T>, D>
{
let is_f = !self.layout.is(CORDER) && self.layout.is(FORDER);
Array::maybe_uninit(self.dimension.clone().set_f(is_f))
}
}

/*
@@ -982,6 +991,42 @@ macro_rules! map_impl {
dimension: self.dimension,
}
}

/// Apply and collect the results into a new array, which has the same size as the
/// inputs.
///
/// If all inputs are c- or f-order respectively, that is preserved in the output.
///
/// Restricted to functions that produce copyable results for technical reasons; other
/// cases are not yet implemented.
pub fn apply_collect<R>(self, f: impl FnMut($($p::Item,)* ) -> R) -> Array<R, D>
where R: Copy,
{
// To support non-Copy elements, implementation of dropping partial array (on
// panic) is needed
let mut output = self.uninitalized_for_current_layout::<R>();
self.apply_assign_into(&mut output, f);
unsafe {
output.assume_init()
}
}

/// Apply and assign the results into the producer `into`, which should have the same
/// size as the other inputs.
///
/// The producer should have assignable items as dictated by the `AssignElem` trait,
/// for example `&mut R`.
pub fn apply_assign_into<R, Q>(self, into: Q, mut f: impl FnMut($($p::Item,)* ) -> R)
where Q: IntoNdProducer<Dim=D>,
Q::Item: AssignElem<R>
{
self.and(into)
.apply(move |$($p, )* output_| {
output_.assign_elem(f($($p ),*));
});
}


);

/// Split the `Zip` evenly in two.
6 changes: 6 additions & 0 deletions src/zip/zipmacro.rs
Original file line number Diff line number Diff line change
@@ -122,6 +122,12 @@ macro_rules! azip {
$(.and($prod))*
.$apply(|$first_pat, $($pat),*| $body)
};

// Unindexed with one or more producer, no loop body
(@build $apply:ident $first_prod:expr $(, $prod:expr)* $(,)?) => {
$crate::Zip::from($first_prod)
$(.and($prod))*
};
// catch-all rule
(@build $($t:tt)*) => { compile_error!("Invalid syntax in azip!()") };
($($t:tt)*) => {
59 changes: 59 additions & 0 deletions tests/azip.rs
Original file line number Diff line number Diff line change
@@ -49,6 +49,65 @@ fn test_azip2_3() {
assert!(a != b);
}

#[test]
#[cfg(feature = "approx")]
fn test_zip_collect() {
use approx::assert_abs_diff_eq;

// test Zip::apply_collect and that it preserves c/f layout.

let b = Array::from_shape_fn((5, 10), |(i, j)| 1. / (i + 2 * j + 1) as f32);
let c = Array::from_shape_fn((5, 10), |(i, j)| f32::exp((i + j) as f32));

{
let a = Zip::from(&b).and(&c).apply_collect(|x, y| x + y);

assert_abs_diff_eq!(a, &b + &c, epsilon = 1e-6);
assert_eq!(a.strides(), b.strides());
}

{
let b = b.t();
let c = c.t();

let a = Zip::from(&b).and(&c).apply_collect(|x, y| x + y);

assert_abs_diff_eq!(a, &b + &c, epsilon = 1e-6);
assert_eq!(a.strides(), b.strides());
}
}

#[test]
#[cfg(feature = "approx")]
fn test_zip_assign_into() {
use approx::assert_abs_diff_eq;

let mut a = Array::<f32, _>::zeros((5, 10));
let b = Array::from_shape_fn((5, 10), |(i, j)| 1. / (i + 2 * j + 1) as f32);
let c = Array::from_shape_fn((5, 10), |(i, j)| f32::exp((i + j) as f32));

Zip::from(&b).and(&c).apply_assign_into(&mut a, |x, y| x + y);

assert_abs_diff_eq!(a, &b + &c, epsilon = 1e-6);
}

#[test]
#[cfg(feature = "approx")]
fn test_zip_assign_into_cell() {
use approx::assert_abs_diff_eq;
use std::cell::Cell;

let a = Array::<Cell<f32>, _>::default((5, 10));
let b = Array::from_shape_fn((5, 10), |(i, j)| 1. / (i + 2 * j + 1) as f32);
let c = Array::from_shape_fn((5, 10), |(i, j)| f32::exp((i + j) as f32));

Zip::from(&b).and(&c).apply_assign_into(&a, |x, y| x + y);
let a2 = a.mapv(|elt| elt.get());

assert_abs_diff_eq!(a2, &b + &c, epsilon = 1e-6);
}


#[test]
fn test_azip_syntax_trailing_comma() {
let mut b = Array::<i32, _>::zeros((5, 5));
42 changes: 42 additions & 0 deletions tests/par_zip.rs
Original file line number Diff line number Diff line change
@@ -70,3 +70,45 @@ fn test_zip_index_4() {
assert_eq!(*elt, j);
}
}

#[test]
#[cfg(feature = "approx")]
fn test_zip_collect() {
use approx::assert_abs_diff_eq;

// test Zip::apply_collect and that it preserves c/f layout.

let b = Array::from_shape_fn((M, N), |(i, j)| 1. / (i + 2 * j + 1) as f32);
let c = Array::from_shape_fn((M, N), |(i, j)| f32::ln((1 + i + j) as f32));

{
let a = Zip::from(&b).and(&c).par_apply_collect(|x, y| x + y);

assert_abs_diff_eq!(a, &b + &c, epsilon = 1e-6);
assert_eq!(a.strides(), b.strides());
}

{
let b = b.t();
let c = c.t();

let a = Zip::from(&b).and(&c).par_apply_collect(|x, y| x + y);

assert_abs_diff_eq!(a, &b + &c, epsilon = 1e-6);
assert_eq!(a.strides(), b.strides());
}
}

#[test]
#[cfg(feature = "approx")]
fn test_zip_assign_into() {
use approx::assert_abs_diff_eq;

let mut a = Array::<f32, _>::zeros((M, N));
let b = Array::from_shape_fn((M, N), |(i, j)| 1. / (i + 2 * j + 1) as f32);
let c = Array::from_shape_fn((M, N), |(i, j)| f32::ln((1 + i + j) as f32));

Zip::from(&b).and(&c).par_apply_assign_into(&mut a, |x, y| x + y);

assert_abs_diff_eq!(a, &b + &c, epsilon = 1e-6);
}