8000 Improve .assume_init() and use Array::maybe_uninit instead of Array::uninitialized by bluss · Pull Request #876 · rust-ndarray/ndarray · GitHub
[go: up one dir, main page]

Skip to content

Improve .assume_init() and use Array::maybe_uninit instead of Array::uninitialized #876

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 5 commits into from
Dec 30, 2020
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

8000 Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
FIX: Update stack/concatenate to use maybe_uninit
  • Loading branch information
bluss committed Dec 29, 2020
commit 1b45140f94622e89382effae4e8fc782594ef04f
3 changes: 2 additions & 1 deletion src/lib.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
// Copyright 2014-2016 bluss and ndarray developers.
// Copyright 2014-2020 bluss and ndarray developers.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
Expand Down Expand Up @@ -187,6 +187,7 @@ mod shape_builder;
mod slice;
mod split_at;
mod stacking;
mod traversal_utils;
#[macro_use]
mod zip;

Expand Down
47 changes: 24 additions & 23 deletions src/stacking.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
// Copyright 2014-2016 bluss and ndarray developers.
// Copyright 2014-2020 bluss and ndarray developers.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
Expand All @@ -8,6 +8,7 @@

use crate::error::{from_kind, ErrorKind, ShapeError};
use crate::imp_prelude::*;
use crate::traversal_utils::assign_to;

/// Stack arrays along the new axis.
///
Expand Down Expand Up @@ -88,25 +89,23 @@ where
let stacked_dim = arrays.iter().fold(0, |acc, a| acc + a.len_of(axis));
res_dim.set_axis(axis, stacked_dim);

// we can safely use uninitialized values here because they are Copy
// and we will only ever write to them
let size = res_dim.size();
let mut v = Vec::with_capacity(size);
unsafe {
v.set_len(size);
}
let mut res = Array::from_shape_vec(res_dim, v)?;
// we can safely use uninitialized values here because we will
// overwrite every one of them.
let mut res = Array::maybe_uninit(res_dim);

{
let mut assign_view = res.view_mut();
for array in arrays {
let len = array.len_of(axis);
let (mut front, rest) = assign_view.split_at(axis, len);
front.assign(array);
let (front, rest) = assign_view.split_at(axis, len);
assign_to(array, front);
assign_view = rest;
}
debug_assert_eq!(assign_view.len(), 0);
}
unsafe {
Ok(res.assume_init())
}
Ok(res)
}

/// Stack arrays along the new axis.
Expand Down Expand Up @@ -158,22 +157,24 @@ where

res_dim.set_axis(axis, arrays.len());

// we can safely use uninitialized values here because they are Copy
// and we will only ever write to them
let size = res_dim.size();
let mut v = Vec::with_capacity(size);
unsafe {
v.set_len(size);
}
let mut res = Array::from_shape_vec(res_dim, v)?;
// we can safely use uninitialized values here because we will
// overwrite every one of them.
let 8000 mut res = Array::maybe_uninit(res_dim);

res.axis_iter_mut(axis)
.zip(arrays.iter())
.for_each(|(mut assign_view, array)| {
assign_view.assign(&array);
.for_each(|(assign_view, array)| {
// assign_view is D::Larger::Smaller which is usually == D
// (but if D is Ix6, we have IxD != Ix6 here; differing types
// but same number of axes).
let assign_view = assign_view.into_dimensionality::<D>()
.expect("same-dimensionality cast");
assign_to(array, assign_view);
});

Ok(res)
unsafe {
Ok(res.assume_init())
}
}

/// Stack arrays along the new axis.
Expand Down
26 changes: 26 additions & 0 deletions src/traversal_utils.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,26 @@
// Copyright 2020 bluss and ndarray developers.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use crate::{
IntoNdProducer,
AssignElem,
Zip,
};

/// Assign values from producer P1 to producer P2
/// P1 and P2 must be of the same shape and dimension
pub(crate) fn assign_to<'a, P1, P2, A>(from: P1, to: P2)
where P1: IntoNdProducer<Item = &'a A>,
P2: IntoNdProducer<Dim = P1::Dim>,
P2::Item: AssignElem<A>,
A: Clone + 'a
{
Zip::from(from)
.apply_assign_into(to, A::clone);
}

0