Merge pull request #624 from robertknight/extend-init

Add `ExtendInit` utility for updating length of destination buffers for vectorized ops

Author: robertknight

rten-vecmath/src/extend_init.rs

@@ -0,0 +1,85 @@
+use std::mem::MaybeUninit;
+
+/// Extend a buffer by incrementally initializing spare capacity.
+///
+/// This is implemented for [`Vec<T>`], where it provides a safe API to
+/// initialize the spare capacity returned by
+/// [`spare_capacity_mut`](Vec::spare_capacity_mut).
+pub trait ExtendInit {
+    /// Element type in the buffer.
+    type Elem;
+
+    /// Extend the buffer by initializing a portion of the buffer's spare
+    /// capacity.
+    ///
+    /// The function `f` is passed the uninitialized portion of the buffer and
+    /// should return the portion that it has initialized. `extend_init` can
+    /// be called many times, until the entire buffer has been initialized.
+    ///
+    /// # Panics
+    ///
+    /// Panics if `f` returns a slice that is not a prefix of the slice that
+    /// was passed to it.
+    fn extend_init<F: Fn(&mut [MaybeUninit<Self::Elem>]) -> &[Self::Elem]>(&mut self, f: F);
+}
+
+impl<T> ExtendInit for Vec<T> {
+    type Elem = T;
+
+    fn extend_init<F: Fn(&mut [MaybeUninit<Self::Elem>]) -> &[Self::Elem]>(&mut self, f: F) {
+        let cap = self.spare_capacity_mut();
+        let cap_ptr = cap.as_ptr();
+        let cap_len = cap.len();
+
+        let initialized = f(cap);
+        assert_eq!(
+            initialized.as_ptr(),
+            cap_ptr as *const T,
+            "returned slice must be a prefix of the input"
+        );
+        assert!(
+            initialized.len() <= cap_len,
+            "initialized slice length {} is longer than input {}",
+            initialized.len(),
+            cap_len
+        );
+        let n_init = initialized.len();
+
+        // Safety: `n_init` elements from the spare capacity have been initialized.
+        unsafe { self.set_len(self.len() + n_init) }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::mem::MaybeUninit;
+
+    use super::ExtendInit;
+
+    // Implementation of `MaybeUninit::fill` from nightly Rust.
+    fn fill<T: Copy>(xs: &mut [MaybeUninit<T>], value: T) -> &mut [T] {
+        for x in xs.iter_mut() {
+            x.write(value);
+        }
+        unsafe { std::mem::transmute::<&mut [MaybeUninit<T>], &mut [T]>(xs) }
+    }
+
+    #[test]
+    fn test_extend_init() {
+        let mut vec = Vec::with_capacity(7);
+
+        vec.extend_init(|uninit| {
+            assert_eq!(uninit.len(), 7);
+            fill(&mut uninit[..3], 1.)
+        });
+        assert_eq!(vec.len(), 3);
+        assert_eq!(vec, &[1., 1., 1.]);
+
+        vec.extend_init(|uninit| {
+            assert_eq!(uninit.len(), 4);
+            fill(uninit, 2.)
+        });
+        assert_eq!(vec.len(), 7);
+        assert_eq!(vec, &[1., 1., 1., 2., 2., 2., 2.]);
+    }
+}

rten-vecmath/src/lib.rs

@@ -8,19 +8,26 @@
 //! The operations are implemented by structs which implement the SIMD operation
 //! traits from [rten-simd](rten_simd). To apply an operation to data, first
 //! construct the operation using the struct from this crate, then use a
-//! dispatch method from the [`SimdOp`](rten_simd::dispatch::SimdOp) or
-//! [`SimdUnaryOp`](rten_simd::dispatch::SimdUnaryOp) traits to execute the
-//! operation using the preferred SIMD instruction set.
+//! dispatch method from the [`SimdOp`](rten_simd::safe::SimdOp) or
+//! [`SimdUnaryOp`](rten_simd::safe::SimdUnaryOp) traits to execute
+//! the operation.
 //!
-//! ## In-place versus mutating operations
+//! ## In-place and non in-place operations
 //!
 //! Some operations support both updating data in place or reading input from
 //! one slice and writing to another. For unary operations this is controlled by
-//! dispatching with either [`map`](rten_simd::dispatch::SimdUnaryOp::map) or
-//! [`map_mut`](rten_simd::dispatch::SimdUnaryOp::map_mut). For other operations
+//! dispatching with either [`map`](rten_simd::safe::SimdUnaryOp::map) or
+//! [`map_mut`](rten_simd::safe::SimdUnaryOp::map_mut). For other operations
 //! this is handled by exposing different constructors for the in-place and
 //! mutating cases, such as [`Softmax::new`] and [`Softmax::new_mut`].
 //!
+//! For operations which use a separate source and destination, the destination
+//! is expected to be an uninitialized slice (`[MaybeUninit<T>]`). This allows
+//! the caller to control allocation of the buffer and avoid the overhead of
+//! initializing elements which the operation will overwrite. The [`ExtendInit`]
+//! trait provides a safe API for the common task of filling a new `Vec` with
+//! the result of the operation.
+//!
 //! ## Examples
 //!
 //! ### Applying a vectorized unary function
@@ -44,6 +51,8 @@
 //!
 //! ### Applying softmax in place
 //!
+//! This example applies the softmax function in-place to a mutable slice.
+//!
 //! ```
 //! use rten_simd::safe::SimdOp;
 //! use rten_vecmath::Softmax;
@@ -55,20 +64,24 @@
 //! ### Applying softmax with separate input and output buffers
 //!
 //! This example reads data from an input and writes to an uninitialized output
-//! buffer. The softmax operation returns the initialized slice.
+//! buffer (`&mut [MaybeUninit<f32>]`), obtained from the uninitialized portion
+//! of a `Vec<f32>`. To update the length of the `Vec<f32>` after it is
+//! initialized, the helper `ExtendInit` trait is used.
 //!
 //! ```
 //! use rten_simd::safe::SimdOp;
-//! use rten_vecmath::Softmax;
+//! use rten_vecmath::{Softmax, ExtendInit};
 //!
 //! let data = [1., 0.5, 2.0];
 //! let mut output = Vec::with_capacity(data.len());
-//! let output_uninit = &mut output.spare_capacity_mut()[..data.len()];
-//! let output_init = Softmax::new(&data, output_uninit).dispatch();
-//!
-//! // Safety: The softmax operation initialized all output elements.
-//! let init_len = output_init.len();
-//! unsafe { output.set_len(init_len) };
+//! output.extend_init(|output_uninit| {
+//!     // `output_uninit` is the uninitialized part of `output`, as returned by
+//!     // `output.spare_capacity_mut()`.
+//!     //
+//!     // The `dispatch` call initializes it and returns the initialized slice.
+//!     Softmax::new(&data, output_uninit).dispatch()
+//! });
+//! assert_eq!(output.len(), 3);
 //! ```
 //!
 //! ### Computing the sum of a list of floats
@@ -96,6 +109,8 @@ mod ulp;
 #[cfg(test)]
 mod testing;
 
+mod extend_init;
+
 // Unary functions.
 pub use erf::{Erf, Gelu};
 pub use exp::{Exp, Sigmoid, Silu, Swish};
@@ -107,3 +122,6 @@ pub use min_max::MinMax;
 pub use normalize::{Normalize, NormalizeOptions};
 pub use softmax::Softmax;
 pub use sum::{Sum, SumSquare, SumSquareSub};
+
+// Utilities
+pub use extend_init::ExtendInit;

src/ops/norm.rs

@@ -5,6 +5,7 @@ use rten_simd::safe::SimdOp;
 use rten_tensor::prelude::*;
 use rten_tensor::{NdTensorView, Tensor, TensorView};
 use rten_vecmath as vecmath;
+use rten_vecmath::ExtendInit;
 
 use crate::ops::static_dims;
 use crate::ops::{resolve_axis, InputList, IntoOpResult, OpError, Operator, Output, OutputList};
@@ -70,23 +71,23 @@ impl Default for NormalizeOptions<'_> {
     }
 }
 
