From 4859ddf04c34bce5522009718a154d28ff2ede3a Mon Sep 17 00:00:00 2001 From: Ben Liepert Date: Wed, 4 Sep 2024 18:30:47 -0400 Subject: [PATCH 1/5] Replace into_raw_vec() with into_raw_vec_and_offset().0 --- Cargo.lock | 23 ++++++++++++++----- Cargo.toml | 4 ++-- crates/store/re_types/src/archetypes/image.rs | 2 +- .../re_types/src/datatypes/tensor_data_ext.rs | 6 +++-- .../all/archetypes/image_send_columns.rs | 2 +- 5 files changed, 25 insertions(+), 12 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 3a9e7208d00d..b90b85b2e833 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3509,22 +3509,24 @@ checksum = "308d96db8debc727c3fd9744aac51751243420e46edf401010908da7f8d5e57c" [[package]] name = "ndarray" -version = "0.15.6" +version = "0.16.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "adb12d4e967ec485a5f71c6311fe28158e9d6f4bc4a447b474184d0f91a8fa32" +checksum = "882ed72dce9365842bf196bdeedf5055305f11fc8c03dee7bb0194a6cad34841" dependencies = [ "matrixmultiply", "num-complex", "num-integer", "num-traits", + "portable-atomic", + "portable-atomic-util", "rawpointer", ] [[package]] name = "ndarray-rand" -version = "0.14.0" +version = "0.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "65608f937acc725f5b164dcf40f4f0bc5d67dc268ab8a649d3002606718c4588" +checksum = "f093b3db6fd194718dcdeea6bd8c829417deae904e3fcc7732dabcd4416d25d8" dependencies = [ "ndarray", "rand", @@ -4328,9 +4330,18 @@ checksum = "22686f4785f02a4fcc856d3b3bb19bf6c8160d103f7a99cc258bddd0251dc7f2" [[package]] name = "portable-atomic" -version = "1.5.1" +version = "1.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da544ee218f0d287a911e9c99a39a8c9bc8fcad3cb8db5959940044ecfc67265" + +[[package]] +name = "portable-atomic-util" +version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3bccab0e7fd7cc19f820a1c8c91720af652d0c88dc9664dd72aef2614f04af3b" +checksum = "fcdd8420072e66d54a407b3316991fe946ce3ab1083a7f575b2463866624704d" +dependencies = [ + "portable-atomic", +] [[package]] name = "powerfmt" diff --git a/Cargo.toml b/Cargo.toml index 7a12442f8d52..857f8705864f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -197,8 +197,8 @@ mime_guess2 = "2.0" # infer MIME type by file extension, and map mime to file ex mint = "0.5.9" mp4 = "0.14.0" natord = "1.0.9" -ndarray = "0.15" -ndarray-rand = "0.14" +ndarray = "0.16" +ndarray-rand = "0.15" never = "0.1" nohash-hasher = "0.2" notify = "6.0" diff --git a/crates/store/re_types/src/archetypes/image.rs b/crates/store/re_types/src/archetypes/image.rs index 119bc41e3929..1e87272285c7 100644 --- a/crates/store/re_types/src/archetypes/image.rs +++ b/crates/store/re_types/src/archetypes/image.rs @@ -99,7 +99,7 @@ use ::re_types_core::{DeserializationError, DeserializationResult}; /// /// // Split up the image data into several components referencing the underlying data. /// let image_size_in_bytes = width * height * 3; -/// let blob = rerun::datatypes::Blob::from(images.into_raw_vec()); +/// let blob = rerun::datatypes::Blob::from(images.into_raw_vec_and_offset().0); /// let image_column = times /// .iter() /// .map(|&t| { diff --git a/crates/store/re_types/src/datatypes/tensor_data_ext.rs b/crates/store/re_types/src/datatypes/tensor_data_ext.rs index 661aa33bfa4c..776592084a08 100644 --- a/crates/store/re_types/src/datatypes/tensor_data_ext.rs +++ b/crates/store/re_types/src/datatypes/tensor_data_ext.rs @@ -182,7 +182,9 @@ macro_rules! tensor_from_ndarray { .is_standard_layout() .then(|| TensorData { shape, - buffer: TensorBuffer::$variant(value.to_owned().into_raw_vec().into()), + buffer: TensorBuffer::$variant( + value.to_owned().into_raw_vec_and_offset().0.into(), + ), }) .ok_or(TensorCastError::NotContiguousStdOrder) } @@ -314,7 +316,7 @@ impl TryFrom<::ndarray::Array> for Tensor Ok(Self { shape, buffer: TensorBuffer::F16( - bytemuck::cast_slice(value.into_raw_vec().as_slice()) + bytemuck::cast_slice(value.into_raw_vec_and_offset().0.as_slice()) .to_vec() .into(), ), diff --git a/docs/snippets/all/archetypes/image_send_columns.rs b/docs/snippets/all/archetypes/image_send_columns.rs index 3baf8ede10c2..f03df013e675 100644 --- a/docs/snippets/all/archetypes/image_send_columns.rs +++ b/docs/snippets/all/archetypes/image_send_columns.rs @@ -31,7 +31,7 @@ fn main() -> Result<(), Box> { // Split up the image data into several components referencing the underlying data. let image_size_in_bytes = width * height * 3; - let blob = rerun::datatypes::Blob::from(images.into_raw_vec()); + let blob = rerun::datatypes::Blob::from(images.into_raw_vec_and_offset().0); let image_column = times .iter() .map(|&t| { From 90b32eb1f0afe3c872608d50a42cf93c8c73ba4d Mon Sep 17 00:00:00 2001 From: Ben Liepert Date: Wed, 4 Sep 2024 20:10:00 -0400 Subject: [PATCH 2/5] Replace into_shape() with into_shape_with_order() --- crates/viewer/re_space_view_tensor/src/space_view_class.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/viewer/re_space_view_tensor/src/space_view_class.rs b/crates/viewer/re_space_view_tensor/src/space_view_class.rs index 729b3b8ff46d..d20aa55090e2 100644 --- a/crates/viewer/re_space_view_tensor/src/space_view_class.rs +++ b/crates/viewer/re_space_view_tensor/src/space_view_class.rs @@ -424,7 +424,7 @@ pub fn selected_tensor_slice<'a, T: Copy>( // This is important for above width/height conversion to work since this assumes at least 2 dimensions. tensor .view() - .into_shape(ndarray::IxDyn(&[tensor.len(), 1])) + .into_shape_with_order(ndarray::IxDyn(&[tensor.len(), 1])) .unwrap() } else { tensor.view() From 7b125e8fdba17932a51ec1facea0367202225f49 Mon Sep 17 00:00:00 2001 From: Ben Liepert Date: Fri, 6 Sep 2024 18:29:06 -0400 Subject: [PATCH 3/5] First pass at updating TensorData macros --- .../re_types/src/datatypes/tensor_data_ext.rs | 50 ++++++++++++------- 1 file changed, 32 insertions(+), 18 deletions(-) diff --git a/crates/store/re_types/src/datatypes/tensor_data_ext.rs b/crates/store/re_types/src/datatypes/tensor_data_ext.rs index 776592084a08..38c26a157c26 100644 --- a/crates/store/re_types/src/datatypes/tensor_data_ext.rs +++ b/crates/store/re_types/src/datatypes/tensor_data_ext.rs @@ -8,9 +8,6 @@ use crate::archetypes::EncodedImage; use super::{TensorBuffer, TensorData, TensorDimension}; -// Much of the following duplicates code from: `crates/re_components/src/tensor.rs`, which -// will eventually go away as the Tensor migration is completed. - // ---------------------------------------------------------------------------- impl TensorData { @@ -169,7 +166,6 @@ macro_rules! tensor_from_ndarray { type Error = TensorCastError; fn try_from(value: ndarray::Array<$type, D>) -> Result { - let value = value.as_standard_layout(); let shape = value .shape() .iter() @@ -178,15 +174,27 @@ macro_rules! tensor_from_ndarray { name: None, }) .collect(); - value - .is_standard_layout() - .then(|| TensorData { - shape, - buffer: TensorBuffer::$variant( - value.to_owned().into_raw_vec_and_offset().0.into(), - ), - }) - .ok_or(TensorCastError::NotContiguousStdOrder) + + let vec = if value.is_standard_layout() { + let (mut vec, offset) = value.into_raw_vec_and_offset(); + // into_raw_vec_and_offset() guarantees that the logical element order (.iter()) matches the internal + // storage order in the returned vector. Therefore, since it's in standard layout, it's contiguous in + // memory and safe to assume all our data is stored starting at the offset returned + if let Some(offset) = offset { + vec.drain(..offset); + vec + } else { + debug_assert!(vec.is_empty()); + vec + } + } else { + value.into_iter().collect::>() + }; + + Ok(Self { + shape, + buffer: TensorBuffer::$variant(vec.into()), + }) } } @@ -313,13 +321,19 @@ impl TryFrom<::ndarray::Array> for Tensor }) .collect(); if value.is_standard_layout() { + let (vec, offset) = value.into_raw_vec_and_offset(); + // into_raw_vec_and_offset() guarantees that the logical element order (.iter()) matches the internal + // storage order in the returned vector. Therefore, since it's in standard layout, it's contiguous in + // memory and it's safe to assume all our data is stored starting at the offset returned + let vec_slice = if let Some(offset) = offset { + &vec[offset..] + } else { + debug_assert!(vec.is_empty()); + &vec + }; Ok(Self { shape, - buffer: TensorBuffer::F16( - bytemuck::cast_slice(value.into_raw_vec_and_offset().0.as_slice()) - .to_vec() - .into(), - ), + buffer: TensorBuffer::F16(Vec::from(bytemuck::cast_slice(vec_slice)).into()), }) } else { Ok(Self { From 773c6d2896d3cc5f4948a0059baba15f2ee52c3e Mon Sep 17 00:00:00 2001 From: Ben Liepert Date: Thu, 12 Sep 2024 18:00:29 -0400 Subject: [PATCH 4/5] Add tests for std/nonstd layouts and zero/nonzero offsets --- .../re_types/src/datatypes/tensor_data_ext.rs | 6 +- crates/store/re_types/tests/types/tensor.rs | 85 ++++++++++++++++++- 2 files changed, 86 insertions(+), 5 deletions(-) diff --git a/crates/store/re_types/src/datatypes/tensor_data_ext.rs b/crates/store/re_types/src/datatypes/tensor_data_ext.rs index 38c26a157c26..d9cf2ff39cd0 100644 --- a/crates/store/re_types/src/datatypes/tensor_data_ext.rs +++ b/crates/store/re_types/src/datatypes/tensor_data_ext.rs @@ -178,8 +178,7 @@ macro_rules! tensor_from_ndarray { let vec = if value.is_standard_layout() { let (mut vec, offset) = value.into_raw_vec_and_offset(); // into_raw_vec_and_offset() guarantees that the logical element order (.iter()) matches the internal - // storage order in the returned vector. Therefore, since it's in standard layout, it's contiguous in - // memory and safe to assume all our data is stored starting at the offset returned + // storage order in the returned vector if the array is in standard layout. if let Some(offset) = offset { vec.drain(..offset); vec @@ -323,8 +322,7 @@ impl TryFrom<::ndarray::Array> for Tensor if value.is_standard_layout() { let (vec, offset) = value.into_raw_vec_and_offset(); // into_raw_vec_and_offset() guarantees that the logical element order (.iter()) matches the internal - // storage order in the returned vector. Therefore, since it's in standard layout, it's contiguous in - // memory and it's safe to assume all our data is stored starting at the offset returned + // storage order in the returned vector if the array is in standard layout. let vec_slice = if let Some(offset) = offset { &vec[offset..] } else { diff --git a/crates/store/re_types/tests/types/tensor.rs b/crates/store/re_types/tests/types/tensor.rs index d995384bcc98..0b47514a9f17 100644 --- a/crates/store/re_types/tests/types/tensor.rs +++ b/crates/store/re_types/tests/types/tensor.rs @@ -106,7 +106,7 @@ fn convert_tensor_to_ndarray_f32() { } #[test] -fn convert_ndarray_u8_to_tensor() { +fn convert_ndarray_f64_to_tensor() { let n = ndarray::array![[1., 2., 3.], [4., 5., 6.]]; let t = TensorData::try_from(n).unwrap(); @@ -125,6 +125,89 @@ fn convert_ndarray_slice_to_tensor() { assert_eq!(t.shape(), &[TensorDimension::unnamed(2)]); } +#[test] +fn convert_ndarray_to_tensor_both_layouts() { + #[rustfmt::skip] + let row_major_vec = vec![ + 1, 2, 3, + 4, 5, 6, + 7, 8, 9 + ]; + #[rustfmt::skip] + let col_major_vec = vec![ + 1, 4, 7, + 2, 5, 8, + 3, 6, 9 + ]; + + let shape = ndarray::Ix2(3, 3); + + let row_major = ndarray::Array::from_vec(row_major_vec) + .into_shape_with_order((shape, ndarray::Order::RowMajor)) + .unwrap(); + + let col_major = ndarray::Array::from_vec(col_major_vec) + .into_shape_with_order((shape, ndarray::Order::ColumnMajor)) + .unwrap(); + + // make sure that the offset is in fact negative, in case ndarray behavior changes + let rm = row_major.clone(); + let cm = col_major.clone(); + let (_, rm_offset) = rm.into_raw_vec_and_offset(); + let (_, cm_offset) = cm.into_raw_vec_and_offset(); + assert_eq!(rm_offset.unwrap(), 0); + assert_eq!(cm_offset.unwrap(), 0); + + let tensor_row_major = TensorData::try_from(row_major).unwrap(); + let tensor_col_major = TensorData::try_from(col_major).unwrap(); + + assert_eq!(tensor_row_major, tensor_col_major); +} + +#[test] +fn convert_ndarray_to_tensor_both_layouts_nonzero_offset() { + #[rustfmt::skip] + let row_major_vec = vec![ + 1, 2, 3, + 4, 5, 6, + 7, 8, 9 + ]; + #[rustfmt::skip] + let col_major_vec = vec![ + 1, 4, 7, + 2, 5, 8, + 3, 6, 9 + ]; + + let shape = ndarray::Ix2(3, 3); + + let row_major = ndarray::Array::from_vec(row_major_vec) + .into_shape_with_order((shape, ndarray::Order::RowMajor)) + .unwrap(); + let row_major_nonzero_offset = row_major.slice_move(ndarray::s![1.., ..]); + + let col_major = ndarray::Array::from_vec(col_major_vec) + .into_shape_with_order((shape, ndarray::Order::ColumnMajor)) + .unwrap(); + let col_major_nonzero_offset = col_major.slice_move(ndarray::s![1.., ..]); + + // make sure that the offset is in fact negative, in case ndarray behavior changes + let rmno = row_major_nonzero_offset.clone(); + let cmno = col_major_nonzero_offset.clone(); + let (_, rm_offset) = rmno.into_raw_vec_and_offset(); + let (_, cm_offset) = cmno.into_raw_vec_and_offset(); + assert!(rm_offset.unwrap() > 0); + assert!(cm_offset.unwrap() > 0); + + let tensor_row_major_nonzero_offset = TensorData::try_from(row_major_nonzero_offset).unwrap(); + let tensor_col_major_nonzero_offset = TensorData::try_from(col_major_nonzero_offset).unwrap(); + + assert_eq!( + tensor_row_major_nonzero_offset, + tensor_col_major_nonzero_offset + ); +} + #[test] fn check_slices() { let t = TensorData::new( From 74a14eb6101952191a273c7b629ce5de6c6809d4 Mon Sep 17 00:00:00 2001 From: Ben Liepert Date: Mon, 16 Sep 2024 19:46:34 -0400 Subject: [PATCH 5/5] Add asserts for std/non-std layout, fix comments --- crates/store/re_types/tests/types/tensor.rs | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/crates/store/re_types/tests/types/tensor.rs b/crates/store/re_types/tests/types/tensor.rs index 0b47514a9f17..8c72dfb4611d 100644 --- a/crates/store/re_types/tests/types/tensor.rs +++ b/crates/store/re_types/tests/types/tensor.rs @@ -150,7 +150,10 @@ fn convert_ndarray_to_tensor_both_layouts() { .into_shape_with_order((shape, ndarray::Order::ColumnMajor)) .unwrap(); - // make sure that the offset is in fact negative, in case ndarray behavior changes + assert!(row_major.is_standard_layout()); + assert!(!col_major.is_standard_layout()); + + // make sure that the offset is in fact zero, in case ndarray behavior changes let rm = row_major.clone(); let cm = col_major.clone(); let (_, rm_offset) = rm.into_raw_vec_and_offset(); @@ -184,14 +187,19 @@ fn convert_ndarray_to_tensor_both_layouts_nonzero_offset() { let row_major = ndarray::Array::from_vec(row_major_vec) .into_shape_with_order((shape, ndarray::Order::RowMajor)) .unwrap(); + assert!(row_major.is_standard_layout()); let row_major_nonzero_offset = row_major.slice_move(ndarray::s![1.., ..]); let col_major = ndarray::Array::from_vec(col_major_vec) .into_shape_with_order((shape, ndarray::Order::ColumnMajor)) .unwrap(); + assert!(!col_major.is_standard_layout()); let col_major_nonzero_offset = col_major.slice_move(ndarray::s![1.., ..]); - // make sure that the offset is in fact negative, in case ndarray behavior changes + assert!(row_major_nonzero_offset.is_standard_layout()); + assert!(!col_major_nonzero_offset.is_standard_layout()); + + // make sure that the offset is in fact non-zero, in case ndarray behavior changes let rmno = row_major_nonzero_offset.clone(); let cmno = col_major_nonzero_offset.clone(); let (_, rm_offset) = rmno.into_raw_vec_and_offset();