Merge pull request #876 from rust-ndarray/uninit

bluss · web-flow · commit d5b61ba9a6a8 · 2020-12-30T11:00:47.000+01:00
Improve .assume_init() and use Array::maybe_uninit instead of Array::uninitialized
diff --git a/benches/bench1.rs b/benches/bench1.rs
@@ -9,6 +9,8 @@
 
 extern crate test;
 
+use std::mem::MaybeUninit;
+
 use ndarray::ShapeBuilder;
 use ndarray::{arr0, arr1, arr2, azip, s};
 use ndarray::{Array, Array1, Array2, Axis, Ix, Zip};
@@ -269,9 +271,9 @@ 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 c.raw_view_mut())
-            std::ptr::write(c, a + b)
+        let mut c = Array::<MaybeUninit<i32>, _>::maybe_uninit(a.dim());
+        azip!((&a in &a, &b in &b, c in c.raw_view_mut().cast::<i32>())
+            c.write(a + b)
         );
         c
     });
diff --git a/src/data_repr.rs b/src/data_repr.rs
@@ -10,7 +10,11 @@ use crate::extension::nonnull;
 /// *Don’t use this type directly—use the type alias
 /// [`Array`](type.Array.html) for the array type!*
 // Like a Vec, but with non-unique ownership semantics
+//
+// repr(C) to make it transmutable OwnedRepr<A> -> OwnedRepr<B> if
+// transmutable A -> B.
 #[derive(Debug)]
