From 90010b1a501c411ef042b36742449aa02cd3ec0e Mon Sep 17 00:00:00 2001 From: Avi Weinstock Date: Sun, 5 Feb 2023 16:58:59 -0500 Subject: [PATCH] Implement HStackLazy and VStackLazy, allowing for the removal of intermediate allocations when building a matrix with a mix of horizontal and vertical stacking. --- nalgebra-macros/tests/tests.rs | 20 +- src/base/mod.rs | 2 +- src/base/stacking.rs | 372 ++++++++++++++++++++------------- 3 files changed, 243 insertions(+), 151 deletions(-) diff --git a/nalgebra-macros/tests/tests.rs b/nalgebra-macros/tests/tests.rs index f58c6733..c669d271 100644 --- a/nalgebra-macros/tests/tests.rs +++ b/nalgebra-macros/tests/tests.rs @@ -308,7 +308,7 @@ fn dvector_arbitrary_expressions() { #[test] fn test_stacking() { - use nalgebra::{hstack, vstack, RowVector3, RowVector4}; + use nalgebra::{hstack, vstack, Const, HStackLazy, RowVector3, RowVector4, VStackLazy}; assert_eq_and_type!( vstack((&RowVector3::new(1, 2, 3), &RowVector3::new(4, 5, 6))), Matrix2x3::new(1, 2, 3, 4, 5, 6) @@ -347,4 +347,22 @@ fn test_stacking() { )), dmatrix![1, 0, 0, 2; 0, 1, 0, 2; 0, 0, 1, 2; 3, 3, 3, 3] ); + assert_eq_and_type!( + hstack::<_, _, Const<4>, _>((&Matrix2::::identity(), &nalgebra::SMatrix::zeros())), + matrix![1, 0, 0, 0; 0, 1, 0, 0] + ); + assert_eq_and_type!( + vstack(( + HStackLazy((&Matrix2::::identity(), &Matrix2::zeros())), + HStackLazy((&Matrix2::identity(), &Matrix2::identity())) + )), + matrix![1, 0, 0, 0; 0, 1, 0, 0; 1, 0, 1, 0; 0, 1, 0, 1] + ); + assert_eq_and_type!( + hstack(( + VStackLazy((&Matrix2::::identity(), &Matrix2::identity())), + VStackLazy((&Matrix2::zeros(), &Matrix2::identity())) + )), + matrix![1, 0, 0, 0; 0, 1, 0, 0; 1, 0, 1, 0; 0, 1, 0, 1] + ); } diff --git a/src/base/mod.rs b/src/base/mod.rs index c2b17d17..2c894b0d 100644 --- a/src/base/mod.rs +++ b/src/base/mod.rs @@ -62,7 +62,7 @@ pub use self::alias_slice::*; pub use self::alias_view::*; pub use self::array_storage::*; pub use self::matrix_view::*; -pub use self::stacking::{hstack, vstack}; +pub use self::stacking::*; pub use self::storage::*; #[cfg(any(feature = "std", feature = "alloc"))] pub use self::vec_storage::*; diff --git a/src/base/stacking.rs b/src/base/stacking.rs index 0585249c..d3444755 100644 --- a/src/base/stacking.rs +++ b/src/base/stacking.rs @@ -2,9 +2,11 @@ use crate::{ base::allocator::Allocator, constraint::{DimEq, SameNumberOfColumns, SameNumberOfRows, ShapeConstraint}, - Const, DefaultAllocator, Dim, DimAdd, DimSum, Dyn, Matrix, RawStorage, RawStorageMut, Scalar, + Const, DefaultAllocator, Dim, DimAdd, DimDiff, DimSub, DimSum, Matrix, RawStorage, + RawStorageMut, Scalar, ViewStorageMut, }; use num_traits::Zero; +use std::marker::PhantomData; /// A visitor for each folding over each element of a tuple. pub trait Visitor { @@ -59,42 +61,136 @@ macro_rules! impl_visit_tuple { impl_visit_tuple!(H, G, F, E, D, C, B, A); +/// Source of data that can populate a block of a matrix. +pub trait Block { + /// The scalar type of the data. + type T: Scalar; + /// The number of rows of the block. + type Rows: Dim; + /// The number of columns of the block. + type Cols: Dim; + /// The shape of the block. + fn shape(self) -> (Self::Rows, Self::Cols); +} + +/// Source of data that can populate a block of a matrix. +/// Separate from Block because it's useful to specify the bound on the storage independently of +/// the other bounds. +pub trait BlockPopulate: Block +where + S: RawStorageMut, +{ + /// Populate a matrix from this block's data. + fn populate(self, m: &mut Matrix); +} + +impl<'a, T: Scalar, R: Dim, C: Dim, S: RawStorage> Block for &'a Matrix { + type T = T; + type Rows = R; + type Cols = C; + #[inline(always)] + fn shape(self) -> (Self::Rows, Self::Cols) { + self.shape_generic() + } +} +impl<'a, T: Scalar, R: Dim, C: Dim, S: RawStorage, S2: RawStorageMut> + BlockPopulate for &'a Matrix +{ + #[inline(always)] + fn populate(self, m: &mut Matrix) { + m.copy_from(self); + } +} + +#[inline] +fn build(x: B) -> Matrix +where + S: RawStorageMut, + DefaultAllocator: Allocator, + B: Block + BlockPopulate, +{ + let (r, c) = x.shape(); + let mut out = Matrix::zeros_generic(r, c); + x.populate(&mut out); + out +} + mod vstack_impl { use super::*; #[derive(Clone, Copy)] pub struct VStackShapeInit; #[derive(Clone, Copy)] - pub struct VStackShape { + pub struct VStackShape { + t: PhantomData, r: R, c: C, } - impl> Visitor<&Matrix> - for VStackShapeInit - { - type Output = VStackShape; - fn visit(self, x: &Matrix) -> Self::Output { - let (r, c) = x.shape_generic(); - VStackShape { r, c } + impl Visitor for VStackShapeInit { + type Output = VStackShape; + fn visit(self, x: B) -> Self::Output { + let (r, c) = x.shape(); + VStackShape { + t: PhantomData, + r, + c, + } } } - impl, C1: Dim, R2: Dim, C2: Dim, S2: RawStorage> - Visitor<&Matrix> for VStackShape + impl, C1: Dim> Visitor for VStackShape where - ShapeConstraint: SameNumberOfColumns, + DimSum: DimSub + DimSub, + ShapeConstraint: SameNumberOfColumns + + SameNumberOfRows, R1>, B::Rows> + + SameNumberOfRows, B::Rows>, R1>, { - type Output = - VStackShape, >::Representative>; - fn visit(self, x: &Matrix) -> Self::Output { - let (r, c) = x.shape_generic(); + type Output = VStackShape< + B::T, + DimSum, + >::Representative, + >; + fn visit(self, x: B) -> Self::Output { + let (r, c) = x.shape(); VStackShape { + t: self.t, r: self.r.add(r), - c: >::Representative::from_usize(c.value()), + c: >::Representative::from_usize(c.value()), } } } + /// Specify vertical stacking as a Block. + #[derive(Copy, Clone)] + pub struct VStackLazy(pub X); + + impl Block for VStackLazy + where + X: Copy + VisitTuple>, + { + type T = T; + type Rows = R; + type Cols = C; + fn shape(self) -> (Self::Rows, Self::Cols) { + let shape = >::visit(VStackShapeInit, self.0); + (shape.r, shape.c) + } + } + impl, X> BlockPopulate for VStackLazy + where + X: Copy + + VisitTuple> + + for<'a> VisitTuple>, Output = VStack<'a, T, R, C, S, R>>, + { + fn populate(self, m: &mut Matrix) { + let vstack_visitor = VStack { + out: m, + current_row: Const, + }; + let _ = >::visit(vstack_visitor, self.0); + } + } + pub struct VStack<'a, T, R, C, S, R2> { out: &'a mut Matrix, current_row: R2, @@ -102,47 +198,32 @@ mod vstack_impl { impl< 'a, - T: Scalar, - R1: Dim + DimAdd>, + B: Copy + + Block + + for<'b> BlockPopulate< + ViewStorageMut< + 'b, + ::T, + ::Rows, + ::Cols, + ::T, R1, C1>>::RStride, + ::T, R1, C1>>::CStride, + >, + >, + R1: Dim + DimAdd, C1: Dim, - S1: RawStorageMut, - C2: Dim, - S2: RawStorage, C2>, - R3: Dim + DimAdd>, - const R2: usize, - > Visitor<&Matrix, C2, S2>> for VStack<'a, T, R1, C1, S1, R3> + S1: RawStorageMut, + R3: Dim + DimAdd, + > Visitor for VStack<'a, B::T, R1, C1, S1, R3> where - ShapeConstraint: SameNumberOfColumns, + B::T: Scalar, + ShapeConstraint: SameNumberOfColumns, { - type Output = VStack<'a, T, R1, C1, S1, DimSum>>; - fn visit(self, x: &Matrix, C2, S2>) -> Self::Output { - let (r2, _) = x.shape_generic(); + type Output = VStack<'a, B::T, R1, C1, S1, DimSum>; + fn visit(self, x: B) -> Self::Output { + let (r2, c2) = x.shape(); let VStack { out, current_row } = self; - out.fixed_rows_mut::<{ R2 }>(current_row.value()) - .copy_from::, C2, S2>(x); - let current_row = current_row.add(r2); - VStack { out, current_row } - } - } - impl< - 'a, - T: Scalar, - R1: Dim + DimAdd, - C1: Dim, - S1: RawStorageMut, - C2: Dim, - S2: RawStorage, - R3: Dim + DimAdd, - > Visitor<&Matrix> for VStack<'a, T, R1, C1, S1, R3> - where - ShapeConstraint: SameNumberOfColumns, - { - type Output = VStack<'a, T, R1, C1, S1, DimSum>; - fn visit(self, x: &Matrix) -> Self::Output { - let (r2, _) = x.shape_generic(); - let VStack { out, current_row } = self; - out.rows_mut(current_row.value(), r2.value()) - .copy_from::(x); + x.populate(&mut out.generic_view_mut((current_row.value(), 0), (r2, c2))); let current_row = current_row.add(r2); VStack { out, current_row } } @@ -151,34 +232,18 @@ mod vstack_impl { /// Stack a tuple of references to matrices with equal column counts vertically, yielding a /// matrix with every row of the input matrices. #[inline] - pub fn vstack< - T: Scalar + Zero, - R: Dim, - C: Dim, - X: Copy - + VisitTuple> - + for<'a> VisitTuple< - VStack<'a, T, R, C, >::Buffer, Const<0>>, - Output = VStack<'a, T, R, C, >::Buffer, R>, - >, - >( + pub fn vstack( x: X, ) -> Matrix>::Buffer> where DefaultAllocator: Allocator, + VStackLazy: Block + + BlockPopulate<>::Buffer>, { - let vstack_shape = VStackShapeInit; - let vstack_shape = >::visit(vstack_shape, x); - let mut out = Matrix::zeros_generic(vstack_shape.r, vstack_shape.c); - let vstack_visitor = VStack { - out: &mut out, - current_row: Const, - }; - let _ = >::visit(vstack_visitor, x); - out + build(VStackLazy(x)) } } -pub use vstack_impl::vstack; +pub use vstack_impl::{vstack, VStackLazy}; mod hstack_impl { use super::*; @@ -186,36 +251,76 @@ mod hstack_impl { pub struct HStackShapeInit; #[derive(Clone, Copy)] - pub struct HStackShape { + pub struct HStackShape { + t: PhantomData, r: R, c: C, } - impl> Visitor<&Matrix> - for HStackShapeInit - { - type Output = HStackShape; - fn visit(self, x: &Matrix) -> Self::Output { - let (r, c) = x.shape_generic(); - HStackShape { r, c } + impl Visitor for HStackShapeInit { + type Output = HStackShape; + fn visit(self, x: B) -> Self::Output { + let (r, c) = x.shape(); + HStackShape { + t: PhantomData, + r, + c, + } } } - impl, R2: Dim, C2: Dim, S2: RawStorage> - Visitor<&Matrix> for HStackShape + impl> Visitor for HStackShape where - ShapeConstraint: SameNumberOfRows, + DimSum: DimSub + DimSub, + ShapeConstraint: SameNumberOfRows + + SameNumberOfColumns, C1>, B::Cols> + + SameNumberOfColumns, B::Cols>, C1>, { - type Output = - HStackShape<>::Representative, DimSum>; - fn visit(self, x: &Matrix) -> Self::Output { - let (r, c) = x.shape_generic(); + type Output = HStackShape< + B::T, + >::Representative, + DimSum, + >; + fn visit(self, x: B) -> Self::Output { + let (r, c) = x.shape(); HStackShape { - r: >::Representative::from_usize(r.value()), + t: self.t, + r: >::Representative::from_usize(r.value()), c: self.c.add(c), } } } + /// Specify horizontal stacking as a Block. + #[derive(Copy, Clone)] + pub struct HStackLazy(pub X); + + impl Block for HStackLazy + where + X: Copy + VisitTuple>, + { + type T = T; + type Rows = R; + type Cols = C; + fn shape(self) -> (Self::Rows, Self::Cols) { + let shape = >::visit(HStackShapeInit, self.0); + (shape.r, shape.c) + } + } + impl, X> BlockPopulate for HStackLazy + where + X: Copy + + VisitTuple> + + for<'a> VisitTuple>, Output = HStack<'a, T, R, C, S, C>>, + { + fn populate(self, m: &mut Matrix) { + let hstack_visitor = HStack { + out: m, + current_col: Const, + }; + let _ = >::visit(hstack_visitor, self.0); + } + } + pub struct HStack<'a, T, R, C, S, C2> { out: &'a mut Matrix, current_col: C2, @@ -223,47 +328,32 @@ mod hstack_impl { impl< 'a, - T: Scalar, + B: Copy + + Block + + for<'b> BlockPopulate< + ViewStorageMut< + 'b, + ::T, + ::Rows, + ::Cols, + ::T, R1, C1>>::RStride, + ::T, R1, C1>>::CStride, + >, + >, R1: Dim, - C1: Dim + DimAdd>, - S1: RawStorageMut, - R2: Dim, - S2: RawStorage>, - C3: Dim + DimAdd>, - const C2: usize, - > Visitor<&Matrix, S2>> for HStack<'a, T, R1, C1, S1, C3> + C1: Dim, + S1: RawStorageMut, + C3: Dim + DimAdd, + > Visitor for HStack<'a, B::T, R1, C1, S1, C3> where - ShapeConstraint: SameNumberOfRows, + B::T: Scalar, + ShapeConstraint: SameNumberOfRows, { - type Output = HStack<'a, T, R1, C1, S1, DimSum>>; - fn visit(self, x: &Matrix, S2>) -> Self::Output { - let (_, c2) = x.shape_generic(); + type Output = HStack<'a, B::T, R1, C1, S1, DimSum>; + fn visit(self, x: B) -> Self::Output { + let (r2, c2) = x.shape(); let HStack { out, current_col } = self; - out.fixed_columns_mut::<{ C2 }>(current_col.value()) - .copy_from::, S2>(x); - let current_col = current_col.add(c2); - HStack { out, current_col } - } - } - impl< - 'a, - T: Scalar, - R1: Dim, - C1: Dim + DimAdd, - S1: RawStorageMut, - R2: Dim, - S2: RawStorage, - C3: Dim + DimAdd, - > Visitor<&Matrix> for HStack<'a, T, R1, C1, S1, C3> - where - ShapeConstraint: SameNumberOfRows, - { - type Output = HStack<'a, T, R1, C1, S1, DimSum>; - fn visit(self, x: &Matrix) -> Self::Output { - let (_, c2) = x.shape_generic(); - let HStack { out, current_col } = self; - out.columns_mut(current_col.value(), c2.value()) - .copy_from::(x); + x.populate(&mut out.generic_view_mut((0, current_col.value()), (r2, c2))); let current_col = current_col.add(c2); HStack { out, current_col } } @@ -272,31 +362,15 @@ mod hstack_impl { /// Stack a tuple of references to matrices with equal row counts horizontally, yielding a /// matrix with every column of the input matrices. #[inline] - pub fn hstack< - T: Scalar + Zero, - R: Dim, - C: Dim, - X: Copy - + VisitTuple> - + for<'a> VisitTuple< - HStack<'a, T, R, C, >::Buffer, Const<0>>, - Output = HStack<'a, T, R, C, >::Buffer, C>, - >, - >( + pub fn hstack( x: X, ) -> Matrix>::Buffer> where DefaultAllocator: Allocator, + HStackLazy: Block + + BlockPopulate<>::Buffer>, { - let hstack_shape = HStackShapeInit; - let hstack_shape = >::visit(hstack_shape, x); - let mut out = Matrix::zeros_generic(hstack_shape.r, hstack_shape.c); - let hstack_visitor = HStack { - out: &mut out, - current_col: Const, - }; - let _ = >::visit(hstack_visitor, x); - out + build(HStackLazy(x)) } } -pub use hstack_impl::hstack; +pub use hstack_impl::{hstack, HStackLazy};