-enum NormalizeData<'a> {
+enum NormalizeData<'src, 'dst> {
     /// Read from a source slice and write normalized data to an output slice
     /// of the same length.
-    SrcDest((&'a [f32], &'a mut [MaybeUninit<f32>])),
+    SrcDest((&'src [f32], &'dst mut [MaybeUninit<f32>])),
 
     /// Normalize elements of a slice in place.
-    InPlace(&'a mut [f32]),
+    InPlace(&'dst mut [f32]),
 }
 
-impl<'a> From<&'a mut [f32]> for NormalizeData<'a> {
-    fn from(val: &'a mut [f32]) -> Self {
+impl<'dst> From<&'dst mut [f32]> for NormalizeData<'dst, 'dst> {
+    fn from(val: &'dst mut [f32]) -> Self {
         NormalizeData::InPlace(val)
     }
 }
 
-impl<'a> From<(&'a [f32], &'a mut [MaybeUninit<f32>])> for NormalizeData<'a> {
-    fn from(val: (&'a [f32], &'a mut [MaybeUninit<f32>])) -> Self {
+impl<'src, 'dst> From<(&'src [f32], &'dst mut [MaybeUninit<f32>])> for NormalizeData<'src, 'dst> {
+    fn from(val: (&'src [f32], &'dst mut [MaybeUninit<f32>])) -> Self {
         NormalizeData::SrcDest(val)
     }
 }
@@ -95,7 +96,10 @@ impl<'a> From<(&'a [f32], &'a mut [MaybeUninit<f32>])> for NormalizeData<'a> {
 /// and bias to the result.
 ///
 /// Returns the normalized elements.
-fn normalize_slice<'a>(data: NormalizeData<'a>, opts: NormalizeOptions<'a>) -> &'a mut [f32] {
+fn normalize_slice<'src, 'dst>(
+    data: NormalizeData<'src, 'dst>,
+    opts: NormalizeOptions<'src>,
+) -> &'dst mut [f32] {
     let NormalizeOptions {
         mean_normalize,
         epsilon,
@@ -469,29 +473,19 @@ fn layer_normalization_impl(
     let scale = scale.to_contiguous_in(pool);
     let scale_data = scale.data().unwrap();
 
-    let mut n_init = 0;
-    for (in_chunk, out_chunk) in input
-        .data()
-        .unwrap()
-        .chunks(chunk_size)
-        .zip(output.spare_capacity_mut().chunks_mut(chunk_size))
-    {
-        normalize_slice(
-            (in_chunk, out_chunk).into(),
-            NormalizeOptions {
-                mean_normalize,
-                epsilon,
-                element_scale: Some(scale_data),
-                element_bias: bias_data,
-                ..Default::default()
-            },
-        );
-        n_init += in_chunk.len();
-    }
-
-    // Safety: We initialized `n_init` elements.
-    unsafe {
-        output.set_len(n_init);
+    for in_chunk in input.data().unwrap().chunks(chunk_size) {
+        output.extend_init(|out_chunk| {
+            normalize_slice(
+                (in_chunk, &mut out_chunk[..chunk_size]).into(),
+                NormalizeOptions {
+                    mean_normalize,
+                    epsilon,
+                    element_scale: Some(scale_data),
+                    element_bias: bias_data,
+                    ..Default::default()
+                },
+            )
+        })
     }
 
     Ok(Tensor::from_data(input.shape(), output))

src/ops/quantize.rs

@@ -4,6 +4,7 @@ use rten_simd::safe::SimdOp;
 use rten_tensor::prelude::*;
 use rten_tensor::{AssumeInit, NdTensor, NdTensorView, Scalar, Tensor, TensorView};
 use rten_vecmath as vecmath;
+use rten_vecmath::ExtendInit;
 
 use crate::ops::{
     resolve_axis, DataType, Input, InputList, IntoOpResult, OpError, Operator, Output, OutputList,
@@ -200,14 +201,9 @@ where
 
             if let Some(data) = input.data() {
                 let mut buf = pool.alloc(data.len());
-                let buf_data = &mut buf.spare_capacity_mut()[..data.len()];
-
-                Quantize::quantize_slice(data, buf_data, inv_scale, zero_point);
-
-                // Safety: `quantize_slice` initialized `data.len()` elements
-                unsafe {
-                    buf.set_len(data.len());
-                }
+                buf.extend_init(|buf_data| {
+                    Quantize::quantize_slice(data, buf_data, inv_scale, zero_point)
+                });
                 Ok(Tensor::from_data(input.shape(), buf))
             } else {
                 Ok(input.map_in(pool, |x| x.quantize(inv_scale, zero_point)))