+#[repr(C)]
 pub struct OwnedRepr<A> {
     ptr: NonNull<A>,
     len: usize,
@@ -50,6 +54,23 @@ impl<A> OwnedRepr<A> {
         self.ptr
     }
 
+    /// Cast self into equivalent repr of other element type
+    ///
+    /// ## Safety
+    ///
+    /// Caller must ensure the two types have the same representation.
+    /// **Panics** if sizes don't match (which is not a sufficient check).
+    pub(crate) unsafe fn data_subst<B>(self) -> OwnedRepr<B> {
+        // necessary but not sufficient check
+        assert_eq!(mem::size_of::<A>(), mem::size_of::<B>());
+        let self_ = ManuallyDrop::new(self);
+        OwnedRepr {
+            ptr: self_.ptr.cast::<B>(),
+            len: self_.len,
+            capacity: self_.capacity,
+        }
+    }
+
     fn take_as_vec(&mut self) -> Vec<A> {
         let capacity = self.capacity;
         let len = self.len;
diff --git a/src/data_traits.rs b/src/data_traits.rs
@@ -526,28 +526,60 @@ unsafe impl<'a, A> DataMut for CowRepr<'a, A> where A: Clone {}
 pub trait RawDataSubst<A>: RawData {
     /// The resulting array storage of the same kind but substituted element type
     type Output: RawData<Elem = A>;
+
+    /// Unsafely translate the data representation from one element
+    /// representation to another.
+    ///
+    /// ## Safety
+    ///
+    /// Caller must ensure the two types have the same representation.
+    unsafe fn data_subst(self) -> Self::Output;
 }
 
 impl<A, B> RawDataSubst<B> for OwnedRepr<A> {
     type Output = OwnedRepr<B>;
+
+    unsafe fn data_subst(self) -> Self::Output {
+        self.data_subst()
+    }
 }
 
 impl<A, B> RawDataSubst<B> for OwnedArcRepr<A> {
     type Output = OwnedArcRepr<B>;
+
+    unsafe fn data_subst(self) -> Self::Output {
+        OwnedArcRepr(Arc::from_raw(Arc::into_raw(self.0) as *const OwnedRepr<B>))
+    }
 }
 
 impl<A, B> RawDataSubst<B> for RawViewRepr<*const A> {
     type Output = RawViewRepr<*const B>;
+
+    unsafe fn data_subst(self) -> Self::Output {
+        RawViewRepr::new()
+    }
 }
 
 impl<A, B> RawDataSubst<B> for RawViewRepr<*mut A> {
     type Output = RawViewRepr<*mut B>;
+
+    unsafe fn data_subst(self) -> Self::Output {
+        RawViewRepr::new()
+    }
 }
 
 impl<'a, A: 'a, B: 'a> RawDataSubst<B> for ViewRepr<&'a A> {
     type Output = ViewRepr<&'a B>;
+
+    unsafe fn data_subst(self) -> Self::Output {
+        ViewRepr::new()
+    }
 }
 
 impl<'a, A: 'a, B: 'a> RawDataSubst<B> for ViewRepr<&'a mut A> {
     type Output = ViewRepr<&'a mut B>;
+
+    unsafe fn data_subst(self) -> Self::Output {
+        ViewRepr::new()
+    }
 }
diff --git a/src/doc/ndarray_for_numpy_users/mod.rs b/src/doc/ndarray_for_numpy_users/mod.rs
@@ -647,7 +647,6 @@
 //! [.index_axis()]: ../../struct.ArrayBase.html#method.index_axis
 //! [.sum_axis()]: ../../struct.ArrayBase.html#method.sum_axis
 //! [.t()]: ../../struct.ArrayBase.html#method.t
-//! [::uninitialized()]: ../../struct.ArrayBase.html#method.uninitialized
 //! [vec-* dot]: ../../struct.ArrayBase.html#method.dot
 //! [.visit()]: ../../struct.ArrayBase.html#method.visit
 //! [::zeros()]: ../../struct.ArrayBase.html#method.zeros
diff --git a/src/impl_special_element_types.rs b/src/impl_special_element_types.rs
@@ -6,8 +6,6 @@
 // option. This file may not be copied, modified, or distributed
 // except according to those terms.
 
-use std::mem::size_of;
-use std::mem::ManuallyDrop;
 use std::mem::MaybeUninit;
 
 use crate::imp_prelude::*;
@@ -37,12 +35,10 @@ where
     /// array's storage; it is for example possible to slice these in place, but that must
     /// only be done after all elements have been initialized.
     pub unsafe fn assume_init(self) -> ArrayBase<<S as RawDataSubst<A>>::Output, D> {
-        // NOTE: Fully initialized includes elements not reachable in current slicing/view.
-
         let ArrayBase { data, ptr, dim, strides } = self;
 
-        // transmute from storage of MaybeUninit<A> to storage of A
-        let data = unlimited_transmute::<S, S::Output>(data);
+        // "transmute" from storage of MaybeUninit<A> to storage of A
+        let data = S::data_subst(data);
         let ptr = ptr.cast::<A>();
 
         ArrayBase {
@@ -53,15 +49,3 @@ where
         }
     }
 }
-
-/// Transmute from A to B.
-///
-/// Like transmute, but does not have the compile-time size check which blocks
-/// using regular transmute for "S to S::Output".
-///
-/// **Panics** if the size of A and B are different.
-unsafe fn unlimited_transmute<A, B>(data: A) -> B {
-    assert_eq!(size_of::<A>(), size_of::<B>());
-    let old_data = ManuallyDrop::new(data);
-    (&*old_data as *const A as *const B).read()
-}
diff --git a/src/lib.rs b/src/lib.rs
@@ -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
@@ -187,6 +187,7 @@ mod shape_builder;
 mod slice;
 mod split_at;
 mod stacking;
+mod traversal_utils;
 #[macro_use]
 mod zip;
 
diff --git a/src/linalg/impl_linalg.rs b/src/linalg/impl_linalg.rs
@@ -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
@@ -325,9 +325,9 @@ where
 
         // Avoid initializing the memory in vec -- set it during iteration
         unsafe {
-            let mut c = Array::uninitialized(m);
-            general_mat_vec_mul(A::one(), self, rhs, A::zero(), &mut c);
-            c
+            let mut c = Array1::maybe_uninit(m);
+            general_mat_vec_mul_impl(A::one(), self, rhs, A::zero(), c.raw_view_mut().cast::<A>());
+            c.assume_init()
         }
     }
 }
@@ -598,6 +598,30 @@ pub fn general_mat_vec_mul<A, S1, S2, S3>(
     S2: Data<Elem = A>,
     S3: DataMut<Elem = A>,
     A: LinalgScalar,
+{
+    unsafe {
+        general_mat_vec_mul_impl(alpha, a, x, beta, y.raw_view_mut())
+    }
+}
+
+/// General matrix-vector multiplication
+///
+/// Use a raw view for the destination vector, so that it can be uninitalized.
+///
+/// ## Safety
+///
+/// The caller must ensure that the raw view is valid for writing.
+/// the destination may be uninitialized iff beta is zero.
+unsafe fn general_mat_vec_mul_impl<A, S1, S2>(
+    alpha: A,
+    a: &ArrayBase<S1, Ix2>,
+    x: &ArrayBase<S2, Ix1>,
+    beta: A,
+    y: RawArrayViewMut<A, Ix1>,
+) where
+    S1: Data<Elem = A>,
+    S2: Data<Elem = A>,
+    A: LinalgScalar,
 {
     let ((m, k), k2) = (a.dim(), x.dim());
     let m2 = y.dim();
@@ -626,22 +650,20 @@ pub fn general_mat_vec_mul<A, S1, S2, S3>(
                         let x_stride = x.strides()[0] as blas_index;
                         let y_stride = y.strides()[0] as blas_index;
 
-                        unsafe {
-                            blas_sys::$gemv(
-                                layout,
-                                a_trans,
-                                m as blas_index,            // m, rows of Op(a)
-                                k as blas_index,            // n, cols of Op(a)
-                                cast_as(&alpha),            // alpha
-                                a.ptr.as_ptr() as *const _, // a
-                                a_stride,                   // lda
-                                x.ptr.as_ptr() as *const _, // x
-                                x_stride,
-                                cast_as(&beta),           // beta
-                                y.ptr.as_ptr() as *mut _, // x
-                                y_stride,
-                            );
-                        }
+                        blas_sys::$gemv(
+                            layout,
+                            a_trans,
+                            m as blas_index,            // m, rows of Op(a)
+                            k as blas_index,            // n, cols of Op(a)
+                            cast_as(&alpha),            // alpha
+                            a.ptr.as_ptr() as *const _, // a
+                            a_stride,                   // lda
+                            x.ptr.as_ptr() as *const _, // x
+                            x_stride,
+                            cast_as(&beta),           // beta
+                            y.ptr.as_ptr() as *mut _, // x
+                            y_stride,
+                        );
                         return;
                     }
                 }
@@ -655,8 +677,9 @@ pub fn general_mat_vec_mul<A, S1, S2, S3>(
         /* general */
 
         if beta.is_zero() {
+            // when beta is zero, c may be uninitialized
             Zip::from(a.outer_iter()).and(y).apply(|row, elt| {
-                *elt = row.dot(x) * alpha;
+                elt.write(row.dot(x) * alpha);
             });
         } else {
             Zip::from(a.outer_iter()).and(y).apply(|row, elt| {
@@ -683,7 +706,7 @@ fn cast_as<A: 'static + Copy, B: 'static + Copy>(a: &A) -> B {
 #[cfg(feature = "blas")]
 fn blas_compat_1d<A, S>(a: &ArrayBase<S, Ix1>) -> bool
 where
-    S: Data,
+    S: RawData,
     A: 'static,
     S::Elem: 'static,
 {
diff --git a/src/stacking.rs b/src/stacking.rs
@@ -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
@@ -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.
 ///
@@ -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.
@@ -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 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.
diff --git a/src/traversal_utils.rs b/src/traversal_utils.rs
@@ -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);
+}
+
