diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml index 780511ea27..bc91f8d22f 100644 --- a/.github/workflows/nightly.yml +++ b/.github/workflows/nightly.yml @@ -132,7 +132,10 @@ jobs: - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2 - name: miri - run: cargo +nightly miri nextest run --locked --package diskann-quantization + run: | + cargo +nightly miri nextest run --locked \ + --package diskann-quantization \ + --package diskann-utils env: MIRIFLAGS: -Zmiri-disable-isolation -Zmiri-strict-provenance diff --git a/agents.md b/agents.md index 6911cf791f..2411a16eb3 100644 --- a/agents.md +++ b/agents.md @@ -47,7 +47,7 @@ The repository uses a Cargo workspace with crates organized into functional tier **Tier 2: Core Libraries** - `diskann-linalg/` - Linear algebra operations -- `diskann-utils/` - Shared utilities (Reborrow, MatrixView traits) +- `diskann-utils/` - Shared utilities (Reborrow, Matrix traits) - `diskann-quantization/` - Vector quantization (PQ, SQ) **Tier 3: Algorithm & Storage** diff --git a/diskann-benchmark-core/src/build/graph/multi.rs b/diskann-benchmark-core/src/build/graph/multi.rs index f47f00fb47..d73e818d3c 100644 --- a/diskann-benchmark-core/src/build/graph/multi.rs +++ b/diskann-benchmark-core/src/build/graph/multi.rs @@ -10,7 +10,10 @@ use diskann::{ graph::{self, glue}, provider, }; -use diskann_utils::{future::AsyncFriendly, views::Matrix}; +use diskann_utils::{ + future::AsyncFriendly, + views::rowmajor::{self, Matrix}, +}; use crate::build::{Build, ids::ToId}; @@ -27,7 +30,7 @@ where DP: provider::DataProvider, { index: Arc>, - data: Arc>, + data: Arc>, strategy: S, to_id: Box>, } @@ -45,7 +48,7 @@ where /// to external IDs. pub fn new( index: Arc>, - data: Arc>, + data: Arc>, strategy: S, to_id: I, ) -> Arc @@ -64,7 +67,7 @@ where impl Build for MultiInsert where DP: provider::DataProvider + for<'a> provider::SetElement<&'a [T]>, - S: glue::MultiInsertStrategy> + Clone + 'static, + S: glue::MultiInsertStrategy> + Clone + 'static, T: AsyncFriendly + Clone, { type Output = (); @@ -101,7 +104,7 @@ where end: range.end, }) })? - .to_owned(); + .to_rowmajor_owned(); let ids: ANNResult> = range.into_iter().map(|i| self.to_id.to_id(i)).collect(); let context = DP::Context::default(); diff --git a/diskann-benchmark-core/src/build/graph/single.rs b/diskann-benchmark-core/src/build/graph/single.rs index 96a367b6a7..deda2da006 100644 --- a/diskann-benchmark-core/src/build/graph/single.rs +++ b/diskann-benchmark-core/src/build/graph/single.rs @@ -10,7 +10,10 @@ use diskann::{ graph::{self, glue}, provider, }; -use diskann_utils::{future::AsyncFriendly, views::Matrix}; +use diskann_utils::{ + future::AsyncFriendly, + views::rowmajor::{self, Matrix}, +}; use crate::build::{Build, ids::ToId}; @@ -23,7 +26,7 @@ where DP: provider::DataProvider, { index: Arc>, - data: Arc>, + data: Arc>, strategy: S, to_id: Box>, } @@ -41,7 +44,7 @@ where /// to external IDs. pub fn new( index: Arc>, - data: Arc>, + data: Arc>, strategy: S, to_id: I, ) -> Arc diff --git a/diskann-benchmark-core/src/recall.rs b/diskann-benchmark-core/src/recall.rs index b464a420db..fecf6fbc6d 100644 --- a/diskann-benchmark-core/src/recall.rs +++ b/diskann-benchmark-core/src/recall.rs @@ -10,7 +10,7 @@ use std::{ use diskann_utils::{ strided::Strided, - views::{Matrix, MatrixView}, + views::rowmajor::{self, Matrix}, }; use thiserror::Error; @@ -81,27 +81,27 @@ pub trait Rows { } } -impl Rows for Matrix { +impl Rows for rowmajor::Owned { fn nrows(&self) -> usize { - Matrix::::nrows(self) + Matrix::nrows(self) } fn row(&self, i: usize) -> &[T] { - Matrix::::row(self, i) + Matrix::row(self, i) } fn ncols(&self) -> Option { - Some(Matrix::::ncols(self)) + Some(Matrix::ncols(self)) } } -impl Rows for MatrixView<'_, T> { +impl Rows for rowmajor::Ref<'_, T> { fn nrows(&self) -> usize { - MatrixView::<'_, T>::nrows(self) + Matrix::nrows(self) } fn row(&self, i: usize) -> &[T] { - MatrixView::<'_, T>::row(self, i) + Matrix::row(self, i) } fn ncols(&self) -> Option { - Some(MatrixView::<'_, T>::ncols(self)) + Some(Matrix::ncols(self)) } } @@ -364,10 +364,10 @@ where #[cfg(test)] mod tests { - use diskann_utils::views::Matrix; - use super::*; + use diskann_utils::views::rowmajor::MatrixMut; + fn test_rows_inner(rows: &dyn Rows, ncols: Option) { assert_eq!(rows.ncols(), ncols); assert_eq!(rows.nrows(), 3); @@ -379,7 +379,7 @@ mod tests { #[test] fn test_rows() { let mut i = 0usize; - let mat = Matrix::from_fn(3, 4, |_| { + let mat = rowmajor::Owned::from_fn(3, 4, |_| { let v = i; i += 1; v @@ -420,7 +420,7 @@ mod tests { #[test] fn test_happy_path() { - let groundtruth = Matrix::try_from( + let groundtruth = rowmajor::Owned::try_from_data( vec![ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, // row 0 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, // row 1 @@ -433,7 +433,7 @@ mod tests { ) .unwrap(); - let distances = Matrix::try_from( + let distances = rowmajor::Owned::try_from_data( vec![ 0.0, 1.0, 2.0, 3.0, 3.0, 3.0, 3.0, 4.0, 5.0, 6.0, // row 0 2.0, 3.0, 3.0, 3.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, // row 1 @@ -447,7 +447,7 @@ mod tests { .unwrap(); // Shift row 0 by one and row 1 by two. - let our_results = Matrix::try_from( + let our_results = rowmajor::Owned::try_from_data( vec![ 100, 0, 1, 2, 5, 6, // row 0 100, 101, 7, 8, 9, 10, // row 1 @@ -480,7 +480,7 @@ mod tests { let epsilon = 1e-6; // Define a small tolerance for (i, expected) in expected_no_ties.iter().enumerate() { - assert_eq!(expected.components.len(), our_results.nrows()); + assert_eq!(expected.components.len(), Matrix::nrows(&our_results)); let recall = knn( &groundtruth, None, @@ -501,7 +501,7 @@ mod tests { i ); - assert_eq!(recall.num_queries, our_results.nrows()); + assert_eq!(recall.num_queries, Matrix::nrows(&our_results)); assert_eq!(recall.recall_k, expected.recall_k); assert_eq!(recall.recall_n, expected.recall_n); } @@ -525,7 +525,7 @@ mod tests { ]; for (i, expected) in expected_with_ties.iter().enumerate() { - assert_eq!(expected.components.len(), our_results.nrows()); + assert_eq!(expected.components.len(), Matrix::nrows(&our_results)); let recall = knn( &groundtruth, Some(distances.as_view().into()), @@ -546,7 +546,7 @@ mod tests { i ); - assert_eq!(recall.num_queries, our_results.nrows()); + assert_eq!(recall.num_queries, Matrix::nrows(&our_results)); assert_eq!(recall.recall_k, expected.recall_k); assert_eq!(recall.recall_n, expected.recall_n); } @@ -554,16 +554,16 @@ mod tests { #[test] fn test_error_recall_k_and_n() { - let groundtruth = Matrix::::from_element(10, 10, 0); - let results = Matrix::::from_element(10, 10, 0); + let groundtruth = rowmajor::Owned::::from_element(10, 10, 0); + let results = rowmajor::Owned::::from_element(10, 10, 0); let err = knn(&groundtruth, None, &results, 11, 10, GroundTruthMode::Fixed).unwrap_err(); assert!(matches!(err, ComputeRecallError::RecallKAndNError(..))); } #[test] fn test_error_rows_mismatch() { - let groundtruth = Matrix::::from_element(11, 10, 0); - let results = Matrix::::from_element(10, 10, 0); + let groundtruth = rowmajor::Owned::::from_element(11, 10, 0); + let results = rowmajor::Owned::::from_element(10, 10, 0); let err = knn(&groundtruth, None, &results, 10, 10, GroundTruthMode::Fixed).unwrap_err(); assert!(matches!(err, ComputeRecallError::RowsMismatch(..))); let err_allow_insufficient_results = @@ -576,8 +576,8 @@ mod tests { #[test] fn test_error_not_enough_groundtruth() { - let groundtruth = Matrix::::from_element(10, 5, 0); - let results = Matrix::::from_element(10, 10, 0); + let groundtruth = rowmajor::Owned::::from_element(10, 5, 0); + let results = rowmajor::Owned::::from_element(10, 10, 0); let err = knn(&groundtruth, None, &results, 10, 10, GroundTruthMode::Fixed).unwrap_err(); assert!(matches!(err, ComputeRecallError::NotEnoughGroundTruth(..))); let err_allow_insufficient_results = @@ -591,7 +591,7 @@ mod tests { #[test] fn test_dynamic_groundtruth_valid() { let groundtruth: Vec<_> = (0..10).map(|_| vec![0u32; 5]).collect(); - let results = Matrix::::from_element(10, 10, 0); + let results = rowmajor::Owned::::from_element(10, 10, 0); // Should succeed: each row uses this_recall_k = min(5, 10) = 5 // Should succeed in Flexible mode, but fail in Fixed mode let recall_flexible = knn( @@ -614,7 +614,7 @@ mod tests { fn test_dynamic_groundtruth_full_match() { let gt_row: Vec = (1..=5).collect(); let groundtruth: Vec<_> = (0..10).map(|_| gt_row.clone()).collect(); - let mut results = Matrix::::from_element(10, 10, 0); + let mut results = rowmajor::Owned::::from_element(10, 10, 0); for i in 0..10 { for (j, v) in (1u32..=10).enumerate() { *results.element_mut(i, j) = v; @@ -637,7 +637,7 @@ mod tests { // groundtruth: [1, 2, 3, 4, 5]; results contain [1, 2, 3, 6, 7, 8, 9, 10, 11, 12] let gt_row: Vec = (1..=5).collect(); let groundtruth: Vec<_> = (0..10).map(|_| gt_row.clone()).collect(); - let mut results = Matrix::::from_element(10, 10, 0); + let mut results = rowmajor::Owned::::from_element(10, 10, 0); let res_row: Vec = vec![1, 2, 3, 6, 7, 8, 9, 10, 11, 12]; for i in 0..10 { for (j, &v) in res_row.iter().enumerate() { @@ -668,7 +668,7 @@ mod tests { groundtruth.push(vec![]); } - let mut results = Matrix::::from_element(10, 10, 0); + let mut results = rowmajor::Owned::::from_element(10, 10, 0); for i in 0..10 { for (j, v) in (1u32..=10).enumerate() { *results.element_mut(i, j) = v; @@ -691,7 +691,7 @@ mod tests { #[test] fn test_dynamic_groundtruth_all_zero() { let groundtruth: Vec> = (0..10).map(|_| vec![]).collect(); - let results = Matrix::::from_element(10, 10, 0); + let results = rowmajor::Owned::::from_element(10, 10, 0); let recall = knn( &groundtruth, @@ -748,9 +748,9 @@ mod tests { #[test] fn test_error_distance_rows_mismatch() { - let groundtruth = Matrix::::from_element(10, 10, 0); - let distances = Matrix::::from_element(9, 10, 0.0); - let results = Matrix::::from_element(10, 10, 0); + let groundtruth = rowmajor::Owned::::from_element(10, 10, 0); + let distances = rowmajor::Owned::::from_element(9, 10, 0.0); + let results = rowmajor::Owned::::from_element(10, 10, 0); let err = knn( &groundtruth, Some(distances.as_view().into()), @@ -765,9 +765,9 @@ mod tests { #[test] fn test_error_distance_cols_mismatch() { - let groundtruth = Matrix::::from_element(10, 10, 0); - let distances = Matrix::::from_element(10, 9, 0.0); - let results = Matrix::::from_element(10, 10, 0); + let groundtruth = rowmajor::Owned::::from_element(10, 10, 0); + let distances = rowmajor::Owned::::from_element(10, 9, 0.0); + let results = rowmajor::Owned::::from_element(10, 10, 0); let err = knn( &groundtruth, Some(distances.as_view().into()), @@ -789,7 +789,7 @@ mod tests { let groundtruth: Vec> = vec![vec![1, 2, 3], vec![4, 5]]; // distances: first row has 2 elements (should be 3), second row has 2 (matches) let distances: Vec> = vec![vec![0.1, 0.2], vec![0.3, 0.4]]; - let distances = Matrix::try_from( + let distances = rowmajor::Owned::try_from_data( distances.into_iter().flatten().collect::>().into(), 2, 2, diff --git a/diskann-benchmark-core/src/search/api.rs b/diskann-benchmark-core/src/search/api.rs index 383cc0f46e..acd5964f31 100644 --- a/diskann-benchmark-core/src/search/api.rs +++ b/diskann-benchmark-core/src/search/api.rs @@ -9,7 +9,7 @@ use diskann::{ANNResult, graph, utils::async_tools}; use diskann_benchmark_runner::utils::MicroSeconds; use diskann_utils::{ future::{AsyncFriendly, boxit}, - views::Matrix, + views::rowmajor::{self, MatrixMut}, }; use crate::{ @@ -517,12 +517,12 @@ where T: Search, { let mut lengths = Vec::with_capacity(range.len()); - let mut ids = Matrix::try_from_fn(range.len(), num_ids.into(), |_| T::Id::default())?; + let mut ids = rowmajor::Owned::try_from_fn(range.len(), num_ids.into(), |_| T::Id::default())?; let mut latencies = Vec::::with_capacity(range.len()); let mut outputs = Vec::::with_capacity(range.len()); - for (ids, index) in std::iter::zip(ids.row_iter_mut(), range) { + for (ids, index) in std::iter::zip(ids.rows_mut(), range) { let mut buffer = internal::buffer::Buffer::slice(ids); let start = std::time::Instant::now(); diff --git a/diskann-benchmark-core/src/search/graph/filtered_range.rs b/diskann-benchmark-core/src/search/graph/filtered_range.rs index a3378c3584..b6bbec526b 100644 --- a/diskann-benchmark-core/src/search/graph/filtered_range.rs +++ b/diskann-benchmark-core/src/search/graph/filtered_range.rs @@ -11,7 +11,10 @@ use diskann::{ provider, }; use diskann_benchmark_runner::utils::{MicroSeconds, percentiles}; -use diskann_utils::{future::AsyncFriendly, views::Matrix}; +use diskann_utils::{ + future::AsyncFriendly, + views::rowmajor::{self, Matrix}, +}; use crate::{ recall, @@ -34,7 +37,7 @@ where DP: provider::DataProvider, { index: Arc>, - queries: Arc>, + queries: Arc>, strategy: Strategy, labels: Arc<[Arc>]>, } @@ -56,7 +59,7 @@ where /// the number of rows in `queries`. pub fn new( index: Arc>, - queries: Arc>, + queries: Arc>, strategy: Strategy, labels: Arc<[Arc>]>, ) -> anyhow::Result> { @@ -223,6 +226,7 @@ mod tests { use super::*; use diskann::graph::{ext::labeled::QueryLabelProvider, test::provider}; + use diskann_utils::views::rowmajor::MatrixMut; #[derive(Debug)] struct NoOdds; @@ -237,7 +241,7 @@ mod tests { fn test_filtered_range() { let index = search::graph::test_grid_provider(); - let mut queries = Matrix::from_element(5, index.provider().dim(), 0.0f32); + let mut queries = rowmajor::Owned::from_element(5, index.provider().dim(), 0.0f32); queries.row_mut(0).copy_from_slice(&[0.0, 0.0, 0.0, 0.0]); queries.row_mut(1).copy_from_slice(&[4.0, 0.0, 0.0, 0.0]); queries.row_mut(2).copy_from_slice(&[0.0, 4.0, 0.0, 0.0]); @@ -330,7 +334,11 @@ mod tests { #[test] fn test_filtered_range_error() { let index = search::graph::test_grid_provider(); - let queries = Arc::new(Matrix::from_element(2, index.provider().dim(), 0.0f32)); + let queries = Arc::new(rowmajor::Owned::from_element( + 2, + index.provider().dim(), + 0.0f32, + )); let labels: Arc<[_]> = (0..queries.nrows() + 1) .map(|_| -> Arc> { Arc::new(NoOdds {}) }) diff --git a/diskann-benchmark-core/src/search/graph/inline.rs b/diskann-benchmark-core/src/search/graph/inline.rs index a13b437bf1..20812dea83 100644 --- a/diskann-benchmark-core/src/search/graph/inline.rs +++ b/diskann-benchmark-core/src/search/graph/inline.rs @@ -10,7 +10,10 @@ use diskann::{ graph::{self, ext::labeled, glue, search::AdaptiveL}, provider, }; -use diskann_utils::{future::AsyncFriendly, views::Matrix}; +use diskann_utils::{ + future::AsyncFriendly, + views::rowmajor::{self, Matrix}, +}; use crate::search::{self, Search, graph::KnnParams, graph::Strategy}; @@ -30,7 +33,7 @@ where DP: provider::DataProvider, { index: Arc>, - queries: Arc>, + queries: Arc>, strategy: Strategy, labels: Arc<[Arc>]>, adaptive_l: Option, @@ -61,7 +64,7 @@ where /// `queries`. pub fn new( index: Arc>, - queries: Arc>, + queries: Arc>, strategy: Strategy, labels: Arc<[Arc>]>, adaptive_l: Option, @@ -150,12 +153,14 @@ where #[cfg(test)] mod tests { + use super::*; + use std::num::NonZeroUsize; - use super::*; + use diskann::graph::{ext::labeled::QueryLabelProvider, test::provider}; + use diskann_utils::views::rowmajor::MatrixMut; use crate::recall::GroundTruthMode; - use diskann::graph::{ext::labeled::QueryLabelProvider, test::provider}; // A simple [`QueryLabelProvider`] that rejects odd indices. #[derive(Debug)] @@ -173,7 +178,7 @@ mod tests { let index = search::graph::test_grid_provider(); - let mut queries = Matrix::from_element(5, index.provider().dim(), 0.0f32); + let mut queries = rowmajor::Owned::from_element(5, index.provider().dim(), 0.0f32); queries.row_mut(0).copy_from_slice(&[0.0, 0.0, 0.0, 0.0]); queries.row_mut(1).copy_from_slice(&[4.0, 0.0, 0.0, 0.0]); queries.row_mut(2).copy_from_slice(&[0.0, 4.0, 0.0, 0.0]); @@ -273,7 +278,11 @@ mod tests { #[test] fn test_inline_error() { let index = search::graph::test_grid_provider(); - let queries = Arc::new(Matrix::from_element(2, index.provider().dim(), 0.0f32)); + let queries = Arc::new(rowmajor::Owned::from_element( + 2, + index.provider().dim(), + 0.0f32, + )); let labels: Arc<[_]> = (0..queries.nrows() + 1) .map(|_| -> Arc> { Arc::new(NoOdds {}) }) diff --git a/diskann-benchmark-core/src/search/graph/knn.rs b/diskann-benchmark-core/src/search/graph/knn.rs index ceedacbef7..6c02873545 100644 --- a/diskann-benchmark-core/src/search/graph/knn.rs +++ b/diskann-benchmark-core/src/search/graph/knn.rs @@ -14,7 +14,10 @@ use diskann::{ provider, }; use diskann_benchmark_runner::utils::{MicroSeconds, percentiles}; -use diskann_utils::{future::AsyncFriendly, views::Matrix}; +use diskann_utils::{ + future::AsyncFriendly, + views::rowmajor::{self, Matrix}, +}; use crate::{ recall, @@ -48,7 +51,7 @@ where DP: provider::DataProvider, { index: Arc>, - queries: Arc>, + queries: Arc>, strategy: Strategy, post_processor: PP, } @@ -70,7 +73,7 @@ where /// the number of rows in `queries`. pub fn new( index: Arc>, - queries: Arc>, + queries: Arc>, strategy: Strategy, ) -> anyhow::Result> { strategy.length_compatible(queries.nrows())?; @@ -96,7 +99,7 @@ where /// the number of rows in `queries`. pub fn with_postprocessor( index: Arc>, - queries: Arc>, + queries: Arc>, strategy: Strategy, post_processor: PP, ) -> anyhow::Result> { @@ -436,11 +439,12 @@ where #[cfg(test)] mod tests { - use std::num::NonZeroUsize; - use super::*; + use std::num::NonZeroUsize; + use diskann::graph::test::provider; + use diskann_utils::views::rowmajor::MatrixMut; #[test] fn test_knn() { @@ -448,7 +452,7 @@ mod tests { let index = search::graph::test_grid_provider(); - let mut queries = Matrix::from_element(5, index.provider().dim(), 0.0f32); + let mut queries = rowmajor::Owned::from_element(5, index.provider().dim(), 0.0f32); queries.row_mut(0).copy_from_slice(&[0.0, 0.0, 0.0, 0.0]); queries.row_mut(1).copy_from_slice(&[4.0, 0.0, 0.0, 0.0]); queries.row_mut(2).copy_from_slice(&[0.0, 4.0, 0.0, 0.0]); @@ -534,7 +538,11 @@ mod tests { fn test_knn_error() { let index = search::graph::test_grid_provider(); - let queries = Arc::new(Matrix::from_element(1, index.provider().dim(), 0.0f32)); + let queries = Arc::new(rowmajor::Owned::from_element( + 1, + index.provider().dim(), + 0.0f32, + )); let strategy = provider::Strategy::new(); let err = KNN::new( diff --git a/diskann-benchmark-core/src/search/graph/multihop.rs b/diskann-benchmark-core/src/search/graph/multihop.rs index f9d372daaf..10197bcf86 100644 --- a/diskann-benchmark-core/src/search/graph/multihop.rs +++ b/diskann-benchmark-core/src/search/graph/multihop.rs @@ -10,7 +10,10 @@ use diskann::{ graph::{self, ext::labeled, glue}, provider, }; -use diskann_utils::{future::AsyncFriendly, views::Matrix}; +use diskann_utils::{ + future::AsyncFriendly, + views::rowmajor::{self, Matrix}, +}; use crate::search::{self, Search, graph::KnnParams, graph::Strategy}; @@ -30,7 +33,7 @@ where DP: provider::DataProvider, { index: Arc>, - queries: Arc>, + queries: Arc>, strategy: Strategy, labels: Arc<[Arc>]>, } @@ -60,7 +63,7 @@ where /// `queries`. pub fn new( index: Arc>, - queries: Arc>, + queries: Arc>, strategy: Strategy, labels: Arc<[Arc>]>, ) -> anyhow::Result> { @@ -146,12 +149,14 @@ where #[cfg(test)] mod tests { + use super::*; + use std::num::NonZeroUsize; - use super::*; + use diskann::graph::{ext::labeled::QueryLabelProvider, test::provider}; + use diskann_utils::views::rowmajor::MatrixMut; use crate::recall::GroundTruthMode; - use diskann::graph::{ext::labeled::QueryLabelProvider, test::provider}; // A simple [`QueryLabelProvider`] that rejects odd indices. #[derive(Debug)] @@ -169,7 +174,7 @@ mod tests { let index = search::graph::test_grid_provider(); - let mut queries = Matrix::from_element(5, index.provider().dim(), 0.0f32); + let mut queries = rowmajor::Owned::from_element(5, index.provider().dim(), 0.0f32); queries.row_mut(0).copy_from_slice(&[0.0, 0.0, 0.0, 0.0]); queries.row_mut(1).copy_from_slice(&[4.0, 0.0, 0.0, 0.0]); queries.row_mut(2).copy_from_slice(&[0.0, 4.0, 0.0, 0.0]); @@ -266,7 +271,11 @@ mod tests { #[test] fn test_multihop_error() { let index = search::graph::test_grid_provider(); - let queries = Arc::new(Matrix::from_element(2, index.provider().dim(), 0.0f32)); + let queries = Arc::new(rowmajor::Owned::from_element( + 2, + index.provider().dim(), + 0.0f32, + )); let labels: Arc<[_]> = (0..queries.nrows() + 1) .map(|_| -> Arc> { Arc::new(NoOdds {}) }) diff --git a/diskann-benchmark-core/src/search/graph/range.rs b/diskann-benchmark-core/src/search/graph/range.rs index 9992f01b50..f33e175d92 100644 --- a/diskann-benchmark-core/src/search/graph/range.rs +++ b/diskann-benchmark-core/src/search/graph/range.rs @@ -11,7 +11,10 @@ use diskann::{ provider, }; use diskann_benchmark_runner::utils::{MicroSeconds, percentiles}; -use diskann_utils::{future::AsyncFriendly, views::Matrix}; +use diskann_utils::{ + future::AsyncFriendly, + views::rowmajor::{self, Matrix}, +}; use crate::{ recall, @@ -33,7 +36,7 @@ where DP: provider::DataProvider, { index: Arc>, - queries: Arc>, + queries: Arc>, strategy: Strategy, } @@ -54,7 +57,7 @@ where /// the number of rows in `queries`. pub fn new( index: Arc>, - queries: Arc>, + queries: Arc>, strategy: Strategy, ) -> anyhow::Result> { strategy.length_compatible(queries.nrows())?; @@ -237,12 +240,13 @@ mod tests { use super::*; use diskann::graph::test::provider; + use diskann_utils::views::rowmajor::MatrixMut; #[test] fn test_range() { let index = search::graph::test_grid_provider(); - let mut queries = Matrix::from_element(5, index.provider().dim(), 0.0f32); + let mut queries = rowmajor::Owned::from_element(5, index.provider().dim(), 0.0f32); queries.row_mut(0).copy_from_slice(&[0.0, 0.0, 0.0, 0.0]); queries.row_mut(1).copy_from_slice(&[4.0, 0.0, 0.0, 0.0]); queries.row_mut(2).copy_from_slice(&[0.0, 4.0, 0.0, 0.0]); @@ -325,7 +329,11 @@ mod tests { fn test_range_error() { let index = search::graph::test_grid_provider(); - let queries = Arc::new(Matrix::from_element(2, index.provider().dim(), 0.0f32)); + let queries = Arc::new(rowmajor::Owned::from_element( + 2, + index.provider().dim(), + 0.0f32, + )); let strategy = provider::Strategy::new(); let err = Range::new(index, queries.clone(), Strategy::collection([strategy])).unwrap_err(); diff --git a/diskann-benchmark-core/src/search/ids.rs b/diskann-benchmark-core/src/search/ids.rs index 421ff005f7..f9ae988e97 100644 --- a/diskann-benchmark-core/src/search/ids.rs +++ b/diskann-benchmark-core/src/search/ids.rs @@ -3,7 +3,7 @@ * Licensed under the MIT license. */ -use diskann_utils::views::Matrix; +use diskann_utils::views::rowmajor::{self, Matrix, MatrixMut}; use crate::recall; @@ -36,7 +36,7 @@ impl ResultIds { /// separately. #[derive(Debug)] pub(crate) struct Bounded { - ids: Matrix, + ids: rowmajor::Owned, // Must have the same length as `matrix.nrows()`. lengths: Vec, } @@ -52,7 +52,7 @@ impl Bounded { /// # Panics /// /// Panics if the number of rows in `ids` does not match the length of `lengths`. - pub(crate) fn new(ids: Matrix, lengths: Vec) -> Self { + pub(crate) fn new(ids: rowmajor::Owned, lengths: Vec) -> Self { assert_eq!( ids.nrows(), lengths.len(), @@ -71,7 +71,7 @@ impl Bounded { /// /// Note that the yielded slices are not guaranteed to have the same length. pub(crate) fn iter(&self) -> impl ExactSizeIterator { - std::iter::zip(self.ids.row_iter(), self.lengths.iter()).map(|(row, len)| { + std::iter::zip(self.ids.rows(), self.lengths.iter()).map(|(row, len)| { match row.get(..*len) { Some(v) => v, None => row, @@ -86,7 +86,7 @@ impl recall::Rows for Bounded { } fn row(&self, index: usize) -> &[I] { let length = self.lengths[index]; - let row = self.ids.row(index); + let row = Matrix::row(&self.ids, index); match row.get(..length) { Some(v) => v, None => row, @@ -231,12 +231,12 @@ where len, num_ids, } => { - let mut dst = Matrix::from_fn(len, num_ids, |_| I::default()); + let mut dst = rowmajor::Owned::from_fn(len, num_ids, |_| I::default()); let mut lengths = Vec::with_capacity(len); let mut output_row = 0; for bounded in matrices { - for row in bounded.ids.row_iter() { + for row in bounded.ids.rows() { dst.row_mut(output_row).clone_from_slice(row); output_row += 1; } @@ -279,10 +279,10 @@ mod tests { let nrows = data.len(); let ncols = data.iter().map(|v| v.len()).max().unwrap_or(0); - let mut matrix = Matrix::from_element(nrows, ncols, 0u32); + let mut matrix = rowmajor::Owned::from_element(nrows, ncols, 0u32); let mut lengths = Vec::with_capacity(nrows); - for (row, row_data) in std::iter::zip(matrix.row_iter_mut(), data.iter()) { + for (row, row_data) in std::iter::zip(matrix.rows_mut(), data.iter()) { let len = std::iter::zip(row.iter_mut(), row_data.iter()) .map(|(dst, src)| { *dst = *src; @@ -296,7 +296,7 @@ mod tests { #[test] fn test_bounded_new_valid() { - let matrix = Matrix::from_element(3, 5, 0u32); + let matrix = rowmajor::Owned::from_element(3, 5, 0u32); let lengths = vec![2, 3, 1]; let bounded = Bounded::new(matrix, lengths); @@ -305,7 +305,7 @@ mod tests { #[test] fn test_bounded_length_clamping() { - let matrix = Matrix::from_element(3, 3, 0u32); + let matrix = rowmajor::Owned::from_element(3, 3, 0u32); let lengths = vec![2, 3, 5]; // Last length exceeds number of columns let bounded = Bounded::new(matrix, lengths); @@ -323,7 +323,7 @@ mod tests { #[test] #[should_panic(expected = "an internal invariant was not upheld")] fn test_bounded_new_mismatched_lengths() { - let matrix = Matrix::from_element(3, 5, 0u32); + let matrix = rowmajor::Owned::from_element(3, 5, 0u32); let lengths = vec![2, 3]; // Only 2 lengths for 3 rows Bounded::new(matrix, lengths); } diff --git a/diskann-benchmark-core/src/streaming/executors/bigann/withdata.rs b/diskann-benchmark-core/src/streaming/executors/bigann/withdata.rs index b302293620..7e26cb929c 100644 --- a/diskann-benchmark-core/src/streaming/executors/bigann/withdata.rs +++ b/diskann-benchmark-core/src/streaming/executors/bigann/withdata.rs @@ -5,7 +5,7 @@ use std::{ops::Range, path::Path, sync::Arc}; -use diskann_utils::views::{Matrix, MatrixView}; +use diskann_utils::views::rowmajor::{self, Matrix}; use crate::{recall, streaming}; @@ -16,8 +16,8 @@ type LoadGroundtruth = dyn FnMut(&Path) -> anyhow::Result { inner: Inner, - dataset: Matrix, - queries: Arc>, + dataset: rowmajor::Owned, + queries: Arc>, load_groundtruth: Box>, } @@ -28,8 +28,8 @@ impl WithData { /// for a given path. pub fn new( inner: Inner, - dataset: Matrix, - queries: Arc>, + dataset: rowmajor::Owned, + queries: Arc>, load_groundtruth: impl FnMut(&Path) -> anyhow::Result>> + 'static, ) -> Self { Self { @@ -55,12 +55,12 @@ where { /// A tuple consisting of the queries for search as well as the corresponding /// groundtruth (as [`recall::Rows`]). - type Search<'a> = (Arc>, &'a dyn recall::Rows); + type Search<'a> = (Arc>, &'a dyn recall::Rows); /// A tuple consisting of the data to insert as well as the external IDs (stored /// as `usize`) for the data. It is assumed that the length of the IDs range matches /// the number of rows in the data matrix. - type Insert<'a> = (MatrixView<'a, T>, Range); + type Insert<'a> = (rowmajor::Ref<'a, T>, Range); /// The external IDs (stored as `usize`) to delete. type Delete<'a> = Range; @@ -68,7 +68,7 @@ where /// A tuple consisting of the data to replace as well as the external IDs (stored /// as `usize`) for the data. It is assumed that the length of the IDs range matches /// the number of rows in the data matrix. - type Replace<'a> = (MatrixView<'a, T>, Range); + type Replace<'a> = (rowmajor::Ref<'a, T>, Range); type Maintain<'a> = (); } diff --git a/diskann-benchmark-core/src/streaming/graph/test.rs b/diskann-benchmark-core/src/streaming/graph/test.rs index e86906959b..ea04b22fe1 100644 --- a/diskann-benchmark-core/src/streaming/graph/test.rs +++ b/diskann-benchmark-core/src/streaming/graph/test.rs @@ -12,6 +12,7 @@ use diskann::{ }, utils::{IntoUsize, ONE}, }; +use diskann_utils::views::rowmajor::Matrix; use crate::build::{self, graph::SingleInsert}; diff --git a/diskann-benchmark-simd/src/lib.rs b/diskann-benchmark-simd/src/lib.rs index c2fec58069..353c8504c5 100644 --- a/diskann-benchmark-simd/src/lib.rs +++ b/diskann-benchmark-simd/src/lib.rs @@ -7,7 +7,7 @@ use std::{io::Write, num::NonZeroUsize}; -use diskann_utils::views::{Matrix, MatrixView}; +use diskann_utils::views::rowmajor::{self, Matrix}; use diskann_vector::distance::simd; use diskann_wide::Architecture; use half::f16; @@ -713,7 +713,7 @@ impl std::fmt::Display for DisplayWrapper<'_, [RunResult]> { } } -fn run_loops(query: &[Q], data: MatrixView, run: &Run, f: F) -> RunResult +fn run_loops(query: &[Q], data: rowmajor::Ref, run: &Run, f: F) -> RunResult where F: Fn(&[Q], &[D]) -> f32, { @@ -723,7 +723,7 @@ where for _ in 0..run.num_measurements.get() { let start = std::time::Instant::now(); for _ in 0..run.loops_per_measurement.get() { - std::iter::zip(dst.iter_mut(), data.row_iter()).for_each(|(d, r)| { + std::iter::zip(dst.iter_mut(), data.rows()).for_each(|(d, r)| { *d = f(query, r); }); std::hint::black_box(&mut dst); @@ -741,7 +741,7 @@ where struct Data { query: Box<[Q]>, - data: Matrix, + data: rowmajor::Owned, } impl Data { @@ -754,7 +754,7 @@ impl Data { let query: Box<[Q]> = (0..run.dim.get()) .map(|_| StandardUniform.sample(&mut rng)) .collect(); - let data = Matrix::::from_fn(run.num_points.get(), run.dim.get(), |_| { + let data = rowmajor::Owned::::from_fn(run.num_points.get(), run.dim.get(), |_| { StandardUniform.sample(&mut rng) }); diff --git a/diskann-benchmark/src/disk_index/search.rs b/diskann-benchmark/src/disk_index/search.rs index 8cd0f5e606..dc55de23d5 100644 --- a/diskann-benchmark/src/disk_index/search.rs +++ b/diskann-benchmark/src/disk_index/search.rs @@ -31,7 +31,7 @@ use diskann_providers::{ utils::{create_thread_pool, ParallelIteratorInPool}, }; use diskann_tools::utils::{search_index_utils, KRecallAtN}; -use diskann_utils::views::Matrix; +use diskann_utils::views::rowmajor::{self, Matrix}; use serde::{Deserialize, Serialize}; use crate::{ @@ -181,7 +181,8 @@ where let mut logger = PerfLogger::new("search_disk_index", true); // Load the query file - let queries: Matrix = datafiles::load_dataset(datafiles::BinFile(&search_params.queries))?; + let queries: rowmajor::Owned = + datafiles::load_dataset(datafiles::BinFile(&search_params.queries))?; let num_queries = queries.nrows(); // Load the vector filters @@ -258,7 +259,7 @@ where }; let zipped = queries - .par_row_iter() + .par_rows() .zip(vector_filters.par_iter()) .zip(result_ids.par_chunks_mut(search_params.recall_at as usize)) .zip(result_dists.par_chunks_mut(search_params.recall_at as usize)) diff --git a/diskann-benchmark/src/exhaustive/algos.rs b/diskann-benchmark/src/exhaustive/algos.rs index 55fc39dbf6..7d27806a19 100644 --- a/diskann-benchmark/src/exhaustive/algos.rs +++ b/diskann-benchmark/src/exhaustive/algos.rs @@ -5,6 +5,7 @@ use diskann::neighbor::{Neighbor, NeighborPriorityQueue}; use diskann_benchmark_runner::utils::MicroSeconds; +use diskann_utils::views::rowmajor::{self, Matrix, MatrixMut}; use diskann_vector::PreprocessedDistanceFunction; use rayon::iter::{IndexedParallelIterator, ParallelIterator}; @@ -33,7 +34,7 @@ where #[derive(Debug, Clone)] pub(super) struct LinearSearch { - pub(super) ids: diskann_utils::views::Matrix, + pub(super) ids: rowmajor::Owned, pub(super) preprocess: Vec, pub(super) search: Vec, pub(super) total: MicroSeconds, @@ -41,7 +42,7 @@ pub(super) struct LinearSearch { pub(super) fn linear_search( store: &Q, - queries: diskann_utils::views::MatrixView, + queries: rowmajor::Ref, builder: &C, results_per_query: usize, progress: &indicatif::ProgressBar, @@ -50,8 +51,7 @@ where Q: QuantStore + Sync, C: CreateQuantComputer + Sync, { - let mut output = - diskann_utils::views::Matrix::from_element(queries.nrows(), results_per_query, u32::MAX); + let mut output = rowmajor::Owned::from_element(queries.nrows(), results_per_query, u32::MAX); struct Times { preprocess: MicroSeconds, @@ -64,8 +64,8 @@ where // this in a properly sized Rayon environment. #[expect(clippy::disallowed_methods)] let times: Vec = output - .par_row_iter_mut() - .zip(queries.par_row_iter()) + .par_rows_mut() + .zip(queries.par_rows()) .map(|(o, q)| -> anyhow::Result { let mut queue = NeighborPriorityQueue::::new(results_per_query); diff --git a/diskann-benchmark/src/exhaustive/minmax.rs b/diskann-benchmark/src/exhaustive/minmax.rs index 8cbac705f2..fa57404e3e 100644 --- a/diskann-benchmark/src/exhaustive/minmax.rs +++ b/diskann-benchmark/src/exhaustive/minmax.rs @@ -54,6 +54,7 @@ mod imp { num::Positive, AsFunctor, CompressInto, }; + use diskann_utils::views::rowmajor::{self, Matrix, MatrixMut}; use diskann_utils::{Reborrow, ReborrowMut}; use diskann_vector::{PreprocessedDistanceFunction, PureDistanceFunction}; use indicatif::{ProgressBar, ProgressStyle}; @@ -370,7 +371,7 @@ mod imp { pub(super) struct Store { // The number of bytes to take from each row. bytes: usize, - data: diskann_utils::views::Matrix, + data: rowmajor::Owned, quantizer: diskann_quantization::minmax::MinMaxQuantizer, } @@ -379,7 +380,7 @@ mod imp { Unsigned: Representation, { fn new( - input: diskann_utils::views::MatrixView, + input: rowmajor::Ref, quantizer: diskann_quantization::minmax::MinMaxQuantizer, progress: &ProgressBar, ) -> anyhow::Result { @@ -389,12 +390,12 @@ mod imp { // The APIs below should correctly handle these variables. let output_dim = quantizer.output_dim(); let bytes = Data::::canonical_bytes(output_dim); - let mut data = diskann_utils::views::Matrix::try_from_element(input.nrows(), bytes, 0)?; + let mut data = rowmajor::Owned::try_from_element(input.nrows(), bytes, 0)?; // Compress the data. // // NOTE: If this gets too slow, we can parallelize it. - std::iter::zip(data.row_iter_mut(), input.row_iter()).try_for_each( + std::iter::zip(data.rows_mut(), input.rows()).try_for_each( |(d, i)| -> anyhow::Result<()> { let c = diskann_quantization::minmax::DataMutRef::::from_canonical_front_mut( &mut d[..bytes], @@ -430,7 +431,7 @@ mod imp { fn iter(&self) -> impl Iterator> { let output_dim = self.quantizer.output_dim(); - self.data.row_iter().map(move |r| { + self.data.rows().map(move |r| { DataRef::::from_canonical_front(&r[..self.bytes], output_dim).unwrap() }) } diff --git a/diskann-benchmark/src/exhaustive/product.rs b/diskann-benchmark/src/exhaustive/product.rs index d64f8ccf31..ca52a0a5dc 100644 --- a/diskann-benchmark/src/exhaustive/product.rs +++ b/diskann-benchmark/src/exhaustive/product.rs @@ -35,6 +35,7 @@ mod imp { Benchmark, Output, }; use diskann_quantization::{product::train::TrainQuantizer, CompressInto}; + use diskann_utils::views::rowmajor::{self, Matrix, MatrixMut}; use indicatif::{ProgressBar, ProgressStyle}; use rayon::iter::{IndexedParallelIterator, ParallelIterator}; use serde::Serialize; @@ -322,31 +323,28 @@ mod imp { /// A store for quantized data. pub(super) struct Store { - data: diskann_utils::views::Matrix, + data: rowmajor::Owned, quantizer: diskann_providers::model::pq::FixedChunkPQTable, } impl Store { fn new( - input: diskann_utils::views::MatrixView, + input: rowmajor::Ref, quantizer: diskann_providers::model::pq::FixedChunkPQTable, progress: &ProgressBar, ) -> anyhow::Result { - let mut data = diskann_utils::views::Matrix::try_from_element( - input.nrows(), - quantizer.get_num_chunks(), - 0, - )?; + let mut data = + rowmajor::Owned::try_from_element(input.nrows(), quantizer.get_num_chunks(), 0)?; // Compress the data. #[expect(clippy::disallowed_methods)] - data.par_row_iter_mut() - .zip(input.par_row_iter()) - .try_for_each(|(d, i)| -> anyhow::Result<()> { + data.par_rows_mut().zip(input.par_rows()).try_for_each( + |(d, i)| -> anyhow::Result<()> { quantizer.compress_into(i, d)?; progress.inc(1); Ok(()) - })?; + }, + )?; Ok(Self { data, quantizer }) } @@ -363,7 +361,7 @@ mod imp { Self: 'a; fn iter(&self) -> impl Iterator> { - self.data.row_iter() + self.data.rows() } } diff --git a/diskann-benchmark/src/exhaustive/spherical.rs b/diskann-benchmark/src/exhaustive/spherical.rs index 012b2d248b..c49ecf2c0c 100644 --- a/diskann-benchmark/src/exhaustive/spherical.rs +++ b/diskann-benchmark/src/exhaustive/spherical.rs @@ -50,6 +50,7 @@ mod imp { spherical::{DataMut, SphericalQuantizer}, CompressIntoWith, }; + use diskann_utils::views::rowmajor::{self, Matrix, MatrixMut}; use indicatif::{ProgressBar, ProgressStyle}; use rand::SeedableRng; use rayon::iter::{IndexedParallelIterator, ParallelIterator}; @@ -394,7 +395,7 @@ mod imp { pub(super) struct Store { // The number of bytes to take from each row. bytes: usize, - data: diskann_utils::views::Matrix, + data: rowmajor::Owned, plan: diskann_quantization::spherical::iface::Impl, } @@ -405,7 +406,7 @@ mod imp { for<'a> CompressIntoWith<&'a [f32], DataMut<'a, NBITS>, ScopedAllocator<'a>>, { fn new( - input: diskann_utils::views::MatrixView, + input: rowmajor::Ref, plan: diskann_quantization::spherical::iface::Impl, progress: &ProgressBar, ) -> anyhow::Result { @@ -416,13 +417,12 @@ mod imp { let output_dim = plan.quantizer().output_dim(); let bytes = diskann_quantization::spherical::DataRef::::canonical_bytes(output_dim); - let mut data = diskann_utils::views::Matrix::try_from_element(input.nrows(), bytes, 0)?; + let mut data = rowmajor::Owned::try_from_element(input.nrows(), bytes, 0)?; // Compress the data. #[expect(clippy::disallowed_methods)] - data.par_row_iter_mut() - .zip(input.par_row_iter()) - .try_for_each(|(d, i)| -> anyhow::Result<()> { + data.par_rows_mut().zip(input.par_rows()).try_for_each( + |(d, i)| -> anyhow::Result<()> { let c = diskann_quantization::spherical::DataMut::::from_canonical_back_mut( &mut d[..bytes], @@ -432,7 +432,8 @@ mod imp { .compress_into_with(i, c, ScopedAllocator::global())?; progress.inc(1); Ok(()) - })?; + }, + )?; Ok(Self { bytes, data, plan }) } @@ -459,7 +460,7 @@ mod imp { fn iter(&self) -> impl Iterator> { self.data - .row_iter() + .rows() .map(move |r| diskann_quantization::spherical::iface::Opaque::new(&r[..self.bytes])) } } diff --git a/diskann-benchmark/src/flat/search.rs b/diskann-benchmark/src/flat/search.rs index 25f93f2687..97b6dbb592 100644 --- a/diskann-benchmark/src/flat/search.rs +++ b/diskann-benchmark/src/flat/search.rs @@ -24,7 +24,10 @@ use diskann_benchmark_runner::{ utils::{datatype::AsDataType, percentiles, MicroSeconds}, Benchmark, Checkpoint, Registry, }; -use diskann_utils::{future::SendFuture, views::Matrix}; +use diskann_utils::{ + future::SendFuture, + views::rowmajor::{self, Matrix}, +}; use diskann_vector::{distance::Metric, PreprocessedDistanceFunction}; use half::f16; use serde::Serialize; @@ -54,7 +57,7 @@ pub(super) fn register_benchmarks(registry: &mut Registry) -> anyhow::Result<()> /// A minimal in-memory provider for flat search benchmarks. struct InMemProvider { - data: Arc>, + data: Arc>, } struct Flat { @@ -102,7 +105,7 @@ where // Load dataset writeln!(output, "Loading dataset...")?; - let data: Matrix = datafiles::load_dataset(datafiles::BinFile(&input.data))?; + let data: rowmajor::Owned = datafiles::load_dataset(datafiles::BinFile(&input.data))?; let nrows = data.nrows(); let ncols = data.ncols(); anyhow::ensure!( @@ -118,7 +121,7 @@ where let provider = InMemProvider { data: data.clone() }; // Load queries and groundtruth - let queries: Matrix = + let queries: rowmajor::Owned = datafiles::load_dataset(datafiles::BinFile(&input.search.queries))?; let groundtruth = datafiles::load_groundtruth( datafiles::BinFile(&input.search.groundtruth), @@ -184,7 +187,7 @@ where /// The visitor that iterates over all vectors in the provider. struct Visitor<'a, T: VectorRepr> { - data: &'a Matrix, + data: &'a rowmajor::Owned, computer: T::QueryDistance, } @@ -209,7 +212,7 @@ impl DistancesUnordered for Visitor<'_, T> { F: Send + FnMut(Self::Id, f32), { async move { - for (i, vector) in self.data.row_iter().enumerate() { + for (i, vector) in self.data.rows().enumerate() { let dist = self.computer.evaluate_similarity(vector); f(i as u32, dist); } @@ -225,7 +228,7 @@ impl DistancesUnordered for Visitor<'_, T> { /// Wraps a flat-search provider and queries to implement [`search::Search`]. struct Searcher { provider: InMemProvider, - queries: Matrix, + queries: rowmajor::Owned, metric: Metric, } @@ -284,12 +287,12 @@ where /// Aggregates results from multiple flat search runs, computing recall metrics. struct Aggregator<'a> { - groundtruth: &'a Matrix, + groundtruth: &'a rowmajor::Owned, recall_k: usize, } impl<'a> Aggregator<'a> { - fn new(groundtruth: &'a Matrix, recall_k: usize) -> Self { + fn new(groundtruth: &'a rowmajor::Owned, recall_k: usize) -> Self { Self { groundtruth, recall_k, diff --git a/diskann-benchmark/src/index/benchmarks.rs b/diskann-benchmark/src/index/benchmarks.rs index 33f5fca649..3aad4d8275 100644 --- a/diskann-benchmark/src/index/benchmarks.rs +++ b/diskann-benchmark/src/index/benchmarks.rs @@ -31,7 +31,7 @@ use diskann_providers::{ }; use diskann_utils::{ future::AsyncFriendly, - views::{Matrix, MatrixView}, + views::rowmajor::{self, Matrix}, }; use half::f16; @@ -364,7 +364,7 @@ type Index = Arc>; pub(crate) fn run_build( input: &IndexBuild, build_strategy: B, - data: Option>>, + data: Option>>, output: &mut dyn Output, create: CF, build: BF, @@ -372,13 +372,13 @@ pub(crate) fn run_build( where DP: DataProvider + for<'a> provider::SetElement<&'a [T]>, - CF: FnOnce(MatrixView) -> anyhow::Result>>, + CF: FnOnce(rowmajor::Ref) -> anyhow::Result>>, T: diskann::graph::SampleableForStart + std::fmt::Debug + Copy + AsyncFriendly + bytemuck::Pod, B: for<'a> glue::SearchStrategy<'a, DP, &'a [T]> + Clone + Send + Sync, BF: FnOnce( Index, B, - Arc>, + Arc>, &IndexBuild, &mut dyn Output, ) -> anyhow::Result, @@ -491,7 +491,7 @@ where ) -> anyhow::Result { let topk = phase.as_topk()?; - let queries: Arc> = + let queries: Arc> = Arc::new(datafiles::load_dataset(datafiles::BinFile(&topk.queries))?); // compute the maximum value of k used in any search @@ -548,7 +548,7 @@ where strategy: &Strategy, ) -> anyhow::Result { let range = phase.as_range()?; - let queries: Arc> = + let queries: Arc> = Arc::new(datafiles::load_dataset(datafiles::BinFile(&range.queries))?); let groundtruth = @@ -599,7 +599,7 @@ where ) -> anyhow::Result { let filtered_range = phase.as_filtered_range()?; - let queries: Arc> = Arc::new(datafiles::load_dataset( + let queries: Arc> = Arc::new(datafiles::load_dataset( datafiles::BinFile(&filtered_range.queries), )?); @@ -665,7 +665,7 @@ where ) -> anyhow::Result { let beta_filter = phase.as_topk_beta_filter()?; - let queries: Arc> = Arc::new(datafiles::load_dataset( + let queries: Arc> = Arc::new(datafiles::load_dataset( datafiles::BinFile(&beta_filter.queries), )?); @@ -731,7 +731,7 @@ where ) -> anyhow::Result { let multihop = phase.as_topk_multihop_filter()?; - let queries: Arc> = Arc::new(datafiles::load_dataset( + let queries: Arc> = Arc::new(datafiles::load_dataset( datafiles::BinFile(&multihop.queries), )?); @@ -793,7 +793,7 @@ where ) -> anyhow::Result { let inline = phase.as_topk_inline_filter()?; - let queries: Arc> = Arc::new(datafiles::load_dataset( + let queries: Arc> = Arc::new(datafiles::load_dataset( datafiles::BinFile(&inline.queries), )?); diff --git a/diskann-benchmark/src/index/bftree/full_precision.rs b/diskann-benchmark/src/index/bftree/full_precision.rs index fcd984887d..7238662a52 100644 --- a/diskann-benchmark/src/index/bftree/full_precision.rs +++ b/diskann-benchmark/src/index/bftree/full_precision.rs @@ -19,6 +19,7 @@ use diskann_providers::{ model::graph::provider::async_::common::FullPrecision, storage::{FileStorageProvider, SaveWith}, }; +use diskann_utils::views::rowmajor::Matrix; use crate::{ index::{ diff --git a/diskann-benchmark/src/index/bftree/full_precision_streaming.rs b/diskann-benchmark/src/index/bftree/full_precision_streaming.rs index 3d26773cd1..c752bf58b7 100644 --- a/diskann-benchmark/src/index/bftree/full_precision_streaming.rs +++ b/diskann-benchmark/src/index/bftree/full_precision_streaming.rs @@ -24,7 +24,7 @@ use diskann_providers::{ model::graph::provider::async_::common::FullPrecision, storage::{FileStorageProvider, SaveWith}, }; -use diskann_utils::views::{Matrix, MatrixView}; +use diskann_utils::views::rowmajor::{self, Matrix}; use crate::{ index::{ @@ -68,10 +68,10 @@ impl BfTreeStream where T: VectorRepr, { - fn insert_(&self, data: MatrixView<'_, T>, slots: &[u32]) -> anyhow::Result { + fn insert_(&self, data: rowmajor::Ref<'_, T>, slots: &[u32]) -> anyhow::Result { let runner = benchmark_core::build::graph::SingleInsert::new( self.index.clone(), - Arc::new(data.to_owned()), + Arc::new(data.to_rowmajor_owned()), FullPrecision, benchmark_core::build::ids::Slice::new(slots.into()), ); @@ -94,7 +94,7 @@ where fn search( &self, - queries: Arc>, + queries: Arc>, groundtruth: &dyn Rows, ) -> anyhow::Result { let knn = benchmark_core::search::graph::KNN::new( @@ -113,11 +113,11 @@ where Ok(StreamStats::Search(results)) } - fn insert(&self, data: MatrixView<'_, T>, slots: &[u32]) -> anyhow::Result { + fn insert(&self, data: rowmajor::Ref<'_, T>, slots: &[u32]) -> anyhow::Result { Ok(StreamStats::Insert(self.insert_(data, slots)?)) } - fn replace(&self, data: MatrixView<'_, T>, slots: &[u32]) -> anyhow::Result { + fn replace(&self, data: rowmajor::Ref<'_, T>, slots: &[u32]) -> anyhow::Result { Ok(StreamStats::Replace(self.insert_(data, slots)?)) } @@ -247,7 +247,7 @@ where ))?); let config = input.try_as_config()?.build()?; - let params = input.bftree_parameters(max_points, data.ncols())?; + let params = input.bftree_parameters(max_points, Matrix::ncols(&data))?; let start_points = input .build() .start_point_strategy() diff --git a/diskann-benchmark/src/index/bftree/spherical.rs b/diskann-benchmark/src/index/bftree/spherical.rs index 3c1939b792..a5d24c15cb 100644 --- a/diskann-benchmark/src/index/bftree/spherical.rs +++ b/diskann-benchmark/src/index/bftree/spherical.rs @@ -17,7 +17,7 @@ use diskann_providers::{ storage::{FileStorageProvider, SaveWith}, }; use diskann_quantization::alloc::GlobalAllocator; -use diskann_utils::views::Matrix; +use diskann_utils::views::rowmajor::{self, Matrix}; use rand::SeedableRng; use crate::{ @@ -112,7 +112,7 @@ impl Benchmark for BfTreeSpherical { writeln!(output, "{}", input)?; let build = input.build(); - let data: Arc> = + let data: Arc> = Arc::new(datafiles::load_dataset(datafiles::BinFile(build.data()))?); // 1. Train the spherical quantizer. diff --git a/diskann-benchmark/src/index/bftree/spherical_streaming.rs b/diskann-benchmark/src/index/bftree/spherical_streaming.rs index 24b74b99f9..dd9d008910 100644 --- a/diskann-benchmark/src/index/bftree/spherical_streaming.rs +++ b/diskann-benchmark/src/index/bftree/spherical_streaming.rs @@ -26,7 +26,7 @@ use diskann_providers::{ storage::{FileStorageProvider, SaveWith}, }; use diskann_quantization::alloc::GlobalAllocator; -use diskann_utils::views::{Matrix, MatrixView}; +use diskann_utils::views::rowmajor::{self, Matrix}; use rand::SeedableRng; use crate::{ @@ -59,10 +59,10 @@ struct BfTreeSQStream { } impl BfTreeSQStream { - fn insert_(&self, data: MatrixView<'_, f32>, slots: &[u32]) -> anyhow::Result { + fn insert_(&self, data: rowmajor::Ref<'_, f32>, slots: &[u32]) -> anyhow::Result { let runner = benchmark_core::build::graph::SingleInsert::new( self.index.clone(), - Arc::new(data.to_owned()), + Arc::new(data.to_rowmajor_owned()), Quantized, benchmark_core::build::ids::Slice::new(slots.into()), ); @@ -82,7 +82,7 @@ impl ManagedStream for BfTreeSQStream { fn search( &self, - queries: Arc>, + queries: Arc>, groundtruth: &dyn Rows, ) -> anyhow::Result { let knn = benchmark_core::search::graph::KNN::new( @@ -101,11 +101,11 @@ impl ManagedStream for BfTreeSQStream { Ok(StreamStats::Search(results)) } - fn insert(&self, data: MatrixView<'_, f32>, slots: &[u32]) -> anyhow::Result { + fn insert(&self, data: rowmajor::Ref<'_, f32>, slots: &[u32]) -> anyhow::Result { Ok(StreamStats::Insert(self.insert_(data, slots)?)) } - fn replace(&self, data: MatrixView<'_, f32>, slots: &[u32]) -> anyhow::Result { + fn replace(&self, data: rowmajor::Ref<'_, f32>, slots: &[u32]) -> anyhow::Result { Ok(StreamStats::Replace(self.insert_(data, slots)?)) } @@ -255,7 +255,7 @@ fn bftree_sq_streaming_impl( }; let config = input.try_as_config()?.build()?; - let params = input.bftree_parameters(max_points, data.ncols())?; + let params = input.bftree_parameters(max_points, Matrix::ncols(&data))?; let start_points = input .build() .start_point_strategy() diff --git a/diskann-benchmark/src/index/build.rs b/diskann-benchmark/src/index/build.rs index b4bf186e68..363952babd 100644 --- a/diskann-benchmark/src/index/build.rs +++ b/diskann-benchmark/src/index/build.rs @@ -22,7 +22,7 @@ use diskann_providers::{ }; use diskann_utils::{ future::AsyncFriendly, - views::{Matrix, MatrixView}, + views::rowmajor::{self, Matrix}, }; use indicatif::{ProgressBar, ProgressStyle}; use serde::Serialize; @@ -35,7 +35,7 @@ use crate::inputs::graph_index::IndexBuild; pub(crate) fn set_start_points( provider: &DP, - data: MatrixView<'_, T>, + data: rowmajor::Ref<'_, T>, start_strategy: StartPointStrategy, ) -> ANNResult<()> where @@ -43,7 +43,7 @@ where T: diskann::graph::SampleableForStart + AsyncFriendly, { let start_points = start_strategy.compute(data).map_err(ANNError::new)?; - provider.set_start_points(start_points.row_iter()) + provider.set_start_points(start_points.rows()) } /////////// @@ -53,7 +53,7 @@ where pub(crate) fn single_or_multi_insert( index: Arc>, strategy: S, - data: Arc>, + data: Arc>, input: &IndexBuild, output: &mut dyn Output, ) -> anyhow::Result @@ -109,7 +109,7 @@ where pub(crate) fn only_single_insert( index: Arc>, strategy: S, - data: Arc>, + data: Arc>, input: &IndexBuild, output: &mut dyn Output, ) -> anyhow::Result diff --git a/diskann-benchmark/src/index/inmem/product.rs b/diskann-benchmark/src/index/inmem/product.rs index 931ca725ce..5956d8ebe3 100644 --- a/diskann-benchmark/src/index/inmem/product.rs +++ b/diskann-benchmark/src/index/inmem/product.rs @@ -49,7 +49,7 @@ mod imp { IndexConfiguration, }, }; - use diskann_utils::views::{Matrix, MatrixView}; + use diskann_utils::views::rowmajor::{self, Matrix}; use diskann_benchmark_runner::{ benchmark::{MatchContext, Score}, @@ -172,12 +172,12 @@ mod imp { (Arc::new(index), None::, MicroSeconds::new(0)) } IndexSource::Build(build) => { - let data: Arc> = + let data: Arc> = Arc::new(datafiles::load_dataset(datafiles::BinFile(build.data()))?); let start = std::time::Instant::now(); let table = { - let train_data = Matrix::try_from( + let train_data = rowmajor::Owned::try_from_data( (&*T::as_f32(data.as_slice())?).into(), data.nrows(), data.ncols(), @@ -192,7 +192,7 @@ mod imp { )? }; - let create_index = |data_view: MatrixView| { + let create_index = |data_view: rowmajor::Ref| { let index = diskann_async::new_quant_index::( input.try_as_config()?.build()?, input.inmem_parameters(data_view.nrows(), data_view.ncols())?, diff --git a/diskann-benchmark/src/index/inmem/scalar.rs b/diskann-benchmark/src/index/inmem/scalar.rs index e0287638b1..da0c8a3c17 100644 --- a/diskann-benchmark/src/index/inmem/scalar.rs +++ b/diskann-benchmark/src/index/inmem/scalar.rs @@ -73,7 +73,7 @@ mod imp { graph::provider::async_::{common, inmem}, }, }; - use diskann_utils::views::{Matrix, MatrixView}; + use diskann_utils::views::rowmajor::{self, Matrix}; use crate::{ index::{ @@ -239,7 +239,7 @@ mod imp { (Arc::new(index), None::, MicroSeconds::new(0)) } IndexSource::Build(build) => { - let data: Arc> = + let data: Arc> = Arc::new(datafiles::load_dataset(datafiles::BinFile(build.data()))?); let start = std::time::Instant::now(); @@ -250,7 +250,7 @@ mod imp { )?, ) .train(data.as_view()); - let create_index = |data_view: MatrixView<$T>| { + let create_index = |data_view: rowmajor::Ref<$T>| { let index = diskann_async::new_quant_index::<$T, _, _>( input.try_as_config()?.build()?, input diff --git a/diskann-benchmark/src/index/inmem/spherical.rs b/diskann-benchmark/src/index/inmem/spherical.rs index 5ee77166a3..07cdad29c3 100644 --- a/diskann-benchmark/src/index/inmem/spherical.rs +++ b/diskann-benchmark/src/index/inmem/spherical.rs @@ -78,7 +78,7 @@ mod imp { model::graph::provider::async_::{common, inmem}, }; use diskann_quantization::alloc::GlobalAllocator; - use diskann_utils::views::Matrix; + use diskann_utils::views::rowmajor::{self, Matrix}; use rand::SeedableRng; use serde::Serialize; use std::{io::Write, sync::Arc}; @@ -259,7 +259,7 @@ mod imp { let build = &input.build; - let data: Arc> = + let data: Arc> = Arc::new(datafiles::load_dataset(datafiles::BinFile(build.data()))?); let start = std::time::Instant::now(); @@ -360,7 +360,7 @@ mod imp { // compute the maximum value of k used in any search let max_k = topk.max_k(); - let queries: Arc> = + let queries: Arc> = Arc::new(datafiles::load_dataset(datafiles::BinFile(&topk.queries))?); let groundtruth = @@ -405,7 +405,7 @@ mod imp { ) -> anyhow::Result { let range = phase.as_range()?; - let queries: Arc> = + let queries: Arc> = Arc::new(datafiles::load_dataset(datafiles::BinFile(&range.queries))?); let groundtruth = @@ -447,9 +447,9 @@ mod imp { ) -> anyhow::Result { let filtered_range = phase.as_filtered_range()?; - let queries: Arc> = Arc::new(datafiles::load_dataset(datafiles::BinFile( - &filtered_range.queries, - ))?); + let queries: Arc> = Arc::new(datafiles::load_dataset( + datafiles::BinFile(&filtered_range.queries), + )?); let groundtruth = datafiles::load_range_groundtruth(datafiles::BinFile(&filtered_range.groundtruth))?; @@ -504,9 +504,9 @@ mod imp { ) -> anyhow::Result { let betafilter = phase.as_topk_beta_filter()?; - let queries: Arc> = Arc::new(datafiles::load_dataset(datafiles::BinFile( - &betafilter.queries, - ))?); + let queries: Arc> = Arc::new(datafiles::load_dataset( + datafiles::BinFile(&betafilter.queries), + )?); let groundtruth = datafiles::load_range_groundtruth(datafiles::BinFile(&betafilter.groundtruth))?; @@ -559,9 +559,9 @@ mod imp { ) -> anyhow::Result { let multihop = phase.as_topk_multihop_filter()?; - let queries: Arc> = Arc::new(datafiles::load_dataset(datafiles::BinFile( - &multihop.queries, - ))?); + let queries: Arc> = Arc::new(datafiles::load_dataset( + datafiles::BinFile(&multihop.queries), + )?); let groundtruth = datafiles::load_range_groundtruth(datafiles::BinFile(&multihop.groundtruth))?; @@ -613,9 +613,9 @@ mod imp { ) -> anyhow::Result { let inline = phase.as_topk_inline_filter()?; - let queries: Arc> = Arc::new(datafiles::load_dataset(datafiles::BinFile( - &inline.queries, - ))?); + let queries: Arc> = Arc::new(datafiles::load_dataset( + datafiles::BinFile(&inline.queries), + )?); let groundtruth = datafiles::load_range_groundtruth(datafiles::BinFile(&inline.groundtruth))?; diff --git a/diskann-benchmark/src/index/inmem2.rs b/diskann-benchmark/src/index/inmem2.rs index b989917e3b..0d574b16d8 100644 --- a/diskann-benchmark/src/index/inmem2.rs +++ b/diskann-benchmark/src/index/inmem2.rs @@ -38,7 +38,7 @@ use diskann_quantization::{ alloc::{GlobalAllocator, Poly}, spherical::iface, }; -use diskann_utils::views::{Matrix, MatrixView}; +use diskann_utils::views::rowmajor::{self, Matrix}; use diskann_vector::distance::Metric; use half::f16; use serde::{Deserialize, Serialize}; @@ -436,7 +436,7 @@ impl Spherical { fn train( &self, - data: MatrixView<'_, f32>, + data: rowmajor::Ref<'_, f32>, metric: Metric, ) -> anyhow::Result> { use diskann_quantization::{algorithms::transforms, spherical}; @@ -680,7 +680,7 @@ where writeln!(output, "{input}\n")?; // Load data. - let data: Arc> = Arc::new(datafiles::load_dataset(datafiles::BinFile( + let data: Arc> = Arc::new(datafiles::load_dataset(datafiles::BinFile( &input.data.data, ))?); @@ -729,9 +729,9 @@ where checkpoint.checkpoint(&total_build_time)?; // Search. - let queries: Arc> = Arc::new(datafiles::load_dataset(datafiles::BinFile( - &input.search.queries, - ))?); + let queries: Arc> = Arc::new(datafiles::load_dataset( + datafiles::BinFile(&input.search.queries), + )?); let max_k = input.search.maximum_recall_k(); let groundtruth = datafiles::load_groundtruth( datafiles::BinFile(&input.search.groundtruth), @@ -827,9 +827,9 @@ impl Benchmark for SphericalBuild { let spherical = input.quantization.as_spherical().unwrap(); // Load data. - let data: Arc> = Arc::new(datafiles::load_dataset(datafiles::BinFile( - &input.data.data, - ))?); + let data: Arc> = Arc::new(datafiles::load_dataset( + datafiles::BinFile(&input.data.data), + )?); let dim = data.ncols(); let num_points = data.nrows(); @@ -879,9 +879,9 @@ impl Benchmark for SphericalBuild { checkpoint.checkpoint(&total_build_time)?; // Search. - let queries: Arc> = Arc::new(datafiles::load_dataset(datafiles::BinFile( - &input.search.queries, - ))?); + let queries: Arc> = Arc::new(datafiles::load_dataset( + datafiles::BinFile(&input.search.queries), + )?); let max_k = input.search.maximum_recall_k(); let groundtruth = datafiles::load_groundtruth( datafiles::BinFile(&input.search.groundtruth), @@ -1204,10 +1204,11 @@ where let max_points = runbook.max_points(); // Load the dataset (consumed by `WithData`) and queries. - let dataset: Matrix = datafiles::load_dataset(datafiles::BinFile(&input.data.data))?; - let queries: Arc> = Arc::new(datafiles::load_dataset(datafiles::BinFile( - &input.search.queries, - ))?); + let dataset: rowmajor::Owned = + datafiles::load_dataset(datafiles::BinFile(&input.data.data))?; + let queries: Arc> = Arc::new(datafiles::load_dataset( + datafiles::BinFile(&input.search.queries), + )?); // Compute the medoid of the dataset as the single start point. let start = StartPointStrategy::Medoid.compute(dataset.as_view())?; @@ -1296,7 +1297,7 @@ where { fn insert_( &mut self, - data: MatrixView<'_, T>, + data: rowmajor::Ref<'_, T>, ids: Range, ) -> anyhow::Result { anyhow::ensure!( @@ -1308,7 +1309,7 @@ where let runner = build_core::graph::SingleInsert::new( self.index.clone(), - Arc::new(data.to_owned()), + Arc::new(data.to_rowmajor_owned()), Strategy, build_core::ids::Range::::new(ids.start as u32..ids.end as u32), ); @@ -1331,7 +1332,7 @@ where fn search( &mut self, - (queries, groundtruth): (Arc>, &dyn recall::Rows), + (queries, groundtruth): (Arc>, &dyn recall::Rows), ) -> anyhow::Result { let knn = benchmark_core::search::graph::KNN::new( self.index.clone(), @@ -1352,7 +1353,7 @@ where fn insert( &mut self, - (data, ids): (MatrixView<'_, T>, Range), + (data, ids): (rowmajor::Ref<'_, T>, Range), ) -> anyhow::Result { self.insert_(data, ids).map(StreamStats::Insert) } @@ -1380,7 +1381,7 @@ where fn replace( &mut self, - (data, ids): (MatrixView<'_, T>, Range), + (data, ids): (rowmajor::Ref<'_, T>, Range), ) -> anyhow::Result { use diskann::provider::Delete; diff --git a/diskann-benchmark/src/index/streaming/full_precision.rs b/diskann-benchmark/src/index/streaming/full_precision.rs index ced09ed3ee..6f402db763 100644 --- a/diskann-benchmark/src/index/streaming/full_precision.rs +++ b/diskann-benchmark/src/index/streaming/full_precision.rs @@ -19,7 +19,7 @@ use diskann_providers::model::graph::provider::async_::{ }; use diskann_utils::{ future::AsyncFriendly, - views::{Matrix, MatrixView}, + views::rowmajor::{self, Matrix}, }; use super::{ @@ -64,10 +64,10 @@ where T: VectorRepr, { // Common code-path for both inserts and replace. - fn insert_(&self, data: MatrixView<'_, T>, slots: &[u32]) -> anyhow::Result { + fn insert_(&self, data: rowmajor::Ref<'_, T>, slots: &[u32]) -> anyhow::Result { let runner = diskann_benchmark_core::build::graph::SingleInsert::new( self.index.clone(), - Arc::new(data.to_owned()), + Arc::new(data.to_rowmajor_owned()), common::FullPrecision, diskann_benchmark_core::build::ids::Slice::new(slots.into()), ); @@ -90,7 +90,7 @@ where fn search( &self, - queries: Arc>, + queries: Arc>, groundtruth: &dyn Rows, ) -> anyhow::Result { let knn = diskann_benchmark_core::search::graph::KNN::new( @@ -109,11 +109,11 @@ where Ok(StreamStats::Search(results)) } - fn insert(&self, data: MatrixView<'_, T>, slots: &[u32]) -> anyhow::Result { + fn insert(&self, data: rowmajor::Ref<'_, T>, slots: &[u32]) -> anyhow::Result { Ok(StreamStats::Insert(self.insert_(data, slots)?)) } - fn replace(&self, data: MatrixView<'_, T>, slots: &[u32]) -> anyhow::Result { + fn replace(&self, data: rowmajor::Ref<'_, T>, slots: &[u32]) -> anyhow::Result { Ok(StreamStats::Replace(self.insert_(data, slots)?)) } diff --git a/diskann-benchmark/src/index/streaming/managed.rs b/diskann-benchmark/src/index/streaming/managed.rs index 2d76f3b4f6..2049d42912 100644 --- a/diskann-benchmark/src/index/streaming/managed.rs +++ b/diskann-benchmark/src/index/streaming/managed.rs @@ -11,7 +11,7 @@ use diskann_benchmark_core::{ streaming::{self, executors}, }; use diskann_benchmark_runner::{timed, utils::MicroSeconds}; -use diskann_utils::views::{Matrix, MatrixView}; +use diskann_utils::views::rowmajor; use crate::utils::streaming::TagSlotManager; @@ -45,15 +45,15 @@ pub(crate) trait ManagedStream { /// See: [`streaming::Stream::search`]. fn search( &self, - queries: Arc>, + queries: Arc>, groundtruth: &dyn Rows, ) -> anyhow::Result; /// See: [`streaming::Stream::insert`]. - fn insert(&self, data: MatrixView<'_, T>, slots: &[u32]) -> anyhow::Result; + fn insert(&self, data: rowmajor::Ref<'_, T>, slots: &[u32]) -> anyhow::Result; /// See: [`streaming::Stream::replace`]. - fn replace(&self, data: MatrixView<'_, T>, slots: &[u32]) -> anyhow::Result; + fn replace(&self, data: rowmajor::Ref<'_, T>, slots: &[u32]) -> anyhow::Result; /// See: [`streaming::Stream::delete`]. fn delete(&self, slots: &[u32]) -> anyhow::Result; @@ -101,7 +101,7 @@ where fn search( &mut self, - (queries, groundtruth): (Arc>, &dyn Rows), + (queries, groundtruth): (Arc>, &dyn Rows), ) -> anyhow::Result { // Translate the groundtruth to the appropriate internal IDs. let (overhead, _): (_, ()) = timed! { @@ -125,7 +125,7 @@ where fn insert( &mut self, - (data, tags): (MatrixView<'_, T>, Range), + (data, tags): (rowmajor::Ref<'_, T>, Range), ) -> anyhow::Result { let (overhead_get, slots) = timed!(self.book_keeping.get_n_empty_slots(tags.len())?); let output = self.stream.insert(data, &slots)?; @@ -136,7 +136,7 @@ where fn replace( &mut self, - (data, tags): (MatrixView<'_, T>, Range), + (data, tags): (rowmajor::Ref<'_, T>, Range), ) -> anyhow::Result { let (overhead, slots) = timed!(self.book_keeping.find_slots_by_tags(tags)?); self.stream diff --git a/diskann-benchmark/src/utils/datafiles.rs b/diskann-benchmark/src/utils/datafiles.rs index ddc507c2c0..4244e506a2 100644 --- a/diskann-benchmark/src/utils/datafiles.rs +++ b/diskann-benchmark/src/utils/datafiles.rs @@ -10,7 +10,7 @@ use bit_set::BitSet; use diskann::utils::IntoUsize; use diskann_benchmark_runner::utils::datatype::DataType; use diskann_providers::storage::StorageReadProvider; -use diskann_utils::views::Matrix; +use diskann_utils::views::rowmajor::{self, Matrix, MatrixMut}; use serde::{Deserialize, Serialize}; pub(crate) struct BinFile<'a>(pub(crate) &'a Path); @@ -18,7 +18,7 @@ pub(crate) struct BinFile<'a>(pub(crate) &'a Path); /// Load a dataset or query set in `.bin` form from disk and return the result as a /// row-major matrix. #[inline(never)] -pub(crate) fn load_dataset(path: BinFile<'_>) -> anyhow::Result> +pub(crate) fn load_dataset(path: BinFile<'_>) -> anyhow::Result> where T: Copy + bytemuck::Pod, { @@ -29,13 +29,13 @@ where Ok(data) } -/// Helper trait to load a `Matrix` from source files that potentially have a different +/// Helper trait to load a `rowmajor::Owned` from source files that potentially have a different /// type. pub(crate) trait ConvertingLoad: Sized { /// Return an error if the provided `data_type` cannot be loaded and converted to `Self`. fn check_converting_load(data_type: DataType) -> anyhow::Result<()>; - /// Attempt to load the data at `path` as a `Matrix` assuming the on-disk + /// Attempt to load the data at `path` as a `rowmajor::Owned` assuming the on-disk /// representation has the encoding specified by `data_type`. /// /// If `data_type` is not compatible with `Self`, return an error. @@ -44,7 +44,10 @@ pub(crate) trait ConvertingLoad: Sized { feature = "minmax-quantization", feature = "product-quantization" ))] - fn converting_load(path: BinFile<'_>, data_type: DataType) -> anyhow::Result>; + fn converting_load( + path: BinFile<'_>, + data_type: DataType, + ) -> anyhow::Result>; } impl ConvertingLoad for f32 { @@ -69,14 +72,20 @@ impl ConvertingLoad for f32 { feature = "minmax-quantization", feature = "product-quantization" ))] - fn converting_load(path: BinFile<'_>, data_type: DataType) -> anyhow::Result> { + fn converting_load( + path: BinFile<'_>, + data_type: DataType, + ) -> anyhow::Result> { #[inline(never)] - fn convert(from: diskann_utils::views::MatrixView) -> anyhow::Result> + fn convert( + from: diskann_utils::views::rowmajor::Ref, + ) -> anyhow::Result> where U: Default + Clone + From, T: Copy, { - let mut to = Matrix::try_from_element(from.nrows(), from.ncols(), U::default())?; + let mut to = + rowmajor::Owned::try_from_element(from.nrows(), from.ncols(), U::default())?; std::iter::zip(to.as_mut_slice().iter_mut(), from.as_slice().iter()) .for_each(|(t, f)| *t = (*f).into()); Ok(to) @@ -95,7 +104,10 @@ impl ConvertingLoad for f32 { } /// Load a groundtruth set from disk and return the result as a row-major matrix. -pub(crate) fn load_groundtruth(path: BinFile<'_>, k: Option) -> anyhow::Result> { +pub(crate) fn load_groundtruth( + path: BinFile<'_>, + k: Option, +) -> anyhow::Result> { let provider = diskann_providers::storage::FileStorageProvider; let mut file = provider .open_reader(&path.0.to_string_lossy()) @@ -111,7 +123,7 @@ pub(crate) fn load_groundtruth(path: BinFile<'_>, k: Option) -> anyhow::R (num_points, dim) }; - let mut groundtruth = Matrix::::from_element(num_points, dim, 0); + let mut groundtruth = rowmajor::Owned::::from_element(num_points, dim, 0); let groundtruth_slice: &mut [u8] = bytemuck::cast_slice_mut(groundtruth.as_mut_slice()); file.read_exact(groundtruth_slice)?; diff --git a/diskann-bftree/src/provider.rs b/diskann-bftree/src/provider.rs index 54dfd0103a..8078770430 100644 --- a/diskann-bftree/src/provider.rs +++ b/diskann-bftree/src/provider.rs @@ -38,7 +38,7 @@ use diskann::{ use diskann_utils::{ future::{AsyncFriendly, SendFuture}, lazy_format, - views::MatrixView, + views::rowmajor::{self, Matrix}, }; use diskann_vector::{distance::Metric, DistanceFunction, PreprocessedDistanceFunction}; @@ -98,7 +98,7 @@ use diskann_providers::storage::{LoadWith, SaveWith, StorageReadProvider, Storag /// }; /// use diskann_bftree::NoStore; /// use diskann_vector::distance::Metric; -/// use diskann_utils::views::Matrix; +/// use diskann_utils::views::rowmajor::{Owned, Matrix}; /// use bf_tree::Config; /// use std::num::NonZeroUsize; /// @@ -116,7 +116,7 @@ use diskann_providers::storage::{LoadWith, SaveWith, StorageReadProvider, Storag /// }; /// /// // Create a table that supports 5 points and 1 start point. -/// let start_points = Matrix::from_element(1, 4, 0.0f32); +/// let start_points = Owned::from_element(1, 4, 0.0f32); /// let provider = BfTreeProvider::::new( /// parameters, /// start_points.as_view(), @@ -135,7 +135,7 @@ use diskann_providers::storage::{LoadWith, SaveWith, StorageReadProvider, Storag /// algorithms::TransformKind, /// spherical::{SphericalQuantizer, SupportedMetric, PreScale}, /// }; -/// use diskann_utils::views::Matrix; +/// use diskann_utils::views::rowmajor::{Owned, Matrix}; /// use diskann_bftree::provider::{ /// BfTreeProvider, BfTreeProviderParameters /// }; @@ -146,7 +146,7 @@ use diskann_providers::storage::{LoadWith, SaveWith, StorageReadProvider, Storag /// use rand::SeedableRng; /// /// let dim = 4; -/// let data = Matrix::from_element(4, dim, 1.0f32); +/// let data = Owned::from_element(4, dim, 1.0f32); /// let mut rng = StdRng::seed_from_u64(42); /// let quantizer = SphericalQuantizer::train( /// data.as_view(), TransformKind::Null, @@ -171,7 +171,7 @@ use diskann_providers::storage::{LoadWith, SaveWith, StorageReadProvider, Storag /// }; /// /// // Create a table that supports 5 points and 1 start point. -/// let start_points = Matrix::from_element(1, 4, 0.0f32); +/// let start_points = Owned::from_element(1, 4, 0.0f32); /// let provider = BfTreeProvider::::new( /// parameters, /// start_points.as_view(), @@ -345,7 +345,7 @@ where /// * `Self: StartPoint` - The provider must implement the `StartPoint` trait. pub fn new( params: BfTreeProviderParameters, - start_points: MatrixView<'_, T>, + start_points: rowmajor::Ref<'_, T>, quant_precursor: TQ, ) -> ANNResult where @@ -735,7 +735,8 @@ pub trait StartPoint { /// This method is internal and should not be called directly by users. /// Use `BfTreeProvider::new` instead. #[doc(hidden)] - fn set_start_points(&self, hidden: Hidden, start_points: MatrixView<'_, T>) -> ANNResult<()>; + fn set_start_points(&self, hidden: Hidden, start_points: rowmajor::Ref<'_, T>) + -> ANNResult<()>; } //////////////////// @@ -751,7 +752,11 @@ where T: VectorRepr, I: BfTreeId, { - fn set_start_points(&self, _hidden: Hidden, start_points: MatrixView<'_, T>) -> ANNResult<()> { + fn set_start_points( + &self, + _hidden: Hidden, + start_points: rowmajor::Ref<'_, T>, + ) -> ANNResult<()> { let start_point_ids: Vec = self.full_vectors.starting_points()?; if start_points.nrows() != start_point_ids.len() { return Err(ANNError::message(format!( @@ -762,7 +767,7 @@ where } let mut scratch = self.neighbor_provider.scratch(&self.locks); - for (id, v) in std::iter::zip(start_point_ids, start_points.row_iter()) { + for (id, v) in std::iter::zip(start_point_ids, start_points.rows()) { // Set the full-precision vector self.full_vectors.set_vector_sync(id.as_index(), v)?; self.quant_vectors.set_vector_sync(id.as_index(), v)?; @@ -783,7 +788,11 @@ where T: VectorRepr, I: BfTreeId, { - fn set_start_points(&self, _hidden: Hidden, start_points: MatrixView<'_, T>) -> ANNResult<()> { + fn set_start_points( + &self, + _hidden: Hidden, + start_points: rowmajor::Ref<'_, T>, + ) -> ANNResult<()> { let start_point_ids: Vec = self.full_vectors.starting_points()?; if start_points.nrows() != start_point_ids.len() { return Err(ANNError::message(format!( @@ -794,7 +803,7 @@ where } let mut scratch = self.neighbor_provider.scratch(&self.locks); - for (id, v) in std::iter::zip(start_point_ids, start_points.row_iter()) { + for (id, v) in std::iter::zip(start_point_ids, start_points.rows()) { // Set the full-precision vector self.full_vectors.set_vector_sync(id.as_index(), v)?; // Initialize empty neighbor list @@ -2129,10 +2138,10 @@ mod tests { neighbor::BackInserter, }; use diskann_providers::storage::FileStorageProvider; - use diskann_utils::views::{Matrix, RowCol}; + use diskann_utils::views::rowmajor::RowCol; fn create_quant_index() -> Arc>> { - let start_point = Matrix::from_element(1, 5, 0.0f32); + let start_point = rowmajor::Owned::from_element(1, 5, 0.0f32); let dim = 5; let logical_max_degree = 6; let physical_max_degree = (logical_max_degree as f32 * 1.3) as u32; @@ -2212,7 +2221,7 @@ mod tests { /// that the `BfTreeProvider<_, _, u64>` path is functional and not merely compilable. #[tokio::test] async fn test_quantized_index_search_u64_ids() { - let start_point = Matrix::from_element(1, 5, 0.0f32); + let start_point = rowmajor::Owned::from_element(1, 5, 0.0f32); let dim = 5; let logical_max_degree = 6; let physical_max_degree = (logical_max_degree as f32 * 1.3) as u32; @@ -2294,7 +2303,7 @@ mod tests { /// full data/quant/neighbor/search stack keys on the complete 8-byte id. #[tokio::test] async fn test_quantized_index_search_u64_high_ids() { - let start_point = Matrix::from_element(1, 5, 0.0f32); + let start_point = rowmajor::Owned::from_element(1, 5, 0.0f32); let dim = 5; let logical_max_degree = 6; let physical_max_degree = (logical_max_degree as f32 * 1.3) as u32; @@ -2392,11 +2401,11 @@ mod tests { let index = create_quant_index(); let ctx = &DefaultContext; - let data = Matrix::from_fn(15, 5, |RowCol { row, .. }| row as f32); + let data = rowmajor::Owned::from_fn(15, 5, |RowCol { row, .. }| row as f32); let ids: Arc<[u32]> = (0u32..15).collect::>().into(); - let batch: Arc> = Arc::new(data); + let batch: Arc> = Arc::new(data); index - .multi_insert::>(Quantized, ctx, batch, ids) + .multi_insert::>(Quantized, ctx, batch, ids) .await .unwrap(); @@ -2475,7 +2484,7 @@ mod tests { } fn create_full_precision_index() -> Arc>> { - let start_point = Matrix::from_element(1, 5, 0.0f32); + let start_point = rowmajor::Owned::from_element(1, 5, 0.0f32); let logical_max_degree = 6; let physical_max_degree = (logical_max_degree as f32 * 1.3) as u32; let metric = Metric::L2; @@ -2632,7 +2641,7 @@ mod tests { let logical_max_degree = 32usize; let physical_max_degree = (logical_max_degree as f32 * 1.3) as u32; let metric = Metric::L2; - let start_point = Matrix::from_element(1, DIM, 0.0f32); + let start_point = rowmajor::Owned::from_element(1, DIM, 0.0f32); let provider: BfTreeProvider = BfTreeProvider::new( BfTreeProviderParameters { @@ -2814,7 +2823,7 @@ mod tests { let ctx = &DefaultContext; let num_start_points = 2; let dim = 5; - let start_points = Matrix::try_from( + let start_points = rowmajor::Owned::try_from_data( vec![0.0f32; dim] .into_iter() .chain(vec![0.5f32; dim]) @@ -2909,7 +2918,7 @@ mod tests { let num_start_points = 2; let dim = 3; - let start_points = Matrix::from_element(num_start_points, dim, 0.0f32); + let start_points = rowmajor::Owned::from_element(num_start_points, dim, 0.0f32); let provider = BfTreeProvider::::new( BfTreeProviderParameters { @@ -3046,7 +3055,7 @@ mod tests { use_snapshot: true, }; - let start_points = Matrix::from_element(num_start_points.into(), dim, 0.0f32); + let start_points = rowmajor::Owned::from_element(num_start_points.into(), dim, 0.0f32); // Create provider let provider = @@ -3178,7 +3187,7 @@ mod tests { use_snapshot: true, }; - let start_points = Matrix::from_element(num_start_points.into(), dim, 0.0f32); + let start_points = rowmajor::Owned::from_element(num_start_points.into(), dim, 0.0f32); // Create provider with quantization let provider = BfTreeProvider::::new( params.clone(), @@ -3297,7 +3306,7 @@ mod tests { let mut neighbor_config = Config::default(); neighbor_config.use_snapshot(true); - let start_points = Matrix::from_element(num_start_points.into(), dim, 0.0f32); + let start_points = rowmajor::Owned::from_element(num_start_points.into(), dim, 0.0f32); // In-memory config (no file path needed) let provider = BfTreeProvider::::new( BfTreeProviderParameters { @@ -3412,7 +3421,7 @@ mod tests { let mut quant_config = Config::default(); quant_config.use_snapshot(true); - let start_points = Matrix::from_element(num_start_points.into(), dim, 0.0f32); + let start_points = rowmajor::Owned::from_element(num_start_points.into(), dim, 0.0f32); let provider = BfTreeProvider::::new( BfTreeProviderParameters { max_points: num_points, @@ -3603,7 +3612,7 @@ mod tests { use_snapshot: true, }; - let start_points = Matrix::from_element(num_start_points.into(), dim, 0.0f32); + let start_points = rowmajor::Owned::from_element(num_start_points.into(), dim, 0.0f32); let provider = BfTreeProvider::::new(params, start_points.as_view(), NoStore) .unwrap(); @@ -3666,7 +3675,7 @@ mod tests { async fn test_new_rejects_capacity_exceeding_id_type() { let dim = 4usize; let num_start_points = NonZeroUsize::new(1).unwrap(); - let start_points = Matrix::from_element(num_start_points.into(), dim, 0.0f32); + let start_points = rowmajor::Owned::from_element(num_start_points.into(), dim, 0.0f32); let params = BfTreeProviderParameters { // Largest index would be u32::MAX + 1, which a u32 id cannot hold. diff --git a/diskann-bftree/src/quant.rs b/diskann-bftree/src/quant.rs index b60d63ff09..01f3620d50 100644 --- a/diskann-bftree/src/quant.rs +++ b/diskann-bftree/src/quant.rs @@ -237,13 +237,13 @@ pub(crate) fn create_test_quantizer(dim: usize) -> Poly { algorithms::TransformKind, spherical::{PreScale, SphericalQuantizer, SupportedMetric}, }; - use diskann_utils::views::Matrix; + use diskann_utils::views::rowmajor::{self, Matrix}; use rand::{rngs::StdRng, SeedableRng}; // Create training data with spread-out values. let nrows = 8; let mut counter = 0.0f32; - let data = Matrix::from_fn(nrows, dim, move |_| { + let data = rowmajor::Owned::from_fn(nrows, dim, move |_| { counter += 0.5; counter }); diff --git a/diskann-disk/src/build/builder/core.rs b/diskann-disk/src/build/builder/core.rs index e28d7d8962..b5114376c2 100644 --- a/diskann-disk/src/build/builder/core.rs +++ b/diskann-disk/src/build/builder/core.rs @@ -20,7 +20,7 @@ use diskann_providers::{ }, }; use diskann_quantization::spherical::DataRef; -use diskann_utils::io::read_bin; +use diskann_utils::{io::read_bin, views::rowmajor::Matrix}; use rand::seq::SliceRandom; use tracing::info; @@ -997,7 +997,7 @@ pub(crate) mod disk_index_builder_tests { // that our simple graph search matches. // // Because this dataset is small, we can expect exact equality. - for (q, query_data) in data.row_iter().enumerate() { + for (q, query_data) in data.rows().enumerate() { let gt = diskann_providers::test_utils::groundtruth(data.as_view(), query_data, |a, b| { distance.evaluate_similarity(a, b) diff --git a/diskann-disk/src/build/builder/quantizer.rs b/diskann-disk/src/build/builder/quantizer.rs index 4bdb9d062f..369f3a6600 100644 --- a/diskann-disk/src/build/builder/quantizer.rs +++ b/diskann-disk/src/build/builder/quantizer.rs @@ -21,7 +21,7 @@ use diskann_quantization::{ scalar::train::ScalarQuantizationParameters, spherical::{PreScale, SphericalQuantizer, SupportedMetric}, }; -use diskann_utils::views::MatrixView; +use diskann_utils::views::rowmajor; use tracing::info; use crate::{ @@ -69,7 +69,8 @@ impl BuildQuantizer { &mut rnd, )?; train_pq( - MatrixView::try_from(&train_data, train_size, train_dim).bridge_err()?, + rowmajor::Ref::try_from_data(&train_data, train_size, train_dim) + .bridge_err()?, num_chunks, &mut rnd, create_thread_pool(index_configuration.num_threads)?.as_ref(), @@ -118,7 +119,8 @@ impl BuildQuantizer { }; let quantizer = quantizer_params.train( - MatrixView::try_from(&train_data_vector, train_size, train_dim).bridge_err()?, + rowmajor::Ref::try_from_data(&train_data_vector, train_size, train_dim) + .bridge_err()?, ); info!("Now quantizer is trained and saving to file"); @@ -141,8 +143,8 @@ impl BuildQuantizer { storage_provider, &mut rnd, )?; - let train_data = - MatrixView::try_from(&train_data, train_size, train_dim).bridge_err()?; + let train_data = rowmajor::Ref::try_from_data(&train_data, train_size, train_dim) + .bridge_err()?; let quantizer = SphericalQuantizer::train( train_data, TransformKind::DoubleHadamard { diff --git a/diskann-disk/src/search/pq/quantizer_preprocess.rs b/diskann-disk/src/search/pq/quantizer_preprocess.rs index ba431f3bed..1bc1e676f9 100644 --- a/diskann-disk/src/search/pq/quantizer_preprocess.rs +++ b/diskann-disk/src/search/pq/quantizer_preprocess.rs @@ -24,7 +24,7 @@ impl PQScratch { fn preprocess_query(&mut self, pq_data: &PQData, metric: Metric) -> ANNResult<()> { let table = pq_data.pq_table(); let expected_len = table.ncenters() * table.nchunks(); - let dst = diskann_utils::views::MutMatrixView::try_from( + let dst = diskann_utils::views::rowmajor::Mut::try_from_data( &mut self.aligned_pqtable_dist_scratch[..expected_len], table.nchunks(), table.ncenters(), diff --git a/diskann-disk/src/search/provider/disk_provider.rs b/diskann-disk/src/search/provider/disk_provider.rs index 4168e5809c..8eabae064d 100644 --- a/diskann-disk/src/search/provider/disk_provider.rs +++ b/diskann-disk/src/search/provider/disk_provider.rs @@ -41,7 +41,7 @@ use diskann_providers::{ use diskann_utils::{ future::SendFuture, object_pool::{ObjectPool, PoolOption, TryAsPooled}, - views::Matrix, + views::rowmajor::{self, Matrix, MatrixMut}, }; use crate::search::pq::{PQData, PQScratch}; @@ -556,7 +556,7 @@ where ensure_vertex_loaded(&mut accessor.scratch.vertex_provider, &candidate_ids)?; let mut candidate_vectors = - Matrix::try_from_element(candidate_ids.len(), query_f32.len(), 0.0)?; + rowmajor::Owned::try_from_element(candidate_ids.len(), query_f32.len(), 0.0)?; let mut candidate_distances = Vec::with_capacity(candidate_ids.len()); let mut associated_data = Vec::with_capacity(candidate_ids.len()); @@ -576,7 +576,7 @@ where } let reranked = determinant_diversity( - candidate_vectors.as_mut_view(), + candidate_vectors.as_view_mut(), &candidate_distances, &query_f32, usize::MAX, @@ -1543,7 +1543,11 @@ mod disk_provider_tests { DynWriteProvider, StorageReadProvider, VirtualStorageProvider, }; use diskann_providers::utils::{create_thread_pool, PQPathNames, ParallelIteratorInPool}; - use diskann_utils::{io::read_bin, test_data_root, views::Matrix}; + use diskann_utils::{ + io::read_bin, + test_data_root, + views::rowmajor::{self, Matrix}, + }; use diskann_vector::distance::Metric; use rayon::prelude::IndexedParallelIterator; use rstest::rstest; @@ -1840,7 +1844,7 @@ mod disk_provider_tests { fn load_source_data( storage_provider: &StorageReader, path: &str, - ) -> Matrix { + ) -> rowmajor::Owned { read_bin(&mut storage_provider.open_reader(path).unwrap()).unwrap() } @@ -1849,7 +1853,7 @@ mod disk_provider_tests { expected_result_count: u32, expected_io_operations: u32, expected_results: impl IntoIterator, - source: Option<&Matrix>, + source: Option<&rowmajor::Owned>, vector_dimension: usize, ) { assert_eq!(indexed.stats.result_count, expected_result_count); @@ -1874,7 +1878,7 @@ mod disk_provider_tests { fn assert_indexed_search_results_match( result: &SearchResult<()>, indexed: &SearchResultWithVectors<(), f32>, - source: &Matrix, + source: &rowmajor::Owned, ) { assert_indexed_results_match( indexed, @@ -1940,7 +1944,7 @@ mod disk_provider_tests { let pool = create_thread_pool(params.thread_num.into_usize()).unwrap(); queries - .par_row_iter() + .par_rows() .enumerate() .for_each_in_pool(pool.as_ref(), |(i, query)| { let mut query_stats = QueryStatistics::default(); @@ -2018,7 +2022,7 @@ mod disk_provider_tests { load_query_result(params.storage_provider, params.truth_result_file_path); let pool = create_thread_pool(params.thread_num.into_usize()).unwrap(); queries - .par_row_iter() + .par_rows() .enumerate() .for_each_in_pool(pool.as_ref(), |(i, query)| { let result = params diff --git a/diskann-disk/src/storage/quant/compressor.rs b/diskann-disk/src/storage/quant/compressor.rs index 27b02a4db1..20a735d779 100644 --- a/diskann-disk/src/storage/quant/compressor.rs +++ b/diskann-disk/src/storage/quant/compressor.rs @@ -4,7 +4,7 @@ */ use diskann::{utils::VectorRepr, ANNResult}; -use diskann_utils::views::{MatrixView, MutMatrixView}; +use diskann_utils::views::rowmajor; /// A quantizer constructed once and shared across compression batches. pub trait QuantCompressor: Sized + Sync @@ -19,6 +19,6 @@ where /// Returns an error if construction fails. fn new(context: &Self::CompressorContext) -> ANNResult; - fn compress(&self, vector: MatrixView, output: MutMatrixView) -> ANNResult<()>; + fn compress(&self, vector: rowmajor::Ref, output: rowmajor::Mut) -> ANNResult<()>; fn compressed_bytes(&self) -> usize; } diff --git a/diskann-disk/src/storage/quant/generator.rs b/diskann-disk/src/storage/quant/generator.rs index 88de1118a1..51748561c9 100644 --- a/diskann-disk/src/storage/quant/generator.rs +++ b/diskann-disk/src/storage/quant/generator.rs @@ -14,7 +14,10 @@ use diskann_providers::{ storage::{StorageReadProvider, StorageWriteProvider}, utils::{load_metadata_from_file, BridgeErr, ParallelIteratorInPool, RayonThreadPoolRef}, }; -use diskann_utils::{io::Metadata, views}; +use diskann_utils::{ + io::Metadata, + views::{self, rowmajor::Matrix, rowmajor::MatrixMut}, +}; use rayon::iter::IndexedParallelIterator; use tracing::info; @@ -133,16 +136,17 @@ where // process `BATCH_SIZE` many dataset vectors at a time. const BATCH_SIZE: usize = 128; - // Wrap the data in `MatrixViews` so we do not need to manually construct view + // Wrap the data in `rowmajor::Mut` so we do not need to manually construct view // in the compression loop. - let mut compressed_block = views::MutMatrixView::try_from( + let mut compressed_block = views::rowmajor::Mut::try_from_data( block_compressed_base, cur_block_size, compressed_size, ) .bridge_err()?; let base_block = - views::MatrixView::try_from(&block_data, cur_block_size, full_dim).bridge_err()?; + views::rowmajor::Ref::try_from_data(&block_data, cur_block_size, full_dim) + .bridge_err()?; base_block .par_window_iter(BATCH_SIZE) .zip_eq(compressed_block.par_window_iter_mut(BATCH_SIZE)) @@ -201,10 +205,7 @@ mod generator_tests { use diskann::utils::read_exact_into; use diskann_providers::storage::VirtualStorageProvider; use diskann_providers::utils::create_thread_pool_for_test; - use diskann_utils::{ - io::{write_bin, Metadata}, - views::MatrixView, - }; + use diskann_utils::io::{write_bin, Metadata}; use rstest::rstest; use vfs::{FileSystem, MemoryFS}; @@ -230,11 +231,11 @@ mod generator_tests { fn compress( &self, - _vector: views::MatrixView, - mut output: views::MutMatrixView, + _vector: views::rowmajor::Ref, + mut output: views::rowmajor::Mut, ) -> ANNResult<()> { output - .row_iter_mut() + .rows_mut() .for_each(|r| r.copy_from_slice(&self.code)); Ok(()) } @@ -272,7 +273,7 @@ mod generator_tests { // Setup test data let data = create_test_data(num_points, dim); - let view = MatrixView::try_from(data.as_slice(), num_points, dim).unwrap(); + let view = views::rowmajor::Ref::try_from_data(data.as_slice(), num_points, dim).unwrap(); write_bin( view, &mut storage_provider.create_for_write(data_path.as_str())?, diff --git a/diskann-disk/src/storage/quant/pq/pq_dataset.rs b/diskann-disk/src/storage/quant/pq/pq_dataset.rs index adf51c433c..f04724eb84 100644 --- a/diskann-disk/src/storage/quant/pq/pq_dataset.rs +++ b/diskann-disk/src/storage/quant/pq/pq_dataset.rs @@ -8,7 +8,7 @@ use core::fmt::Debug; use diskann::ANNResult; use diskann_providers::model::FixedChunkPQTable; use diskann_quantization::{error::Format, product::TransposedTable}; -use diskann_utils::views::Matrix; +use diskann_utils::views::rowmajor::{self, Matrix}; use crate::error::{diskann_error, ErrorKind}; @@ -18,13 +18,13 @@ pub struct PQData { pq_pivot_table: TransposedTable, // pq compressed vectors, shape `num_points × num_pq_chunks`. - pq_compressed_data: Matrix, + pq_compressed_data: rowmajor::Owned, } impl PQData { pub fn new( pq_pivot_table: FixedChunkPQTable, - pq_compressed_data: Matrix, + pq_compressed_data: rowmajor::Owned, ) -> ANNResult { let pq_pivot_table = TransposedTable::from_parts( pq_pivot_table.view_pivots(), @@ -59,7 +59,7 @@ impl PQData { } /// Get pq_compressed_data - pub fn pq_compressed_data(&self) -> &Matrix { + pub fn pq_compressed_data(&self) -> &rowmajor::Owned { &self.pq_compressed_data } @@ -84,8 +84,9 @@ mod tests { let pq_pivot_table = FixedChunkPQTable::new(dim, Box::new([0.0, 0.0, 1.0, 1.0]), Box::new([0, 2])).unwrap(); - let pq_compressed_data = Matrix::try_from(Box::new([123u8, 111, 255]) as Box<[u8]>, 3, 1) - .expect("valid matrix shape"); + let pq_compressed_data = + rowmajor::Owned::try_from_data(Box::new([123u8, 111, 255]) as Box<[u8]>, 3, 1) + .expect("valid matrix shape"); PQData::new(pq_pivot_table, pq_compressed_data) } diff --git a/diskann-disk/src/storage/quant/pq/pq_generation.rs b/diskann-disk/src/storage/quant/pq/pq_generation.rs index 8f5def5f75..eae52f265f 100644 --- a/diskann-disk/src/storage/quant/pq/pq_generation.rs +++ b/diskann-disk/src/storage/quant/pq/pq_generation.rs @@ -13,7 +13,7 @@ use diskann_providers::{ utils::RayonThreadPoolRef, }; use diskann_quantization::{error::Format, product::TransposedTable, CompressInto}; -use diskann_utils::views::MatrixBase; +use diskann_utils::views::rowmajor; use diskann_vector::distance::Metric; use tracing::info; @@ -140,8 +140,8 @@ where fn compress( &self, - vector: MatrixBase<&[f32]>, - output: MatrixBase<&mut [u8]>, + vector: rowmajor::Ref<'_, f32>, + output: rowmajor::Mut<'_, u8>, ) -> Result<(), diskann::ANNError> { self.table .compress_into(vector, output) @@ -171,7 +171,7 @@ mod pq_generation_tests { use diskann_utils::{ io::{read_bin, write_bin}, test_data_root, - views::{MatrixView, MutMatrixView}, + views::rowmajor::{self, Matrix}, }; use diskann_vector::distance::Metric; use rstest::rstest; @@ -261,7 +261,7 @@ mod pq_generation_tests { let (ndata, dim, num_centers, num_chunks, max_k_means_reps) = (5, 8, 2, 2, 5); write_bin( - MatrixView::try_from(VALIDATION_DATA.as_slice(), ndata, dim).unwrap(), + rowmajor::Ref::try_from_data(VALIDATION_DATA.as_slice(), ndata, dim).unwrap(), &mut storage_provider.create_for_write(data_path).unwrap(), ) .unwrap(); @@ -299,7 +299,7 @@ mod pq_generation_tests { let updated_data: Vec = VALIDATION_DATA.iter().map(|x| x + 10.0).collect(); write_bin( - MatrixView::try_from(updated_data.as_slice(), ndata, dim).unwrap(), + rowmajor::Ref::try_from_data(updated_data.as_slice(), ndata, dim).unwrap(), &mut storage_provider.create_for_write(data_path).unwrap(), ) .unwrap(); @@ -336,8 +336,8 @@ mod pq_generation_tests { let mut expected_codes = vec![0; ndata * num_chunks]; compressor .compress( - MatrixView::try_from(updated_data.as_slice(), ndata, dim).unwrap(), - MutMatrixView::try_from(&mut expected_codes, ndata, num_chunks).unwrap(), + rowmajor::Ref::try_from_data(updated_data.as_slice(), ndata, dim).unwrap(), + rowmajor::Mut::try_from_data(&mut expected_codes, ndata, num_chunks).unwrap(), ) .unwrap(); let codes = @@ -371,7 +371,7 @@ mod pq_generation_tests { let pivots_path = "/pivots.bin"; let codes_path = "/codes.bin"; write_bin( - MatrixView::try_from(VALIDATION_DATA.as_slice(), 5, 8).unwrap(), + rowmajor::Ref::try_from_data(VALIDATION_DATA.as_slice(), 5, 8).unwrap(), &mut storage_provider.create_for_write(data_path).unwrap(), ) .unwrap(); @@ -431,7 +431,7 @@ mod pq_generation_tests { let mut train_data: Vec = VALIDATION_DATA.to_vec(); write_bin( - MatrixView::try_from(train_data.as_slice(), ndata, dim).unwrap(), + rowmajor::Ref::try_from_data(train_data.as_slice(), ndata, dim).unwrap(), &mut storage_provider.create_for_write(data_path).unwrap(), ) .unwrap(); @@ -502,7 +502,7 @@ mod pq_generation_tests { use diskann_quantization::CompressInto; let result = table.compress_into( data_matrix.as_view(), - MutMatrixView::try_from(&mut compressed_mat, npts, num_chunks).unwrap(), + rowmajor::Mut::try_from_data(&mut compressed_mat, npts, num_chunks).unwrap(), ); assert!(result.is_ok()); diff --git a/diskann-garnet/src/provider.rs b/diskann-garnet/src/provider.rs index 43a38a2e86..556b6505e1 100644 --- a/diskann-garnet/src/provider.rs +++ b/diskann-garnet/src/provider.rs @@ -24,10 +24,9 @@ use diskann::{ utils::VectorRepr, }; use diskann_quantization::alloc::{AllocatorError, Poly}; -use diskann_utils::views::Matrix; use diskann_utils::{ object_pool::{AsPooled, ObjectPool, PooledRef, Undef}, - views::MatrixView, + views::rowmajor::{self, Matrix, MatrixMut}, }; use diskann_vector::{ DistanceFunction, PreprocessedDistanceFunction, contains::ContainsSimd, distance::Metric, @@ -514,7 +513,7 @@ impl GarnetProvider { }; let rows = quantizer.required_vectors(); - let mut data = Matrix::from_element(rows, self.dim, T::default()); + let mut data = rowmajor::Owned::from_element(rows, self.dim, T::default()); let mut row_idx = 0usize; if self @@ -564,7 +563,7 @@ impl GarnetProvider { Ok(v) => v, Err(_) => return false, }; - let view = match MatrixView::try_from(&*converted, view.nrows(), view.ncols()) { + let view = match rowmajor::Ref::try_from_data(&converted, view.nrows(), view.ncols()) { Ok(v) => v, Err(_) => return false, }; diff --git a/diskann-garnet/src/quantization.rs b/diskann-garnet/src/quantization.rs index a85c2a2136..13556352c6 100644 --- a/diskann-garnet/src/quantization.rs +++ b/diskann-garnet/src/quantization.rs @@ -17,7 +17,7 @@ use diskann_quantization::{ iface::{self, Opaque, OpaqueMut, Quantizer}, }, }; -use diskann_utils::views::MatrixView; +use diskann_utils::views::rowmajor::{self, Matrix}; use diskann_vector::{DistanceFunction, PreprocessedDistanceFunction, distance::Metric}; use thiserror::Error; @@ -57,7 +57,7 @@ pub(crate) trait GarnetQuantizer: Send + Sync { /// Each row of the matrix will be a vector. /// Returns a lock guard for purposes of synchronization; after the guard is released, the /// quantizer will be accessible to all threads. - fn train(&self, metric: Metric, data: MatrixView) -> Result<(), GarnetQuantizerError>; + fn train(&self, metric: Metric, data: rowmajor::Ref) -> Result<(), GarnetQuantizerError>; /// Quantize a vector fn compress(&self, v: &[f32], into: &mut [u8]) -> Result<(), GarnetQuantizerError>; /// Returns a distance computer for comparing quantized vectors @@ -115,7 +115,7 @@ impl GarnetQuantizer for Spherical1Bit { fn train( &self, metric_type: Metric, - data: MatrixView, + data: rowmajor::Ref, ) -> Result<(), GarnetQuantizerError> { let mut rng = rand::rng(); let quantizer = SphericalQuantizer::train( @@ -287,7 +287,11 @@ impl GarnetQuantizer for MinMax8Bit { true } - fn train(&self, _metric: Metric, _data: MatrixView) -> Result<(), GarnetQuantizerError> { + fn train( + &self, + _metric: Metric, + _data: rowmajor::Ref, + ) -> Result<(), GarnetQuantizerError> { Ok(()) } @@ -374,7 +378,7 @@ impl DynQueryComputer for MinMax8BitQueryComputer { #[cfg(test)] mod tests { - use diskann_utils::views::Matrix; + use diskann_utils::views::rowmajor::{self, Matrix, MatrixMut}; use diskann_vector::{DistanceFunction, PreprocessedDistanceFunction, distance::Metric}; use crate::quantization::{GarnetQuantizer, GarnetQuantizerError, MinMax8Bit, Spherical1Bit}; @@ -403,7 +407,7 @@ mod tests { Err(GarnetQuantizerError::NoQuantizer) )); - let mut test_data = Matrix::from_element(1000, 2, 0.0f32); + let mut test_data = rowmajor::Owned::from_element(1000, 2, 0.0f32); for i in 0..1000 { test_data .row_mut(i) @@ -441,7 +445,7 @@ mod tests { let test_v = [0.5f32, 0.5]; let mut test_q = vec![0u8; quantizer.bytes()]; - let mut test_data = Matrix::from_element(1, 2, 0.0f32); + let mut test_data = rowmajor::Owned::from_element(1, 2, 0.0f32); test_data.row_mut(0).copy_from_slice(&[1.0f32, 1.0]); // Training is a no-op, but succeeds. diff --git a/diskann-inmem/integration/index/runner.rs b/diskann-inmem/integration/index/runner.rs index f65e80ea4f..0e5d1e95a8 100644 --- a/diskann-inmem/integration/index/runner.rs +++ b/diskann-inmem/integration/index/runner.rs @@ -13,7 +13,7 @@ use diskann_benchmark_runner::{ files::InputFile, utils::fmt::Indent, }; -use diskann_utils::views::Matrix; +use diskann_utils::views::rowmajor::{self, Matrix}; use diskann_vector::distance::Metric; use serde::{Deserialize, Serialize}; @@ -256,7 +256,7 @@ impl Data { struct Bundle { data: Dataset, queries: Dataset, - groundtruth: Matrix, + groundtruth: rowmajor::Owned, } mod spherical { @@ -564,7 +564,7 @@ impl Test { spherical::Bits::Four => quantizer.as_quantizer::<4>()?, }; - let start_point = Matrix::row_vector(Box::from( + let start_point = rowmajor::Owned::row_vector(Box::from( ::compute_medoid(data), )); diff --git a/diskann-inmem/integration/index/tests.rs b/diskann-inmem/integration/index/tests.rs index 95f7da9251..52c2e6e6cd 100644 --- a/diskann-inmem/integration/index/tests.rs +++ b/diskann-inmem/integration/index/tests.rs @@ -6,7 +6,7 @@ use diskann::graph::search::Knn; use diskann_benchmark_core::recall::{RecallMetrics, Rows}; use diskann_benchmark_runner::utils::fmt::KeyValue; -use diskann_utils::views::Matrix; +use diskann_utils::views::rowmajor::{self, Matrix, MatrixMut}; use serde::{Deserialize, Serialize}; use crate::{ @@ -44,12 +44,12 @@ pub(super) fn knn( groundtruth.nrows(), ); - let mut ids = Matrix::from_element(queries.nrows(), k, u64::MAX); + let mut ids = rowmajor::Owned::from_element(queries.nrows(), k, u64::MAX); let before = index.counters(); let mut misc = KnnSearch::new(); let mut neighbors = Vec::new(); - for (out, query) in std::iter::zip(ids.row_iter_mut(), queries.iter()) { + for (out, query) in std::iter::zip(ids.rows_mut(), queries.iter()) { neighbors.clear(); let stats = rt.block_on(index.search(query, knn, &mut neighbors))?; diff --git a/diskann-inmem/integration/support/datatype.rs b/diskann-inmem/integration/support/datatype.rs index 729d91edbd..b420453317 100644 --- a/diskann-inmem/integration/support/datatype.rs +++ b/diskann-inmem/integration/support/datatype.rs @@ -5,7 +5,7 @@ use diskann_utils::{ sampling::medoid::ComputeMedoid, - views::{Matrix, MatrixView, MutMatrixView}, + views::rowmajor::{self, Matrix, MatrixMut}, }; use diskann_wide::{cast_f16_to_f32, cast_f32_to_f16}; use half::f16; @@ -223,10 +223,10 @@ impl<'a> SliceMut<'a> { #[derive(Debug)] pub(crate) enum Dataset { - F32(Matrix), - F16(Matrix), - U8(Matrix), - I8(Matrix), + F32(rowmajor::Owned), + F16(rowmajor::Owned), + U8(rowmajor::Owned), + I8(rowmajor::Owned), } impl Dataset { @@ -258,10 +258,10 @@ impl Dataset { pub(crate) fn preprocess(&mut self, op: &Preprocess) { match self { - Self::F32(m) => op.apply(m.as_mut_view()), - Self::F16(m) => op.apply(m.as_mut_view()), - Self::U8(m) => op.apply(m.as_mut_view()), - Self::I8(m) => op.apply(m.as_mut_view()), + Self::F32(m) => op.apply(m.as_view_mut()), + Self::F16(m) => op.apply(m.as_view_mut()), + Self::U8(m) => op.apply(m.as_view_mut()), + Self::I8(m) => op.apply(m.as_view_mut()), } } } @@ -278,11 +278,11 @@ pub(crate) enum Preprocess { } trait Apply { - fn apply(&self, m: MutMatrixView<'_, T>); + fn apply(&self, m: rowmajor::Mut<'_, T>); } impl Apply for Preprocess { - fn apply(&self, mut m: MutMatrixView<'_, f32>) { + fn apply(&self, mut m: rowmajor::Mut<'_, f32>) { match self { Self::Halve => m.as_mut_slice().iter_mut().for_each(|v| *v *= 0.5), Self::Floor => m.as_mut_slice().iter_mut().for_each(|v| *v = v.floor()), @@ -291,7 +291,7 @@ impl Apply for Preprocess { } impl Apply for Preprocess { - fn apply(&self, mut m: MutMatrixView<'_, f16>) { + fn apply(&self, mut m: rowmajor::Mut<'_, f16>) { match self { Self::Halve => m.as_mut_slice().iter_mut().for_each(|v| { *v = f16::from_f32(f32::from(*v) * 0.5); @@ -304,7 +304,7 @@ impl Apply for Preprocess { } impl Apply for Preprocess { - fn apply(&self, mut m: MutMatrixView<'_, u8>) { + fn apply(&self, mut m: rowmajor::Mut<'_, u8>) { match self { Self::Halve => m.as_mut_slice().iter_mut().for_each(|v| *v /= 2), Self::Floor => {} @@ -313,7 +313,7 @@ impl Apply for Preprocess { } impl Apply for Preprocess { - fn apply(&self, mut m: MutMatrixView<'_, i8>) { + fn apply(&self, mut m: rowmajor::Mut<'_, i8>) { match self { Self::Halve => m.as_mut_slice().iter_mut().for_each(|v| *v /= 2), Self::Floor => {} @@ -327,10 +327,10 @@ impl Apply for Preprocess { #[derive(Debug, Clone, Copy)] pub(crate) enum DatasetView<'a> { - F32(MatrixView<'a, f32>), - F16(MatrixView<'a, f16>), - U8(MatrixView<'a, u8>), - I8(MatrixView<'a, i8>), + F32(rowmajor::Ref<'a, f32>), + F16(rowmajor::Ref<'a, f16>), + U8(rowmajor::Ref<'a, u8>), + I8(rowmajor::Ref<'a, i8>), } impl<'a> DatasetView<'a> { @@ -372,10 +372,10 @@ impl<'a> DatasetView<'a> { pub(crate) fn medoid(&self) -> Dataset { match self { - Self::F32(v) => Matrix::row_vector(Box::from(f32::compute_medoid(*v))).into(), - Self::F16(v) => Matrix::row_vector(Box::from(f16::compute_medoid(*v))).into(), - Self::U8(v) => Matrix::row_vector(Box::from(u8::compute_medoid(*v))).into(), - Self::I8(v) => Matrix::row_vector(Box::from(i8::compute_medoid(*v))).into(), + Self::F32(v) => rowmajor::Owned::row_vector(Box::from(f32::compute_medoid(*v))).into(), + Self::F16(v) => rowmajor::Owned::row_vector(Box::from(f16::compute_medoid(*v))).into(), + Self::U8(v) => rowmajor::Owned::row_vector(Box::from(u8::compute_medoid(*v))).into(), + Self::I8(v) => rowmajor::Owned::row_vector(Box::from(i8::compute_medoid(*v))).into(), } } @@ -437,8 +437,8 @@ macro_rules! define { } } - impl From> for Dataset { - fn from(m: Matrix<$T>) -> Self { + impl From> for Dataset { + fn from(m: rowmajor::Owned<$T>) -> Self { Self::$variant(m) } } @@ -458,11 +458,11 @@ define!(i8, I8); mod tests { use super::*; - fn matrix(data: &[T], nrows: usize, ncols: usize) -> Matrix + fn matrix(data: &[T], nrows: usize, ncols: usize) -> rowmajor::Owned where T: Copy, { - Matrix::try_from(Box::from(data), nrows, ncols).unwrap() + rowmajor::Owned::try_from_data(Box::from(data), nrows, ncols).unwrap() } //----------// diff --git a/diskann-inmem/integration/support/io.rs b/diskann-inmem/integration/support/io.rs index 21f6d89858..be3f9846aa 100644 --- a/diskann-inmem/integration/support/io.rs +++ b/diskann-inmem/integration/support/io.rs @@ -3,7 +3,10 @@ * Licensed under the MIT license. */ -use diskann_utils::{io::read_bin, views::Matrix}; +use diskann_utils::{ + io::read_bin, + views::rowmajor::{self, MatrixMut}, +}; use half::f16; use super::datatype::{DataType, Dataset, Preprocess, SliceMut}; @@ -34,22 +37,22 @@ where let dst = match target { DataType::F32 => { - let mut dst = Matrix::from_element(data.nrows(), data.ncols(), f32::default()); + let mut dst = rowmajor::Owned::from_element(data.nrows(), data.ncols(), f32::default()); SliceMut::from(dst.as_mut_slice()).convert_lossless(data.as_slice())?; Dataset::from(dst) } DataType::F16 => { - let mut dst = Matrix::from_element(data.nrows(), data.ncols(), f16::default()); + let mut dst = rowmajor::Owned::from_element(data.nrows(), data.ncols(), f16::default()); SliceMut::from(dst.as_mut_slice()).convert_lossless(data.as_slice())?; Dataset::from(dst) } DataType::U8 => { - let mut dst = Matrix::from_element(data.nrows(), data.ncols(), u8::default()); + let mut dst = rowmajor::Owned::from_element(data.nrows(), data.ncols(), u8::default()); SliceMut::from(dst.as_mut_slice()).convert_lossless(data.as_slice())?; Dataset::from(dst) } DataType::I8 => { - let mut dst = Matrix::from_element(data.nrows(), data.ncols(), i8::default()); + let mut dst = rowmajor::Owned::from_element(data.nrows(), data.ncols(), i8::default()); SliceMut::from(dst.as_mut_slice()).convert_lossless(data.as_slice())?; Dataset::from(dst) } diff --git a/diskann-inmem/src/provider.rs b/diskann-inmem/src/provider.rs index bebfa2fd8d..452da17599 100644 --- a/diskann-inmem/src/provider.rs +++ b/diskann-inmem/src/provider.rs @@ -819,7 +819,7 @@ mod tests { neighbor::Neighbor, provider::{DataProvider, Delete}, }; - use diskann_utils::views::Matrix; + use diskann_utils::views::rowmajor::{self, Matrix}; use diskann_vector::distance::Metric; use crate::num::Capacity; @@ -854,7 +854,7 @@ mod tests { Capacity::new(grid.num_points(size)), MaxDegree::new(degree), Metric::L2, - Matrix::row_vector(start.into()), + rowmajor::Owned::row_vector(start.into()), ) .unwrap(); @@ -876,7 +876,7 @@ mod tests { let index = DiskANNIndex::new(config, provider, None); - for (i, data) in data.row_iter().enumerate() { + for (i, data) in data.rows().enumerate() { index .insert(&Strategy, &Context, &((10 * i + 1) as u64), data) .await diff --git a/diskann-inmem/src/repr/full.rs b/diskann-inmem/src/repr/full.rs index d6f768d82b..0c265a0e75 100644 --- a/diskann-inmem/src/repr/full.rs +++ b/diskann-inmem/src/repr/full.rs @@ -38,7 +38,7 @@ mod internal_docs { use std::{fmt::Debug, marker::PhantomData, num::NonZeroUsize}; use diskann::{ANNError, ANNResult, utils::IntoUsize}; -use diskann_utils::views::Matrix; +use diskann_utils::views::rowmajor::{self, Matrix}; use diskann_vector::{ UnalignedSlice, conversion::SliceCast, @@ -114,7 +114,7 @@ pub trait FullPrecision: bytemuck::Pod + std::fmt::Debug + Send + Sync { pub struct Config { layout: store::Layout, metric: Metric, - start_points: Matrix, + start_points: rowmajor::Owned, store: store::Config, lookahead: Option, } @@ -140,7 +140,7 @@ impl Config { capacity: Capacity, max_degree: MaxDegree, metric: Metric, - start_points: Matrix, + start_points: rowmajor::Owned, ) -> Result { let num_start_points: u32 = match start_points.nrows().try_into() { Ok(points) => points, @@ -261,7 +261,7 @@ where capacity: Capacity, max_degree: MaxDegree, metric: Metric, - start_points: Matrix, + start_points: rowmajor::Owned, ) -> Result, ConfigError> { Config::new(capacity, max_degree, metric, start_points) } @@ -286,7 +286,7 @@ where let store = Store::new(layout, store, intrusive)?; // Initialize start points. - for (i, row) in std::iter::zip(store.frozen(), start_points.row_iter()) { + for (i, row) in std::iter::zip(store.frozen(), start_points.rows()) { #[expect( clippy::expect_used, reason = "failing this is an internal, unrecoverable bug" @@ -887,7 +887,7 @@ mod tests { capacity, MaxDegree::new(0), Metric::L2, - Matrix::column_vector(Box::new(start_points)), + rowmajor::Owned::column_vector(Box::new(start_points)), ) .unwrap(), ) @@ -1231,7 +1231,7 @@ mod tests { Capacity::new(1), MaxDegree::new(0), metric, - Matrix::::row_vector(start_point.clone().into()), + rowmajor::Owned::::row_vector(start_point.clone().into()), ) .unwrap(), ) diff --git a/diskann-inmem/src/repr/spherical.rs b/diskann-inmem/src/repr/spherical.rs index 6fd8f5cfda..dde4ae6946 100644 --- a/diskann-inmem/src/repr/spherical.rs +++ b/diskann-inmem/src/repr/spherical.rs @@ -12,7 +12,10 @@ use diskann_quantization::{ alloc::{GlobalAllocator, Poly, ScopedAllocator}, spherical::{SupportedMetric, iface}, }; -use diskann_utils::{lazy_format, views::Matrix}; +use diskann_utils::{ + lazy_format, + views::rowmajor::{self, Matrix}, +}; use diskann_vector::distance::{Distance, DistanceProvider}; use half::f16; use thiserror::Error; @@ -38,7 +41,7 @@ pub struct Config { /// The underlying quantizer for the compressed store. quantizer: Poly, /// The start points. These must have dimensions equal to `quantizer.full_dim()`. - start_points: Matrix, + start_points: rowmajor::Owned, layout: store::Layout, store: store::Config, lookahead: Option, @@ -76,7 +79,7 @@ impl Config { quantizer: Poly, capacity: Capacity, max_degree: MaxDegree, - start_points: Matrix, + start_points: rowmajor::Owned, rerank: Rerank, ) -> Result { let quantizer_dim = quantizer.full_dim(); @@ -332,7 +335,7 @@ impl Spherical { quantizer: Poly, capacity: Capacity, max_degree: MaxDegree, - start_points: Matrix, + start_points: rowmajor::Owned, rerank: Rerank, ) -> Result { Config::new(quantizer, capacity, max_degree, start_points, rerank) @@ -372,7 +375,7 @@ impl Spherical { // Initialize start points. let num_start_points = start_points.nrows(); - for (i, row) in std::iter::zip(this.store.frozen(), start_points.row_iter()) { + for (i, row) in std::iter::zip(this.store.frozen(), start_points.rows()) { #[expect( clippy::expect_used, reason = "failing this is an internal, unrecoverable bug" @@ -637,7 +640,7 @@ mod tests { use super::*; use diskann::{graph::test::synthetic::Grid, neighbor::Neighbor}; - use diskann_utils::{assert_contains, views::MatrixView}; + use diskann_utils::{assert_contains, views::rowmajor::MatrixMut}; use hashbrown::HashMap; use crate::{ @@ -654,7 +657,7 @@ mod tests { } fn train_quantizer( - data: MatrixView<'_, f32>, + data: rowmajor::Ref<'_, f32>, metric: SupportedMetric, bits: Bits, ) -> Poly { @@ -716,7 +719,7 @@ mod tests { let quantizer = train_quantizer(data.as_view(), metric, bits); - let mut start_points = Matrix::from_element(2, data.ncols(), 0.0); + let mut start_points = rowmajor::Owned::from_element(2, data.ncols(), 0.0); start_points.row_mut(0).fill(-2.0); start_points.row_mut(1).fill(2.0); @@ -737,7 +740,7 @@ mod tests { let mut reference = Reference::new(grid.dim().into()); if fill { - for (i, row) in data.row_iter().enumerate() { + for (i, row) in data.rows().enumerate() { let guard = repr::Set::set(&spherical, row).unwrap(); let id = repr::Guard::id(&guard); @@ -747,7 +750,7 @@ mod tests { } // Insert frozen points. - for (slot, point) in spherical.store.frozen().zip(start_points.row_iter()) { + for (slot, point) in spherical.store.frozen().zip(start_points.rows()) { reference.insert(LogicalId(slot.into_usize()), SlotId(slot), point); } @@ -1193,10 +1196,10 @@ mod tests { #[test] fn test_config_dim_mismatch() { - let data = Matrix::from_element(2, 5, 1.0f32); + let data = rowmajor::Owned::from_element(2, 5, 1.0f32); let quantizer = train_quantizer(data.as_view(), SupportedMetric::SquaredL2, Bits::One); - let start_points = Matrix::from_element(1, 6, 0.0f32); // Wrong number of columns + let start_points = rowmajor::Owned::from_element(1, 6, 0.0f32); // Wrong number of columns let err = Spherical::config( quantizer, Capacity::new(10), @@ -1215,10 +1218,10 @@ mod tests { #[test] fn test_empty_start_points() { - let data = Matrix::from_element(2, 5, 1.0f32); + let data = rowmajor::Owned::from_element(2, 5, 1.0f32); let quantizer = train_quantizer(data.as_view(), SupportedMetric::SquaredL2, Bits::One); - let start_points = Matrix::from_element(0, 5, 0.0f32); // Empty + let start_points = rowmajor::Owned::from_element(0, 5, 0.0f32); // Empty let err = Spherical::config( quantizer, Capacity::new(10), @@ -1234,10 +1237,10 @@ mod tests { #[test] fn test_build_error_uncompressible_query() { - let data = Matrix::from_element(2, 5, 1.0f32); + let data = rowmajor::Owned::from_element(2, 5, 1.0f32); let quantizer = train_quantizer(data.as_view(), SupportedMetric::SquaredL2, Bits::One); - let start_points = Matrix::from_element(1, 5, f32::INFINITY); // Wrong number of columns + let start_points = rowmajor::Owned::from_element(1, 5, f32::INFINITY); // Wrong number of columns let config = Spherical::config( quantizer, Capacity::new(10), diff --git a/diskann-providers/src/index/diskann_async.rs b/diskann-providers/src/index/diskann_async.rs index d9ddd91a25..7a50f5c562 100644 --- a/diskann-providers/src/index/diskann_async.rs +++ b/diskann-providers/src/index/diskann_async.rs @@ -10,7 +10,7 @@ use diskann::{ graph::{Config, DiskANNIndex}, utils::VectorRepr, }; -use diskann_utils::future::AsyncFriendly; +use diskann_utils::{future::AsyncFriendly, views::rowmajor::Matrix}; use crate::model::{ self, @@ -59,7 +59,7 @@ pub(crate) fn simplified_builder( } pub fn train_pq( - data: diskann_utils::views::MatrixView, + data: diskann_utils::views::rowmajor::Ref, num_pq_chunks: usize, rng: &mut dyn rand::RngCore, pool: crate::utils::RayonThreadPoolRef<'_>, @@ -180,7 +180,10 @@ pub(crate) mod tests { utils::{IntoUsize, ONE}, }; use diskann_quantization::scalar::train::ScalarQuantizationParameters; - use diskann_utils::{test_data_root, views::Matrix}; + use diskann_utils::{ + test_data_root, + views::rowmajor::{self, Matrix, MatrixMut}, + }; use diskann_vector::{ DistanceFunction, PureDistanceFunction, distance::{Metric, SquaredL2}, @@ -212,9 +215,12 @@ pub(crate) mod tests { // Tests from the original async index // ///////////////////////////////////////// - /// Convert an iterator of vectors into a single Matrix. All elements in `data` must + /// Convert an iterator of vectors into a single rowmajor::Owned. All elements in `data` must /// have the same length, otherwise this function panics. - pub(crate) fn squish<'a, To, T, Itr>(data: Itr, dim: usize) -> diskann_utils::views::Matrix + pub(crate) fn squish<'a, To, T, Itr>( + data: Itr, + dim: usize, + ) -> diskann_utils::views::rowmajor::Owned where To: Clone + Default, T: Clone + Into + 'a, @@ -222,8 +228,9 @@ pub(crate) mod tests { { // Assume that all the vectors in `data` have the same length. // If they don't, `copy_from_slice` will panic, so we're double checking. - let mut mat = diskann_utils::views::Matrix::from_element(data.len(), dim, To::default()); - std::iter::zip(mat.row_iter_mut(), data).for_each(|(output, input)| { + let mut mat = + diskann_utils::views::rowmajor::Owned::from_element(data.len(), dim, To::default()); + std::iter::zip(mat.rows_mut(), data).for_each(|(output, input)| { assert_eq!( input.len(), dim, @@ -278,7 +285,7 @@ pub(crate) mod tests { .unwrap_or_else(|| panic!("{dim}-dimensions is not supported for grid-generation")) } - fn grid_to_vecs(matrix: &Matrix) -> Vec> { + fn grid_to_vecs(matrix: &rowmajor::Owned) -> Vec> { (0..matrix.nrows()) .map(|i| matrix.row(i).to_vec()) .collect() @@ -621,7 +628,7 @@ pub(crate) mod tests { // // So, when we compute the corpus used during groundtruth generation, we take all // but this last point. - let corpus: diskann_utils::views::Matrix = + let corpus: diskann_utils::views::rowmajor::Owned = squish(vectors.iter().take(num_points), dim); let mut paged_tests = Vec::new(); @@ -698,7 +705,7 @@ pub(crate) mod tests { ]); // A matrix view of all vectors (including the start point at the end). - let matrix: Matrix = squish::(vectors.iter(), dim); + let matrix: rowmajor::Owned = squish::(vectors.iter(), dim); let table = train_pq( matrix.map(|i| (*i).into()).as_view(), @@ -728,7 +735,7 @@ pub(crate) mod tests { { let index = init_index(); let ctx = Default::default(); - for (i, v) in matrix.row_iter().take(num_points).enumerate() { + for (i, v) in matrix.rows().take(num_points).enumerate() { index .insert(&FullPrecision, &ctx, &(i as u32), v) .await @@ -742,7 +749,7 @@ pub(crate) mod tests { { let index = init_index(); let ctx = Default::default(); - for (i, v) in matrix.row_iter().take(num_points).enumerate() { + for (i, v) in matrix.rows().take(num_points).enumerate() { index.insert(&hybrid, &ctx, &(i as u32), v).await.unwrap(); } @@ -759,17 +766,22 @@ pub(crate) mod tests { for (batch, batch_data) in matrix .subview(0..num_points) .unwrap() - .window_iter(chunk_size) + .window_iter(NonZeroUsize::new(chunk_size).unwrap()) .enumerate() { - let batch_data = Arc::new(batch_data.to_owned()); + let batch_data = Arc::new(batch_data.to_rowmajor_owned()); let start = batch * chunk_size; let batch_ids: Arc<[u32]> = (start..start + batch_data.nrows()) .map(|i| i as u32) .collect(); index - .multi_insert::<_, Matrix>(FullPrecision, &ctx, batch_data, batch_ids) + .multi_insert::<_, rowmajor::Owned>( + FullPrecision, + &ctx, + batch_data, + batch_ids, + ) .await .unwrap(); } @@ -781,11 +793,11 @@ pub(crate) mod tests { { let index = init_index(); let ctx = Default::default(); - let batch = Arc::new(matrix.subview(0..num_points).unwrap().to_owned()); + let batch = Arc::new(matrix.subview(0..num_points).unwrap().to_rowmajor_owned()); let batch_ids: Arc<[u32]> = (0..num_points as u32).collect(); index - .multi_insert::<_, Matrix>(hybrid, &ctx, batch, batch_ids) + .multi_insert::<_, rowmajor::Owned>(hybrid, &ctx, batch, batch_ids) .await .unwrap(); @@ -879,7 +891,7 @@ pub(crate) mod tests { let data = T::generate_spherical(num, dim, radius, rng); let table = { - let train_data: diskann_utils::views::Matrix = squish(data.iter(), dim); + let train_data: diskann_utils::views::rowmajor::Owned = squish(data.iter(), dim); train_pq( train_data.as_view(), 2.min(dim), @@ -1074,7 +1086,7 @@ pub(crate) mod tests { let beta = 0.5; - let corpus: diskann_utils::views::Matrix = + let corpus: diskann_utils::views::rowmajor::Owned = squish(vectors.iter().take(num_points), dim); let query = vec![grid_size as f32; dim]; @@ -1395,7 +1407,7 @@ pub(crate) mod tests { #[values(1, 10)] batchsize: usize, ) where S: for<'a> InsertStrategy<'a, TestProvider, &'a [f32]> - + MultiInsertStrategy> + + MultiInsertStrategy> + Clone, { let ctx = &DefaultContext; @@ -1441,7 +1453,7 @@ pub(crate) mod tests { // that our simple graph search matches. // // Because this dataset is small, we can expect exact equality. - for (q, query) in data.row_iter().enumerate() { + for (q, query) in data.rows().enumerate() { let gt = groundtruth(data.as_view(), query, |a, b| SquaredL2::evaluate(a, b)); { let mut result_output_buffer = @@ -1489,7 +1501,7 @@ pub(crate) mod tests { #[values((-2.0,-1.0), (-1.0, 0.0), (40000.0,50000.0), (50000.0,75000.0))] radii: (f32, f32), ) where S: for<'a> InsertStrategy<'a, TestProvider, &'a [f32]> - + MultiInsertStrategy> + + MultiInsertStrategy> + Clone, { let ctx = &DefaultContext; @@ -1526,7 +1538,7 @@ pub(crate) mod tests { // Because this dataset is small, we can expect exact equality expect for the // case where we use a lower initial beam, which will trigger more two-round searches. - for (q, query) in data.row_iter().enumerate() { + for (q, query) in data.rows().enumerate() { let gt = groundtruth(data.as_view(), query, |a, b| SquaredL2::evaluate(a, b)); { // Full Precision Search. @@ -1599,10 +1611,10 @@ pub(crate) mod tests { file: &str, create_fn: C, build_fn: B, - ) -> (Arc>, Arc>) + ) -> (Arc>, Arc>) where - C: FnOnce(Arc>, &[f32]) -> Arc>, - B: AsyncFnOnce(Arc>, Arc>), + C: FnOnce(Arc>, &[f32]) -> Arc>, + B: AsyncFnOnce(Arc>, Arc>), DP: DataProvider + for<'a> diskann::provider::SetElement<&'a [f32]>, { @@ -1620,14 +1632,16 @@ pub(crate) mod tests { (index, data) } - async fn build_using_single_insert(index: Arc>, data: Arc>) - where + async fn build_using_single_insert( + index: Arc>, + data: Arc>, + ) where DP: DataProvider + for<'a> diskann::provider::SetElement<&'a [f32]>, Quantized: for<'a> InsertStrategy<'a, DP, &'a [f32]> + Clone + Send + Sync, { let ctx = &DefaultContext; - for (i, vector) in data.row_iter().enumerate() { + for (i, vector) in data.rows().enumerate() { index .insert(&Quantized, ctx, &(i as u32), vector) .await @@ -1647,7 +1661,7 @@ pub(crate) mod tests { batchsize: NonZeroUsize::new(1).unwrap(), }; - let create_fn = |data: Arc>, start_point: &[f32]| { + let create_fn = |data: Arc>, start_point: &[f32]| { let quantizer = ScalarQuantizationParameters::default().train(data.as_view()); let (config, params) = parameters.materialize(data.nrows(), data.ncols()).unwrap(); @@ -1693,7 +1707,7 @@ pub(crate) mod tests { // that our simple graph search matches. // // Because this dataset is small, we can expect exact equality. - for (q, query) in data.row_iter().enumerate() { + for (q, query) in data.rows().enumerate() { let gt = groundtruth(data.as_view(), query, |a, b| SquaredL2::evaluate(a, b)); { let mut result_output_buffer = @@ -1752,7 +1766,7 @@ pub(crate) mod tests { batchsize: NonZeroUsize::new(1).unwrap(), }; - let create_fn = |data: Arc>, start_point: &[f32]| { + let create_fn = |data: Arc>, start_point: &[f32]| { let quantizer = ScalarQuantizationParameters::default().train(data.as_view()); let (config, params) = parameters.materialize(data.nrows(), data.ncols()).unwrap(); @@ -1799,7 +1813,7 @@ pub(crate) mod tests { // that our simple graph search matches. // // Because this dataset is small, we can expect exact equality. - for (q, query) in data.row_iter().enumerate() { + for (q, query) in data.rows().enumerate() { { let mut result_output_buffer = search_output_buffer::IdDistance::new(&mut ids, &mut distances); @@ -1845,7 +1859,7 @@ pub(crate) mod tests { let rng = &mut create_rnd_from_seed_in_tests(0x56870bccb0c44b66); - let create_fn = |data: Arc>, start_point: &[f32]| { + let create_fn = |data: Arc>, start_point: &[f32]| { let quantizer = diskann_quantization::spherical::SphericalQuantizer::train( data.as_view(), diskann_quantization::algorithms::transforms::TransformKind::PaddingHadamard { @@ -1875,10 +1889,10 @@ pub(crate) mod tests { index }; - let build_fn = async |index: Arc>, data: Arc>| { + let build_fn = async |index: Arc>, data: Arc>| { let ctx = &DefaultContext; let strategy = inmem::spherical::Quantized::build(); - for (i, vector) in data.row_iter().enumerate() { + for (i, vector) in data.rows().enumerate() { index .insert(&strategy, ctx, &(i as u32), vector) .await @@ -1912,7 +1926,7 @@ pub(crate) mod tests { // that our simple graph search matches. // // Because this dataset is small, we can expect exact equality. - for (q, query) in data.row_iter().enumerate() { + for (q, query) in data.rows().enumerate() { let gt = groundtruth(data.as_view(), query, |a, b| SquaredL2::evaluate(a, b)); // Full Precision Search. @@ -1958,7 +1972,7 @@ pub(crate) mod tests { let ctx = &DefaultContext; let rng = &mut create_rnd_from_seed_in_tests(0x56870bccb0c44b66); - let create_fn = |data: Arc>, start_points: &[f32]| { + let create_fn = |data: Arc>, start_points: &[f32]| { let quantizer = diskann_quantization::spherical::SphericalQuantizer::train( data.as_view(), diskann_quantization::algorithms::transforms::TransformKind::PaddingHadamard { @@ -1989,10 +2003,10 @@ pub(crate) mod tests { Arc::new(index) }; - let build_fn = async |index: Arc>, data: Arc>| { + let build_fn = async |index: Arc>, data: Arc>| { let ctx = &DefaultContext; let strategy = inmem::spherical::Quantized::build(); - for (i, vector) in data.row_iter().enumerate() { + for (i, vector) in data.rows().enumerate() { index .insert(&strategy, ctx, &(i as u32), vector) .await @@ -2027,7 +2041,7 @@ pub(crate) mod tests { // that our simple graph search matches. // // Because this dataset is small, we can expect exact equality. - for (q, query) in data.row_iter().enumerate() { + for (q, query) in data.rows().enumerate() { // Quantized Search let mut output = search_output_buffer::IdDistance::new(&mut ids, &mut distances); let strategy = inmem::spherical::Quantized::search( @@ -2068,7 +2082,7 @@ pub(crate) mod tests { #[tokio::test] async fn test_sift_pq_only_build_and_search(#[case] metric: Metric, #[case] file: &str) { let ctx = &DefaultContext; - let create_fn = |data: Arc>, start_points: &[f32]| { + let create_fn = |data: Arc>, start_points: &[f32]| { let pq_table = train_pq( data.as_view(), 32, @@ -2116,7 +2130,7 @@ pub(crate) mod tests { // that our simple graph search matches. // // Because this dataset is small, we can expect exact equality. - for (q, query) in data.row_iter().enumerate() { + for (q, query) in data.rows().enumerate() { let gt = groundtruth(data.as_view(), query, |a, b| SquaredL2::evaluate(a, b)); let mut result_output_buffer = @@ -2211,13 +2225,13 @@ pub(crate) mod tests { parameters: InitParams, file: &str, start_strategy: StartPointStrategy, - train_data: diskann_utils::views::MatrixView<'_, f32>, + train_data: diskann_utils::views::rowmajor::Ref<'_, f32>, ) where DefaultProvider: DataProvider + for<'a> SetElement<&'a [f32]> + SetStartPoints<[f32]>, S: for<'a> InsertStrategy<'a, DefaultProvider, &'a [f32]> - + MultiInsertStrategy, Matrix> + + MultiInsertStrategy, rowmajor::Owned> + Clone, { let ctx = &DefaultContext; @@ -2225,11 +2239,11 @@ pub(crate) mod tests { let mut iter = VectorDataIterator::<_, f32>::new(file, None, &storage).unwrap(); - let start_vectors: Matrix = start_strategy.compute(train_data).unwrap(); + let start_vectors: rowmajor::Owned = start_strategy.compute(train_data).unwrap(); index .provider() - .set_start_points(start_vectors.row_iter()) + .set_start_points(start_vectors.rows()) .unwrap(); let batchsize: usize = parameters.batchsize.into(); @@ -2243,8 +2257,9 @@ pub(crate) mod tests { } else { let mut i: u32 = 0; while let Some(data) = iter.next_n(batchsize) { - let mut vectors = Matrix::from_element(data.len(), start_vectors.ncols(), 0.0f32); - let ids: Arc<[_]> = std::iter::zip(vectors.row_iter_mut(), data.iter()) + let mut vectors = + rowmajor::Owned::from_element(data.len(), start_vectors.ncols(), 0.0f32); + let ids: Arc<[_]> = std::iter::zip(vectors.rows_mut(), data.iter()) .map(|(dst, (v, _))| { dst.copy_from_slice(v); let id = i; @@ -2267,10 +2282,10 @@ pub(crate) mod tests { file: &str, num_pq_chunks: usize, startpoint: StartPointStrategy, - ) -> (Arc, diskann_utils::views::Matrix) + ) -> (Arc, diskann_utils::views::rowmajor::Owned) where S: for<'a> InsertStrategy<'a, TestProvider, &'a [f32]> - + MultiInsertStrategy> + + MultiInsertStrategy> + Clone, { let storage = VirtualStorageProvider::new_overlay(test_data_root()); @@ -2317,7 +2332,7 @@ pub(crate) mod tests { S: for<'a> InsertStrategy<'a, TestProvider, &'a [f32]> + for<'a> SearchStrategy<'a, TestProvider, &'a [f32]> + for<'a> InplaceDeleteStrategy = &'a [f32]> - + MultiInsertStrategy> + + MultiInsertStrategy> + Clone, { let ctx = &DefaultContext; @@ -2417,7 +2432,7 @@ pub(crate) mod tests { S: for<'a> InsertStrategy<'a, TestProvider, &'a [f32]> + for<'a> SearchStrategy<'a, TestProvider, &'a [f32]> + for<'a> InplaceDeleteStrategy = &'a [f32]> - + MultiInsertStrategy> + + MultiInsertStrategy> + Clone, { let ctx = &DefaultContext; @@ -2666,12 +2681,12 @@ pub(crate) mod tests { let rng = &mut create_rnd_from_seed_in_tests(0x7dc205fcda38d3a3); indices.shuffle(rng); let mut queries = - diskann_utils::views::Matrix::from_element(data.nrows(), data.ncols(), 0.0); - std::iter::zip(queries.row_iter_mut(), indices.iter()).for_each(|(row, i)| { + diskann_utils::views::rowmajor::Owned::from_element(data.nrows(), data.ncols(), 0.0); + std::iter::zip(queries.rows_mut(), indices.iter()).for_each(|(row, i)| { row.copy_from_slice(data.row(*i)); }); - for (pos, query) in queries.row_iter().enumerate() { + for (pos, query) in queries.rows().enumerate() { index .insert( &Hybrid::new(max_fp_vecs_per_prune), @@ -2701,7 +2716,7 @@ pub(crate) mod tests { let mut ids = vec![0; top_k]; let mut distances = vec![0.0; top_k]; - for (q, query) in queries.row_iter().enumerate() { + for (q, query) in queries.rows().enumerate() { let gt = groundtruth(queries.as_view(), query, |a, b| SquaredL2::evaluate(a, b)); let mut result_output_buffer = search_output_buffer::IdDistance::new(&mut ids, &mut distances); diff --git a/diskann-providers/src/index/wrapped_async.rs b/diskann-providers/src/index/wrapped_async.rs index 9c8242f5a4..74f4fdfdbe 100644 --- a/diskann-providers/src/index/wrapped_async.rs +++ b/diskann-providers/src/index/wrapped_async.rs @@ -664,7 +664,7 @@ mod tests { provider::DefaultContext, utils::ONE, }; - use diskann_utils::test_data_root; + use diskann_utils::{test_data_root, views::rowmajor::Matrix}; use diskann_vector::distance::Metric; use super::DiskANNIndex; @@ -723,7 +723,7 @@ mod tests { let storage = VirtualStorageProvider::new_memory(); let ctx = DefaultContext; - for (i, v) in train_data.row_iter().enumerate() { + for (i, v) in train_data.rows().enumerate() { index.insert(&FullPrecision, &ctx, &(i as u32), v).unwrap(); } diff --git a/diskann-providers/src/model/graph/provider/async_/inmem/full_precision.rs b/diskann-providers/src/model/graph/provider/async_/inmem/full_precision.rs index 34e72cb340..f9cd9f07d1 100644 --- a/diskann-providers/src/model/graph/provider/async_/inmem/full_precision.rs +++ b/diskann-providers/src/model/graph/provider/async_/inmem/full_precision.rs @@ -23,7 +23,7 @@ use diskann::{ }; use diskann_utils::future::AsyncFriendly; -use diskann_utils::views::Matrix; +use diskann_utils::views::rowmajor::{self, MatrixMut}; use diskann_vector::{DistanceFunction, PreprocessedDistanceFunction, distance::Metric}; use crate::model::graph::provider::async_::{ @@ -434,10 +434,11 @@ where let candidates: Vec> = candidates.collect(); let candidate_count = candidates.len(); let store: &FullPrecisionStore = accessor.as_full_precision(); - let mut vectors = match Matrix::try_from_element(candidate_count, query.len(), 0.0f32) { - Ok(vectors) => vectors, - Err(error) => return std::future::ready(Err(error.into())), - }; + let mut vectors = + match rowmajor::Owned::try_from_element(candidate_count, query.len(), 0.0f32) { + Ok(vectors) => vectors, + Err(error) => return std::future::ready(Err(error.into())), + }; let mut ids = Vec::with_capacity(candidate_count); let mut distances = Vec::with_capacity(candidate_count); @@ -451,7 +452,7 @@ where } let indices = match determinant_diversity( - vectors.as_mut_view(), + vectors.as_view_mut(), &distances, query, candidate_count, diff --git a/diskann-providers/src/model/graph/provider/async_/inmem/scalar.rs b/diskann-providers/src/model/graph/provider/async_/inmem/scalar.rs index a0b0e3077f..273470033a 100644 --- a/diskann-providers/src/model/graph/provider/async_/inmem/scalar.rs +++ b/diskann-providers/src/model/graph/provider/async_/inmem/scalar.rs @@ -900,7 +900,7 @@ mod tests { use crate::storage::VirtualStorageProvider; use diskann::utils::ONE; use diskann_quantization::scalar::train::ScalarQuantizationParameters; - use diskann_utils::views::MatrixView; + use diskann_utils::views::rowmajor; use diskann_vector::distance::Metric; use rstest::rstest; @@ -918,7 +918,7 @@ mod tests { fn make_store(metric: Metric) -> SQStore { let quantizer = ScalarQuantizationParameters::default() - .train(MatrixView::try_from(&DATA, NPTS, DIM).unwrap()); + .train(rowmajor::Ref::try_from_data(&DATA, NPTS, DIM).unwrap()); SQStore::new(quantizer, /* capacity */ 5, metric, None) } diff --git a/diskann-providers/src/model/graph/provider/async_/inmem/spherical.rs b/diskann-providers/src/model/graph/provider/async_/inmem/spherical.rs index 596c86b3e2..be13022c00 100644 --- a/diskann-providers/src/model/graph/provider/async_/inmem/spherical.rs +++ b/diskann-providers/src/model/graph/provider/async_/inmem/spherical.rs @@ -703,7 +703,7 @@ mod tests { alloc::GlobalAllocator, spherical::{SphericalQuantizer, SupportedMetric}, }; - use diskann_utils::views::{Matrix, MatrixView}; + use diskann_utils::views::rowmajor::{self, Matrix}; use diskann_vector::{ DistanceFunction, PreprocessedDistanceFunction, PureDistanceFunction, distance::{InnerProduct, Metric, SquaredL2}, @@ -718,7 +718,7 @@ mod tests { //////////////// fn make_store( - data: MatrixView, + data: rowmajor::Ref, metric: SupportedMetric, rng: &mut StdRng, ) -> SphericalStore @@ -745,8 +745,8 @@ mod tests { ) } - fn dataset(nrows: usize, ncols: usize, rng: &mut StdRng) -> Matrix { - Matrix::from_fn(nrows, ncols, |_| StandardNormal {}.sample(rng)) + fn dataset(nrows: usize, ncols: usize, rng: &mut StdRng) -> rowmajor::Owned { + rowmajor::Owned::from_fn(nrows, ncols, |_| StandardNormal {}.sample(rng)) } #[test] @@ -853,12 +853,12 @@ mod tests { let max_relative_error = 0.25; - for (i, r) in data.row_iter().enumerate() { + for (i, r) in data.rows().enumerate() { store.set_vector(i, r).unwrap(); } - for (i, a) in data.row_iter().enumerate() { - for (j, b) in data.row_iter().enumerate().skip(i + 1) { + for (i, a) in data.rows().enumerate() { + for (j, b) in data.rows().enumerate().skip(i + 1) { let expected: f32 = SquaredL2::evaluate(a, b); let got: f32 = computer .evaluate_similarity( @@ -888,15 +888,15 @@ mod tests { let store = make_store::<1>(data.as_view(), SupportedMetric::InnerProduct, &mut rng); let computer = store.distance_computer().unwrap(); - for (i, r) in data.row_iter().enumerate() { + for (i, r) in data.rows().enumerate() { store.set_vector(i, r).unwrap(); } let mut signs_match = 0; let mut total = 0; - for (i, a) in data.row_iter().enumerate() { - for (j, b) in data.row_iter().enumerate().skip(i + 1) { + for (i, a) in data.rows().enumerate() { + for (j, b) in data.rows().enumerate().skip(i + 1) { total += 1; let expected: f32 = InnerProduct::evaluate(a, b); let got: f32 = computer @@ -933,15 +933,15 @@ mod tests { let store = make_store::<1>(data.as_view(), SupportedMetric::SquaredL2, &mut rng); let max_relative_error = 0.2; - for (i, r) in data.row_iter().enumerate() { + for (i, r) in data.rows().enumerate() { store.set_vector(i, r).unwrap(); } - for (i, a) in data.row_iter().enumerate() { + for (i, a) in data.rows().enumerate() { let computer = store .query_computer(a, spherical::iface::QueryLayout::FourBitTransposed, false) .unwrap(); - for (j, b) in data.row_iter().enumerate() { + for (j, b) in data.rows().enumerate() { if i == j { continue; } @@ -974,18 +974,18 @@ mod tests { { let store = make_store::<1>(data.as_view(), SupportedMetric::InnerProduct, &mut rng); - for (i, r) in data.row_iter().enumerate() { + for (i, r) in data.rows().enumerate() { store.set_vector(i, r).unwrap(); } let mut signs_match = 0; let mut total = 0; - for (i, a) in data.row_iter().enumerate() { + for (i, a) in data.rows().enumerate() { let computer = store .query_computer(a, spherical::iface::QueryLayout::FourBitTransposed, true) .unwrap(); - for (j, b) in data.row_iter().enumerate() { + for (j, b) in data.rows().enumerate() { if i == j { continue; } @@ -1027,7 +1027,7 @@ mod tests { assert!(matches!(err, RQError::CompressionError(..))); } - fn test_dataset() -> Matrix { + fn test_dataset() -> rowmajor::Owned { let data = vec![ 0.28657, -0.0318168, @@ -1159,6 +1159,6 @@ mod tests { -0.324718, // row 15 ]; - Matrix::try_from(data.into(), 16, 8).unwrap() + rowmajor::Owned::try_from_data(data.into(), 16, 8).unwrap() } } diff --git a/diskann-providers/src/model/graph/provider/determinant_diversity.rs b/diskann-providers/src/model/graph/provider/determinant_diversity.rs index 1f098931f4..fb4a9a272a 100644 --- a/diskann-providers/src/model/graph/provider/determinant_diversity.rs +++ b/diskann-providers/src/model/graph/provider/determinant_diversity.rs @@ -55,7 +55,7 @@ use std::fmt; -use diskann_utils::views::MutMatrixView; +use diskann_utils::views::rowmajor::{self, Matrix, MatrixMut}; use diskann_vector::{MathematicalValue, PureDistanceFunction, distance::InnerProduct}; /// Parameters for Determinant-Diversity post-processor with validation. @@ -187,7 +187,7 @@ struct DistanceRange { /// An empty candidate set, a `k` of zero, or zero-dimensional vectors yield an /// empty result. pub fn determinant_diversity( - candidates: MutMatrixView<'_, f32>, + candidates: rowmajor::Mut<'_, f32>, distances: &[f32], query: &[f32], k: usize, @@ -326,7 +326,7 @@ pub fn determinant_diversity( /// O(n * k * dim) -- for each of k pivots we touch all n residual rows of /// length `dim`. Memory is O(n * dim) for the contiguous residual matrix. fn greedy_orthogonal_select( - mut candidates: MutMatrixView<'_, f32>, + mut candidates: rowmajor::Mut<'_, f32>, distances: &[f32], k: usize, power: f32, @@ -469,7 +469,6 @@ fn dot_product(a: &[f32], b: &[f32]) -> f32 { mod tests { use super::*; use diskann_quantization::num::Positive; - use diskann_utils::views::Matrix; #[test] fn test_valid_params() { @@ -518,7 +517,7 @@ mod tests { } let dim = candidates[0].2.len(); - let mut matrix = Matrix::from_element(candidates.len(), dim, 0.0f32); + let mut matrix = rowmajor::Owned::from_element(candidates.len(), dim, 0.0f32); let mut ids = Vec::with_capacity(candidates.len()); let mut distances = Vec::with_capacity(candidates.len()); @@ -529,7 +528,7 @@ mod tests { } let params = DeterminantDiversityParams::new(power.into_inner(), eta).unwrap(); - determinant_diversity(matrix.as_mut_view(), &distances, query, k, ¶ms) + determinant_diversity(matrix.as_view_mut(), &distances, query, k, ¶ms) .expect("valid determinant-diversity inputs") .into_iter() .map(|idx| (ids[idx], distances[idx])) @@ -552,11 +551,11 @@ mod tests { // A zero-length query against non-empty candidates is a structural // mismatch (candidate columns != query dimension), not a valid request // that trivially returns nothing. - let mut matrix = Matrix::from_element(1, 2, 0.0f32); + let mut matrix = rowmajor::Owned::from_element(1, 2, 0.0f32); matrix.row_mut(0).copy_from_slice(&[1.0, 2.0]); let params = DeterminantDiversityParams::new(1.0, 0.5).unwrap(); - let result = determinant_diversity(matrix.as_mut_view(), &[0.5], &[], 5, ¶ms); + let result = determinant_diversity(matrix.as_view_mut(), &[0.5], &[], 5, ¶ms); assert!(matches!( result, Err(DeterminantDiversityError::QueryDimensionMismatch { @@ -570,12 +569,12 @@ mod tests { fn test_mismatched_dimensions_errors() { // Candidate vectors are 2-D, but the query is 3-D, so // `determinant_diversity` should report a dimension mismatch. - let mut matrix = Matrix::from_element(1, 2, 0.0f32); + let mut matrix = rowmajor::Owned::from_element(1, 2, 0.0f32); matrix.row_mut(0).copy_from_slice(&[1.0, 2.0]); let params = DeterminantDiversityParams::new(1.0, 0.5).unwrap(); let result = - determinant_diversity(matrix.as_mut_view(), &[0.5], &[1.0, 2.0, 3.0], 5, ¶ms); + determinant_diversity(matrix.as_view_mut(), &[0.5], &[1.0, 2.0, 3.0], 5, ¶ms); assert!(matches!( result, Err(DeterminantDiversityError::QueryDimensionMismatch { @@ -588,12 +587,12 @@ mod tests { #[test] fn test_mismatched_distances_errors() { // Two candidate rows but only one distance is a structural mismatch. - let mut matrix = Matrix::from_element(2, 2, 0.0f32); + let mut matrix = rowmajor::Owned::from_element(2, 2, 0.0f32); matrix.row_mut(0).copy_from_slice(&[1.0, 0.0]); matrix.row_mut(1).copy_from_slice(&[0.0, 1.0]); let params = DeterminantDiversityParams::new(1.0, 0.5).unwrap(); - let result = determinant_diversity(matrix.as_mut_view(), &[0.5], &[1.0, 1.0], 2, ¶ms); + let result = determinant_diversity(matrix.as_view_mut(), &[0.5], &[1.0, 1.0], 2, ¶ms); assert!(matches!( result, Err(DeterminantDiversityError::DistanceCountMismatch { diff --git a/diskann-providers/src/model/pq/debug.rs b/diskann-providers/src/model/pq/debug.rs index 5ecbe8c800..438bc98306 100644 --- a/diskann-providers/src/model/pq/debug.rs +++ b/diskann-providers/src/model/pq/debug.rs @@ -4,7 +4,7 @@ */ use diskann::utils::IntoUsize; -use diskann_utils::views; +use diskann_utils::views::{self, rowmajor::Matrix}; use diskann_vector::{PureDistanceFunction, distance::SquaredL2}; pub struct MismatchRecord { @@ -48,18 +48,18 @@ impl std::fmt::Display for MismatchRecord { /// on the sizes of the provided arguments, but can be helpful for writing test routines /// and as such is still marked as public. pub fn compare_pq( - data: views::MatrixView<'_, T>, + data: views::rowmajor::Ref<'_, T>, schema: diskann_quantization::views::ChunkOffsetsView<'_>, - pivots: views::MatrixView<'_, f32>, + pivots: views::rowmajor::Ref<'_, f32>, center: &[f32], - a: views::MatrixView<'_, U>, - b: views::MatrixView<'_, U>, + a: views::rowmajor::Ref<'_, U>, + b: views::rowmajor::Ref<'_, U>, ) -> Vec where T: Copy + Into, U: Copy + IntoUsize, { - std::iter::zip(a.row_iter(), b.row_iter()) + std::iter::zip(a.rows(), b.rows()) .enumerate() .flat_map(|(row, (a_row, b_row))| { std::iter::zip(a_row.iter(), b_row.iter()) diff --git a/diskann-providers/src/model/pq/fixed_chunk_pq_table.rs b/diskann-providers/src/model/pq/fixed_chunk_pq_table.rs index bfaba66853..511b009529 100644 --- a/diskann-providers/src/model/pq/fixed_chunk_pq_table.rs +++ b/diskann-providers/src/model/pq/fixed_chunk_pq_table.rs @@ -11,7 +11,10 @@ use diskann_quantization::{ }; use diskann_utils::{ lazy_format, - views::{self, MatrixBase, MatrixView}, + views::{ + self, + rowmajor::{self, Matrix}, + }, }; use diskann_vector::{PureDistanceFunction, distance}; use diskann_wide::ARCH; @@ -136,7 +139,7 @@ impl FixedChunkPQTable { pub fn new(dim: usize, pq_table: Box<[f32]>, chunk_offsets: Box<[usize]>) -> ANNResult { let len = pq_table.len(); let table = BasicTable::new( - MatrixBase::try_from(pq_table, len / dim, dim).bridge_err()?, + rowmajor::Owned::try_from_data(pq_table, len / dim, dim).bridge_err()?, ChunkOffsetsBase::new(chunk_offsets).bridge_err()?, ) .map_err(ANNError::new)?; @@ -168,7 +171,7 @@ impl FixedChunkPQTable { } let offsets: &[usize] = self.table.view_offsets().into(); - let table: &[f32] = self.table.view_pivots().into(); + let table: &[f32] = self.table.view_pivots().into_slice(); for centroid_index in 0..num_centers { let table_start = dim * centroid_index; @@ -301,7 +304,7 @@ impl FixedChunkPQTable { let mut accumulator = distance::simd::Resumable::new(T::init(ARCH)); - let pq_table: &[f32] = self.table.view_pivots().into(); + let pq_table: &[f32] = self.table.view_pivots().into_slice(); let chunk_offsets: &[usize] = self.table.view_offsets().into(); let mut start = chunk_offsets[0]; @@ -386,7 +389,7 @@ impl FixedChunkPQTable { assert_eq!(base_vec.len(), self.get_num_chunks()); assert_eq!(out.len(), self.get_dim()); let chunk_offsets: &[usize] = self.table.view_offsets().into(); - let pq_table: &[f32] = self.table.view_pivots().into(); + let pq_table: &[f32] = self.table.view_pivots().into_slice(); let dim = self.get_dim(); base_vec.iter().enumerate().for_each(|(i, b)| { @@ -406,7 +409,7 @@ impl FixedChunkPQTable { /// Returns an immutable reference to the `pq_table`. pub fn get_pq_table(&self) -> &[f32] { - self.table.view_pivots().into() + self.table.view_pivots().into_slice() } /// Returns an immutable reference to the `chunk_offsets`. @@ -419,8 +422,8 @@ impl FixedChunkPQTable { self.table.dim() } - /// Return the pivots as a `MatrixView`. - pub fn view_pivots(&self) -> views::MatrixView<'_, f32> { + /// Return the pivots as a `rowmajor::Ref`. + pub fn view_pivots(&self) -> views::rowmajor::Ref<'_, f32> { self.table.view_pivots() } @@ -497,8 +500,9 @@ fn pq_dist_lookup( ); } - let coordinates = MatrixView::::try_from(pq_coordinates, n_pts, pq_nchunks).bridge_err()?; - let distances = MatrixView::try_from( + let coordinates = + rowmajor::Ref::::try_from_data(pq_coordinates, n_pts, pq_nchunks).bridge_err()?; + let distances = rowmajor::Ref::try_from_data( &pq_dists[..NUM_PQ_CENTROIDS * pq_nchunks], pq_nchunks, NUM_PQ_CENTROIDS, @@ -589,8 +593,8 @@ unsafe fn add_distance_for_a_tile( tile_size: usize, cur_tile_size: usize, dists_out: &mut [f32], - coordinates: MatrixView<'_, u8>, - distances: MatrixView<'_, f32>, + coordinates: rowmajor::Ref<'_, u8>, + distances: rowmajor::Ref<'_, f32>, ) { dists_out.iter_mut().enumerate().for_each(|(point, d)| { for offset in 0..cur_tile_size { @@ -603,7 +607,7 @@ unsafe fn add_distance_for_a_tile( // SAFETY: From above, `chunk` is less than `coordinatges.ncols()`, which must // be equal to `distances.nrows()` by the pre-conditions for this function. - let row = unsafe { distances.get_row_unchecked(chunk) }; + let row = unsafe { distances.row_unchecked(chunk) }; // SAFETY: It's safe to query `row` with `centroid` since // it's less than 256(`NUM_PQ_CENTROIDS`) given that it's a u8. @@ -777,7 +781,7 @@ mod fixed_chunk_pq_table_test { fn conversion_rejects_too_many_centers() { let dim = 5; let table = BasicTable::new( - MatrixBase::try_from( + rowmajor::Owned::try_from_data( vec![0.0; dim * (NUM_PQ_CENTROIDS + 1)].into_boxed_slice(), NUM_PQ_CENTROIDS + 1, dim, @@ -831,10 +835,13 @@ mod fixed_chunk_pq_table_test { .unwrap(); // Calculate the expected output naively - let pq_data = MatrixView::try_from(&pq_data, n_pts, num_pq_chunks).unwrap(); - let distances = - MatrixView::try_from(&query_centroid_l2_distance, num_pq_chunks, NUM_PQ_CENTROIDS) - .unwrap(); + let pq_data = rowmajor::Ref::try_from_data(&pq_data, n_pts, num_pq_chunks).unwrap(); + let distances = rowmajor::Ref::try_from_data( + &query_centroid_l2_distance, + num_pq_chunks, + NUM_PQ_CENTROIDS, + ) + .unwrap(); let mut expected_pd_distance = vec![0.0; n_nbrs]; expected_pd_distance .iter_mut() diff --git a/diskann-providers/src/model/pq/pq_construction.rs b/diskann-providers/src/model/pq/pq_construction.rs index 5b70b76df6..f6202810e0 100644 --- a/diskann-providers/src/model/pq/pq_construction.rs +++ b/diskann-providers/src/model/pq/pq_construction.rs @@ -25,7 +25,7 @@ use diskann_quantization::{ }; use diskann_utils::{ io::Metadata, - views::{MatrixView, MutMatrixView}, + views::rowmajor::{self, Matrix, MatrixMut}, }; use rand::{Rng, distr::Distribution}; use rayon::prelude::*; @@ -89,7 +89,8 @@ where Storage: StorageWriteProvider, Random: Rng, { - MatrixView::try_from(&*train_data, parameters.num_train(), parameters.dim()).bridge_err()?; + rowmajor::Ref::try_from_data(&*train_data, parameters.num_train(), parameters.dim()) + .bridge_err()?; let centroid = if legacy_center_data { let mut centroid: Vec = vec![0.0; parameters.dim()]; @@ -118,7 +119,7 @@ where let full_pivot_data = pool.install(|| -> Result, ANNError> { let result = trainer .train( - MatrixView::try_from(train_data, parameters.num_train(), parameters.dim()) + rowmajor::Ref::try_from_data(train_data, parameters.num_train(), parameters.dim()) .bridge_err()?, chunk_offsets.as_view(), diskann_quantization::Parallelism::Rayon, @@ -217,7 +218,7 @@ pub fn generate_pq_pivots_from_membuf>( let result = trainer .train( - MatrixView::try_from( + rowmajor::Ref::try_from_data( train_data.as_slice(), parameters.num_train(), parameters.dim(), @@ -288,12 +289,12 @@ pub fn move_train_data_by_centroid( /// # Panics /// /// Panics if `y.len() != x.ncols()`. -pub fn accum_row_inplace(mut x: MutMatrixView, y: &[T]) +pub fn accum_row_inplace(mut x: rowmajor::Mut, y: &[T]) where T: Copy + std::ops::AddAssign, { assert_eq!(x.ncols(), y.len()); - x.row_iter_mut().for_each(|row| { + x.rows_mut().for_each(|row| { std::iter::zip(row.iter_mut(), y.iter()).for_each(|(a, b)| { *a += *b; }); @@ -416,12 +417,13 @@ where // process `BATCH_SIZE` many dataset vectors at a time. const BATCH_SIZE: usize = 128; - // Wrap the data in `MatrixViews` so we do not need to manually construct view + // Wrap the data in `rowmajor::Refs` so we do not need to manually construct view // in the compression loop. let mut compressed_block = - MutMatrixView::try_from(&mut block_compressed_base, cur_block_size, num_pq_chunks) + rowmajor::Mut::try_from_data(&mut block_compressed_base, cur_block_size, num_pq_chunks) .bridge_err()?; - let base_block = MatrixView::try_from(block_data, cur_block_size, full_dim).bridge_err()?; + let base_block = + rowmajor::Ref::try_from_data(block_data, cur_block_size, full_dim).bridge_err()?; base_block .par_window_iter(BATCH_SIZE) @@ -474,7 +476,7 @@ pub fn generate_pq_data_from_pivots_from_membuf_batch( } let table = BasicTableView::new( - MatrixView::try_from(pivot_data, parameters.num_centers(), dim).bridge_err()?, + rowmajor::Ref::try_from_data(pivot_data, parameters.num_centers(), dim).bridge_err()?, ChunkOffsetsView::new(offsets).bridge_err()?, ) .map_err(|err| ANNError::message(diskann_quantization::error::format(&err)))?; @@ -947,10 +949,12 @@ mod pq_test { .unwrap(); let membuf_view = - MatrixView::try_from(membuf_pq_data.as_slice(), num_train, num_pq_chunks).unwrap(); + rowmajor::Ref::try_from_data(membuf_pq_data.as_slice(), num_train, num_pq_chunks) + .unwrap(); let original_view = - MatrixView::try_from(original_pq_data.as_slice(), num_train, num_pq_chunks).unwrap(); + rowmajor::Ref::try_from_data(original_pq_data.as_slice(), num_train, num_pq_chunks) + .unwrap(); // Pre-emptively construct an offset view to compare mismatched slices. // We want to check that the difference in the mismatched chunks is small. @@ -961,7 +965,8 @@ mod pq_test { .unwrap(); let offset_view = chunk_offsets.as_view(); let full_data = - MatrixView::try_from(full_data_vector.as_slice(), num_train, train_dim).unwrap(); + rowmajor::Ref::try_from_data(full_data_vector.as_slice(), num_train, train_dim) + .unwrap(); let pivot_view = table.view_pivots(); let centroid = vec![0.0; train_dim]; diff --git a/diskann-providers/src/model/pq/views.rs b/diskann-providers/src/model/pq/views.rs index d9c757d743..b1e3dbbac5 100644 --- a/diskann-providers/src/model/pq/views.rs +++ b/diskann-providers/src/model/pq/views.rs @@ -13,9 +13,9 @@ convert_error!(Bridge); convert_error!(Bridge); convert_error!(Bridge); -impl From>> for ANNError { +impl From>> for ANNError { #[track_caller] - fn from(value: Bridge>) -> Self { + fn from(value: Bridge>) -> Self { ANNError::new(value.into_inner().as_static()) } } @@ -75,7 +75,7 @@ mod tests { let data = vec![0; ncols * nrows]; test_error(|| { - views::MatrixView::try_from(&*data, nrows, ncols + 1) + views::rowmajor::Ref::try_from_data(&data, nrows, ncols + 1) .bridge_err() .unwrap_err() }); diff --git a/diskann-providers/src/storage/index_storage.rs b/diskann-providers/src/storage/index_storage.rs index adb89cf62c..4eb7b30967 100644 --- a/diskann-providers/src/storage/index_storage.rs +++ b/diskann-providers/src/storage/index_storage.rs @@ -223,7 +223,10 @@ mod tests { provider::SetElement, utils::{IntoUsize, ONE}, }; - use diskann_utils::{test_data_root, views::MatrixView}; + use diskann_utils::{ + test_data_root, + views::rowmajor::{self, Matrix}, + }; use diskann_vector::distance::Metric; use super::*; @@ -239,14 +242,14 @@ mod tests { async fn build_index( index: &Arc>, strategy: S, - data: MatrixView<'_, f32>, + data: rowmajor::Ref<'_, f32>, ) where DP: DataProvider + for<'a> SetElement<&'a [f32]>, DP::Context: Default, S: for<'a> InsertStrategy<'a, DP, &'a [f32]> + Clone, { let ctx = &DP::Context::default(); - for (i, v) in data.row_iter().enumerate() { + for (i, v) in data.rows().enumerate() { index.insert(&strategy, ctx, &(i as u32), v).await.unwrap(); } } diff --git a/diskann-providers/src/storage/pq_storage.rs b/diskann-providers/src/storage/pq_storage.rs index 63ca583515..201c281b3f 100644 --- a/diskann-providers/src/storage/pq_storage.rs +++ b/diskann-providers/src/storage/pq_storage.rs @@ -12,7 +12,7 @@ use diskann::{ use diskann_quantization::{product::BasicTable, views::ChunkOffsetsBase}; use diskann_utils::{ io::{Metadata, read_bin, write_bin}, - views::{Matrix, MatrixView}, + views::rowmajor::{self, Matrix, MatrixMut}, }; use rand::Rng; use tracing::info; @@ -107,13 +107,16 @@ impl PQStorage { writer.seek(SeekFrom::Start(cumul_bytes[0] as u64))?; // Write PQ centroid vectors - let pivot_view = MatrixView::try_from(full_pivot_data, num_centers, dim)?; + let pivot_view = rowmajor::Ref::try_from_data(full_pivot_data, num_centers, dim)?; cumul_bytes[1] = cumul_bytes[0] + write_bin(pivot_view, writer)?; // Write the centroid of PQ centroid vectors let centroid_bytes = match centroid { - Some(centroid) => write_bin(MatrixView::column_vector(centroid), writer)?, - None => write_bin(Matrix::::from_element(dim, 1, 0.0).as_view(), writer)?, + Some(centroid) => write_bin(rowmajor::Ref::column_vector(centroid), writer)?, + None => write_bin( + rowmajor::Owned::::from_element(dim, 1, 0.0).as_view(), + writer, + )?, }; cumul_bytes[2] = cumul_bytes[1] + centroid_bytes; @@ -121,14 +124,14 @@ impl PQStorage { let chunk_offsets_u32: Vec = chunk_offsets.iter().map(|&x| x as u32).collect(); cumul_bytes[3] = cumul_bytes[2] + write_bin( - MatrixView::column_vector(chunk_offsets_u32.as_slice()), + rowmajor::Ref::column_vector(chunk_offsets_u32.as_slice()), writer, )?; // Seek back to offset 0 and write the offset table. let cumul_bytes_u64: Vec = cumul_bytes.iter().map(|&x| x as u64).collect(); write_bin_from( - MatrixView::column_vector(cumul_bytes_u64.as_slice()), + rowmajor::Ref::column_vector(cumul_bytes_u64.as_slice()), writer, 0, )?; @@ -170,7 +173,7 @@ impl PQStorage { self.read_pivot_file(&self.pivot_data_path, storage_provider)?; if centroid.as_slice().iter().any(|c| *c != 0.0) { - accum_row_inplace(pivots.as_mut_view(), centroid.as_slice()) + accum_row_inplace(pivots.as_view_mut(), centroid.as_slice()) } Self::pivot_data_into_basic_table(&self.pivot_data_path, pivots, chunk_offsets) @@ -184,7 +187,7 @@ impl PQStorage { num_points_to_load: usize, num_pq_chunks: usize, storage_provider: &Storage, - ) -> ANNResult> { + ) -> ANNResult> { info!( "Loading compressed from pq compressed data file {}...", pq_compressed_data, @@ -215,7 +218,11 @@ impl PQStorage { &self, pq_pivots: &str, storage_provider: &Storage, - ) -> ANNResult<(Matrix, Matrix, Matrix)> { + ) -> ANNResult<( + rowmajor::Owned, + rowmajor::Owned, + rowmajor::Owned, + )> { if !storage_provider.exists(pq_pivots) { return Err(ANNError::message(format!( "ERROR: PQ k-means pivot file not found: {pq_pivots}." @@ -282,8 +289,8 @@ impl PQStorage { fn pivot_data_into_basic_table( pq_pivots: &str, - pivots: Matrix, - chunk_offsets: Matrix, + pivots: rowmajor::Owned, + chunk_offsets: rowmajor::Owned, ) -> ANNResult { let offsets = ChunkOffsetsBase::new(chunk_offsets.into_inner()).map_err(|err| { ANNError::message(format!( @@ -504,26 +511,30 @@ mod pq_storage_tests { let pivots = [0.0, 1.0, 2.0, 3.0]; cumul_bytes[1] = cumul_bytes[0] + write_bin( - MatrixView::try_from(pivots.as_slice(), 2, 2).unwrap(), + rowmajor::Ref::try_from_data(pivots.as_slice(), 2, 2).unwrap(), &mut writer, ) .unwrap(); let centroid = [0.0, 0.0]; cumul_bytes[2] = cumul_bytes[1] - + write_bin(MatrixView::column_vector(centroid.as_slice()), &mut writer).unwrap(); + + write_bin( + rowmajor::Ref::column_vector(centroid.as_slice()), + &mut writer, + ) + .unwrap(); let chunk_offsets = [0_u32, 2_u32]; cumul_bytes[3] = cumul_bytes[2] + write_bin( - MatrixView::try_from(chunk_offsets.as_slice(), 1, 2).unwrap(), + rowmajor::Ref::try_from_data(chunk_offsets.as_slice(), 1, 2).unwrap(), &mut writer, ) .unwrap(); let offsets: Vec = cumul_bytes.iter().map(|&offset| offset as u64).collect(); write_bin_from( - MatrixView::column_vector(offsets.as_slice()), + rowmajor::Ref::column_vector(offsets.as_slice()), &mut writer, 0, ) @@ -568,7 +579,11 @@ mod pq_storage_tests { { let mut writer = storage_provider.create_for_write(pivot_path).unwrap(); let offsets = [METADATA_SIZE as u64, 0, 0]; - write_bin(MatrixView::column_vector(offsets.as_slice()), &mut writer).unwrap(); + write_bin( + rowmajor::Ref::column_vector(offsets.as_slice()), + &mut writer, + ) + .unwrap(); } assert!( diff --git a/diskann-providers/src/test_utils/search_utils.rs b/diskann-providers/src/test_utils/search_utils.rs index 6265b3cc69..dfa9482552 100644 --- a/diskann-providers/src/test_utils/search_utils.rs +++ b/diskann-providers/src/test_utils/search_utils.rs @@ -4,7 +4,7 @@ */ use diskann::neighbor::{self, Neighbor}; -use diskann_utils::views::MatrixView; +use diskann_utils::views::rowmajor::{self, Matrix}; /// Compute the ground truth for a small dataset. /// @@ -13,12 +13,12 @@ use diskann_utils::views::MatrixView; /// /// This allows filtering by `is_match` to be much more efficient because it decreases /// the number of elements that have to be moved. -pub fn groundtruth(data: MatrixView, query: &[T], f: F) -> Vec> +pub fn groundtruth(data: rowmajor::Ref, query: &[T], f: F) -> Vec> where F: Fn(&[T], &[T]) -> f32, { let mut results: Vec<_> = data - .row_iter() + .rows() .enumerate() .map(|(i, row)| Neighbor::new(i as u32, f(row, query))) .collect(); diff --git a/diskann-providers/src/utils/file_util.rs b/diskann-providers/src/utils/file_util.rs index 2d577de30e..abf45404cd 100644 --- a/diskann-providers/src/utils/file_util.rs +++ b/diskann-providers/src/utils/file_util.rs @@ -10,7 +10,11 @@ use std::io::Read; use crate::storage::StorageReadProvider; use diskann::{ANNError, ANNResult, utils::IntoUsize}; -use diskann_utils::{io::Metadata, lazy_format, views::Matrix}; +use diskann_utils::{ + io::Metadata, + lazy_format, + views::rowmajor::{self, MatrixMut}, +}; /// Read metadata of data file. pub fn load_metadata_from_file( @@ -42,7 +46,7 @@ pub fn file_exists( pub fn load_multivec_bin( storage_read_provider: &StorageReader, bin_file: &str, -) -> ANNResult<(Vec>, usize, usize, usize)> { +) -> ANNResult<(Vec>, usize, usize, usize)> { let mut reader = storage_read_provider.open_reader(bin_file)?; let (num_points, dimension, total_results) = { @@ -85,10 +89,10 @@ pub fn load_multivec_bin> = Vec::with_capacity(num_points); + let mut all_vectors: Vec> = Vec::with_capacity(num_points); for &length in &vec_lengths { - let mut vectors = Matrix::from_element(length as usize, dimension, T::default()); + let mut vectors = rowmajor::Owned::from_element(length as usize, dimension, T::default()); reader.read_exact(bytemuck::must_cast_slice_mut::( vectors.as_mut_slice(), ))?; @@ -102,6 +106,7 @@ pub fn load_multivec_bin(reader: &mut Reader) -> std::io::Result<(us pub fn read_bin_from( reader: &mut (impl Read + Seek), offset: usize, -) -> Result, ReadBinError> { +) -> Result, ReadBinError> { reader.seek(std::io::SeekFrom::Start(offset as u64))?; read_bin(reader) } /// Write a matrix at the given byte offset. pub fn write_bin_from( - data: MatrixView<'_, T>, + data: rowmajor::Ref<'_, T>, writer: &mut (impl Write + Seek), offset: usize, ) -> Result { @@ -108,6 +108,7 @@ pub fn save_data_in_base_dimensions Matrix { + /// Retrieve the 8x8 hadamard matrix as a `rowmajor::Owned`. + fn get_hadamard_8() -> rowmajor::Owned { let v: Box<[f32]> = HADAMARD_8.iter().flatten().copied().collect(); - Matrix::try_from(v, 8, 8).unwrap() + rowmajor::Owned::try_from_data(v, 8, 8).unwrap() } - fn hadamard_by_sylvester(dim: usize) -> Matrix { + fn hadamard_by_sylvester(dim: usize) -> rowmajor::Owned { assert_ne!(dim, 0); // Base case. if dim == 1 { - Matrix::from_element(dim, dim, 1.0) + rowmajor::Owned::from_element(dim, dim, 1.0) } else { let half = dim / 2; let sub = hadamard_by_sylvester(half); - let mut m = Matrix::::from_element(dim, dim, 0.0); + let mut m = rowmajor::Owned::::from_element(dim, dim, 0.0); for c in 0..m.ncols() { for r in 0..m.nrows() { @@ -423,9 +423,9 @@ mod tests { } // A naive reference implementation. - fn matmul(a: MatrixView, b: MatrixView) -> Matrix { + fn matmul(a: rowmajor::Ref, b: rowmajor::Ref) -> rowmajor::Owned { assert_eq!(a.ncols(), b.nrows()); - let mut c = Matrix::from_element(a.nrows(), b.ncols(), 0.0); + let mut c = rowmajor::Owned::from_element(a.nrows(), b.ncols(), 0.0); for i in 0..c.nrows() { for j in 0..c.ncols() { @@ -443,7 +443,7 @@ mod tests { fn test_micro_kernel_64() { let mut src = { let mut rng = StdRng::seed_from_u64(0xde1936d651285fc8); - Matrix::from_fn(64, 1, |_| StandardUniform {}.sample(&mut rng)) + rowmajor::Owned::from_fn(64, 1, |_| StandardUniform {}.sample(&mut rng)) }; let h = hadamard_by_sylvester(64); @@ -475,7 +475,7 @@ mod tests { fn test_hadamard_transform(dim: usize, seed: u64) { let src = { let mut rng = StdRng::seed_from_u64(seed); - Matrix::from_fn(dim, 1, |_| StandardUniform {}.sample(&mut rng)) + rowmajor::Owned::from_fn(dim, 1, |_| StandardUniform {}.sample(&mut rng)) }; let h = hadamard_by_sylvester(dim); diff --git a/diskann-quantization/src/algorithms/kmeans/lloyds.rs b/diskann-quantization/src/algorithms/kmeans/lloyds.rs index 53818c7f49..864f0caa34 100644 --- a/diskann-quantization/src/algorithms/kmeans/lloyds.rs +++ b/diskann-quantization/src/algorithms/kmeans/lloyds.rs @@ -9,7 +9,7 @@ use super::common::square_norm; use crate::multi_vector::{BlockTransposed, BlockTransposedRef}; use diskann_utils::{ strided::Strided, - views::{Matrix, MatrixView, MutMatrixView}, + views::rowmajor::{self, Matrix, MatrixMut}, }; //////////////////////////////// @@ -28,7 +28,7 @@ diskann_wide::alias!(u32s = u32x8); pub fn distances_in_place( dataset: BlockTransposedRef<'_, f32, 16>, data_norms: &[f32], - centers: MatrixView<'_, f32>, + centers: rowmajor::Ref<'_, f32>, center_norms: &[f32], nearest: &mut [u32], ) -> f32 { @@ -342,8 +342,8 @@ fn update((d0, i0): (f32s, u32s), (d1, i1): (f32s, u32s)) -> (f32s, u32s) { // Update Step // ///////////////// -fn update_centroids(mut centers: MutMatrixView<'_, f32>, data: Strided<'_, f32>, map: &[u32]) { - let mut sums = Matrix::::from_element(centers.nrows(), centers.ncols(), 0.0); +fn update_centroids(mut centers: rowmajor::Mut<'_, f32>, data: Strided<'_, f32>, map: &[u32]) { + let mut sums = rowmajor::Owned::::from_element(centers.nrows(), centers.ncols(), 0.0); let mut counts: Vec = vec![0; centers.nrows()]; data.rows().zip(map.iter()).for_each(|(row, ¢er)| { counts[center as usize] += 1; @@ -353,8 +353,8 @@ fn update_centroids(mut centers: MutMatrixView<'_, f32>, data: Strided<'_, f32>, }); }); - std::iter::zip(counts.iter(), sums.row_iter()) - .zip(centers.row_iter_mut()) + std::iter::zip(counts.iter(), sums.rows()) + .zip(centers.rows_mut()) .for_each(|((count, sum), center)| { // If the count is zero - we do not want to divide by it because that will // result in `NaN`. @@ -373,7 +373,7 @@ pub(crate) fn lloyds_inner( data: Strided<'_, f32>, square_norms: &[f32], transpose: BlockTransposedRef<'_, f32, 16>, - mut centers: MutMatrixView<'_, f32>, + mut centers: rowmajor::Mut<'_, f32>, max_reps: usize, ) -> (Vec, f32) { // Check our requirements. @@ -401,7 +401,7 @@ pub(crate) fn lloyds_inner( "data and centers should have the same dimensions" ); - let mut center_square_norms: Vec = centers.row_iter().map(square_norm).collect(); + let mut center_square_norms: Vec = centers.rows().map(square_norm).collect(); let mut assignments: Vec = vec![0; num_data]; let mut residual = 0.0; @@ -413,9 +413,9 @@ pub(crate) fn lloyds_inner( ¢er_square_norms, &mut assignments, ); - update_centroids(centers.as_mut_view(), data, &assignments); + update_centroids(centers.as_view_mut(), data, &assignments); if i != max_reps - 1 { - std::iter::zip(center_square_norms.iter_mut(), centers.row_iter()).for_each( + std::iter::zip(center_square_norms.iter_mut(), centers.rows()).for_each( |(c, center)| { *c = square_norm(center); }, @@ -439,8 +439,8 @@ pub(crate) fn lloyds_inner( /// Panics if `data.ncols() != centers.ncols()`. The data and centers must have the same /// dimension. pub fn lloyds( - data: MatrixView<'_, f32>, - centers: MutMatrixView<'_, f32>, + data: rowmajor::Ref<'_, f32>, + centers: rowmajor::Mut<'_, f32>, max_reps: usize, ) -> (Vec, f32) { assert_eq!( @@ -450,7 +450,7 @@ pub fn lloyds( ); let transpose = BlockTransposed::::from_matrix_view(data); - let square_norms: Vec = data.row_iter().map(square_norm).collect(); + let square_norms: Vec = data.rows().map(square_norm).collect(); lloyds_inner( data.into(), &square_norms, @@ -464,7 +464,6 @@ pub fn lloyds( mod tests { #[cfg(not(miri))] use diskann_utils::lazy_format; - use diskann_utils::views::Matrix; use diskann_vector::{PureDistanceFunction, distance::SquaredL2}; use rand::{Rng, SeedableRng, rngs::StdRng, seq::SliceRandom}; #[cfg(not(miri))] @@ -494,21 +493,21 @@ mod tests { ) { let context = lazy_format!("ncenters = {}, ndata = {}, dim = {}", ncenters, ndata, dim,); - let mut centers = Matrix::from_element(ncenters, dim, 0.0); - let mut data = Matrix::from_element(ndata, dim, 0.0); + let mut centers = rowmajor::Owned::from_element(ncenters, dim, 0.0); + let mut data = rowmajor::Owned::from_element(ndata, dim, 0.0); // A list of random "nice" offsets that get applied to each center and data point // to ensure proper visitation during computation. let offsets = [-0.125, -0.0625, -0.03125, 0.03125, 0.0625, 0.125]; // Initialize `centers` uniformly but with random offsets applied to each dimension. - for (i, row) in centers.row_iter_mut().enumerate() { + for (i, row) in centers.rows_mut().enumerate() { for c in row { *c = (i as f32) + *offsets.choose(rng).unwrap(); } } - let center_norms: Vec = centers.row_iter().map(square_norm).collect(); + let center_norms: Vec = centers.rows().map(square_norm).collect(); // This is the distribution of how we assign data points to centers. let assignment_distribution = Uniform::::new(0, centers.nrows()).unwrap(); @@ -518,13 +517,13 @@ mod tests { .map(|_| assignment_distribution.sample(rng)) .collect(); - for (assignment, row) in std::iter::zip(assignments.iter(), data.row_iter_mut()) { + for (assignment, row) in std::iter::zip(assignments.iter(), data.rows_mut()) { for c in row.iter_mut() { *c = (*assignment as f32) + offsets.choose(rng).unwrap() } } - let data_norms: Vec = data.row_iter().map(square_norm).collect(); + let data_norms: Vec = data.rows().map(square_norm).collect(); let residual = distances_in_place( BlockTransposed::::from_matrix_view(data.as_view()).as_view(), @@ -556,7 +555,7 @@ mod tests { // Check that the residual computation is correct. let mut sum: f32 = 0.0; - for (a, row) in std::iter::zip(assignments.iter(), data.row_iter()) { + for (a, row) in std::iter::zip(assignments.iter(), data.rows()) { let distance: f32 = SquaredL2::evaluate(row, centers.row(*a)); sum += distance; } @@ -583,8 +582,8 @@ mod tests { // We do not perform any value-dependent control-flow for memory accesses. // Therefore, the miri tests don't require any setup (this helps everything run faseter). fn test_miri_distances_in_place_impl(ndata: usize, ncenters: usize, dim: usize) { - let centers = Matrix::from_element(ncenters, dim, 0.0); - let data = Matrix::from_element(ndata, dim, 0.0); + let centers = rowmajor::Owned::from_element(ncenters, dim, 0.0); + let data = rowmajor::Owned::from_element(ndata, dim, 0.0); let data_norms = vec![0.0; ndata]; let center_norms = vec![0.0; ncenters]; let mut nearest = vec![0; ndata]; @@ -659,8 +658,8 @@ mod tests { let mut center_order: Vec = (0..setup.ncenters).collect(); let mut data = - Matrix::from_element(setup.ncenters * setup.data_per_center, setup.ndim, 0.0); - let mut centers = Matrix::from_element(setup.ncenters, setup.ndim, 0.0); + rowmajor::Owned::from_element(setup.ncenters * setup.data_per_center, setup.ndim, 0.0); + let mut centers = rowmajor::Owned::from_element(setup.ncenters, setup.ndim, 0.0); for trial in 0..setup.ntrials { values.shuffle(rng); @@ -668,20 +667,20 @@ mod tests { // Populate centers assert_eq!(center_order.len(), centers.nrows()); - for (c, row) in std::iter::zip(center_order.iter(), centers.row_iter_mut()) { + for (c, row) in std::iter::zip(center_order.iter(), centers.rows_mut()) { row.fill((setup.step_between_clusters * c) as f32 - 1.0); } // Populate data. assert_eq!(values.len(), data.nrows()); - for (d, row) in std::iter::zip(values.iter(), data.row_iter_mut()) { + for (d, row) in std::iter::zip(values.iter(), data.rows_mut()) { row.fill(*d as f32); } // Run 2 iteration of lloyds. // The second iteration ensures that we recompute norms properly. let lloyds_iter = 2; - let (assignments, loss) = lloyds(data.as_view(), centers.as_mut_view(), lloyds_iter); + let (assignments, loss) = lloyds(data.as_view(), centers.as_view_mut(), lloyds_iter); // Make sure all the assignments are returned correctly. assert_eq!(assignments.len(), values.len()); @@ -711,7 +710,7 @@ mod tests { }); // Verify the loss is correct. - let expected_loss: f32 = std::iter::zip(assignments.iter(), data.row_iter()) + let expected_loss: f32 = std::iter::zip(assignments.iter(), data.rows()) .map(|(a, row)| -> f32 { let c = centers.row(*a as usize); SquaredL2::evaluate(row, c) @@ -752,9 +751,9 @@ mod tests { #[test] #[should_panic(expected = "dataset and data norms should have the same length")] fn distances_in_place_panics_data_norms() { - let data = Matrix::from_element(5, 8, 0.0); + let data = rowmajor::Owned::from_element(5, 8, 0.0); let data_norms = vec![0.0; data.nrows() + 1]; // Incorrect - let centers = Matrix::from_element(2, 8, 0.0); + let centers = rowmajor::Owned::from_element(2, 8, 0.0); let center_norms = vec![0.0; centers.nrows()]; let mut nearest = vec![0; data.nrows()]; distances_in_place( @@ -769,9 +768,9 @@ mod tests { #[test] #[should_panic(expected = "dataset and centers should have the same dimension")] fn distances_in_place_panics_different_dim() { - let data = Matrix::from_element(5, 8, 0.0); + let data = rowmajor::Owned::from_element(5, 8, 0.0); let data_norms = vec![0.0; data.nrows()]; - let centers = Matrix::from_element(2, 9, 0.0); // Incorrect + let centers = rowmajor::Owned::from_element(2, 9, 0.0); // Incorrect let center_norms = vec![0.0; centers.nrows()]; let mut nearest = vec![0; data.nrows()]; distances_in_place( @@ -786,9 +785,9 @@ mod tests { #[test] #[should_panic(expected = "centers and center norms should have the same length")] fn distances_in_place_panics_center_norms() { - let data = Matrix::from_element(5, 8, 0.0); + let data = rowmajor::Owned::from_element(5, 8, 0.0); let data_norms = vec![0.0; data.nrows()]; - let centers = Matrix::from_element(2, 8, 0.0); + let centers = rowmajor::Owned::from_element(2, 8, 0.0); let center_norms = vec![0.0; centers.nrows() + 1]; // Incorrect let mut nearest = vec![0; data.nrows()]; distances_in_place( @@ -803,9 +802,9 @@ mod tests { #[test] #[should_panic(expected = "dataset and nearest-buffer should have the same length")] fn distances_in_place_panics_nearest() { - let data = Matrix::from_element(5, 8, 0.0); + let data = rowmajor::Owned::from_element(5, 8, 0.0); let data_norms = vec![0.0; data.nrows()]; - let centers = Matrix::from_element(2, 8, 0.0); + let centers = rowmajor::Owned::from_element(2, 8, 0.0); let center_norms = vec![0.0; centers.nrows()]; let mut nearest = vec![0; data.nrows() + 1]; // Incorrect distances_in_place( @@ -824,14 +823,14 @@ mod tests { #[test] #[should_panic(expected = "data and norms should have the same length")] fn lloyds_inner_panics_norms_length() { - let data = Matrix::from_element(5, 8, 0.0); + let data = rowmajor::Owned::from_element(5, 8, 0.0); let square_norms = vec![0.0; data.nrows() + 1]; // Incorrect - let mut centers = Matrix::from_element(2, 8, 0.0); + let mut centers = rowmajor::Owned::from_element(2, 8, 0.0); lloyds_inner( data.as_view().into(), &square_norms, BlockTransposed::::from_matrix_view(data.as_view()).as_view(), - centers.as_mut_view(), + centers.as_view_mut(), 1, ); } @@ -839,15 +838,15 @@ mod tests { #[test] #[should_panic(expected = "data and transpose should have the same length")] fn lloyds_inner_panics_transpose_length() { - let data = Matrix::from_element(5, 8, 0.0); - let data_incorrect = Matrix::from_element(5 + 1, 8, 0.0); // Incorrect + let data = rowmajor::Owned::from_element(5, 8, 0.0); + let data_incorrect = rowmajor::Owned::from_element(5 + 1, 8, 0.0); // Incorrect let square_norms = vec![0.0; data.nrows()]; - let mut centers = Matrix::from_element(2, 8, 0.0); + let mut centers = rowmajor::Owned::from_element(2, 8, 0.0); lloyds_inner( data.as_view().into(), &square_norms, BlockTransposed::::from_matrix_view(data_incorrect.as_view()).as_view(), - centers.as_mut_view(), + centers.as_view_mut(), 1, ); } @@ -855,15 +854,15 @@ mod tests { #[test] #[should_panic(expected = "data and transpose should have the same dimensions")] fn lloyds_inner_panics_transpose_dim() { - let data = Matrix::from_element(5, 8, 0.0); - let data_incorrect = Matrix::from_element(5, 8 + 1, 0.0); // Incorrect + let data = rowmajor::Owned::from_element(5, 8, 0.0); + let data_incorrect = rowmajor::Owned::from_element(5, 8 + 1, 0.0); // Incorrect let square_norms = vec![0.0; data.nrows()]; - let mut centers = Matrix::from_element(2, 8, 0.0); + let mut centers = rowmajor::Owned::from_element(2, 8, 0.0); lloyds_inner( data.as_view().into(), &square_norms, BlockTransposed::::from_matrix_view(data_incorrect.as_view()).as_view(), // Incorrect - centers.as_mut_view(), + centers.as_view_mut(), 1, ); } @@ -871,14 +870,14 @@ mod tests { #[test] #[should_panic(expected = "data and centers should have the same dimensions")] fn lloyds_inner_panics_centers_dim() { - let data = Matrix::from_element(5, 8, 0.0); + let data = rowmajor::Owned::from_element(5, 8, 0.0); let square_norms = vec![0.0; data.nrows()]; - let mut centers = Matrix::from_element(2, 8 + 1, 0.0); // Incorrect + let mut centers = rowmajor::Owned::from_element(2, 8 + 1, 0.0); // Incorrect lloyds_inner( data.as_view().into(), &square_norms, BlockTransposed::::from_matrix_view(data.as_view()).as_view(), - centers.as_mut_view(), + centers.as_view_mut(), 1, ); } @@ -890,8 +889,8 @@ mod tests { #[test] #[should_panic(expected = "data and centers must have the same dimension")] fn lloyds_panics_dim_mismatch() { - let data = Matrix::from_element(5, 8, 0.0); - let mut centers = Matrix::from_element(5, 8 + 1, 0.0); // Incorrect - lloyds(data.as_view(), centers.as_mut_view(), 1); + let data = rowmajor::Owned::from_element(5, 8, 0.0); + let mut centers = rowmajor::Owned::from_element(5, 8 + 1, 0.0); // Incorrect + lloyds(data.as_view(), centers.as_view_mut(), 1); } } diff --git a/diskann-quantization/src/algorithms/kmeans/plusplus.rs b/diskann-quantization/src/algorithms/kmeans/plusplus.rs index ef4f3ea208..2a4211be64 100644 --- a/diskann-quantization/src/algorithms/kmeans/plusplus.rs +++ b/diskann-quantization/src/algorithms/kmeans/plusplus.rs @@ -7,7 +7,7 @@ use std::{collections::HashSet, fmt}; use diskann_utils::{ strided::Strided, - views::{MatrixView, MutMatrixView}, + views::rowmajor::{self, Matrix, MatrixMut}, }; use diskann_wide::{SIMDMulAdd, SIMDPartialOrd, SIMDSelect, SIMDVector}; use rand::{ @@ -379,7 +379,7 @@ impl KMeansPlusPlusError { } pub(crate) fn kmeans_plusplus_into_inner( - mut points: MutMatrixView<'_, f32>, + mut points: rowmajor::Mut<'_, f32>, data: Strided<'_, f32>, transpose: BlockTransposedRef<'_, f32, N>, norms: &[f32], @@ -498,8 +498,8 @@ where } pub fn kmeans_plusplus_into( - centers: MutMatrixView<'_, f32>, - data: MatrixView<'_, f32>, + centers: rowmajor::Mut<'_, f32>, + data: rowmajor::Ref<'_, f32>, rng: &mut dyn RngCore, ) -> Result<(), KMeansPlusPlusError> { assert_eq!( @@ -511,7 +511,7 @@ pub fn kmeans_plusplus_into( const GROUPSIZE: usize = 16; let mut norms: Vec = vec![0.0; data.nrows()]; - for (n, d) in std::iter::zip(norms.iter_mut(), data.row_iter()) { + for (n, d) in std::iter::zip(norms.iter_mut(), data.rows()) { *n = square_norm(d); } @@ -521,21 +521,21 @@ pub fn kmeans_plusplus_into( #[cfg(test)] mod tests { - use diskann_utils::{lazy_format, views::Matrix}; + use diskann_utils::lazy_format; use diskann_vector::{PureDistanceFunction, distance::SquaredL2}; use rand::{Rng, SeedableRng, rngs::StdRng, seq::SliceRandom}; use super::*; use crate::utils; - fn is_in(needle: &[f32], haystack: MatrixView<'_, f32>) -> bool { + fn is_in(needle: &[f32], haystack: rowmajor::Ref<'_, f32>) -> bool { assert_eq!(needle.len(), haystack.ncols()); - haystack.row_iter().any(|row| row == needle) + haystack.rows().any(|row| row == needle) } fn check_post_conditions( - centers: MatrixView<'_, f32>, - data: MatrixView<'_, f32>, + centers: rowmajor::Ref<'_, f32>, + data: rowmajor::Ref<'_, f32>, err: &KMeansPlusPlusError, ) { assert_eq!(err.expected, centers.nrows()); @@ -581,8 +581,8 @@ mod tests { /// ... /// K-1, K, K+1, K+3 ... N+K-2 /// ``` - fn set_default_values(mut x: MutMatrixView<'_, f32>) { - for (i, row) in x.row_iter_mut().enumerate() { + fn set_default_values(mut x: rowmajor::Mut<'_, f32>) { + for (i, row) in x.rows_mut().enumerate() { for (j, r) in row.iter_mut().enumerate() { *r = (i + j) as f32; } @@ -609,14 +609,14 @@ mod tests { dim ); - let mut data = Matrix::::from_element(num_points, dim, 0.0); - set_default_values(data.as_mut_view()); + let mut data = rowmajor::Owned::::from_element(num_points, dim, 0.0); + set_default_values(data.as_view_mut()); - let square_norms: Vec = data.row_iter().map(square_norm).collect(); + let square_norms: Vec = data.rows().map(square_norm).collect(); // The sample points we are computing the distances against. let num_samples = 3; - let mut samples = Matrix::::from_element(num_samples, dim, 0.0); + let mut samples = rowmajor::Owned::::from_element(num_samples, dim, 0.0); let mut distances = vec![f32::INFINITY; num_points]; let distribution = Uniform::::new(0, (num_points + dim) as u32).unwrap(); let transpose = BlockTransposed::::from_matrix_view(data.as_view()); @@ -642,7 +642,7 @@ mod tests { ); // Make sure all the distances are correct. - for (n, (d, data)) in std::iter::zip(distances.iter(), data.row_iter()).enumerate() { + for (n, (d, data)) in std::iter::zip(distances.iter(), data.rows()).enumerate() { let mut min_distance = f32::INFINITY; for j in 0..=i { let distance = SquaredL2::evaluate(samples.row(j), data); @@ -725,17 +725,17 @@ mod tests { assert_eq!(values.len(), ndata); values.shuffle(rng); - let mut data = Matrix::from_element(ndata, dim, 0.0); - for (r, v) in std::iter::zip(data.row_iter_mut(), values.iter()) { + let mut data = rowmajor::Owned::from_element(ndata, dim, 0.0); + for (r, v) in std::iter::zip(data.rows_mut(), values.iter()) { r.fill(*v); } - let mut centers = Matrix::from_element(ncenters, dim, f32::INFINITY); - kmeans_plusplus_into(centers.as_mut_view(), data.as_view(), rng).unwrap(); + let mut centers = rowmajor::Owned::from_element(ncenters, dim, f32::INFINITY); + kmeans_plusplus_into(centers.as_view_mut(), data.as_view(), rng).unwrap(); // Make sure that each value was selected for a center. let mut seen = HashSet::::new(); - for c in centers.row_iter() { + for c in centers.rows() { let first = c[0]; assert!(c.iter().all(|i| *i == first)); @@ -803,17 +803,17 @@ mod tests { assert_eq!(values.len(), ndata); values.shuffle(rng); - let mut data = Matrix::from_element(ndata, dim, 0.0); - for (r, v) in std::iter::zip(data.row_iter_mut(), values.iter()) { + let mut data = rowmajor::Owned::from_element(ndata, dim, 0.0); + for (r, v) in std::iter::zip(data.rows_mut(), values.iter()) { r.fill(*v); } - let mut centers = Matrix::from_element(ncenters, dim, f32::INFINITY); - kmeans_plusplus_into(centers.as_mut_view(), data.as_view(), rng).unwrap(); + let mut centers = rowmajor::Owned::from_element(ncenters, dim, f32::INFINITY); + kmeans_plusplus_into(centers.as_view_mut(), data.as_view(), rng).unwrap(); // Make sure that each value was selected for a center. let mut seen = HashSet::::new(); - for (i, c) in centers.row_iter().enumerate() { + for (i, c) in centers.rows().enumerate() { let first = c[0]; let v: usize = first.round() as usize; assert_eq!( @@ -828,7 +828,7 @@ mod tests { // Make sure the center is equal to one of the data points. let mut found = false; - for r in data.row_iter() { + for r in data.rows() { if r == c { found = true; break; @@ -868,12 +868,12 @@ mod tests { // Failure modes #[test] fn fail_empty_dataset() { - let data = Matrix::from_element(0, 5, 0.0); - let mut centers = Matrix::from_element(10, data.ncols(), 0.0); + let data = rowmajor::Owned::from_element(0, 5, 0.0); + let mut centers = rowmajor::Owned::from_element(10, data.ncols(), 0.0); let mut rng = StdRng::seed_from_u64(0xa9eae150d30845a1); - let result = kmeans_plusplus_into(centers.as_mut_view(), data.as_view(), &mut rng); + let result = kmeans_plusplus_into(centers.as_view_mut(), data.as_view(), &mut rng); assert!( result.is_err(), "kmeans++ on an empty dataset with non-empty centers should be an error" @@ -889,10 +889,10 @@ mod tests { #[test] fn both_empty_is_okay() { - let data = Matrix::from_element(0, 5, 0.0); - let mut centers = Matrix::from_element(0, data.ncols(), 0.0); + let data = rowmajor::Owned::from_element(0, 5, 0.0); + let mut centers = rowmajor::Owned::from_element(0, data.ncols(), 0.0); let mut rng = StdRng::seed_from_u64(0x6f7031afd9b5aa18); - let result = kmeans_plusplus_into(centers.as_mut_view(), data.as_view(), &mut rng); + let result = kmeans_plusplus_into(centers.as_view_mut(), data.as_view(), &mut rng); assert!( result.is_ok(), "selecting 0 points from an empty dataset is okay" @@ -905,13 +905,13 @@ mod tests { let ncenters = 10; let dim = 5; - let mut data = Matrix::from_element(ndata, dim, 0.0); - set_default_values(data.as_mut_view()); - let mut centers = Matrix::from_element(ncenters, data.ncols(), f32::INFINITY); + let mut data = rowmajor::Owned::from_element(ndata, dim, 0.0); + set_default_values(data.as_view_mut()); + let mut centers = rowmajor::Owned::from_element(ncenters, data.ncols(), f32::INFINITY); let mut rng = StdRng::seed_from_u64(0xa9eae150d30845a1); - let result = kmeans_plusplus_into(centers.as_mut_view(), data.as_view(), &mut rng); + let result = kmeans_plusplus_into(centers.as_view_mut(), data.as_view(), &mut rng); assert!( result.is_err(), "kmeans++ on an empty dataset with non-empty centers should be an error" @@ -943,13 +943,13 @@ mod tests { assert!(values.len() >= ndata); values.shuffle(&mut rng); - let mut data = Matrix::from_element(ndata, dim, 0.0); - for (r, v) in std::iter::zip(data.row_iter_mut(), values.iter()) { + let mut data = rowmajor::Owned::from_element(ndata, dim, 0.0); + for (r, v) in std::iter::zip(data.rows_mut(), values.iter()) { r.fill(*v); } - let mut centers = Matrix::from_element(ncenters, dim, f32::INFINITY); - let result = kmeans_plusplus_into(centers.as_mut_view(), data.as_view(), &mut rng); + let mut centers = rowmajor::Owned::from_element(ncenters, dim, f32::INFINITY); + let result = kmeans_plusplus_into(centers.as_view_mut(), data.as_view(), &mut rng); assert!( result.is_err(), "dataset should not have enough unique points" @@ -965,16 +965,16 @@ mod tests { #[test] fn fail_intinity_check() { - let mut data = Matrix::from_element(10, 1, 0.0); - set_default_values(data.as_mut_view()); + let mut data = rowmajor::Owned::from_element(10, 1, 0.0); + set_default_values(data.as_view_mut()); // A very large value that will overflow to infinity when computing the norm. *data.element_mut(6, 0) = -3.4028235e38; - let mut centers = Matrix::from_element(2, 1, 0.0); + let mut centers = rowmajor::Owned::from_element(2, 1, 0.0); let mut rng = StdRng::seed_from_u64(0xc0449b2aa4e12f05); - let result = kmeans_plusplus_into(centers.as_mut_view(), data.as_view(), &mut rng); + let result = kmeans_plusplus_into(centers.as_view_mut(), data.as_view(), &mut rng); assert!(result.is_err(), "result should complain about infinity"); let err = result.unwrap_err(); assert_eq!(err.selected, 1); @@ -987,16 +987,16 @@ mod tests { #[test] fn fail_nan_check() { - let mut data = Matrix::from_element(10, 1, 0.0); - set_default_values(data.as_mut_view()); + let mut data = rowmajor::Owned::from_element(10, 1, 0.0); + set_default_values(data.as_view_mut()); // A very large value that will overflow to infinity when computing the norm. *data.element_mut(6, 0) = f32::NAN; - let mut centers = Matrix::from_element(2, 1, 0.0); + let mut centers = rowmajor::Owned::from_element(2, 1, 0.0); let mut rng = StdRng::seed_from_u64(0x55808c6c728c8473); - let result = kmeans_plusplus_into(centers.as_mut_view(), data.as_view(), &mut rng); + let result = kmeans_plusplus_into(centers.as_view_mut(), data.as_view(), &mut rng); assert!(result.is_err(), "result should complain about NaN"); let err = result.unwrap_err(); assert_eq!(err.selected, 1); @@ -1017,7 +1017,7 @@ mod tests { let npoints = 5; let dim = 8; let mut square_distances = vec![0.0; npoints]; - let data = Matrix::from_element(npoints, dim, 0.0); + let data = rowmajor::Owned::from_element(npoints, dim, 0.0); let norms = vec![0.0; npoints]; let this = vec![0.0; dim + 1]; // Incorrect let this_square_norm = 0.0; @@ -1036,7 +1036,7 @@ mod tests { let npoints = 5; let dim = 8; let mut square_distances = vec![0.0; npoints + 1]; // Incorrect - let data = Matrix::from_element(npoints, dim, 0.0); + let data = rowmajor::Owned::from_element(npoints, dim, 0.0); let norms = vec![0.0; npoints]; let this = vec![0.0; dim]; let this_square_norm = 0.0; @@ -1055,7 +1055,7 @@ mod tests { let npoints = 5; let dim = 8; let mut square_distances = vec![0.0; npoints]; - let data = Matrix::from_element(npoints, dim, 0.0); + let data = rowmajor::Owned::from_element(npoints, dim, 0.0); let norms = vec![0.0; npoints + 1]; // Incorrect let this = vec![0.0; dim]; let this_square_norm = 0.0; @@ -1077,10 +1077,10 @@ mod tests { expected = "centers output matrix should have the same dimensionality as the dataset" )] fn kmeans_plusplus_into_panics_dim_mismatch() { - let mut centers = Matrix::from_element(2, 10, 0.0); - let data = Matrix::from_element(2, 9, 0.0); + let mut centers = rowmajor::Owned::from_element(2, 10, 0.0); + let data = rowmajor::Owned::from_element(2, 9, 0.0); kmeans_plusplus_into( - centers.as_mut_view(), + centers.as_view_mut(), data.as_view(), &mut rand::rngs::ThreadRng::default(), ) diff --git a/diskann-quantization/src/algorithms/transforms/random_rotation.rs b/diskann-quantization/src/algorithms/transforms/random_rotation.rs index 1f0f93da93..19ca98483b 100644 --- a/diskann-quantization/src/algorithms/transforms/random_rotation.rs +++ b/diskann-quantization/src/algorithms/transforms/random_rotation.rs @@ -6,6 +6,7 @@ use std::num::NonZeroUsize; use diskann_linalg::{self, Transpose}; +use diskann_utils::views::rowmajor::{Matrix, MatrixMut}; #[cfg(feature = "flatbuffers")] use flatbuffers::{FlatBufferBuilder, WIPOffset}; use rand::Rng; @@ -31,7 +32,7 @@ use crate::flatbuffers as fb; #[cfg_attr(test, derive(PartialEq))] pub struct RandomRotation { /// This data structure maintains the invariant that this **must** be a square matrix. - transform: diskann_utils::views::Matrix, + transform: diskann_utils::views::rowmajor::Owned, } impl RandomRotation { @@ -88,7 +89,7 @@ impl RandomRotation { // Lint: By construction, the matrix returned from // `diskann_linalg::random_distance_preserving_matrix` will by `matrix_dim x matrix_dim`. #[expect(clippy::unwrap_used)] - let initial = diskann_utils::views::Matrix::try_from( + let initial = diskann_utils::views::rowmajor::Owned::try_from_data( diskann_linalg::random_distance_preserving_matrix(matrix_dim, rng).into(), matrix_dim, matrix_dim, @@ -103,8 +104,8 @@ impl RandomRotation { let scaling = (dim as f32 / target_dim as f32).sqrt(); let mut transform = - diskann_utils::views::Matrix::from_element(target_dim, dim, 0.0f32); - std::iter::zip(transform.row_iter_mut(), indices.iter()).for_each(|(ro, ri)| { + diskann_utils::views::rowmajor::Owned::from_element(target_dim, dim, 0.0f32); + std::iter::zip(transform.rows_mut(), indices.iter()).for_each(|(ro, ri)| { std::iter::zip(ro.iter_mut(), initial.row(ri).iter()).for_each(|(o, i)| { *o = scaling * (*i); }) @@ -113,8 +114,8 @@ impl RandomRotation { } std::cmp::Ordering::Greater => { let mut transform = - diskann_utils::views::Matrix::from_element(target_dim, dim, 0.0f32); - std::iter::zip(transform.row_iter_mut(), initial.row_iter()) + diskann_utils::views::rowmajor::Owned::from_element(target_dim, dim, 0.0f32); + std::iter::zip(transform.rows_mut(), initial.rows()) .for_each(|(o, i)| o.copy_from_slice(&i[..dim])); transform } @@ -221,9 +222,12 @@ impl RandomRotation { } let data = proto.data().into_iter().collect(); - let transform = - diskann_utils::views::Matrix::try_from(data, nrows as usize, ncols as usize) - .map_err(|_| RandomRotationError::IncorrectDim)?; + let transform = diskann_utils::views::rowmajor::Owned::try_from_data( + data, + nrows as usize, + ncols as usize, + ) + .map_err(|_| RandomRotationError::IncorrectDim)?; Ok(Self { transform }) } diff --git a/diskann-quantization/src/binary/quantizer.rs b/diskann-quantization/src/binary/quantizer.rs index 05593b38c3..11b7bf03c8 100644 --- a/diskann-quantization/src/binary/quantizer.rs +++ b/diskann-quantization/src/binary/quantizer.rs @@ -129,7 +129,10 @@ impl AsFunctor for BinaryQuantizer { #[cfg(test)] mod tests { - use diskann_utils::{ReborrowMut, views::Matrix}; + use diskann_utils::{ + ReborrowMut, + views::rowmajor::{self, Matrix, MatrixMut}, + }; use rand::{SeedableRng, rngs::StdRng, seq::SliceRandom}; use super::*; @@ -137,19 +140,19 @@ mod tests { fn test_compression_impl(len: usize, rng: &mut StdRng) { let mut domain = [-10, -1, 0, 1, 10]; - let mut test_pattern = Matrix::::from_element(domain.len(), len, 0); + let mut test_pattern = rowmajor::Owned::::from_element(domain.len(), len, 0); // Fill the test patterns randomly. for col in 0..len { domain.shuffle(rng); - for (row, d) in std::iter::zip(test_pattern.row_iter_mut(), domain) { + for (row, d) in std::iter::zip(test_pattern.rows_mut(), domain) { row[col] = d; } } let quantizer = BinaryQuantizer; let mut binary = BoxedBitSlice::<1, Binary>::new_boxed(len); - for row in test_pattern.row_iter() { + for row in test_pattern.rows() { quantizer.compress_into(row, binary.reborrow_mut()).unwrap(); // Check the compression. diff --git a/diskann-quantization/src/matrix_kernels/blocks/packed.rs b/diskann-quantization/src/matrix_kernels/blocks/packed.rs index 402858a022..2603b8c52d 100644 --- a/diskann-quantization/src/matrix_kernels/blocks/packed.rs +++ b/diskann-quantization/src/matrix_kernels/blocks/packed.rs @@ -289,7 +289,7 @@ mod tests { use diskann_utils::{ assert_contains, - views::{Matrix, MatrixView}, + views::rowmajor::{self, Matrix}, }; use crate::matrix_kernels::test_util::panic_message_for; @@ -408,7 +408,7 @@ mod tests { fn assert_panel( panel: Panel<'_, f32, SZ>, - reference: MatrixView<'_, f32>, + reference: rowmajor::Ref<'_, f32>, block: usize, ctx: std::fmt::Arguments<'_>, ) { @@ -486,16 +486,16 @@ mod tests { assert_contains!(message, "equal to 2"); } - fn test_matrix(nrows: usize, ncols: usize) -> Matrix { + fn test_matrix(nrows: usize, ncols: usize) -> rowmajor::Owned { let mut value = 0.0; - Matrix::from_fn(nrows, ncols, |_| { + rowmajor::Owned::from_fn(nrows, ncols, |_| { let current = value; value += 1.0; current }) } - fn pack(matrix: MatrixView<'_, f32>) -> Vec { + fn pack(matrix: rowmajor::Ref<'_, f32>) -> Vec { assert!(matrix.nrows().is_multiple_of(SZ)); let mut packed = Vec::with_capacity(matrix.as_slice().len()); diff --git a/diskann-quantization/src/matrix_kernels/blocks/unpacked.rs b/diskann-quantization/src/matrix_kernels/blocks/unpacked.rs index 0ae7a28a66..55d091948c 100644 --- a/diskann-quantization/src/matrix_kernels/blocks/unpacked.rs +++ b/diskann-quantization/src/matrix_kernels/blocks/unpacked.rs @@ -5,7 +5,7 @@ use std::num::NonZeroUsize; -use diskann_utils::views::MatrixView; +use diskann_utils::views::rowmajor::{self, Matrix}; use crate::matrix_kernels::{ bounds::{self, Bound}, @@ -35,18 +35,18 @@ pub(crate) struct View<'a, T> { } impl<'a, T> View<'a, T> { - /// Construct a [`View`] from a [`MatrixView`]. + /// Construct a [`View`] from a [`rowmajor::Ref`]. /// - /// Since [`MatrixView`]s are interpreted as "row-major", the value `k` will be derived + /// Since [`rowmajor::Ref`]s are interpreted as "row-major", the value `k` will be derived /// from `v.ncols()` and the extent will be taken from `v.nrows()`. /// /// Returns `None` if either dimension is zero. - pub(crate) fn from_matrix_view(v: MatrixView<'a, T>) -> Option { + pub(crate) fn from_matrix_view(v: rowmajor::Ref<'a, T>) -> Option { let extent = NonZeroUsize::new(v.nrows())?; let k = DimK::new(NonZeroUsize::new(v.ncols())?); - // SAFETY: The `MatrixView` ensures that the inner slice has size `extent * k`. - Some(unsafe { Self::new(Slice::new(v.into_inner()), extent, k) }) + // SAFETY: The `rowmajor::Ref` ensures that the inner slice has size `extent * k`. + Some(unsafe { Self::new(Slice::new(v.into_slice()), extent, k) }) } /// Construct a new [`View`] over `ptr`. @@ -419,7 +419,7 @@ impl<'a, T, const CAPACITY: usize> Remainder<'a, T, CAPACITY> { #[cfg(test)] mod test { use super::*; - use diskann_utils::{assert_contains, views::Matrix}; + use diskann_utils::assert_contains; use crate::matrix_kernels::test_util::panic_message_for; @@ -449,7 +449,7 @@ mod test { v }; - Matrix::from_fn(nrows.get(), ncols.get(), init) + rowmajor::Owned::from_fn(nrows.get(), ncols.get(), init) }; let view = View::from_matrix_view(mat.as_view()).unwrap(); @@ -471,7 +471,7 @@ mod test { fn visit_panels( dut: View<'_, f32>, - reference: MatrixView<'_, f32>, + reference: rowmajor::Ref<'_, f32>, ctx: std::fmt::Arguments<'_>, ) { let mut count = 0; @@ -553,7 +553,7 @@ mod test { v }; - Matrix::from_fn(nrows.get(), ncols.get(), init) + rowmajor::Owned::from_fn(nrows.get(), ncols.get(), init) }; let view = View::from_matrix_view(mat.as_view()).unwrap(); diff --git a/diskann-quantization/src/matrix_kernels/maxsim/packed_f32_x_unpacked_f16.rs b/diskann-quantization/src/matrix_kernels/maxsim/packed_f32_x_unpacked_f16.rs index 79d48184fc..832f874ae4 100644 --- a/diskann-quantization/src/matrix_kernels/maxsim/packed_f32_x_unpacked_f16.rs +++ b/diskann-quantization/src/matrix_kernels/maxsim/packed_f32_x_unpacked_f16.rs @@ -237,6 +237,7 @@ mod tests { use std::num::NonZeroUsize; + use diskann_utils::views::rowmajor::Matrix; use diskann_wide::arch::Scalar; use rand::{SeedableRng, rngs::StdRng}; diff --git a/diskann-quantization/src/matrix_kernels/maxsim/packed_f32_x_unpacked_f32.rs b/diskann-quantization/src/matrix_kernels/maxsim/packed_f32_x_unpacked_f32.rs index 3de91f7a56..6c025fd45b 100644 --- a/diskann-quantization/src/matrix_kernels/maxsim/packed_f32_x_unpacked_f32.rs +++ b/diskann-quantization/src/matrix_kernels/maxsim/packed_f32_x_unpacked_f32.rs @@ -781,6 +781,7 @@ mod aarch64 { mod tests { use super::*; + use diskann_utils::views::rowmajor::Matrix; use rand::{SeedableRng, rngs::StdRng}; #[cfg(target_arch = "x86_64")] diff --git a/diskann-quantization/src/matrix_kernels/maxsim/test.rs b/diskann-quantization/src/matrix_kernels/maxsim/test.rs index b087d56df4..8b5dabc8b0 100644 --- a/diskann-quantization/src/matrix_kernels/maxsim/test.rs +++ b/diskann-quantization/src/matrix_kernels/maxsim/test.rs @@ -3,7 +3,7 @@ * Licensed under the MIT license. */ -use diskann_utils::views::Matrix; +use diskann_utils::views::rowmajor::{self, Matrix}; use crate::matrix_kernels::test_util::TestDistr; @@ -13,12 +13,12 @@ pub(super) fn generate( k: usize, n: usize, rng: &mut impl rand::Rng, -) -> (Matrix, Matrix, Vec) { +) -> (rowmajor::Owned, rowmajor::Owned, Vec) { let ref_a = TestDistr::matrix::(m, k, rng); let ref_b = TestDistr::matrix::(k, n, rng); let ref_c: Vec = ref_a - .row_iter() + .rows() .map(|a_row| { let mut max_ip = f32::NEG_INFINITY; for b_col in 0..n { diff --git a/diskann-quantization/src/matrix_kernels/test_util.rs b/diskann-quantization/src/matrix_kernels/test_util.rs index 3bff0cef05..3d8ef67df1 100644 --- a/diskann-quantization/src/matrix_kernels/test_util.rs +++ b/diskann-quantization/src/matrix_kernels/test_util.rs @@ -3,7 +3,7 @@ * Licensed under the MIT license. */ -use diskann_utils::views::Matrix; +use diskann_utils::views::rowmajor; use half::f16; use rand::{Rng, distr::Distribution}; @@ -32,11 +32,15 @@ where pub(super) struct TestDistr; impl TestDistr { - pub(super) fn matrix(nrows: usize, ncols: usize, rng: &mut impl rand::Rng) -> Matrix + pub(super) fn matrix( + nrows: usize, + ncols: usize, + rng: &mut impl rand::Rng, + ) -> rowmajor::Owned where Self: Distribution, { - Matrix::from_fn(nrows, ncols, |_| (Self).sample(rng)) + rowmajor::Owned::from_fn(nrows, ncols, |_| (Self).sample(rng)) } } diff --git a/diskann-quantization/src/multi_vector/block_transposed.rs b/diskann-quantization/src/multi_vector/block_transposed.rs index 47c05d3da8..7de07b9262 100644 --- a/diskann-quantization/src/multi_vector/block_transposed.rs +++ b/diskann-quantization/src/multi_vector/block_transposed.rs @@ -79,11 +79,7 @@ use std::{alloc::Layout, marker::PhantomData, ptr::NonNull}; -use diskann_utils::{ - Reborrow, ReborrowMut, - strided::Strided, - views::{MatrixView, MutMatrixView}, -}; +use diskann_utils::{Reborrow, ReborrowMut, strided::Strided, views::rowmajor}; use super::matrix::{ Defaulted, LayoutError, Mat, MatMut, MatRef, NewCloned, NewMut, NewOwned, NewRef, Overflow, @@ -812,7 +808,7 @@ impl<'a, T: Copy, const GROUP: usize, const PACK: usize> BlockTransposedRef<'a, unsafe { self.as_ptr().add(self.data.repr().block_offset(block)) } } - /// Return a view over a full block as a [`MatrixView`]. + /// Return a view over a full block as a [`rowmajor::Ref`]. /// /// The returned view has `padded_ncols / PACK` rows and `GROUP * PACK` /// columns. For `PACK == 1` this simplifies to `ncols` rows and `GROUP` @@ -822,14 +818,14 @@ impl<'a, T: Copy, const GROUP: usize, const PACK: usize> BlockTransposedRef<'a, /// /// Panics if `block >= self.full_blocks()`. #[expect(clippy::expect_used)] - pub fn block(&self, block: usize) -> MatrixView<'a, T> { + pub fn block(&self, block: usize) -> rowmajor::Ref<'a, T> { assert!(block < self.full_blocks()); let offset = self.data.repr().block_offset(block); let stride = self.data.repr().block_stride(); // SAFETY: `block < full_blocks()` (asserted above) guarantees // `offset + stride` is within the backing allocation. let data: &[T] = unsafe { std::slice::from_raw_parts(self.as_ptr().add(offset), stride) }; - MatrixView::try_from(data, self.padded_ncols() / PACK, GROUP * PACK) + rowmajor::Ref::try_from_data(data, self.padded_ncols() / PACK, GROUP * PACK) .expect("base data should have been sized correctly") } @@ -839,7 +835,7 @@ impl<'a, T: Copy, const GROUP: usize, const PACK: usize> BlockTransposedRef<'a, /// The returned view has the same dimensions as [`block()`](Self::block): /// `padded_ncols / PACK` rows and `GROUP * PACK` columns. #[expect(clippy::expect_used)] - pub fn remainder_block(&self) -> Option> { + pub fn remainder_block(&self) -> Option> { if self.remainder() == 0 { None } else { @@ -850,7 +846,7 @@ impl<'a, T: Copy, const GROUP: usize, const PACK: usize> BlockTransposedRef<'a, let data: &[T] = unsafe { std::slice::from_raw_parts(self.as_ptr().add(offset), stride) }; Some( - MatrixView::try_from(data, self.padded_ncols() / PACK, GROUP * PACK) + rowmajor::Ref::try_from_data(data, self.padded_ncols() / PACK, GROUP * PACK) .expect("base data should have been sized correctly"), ) } @@ -910,8 +906,8 @@ impl<'a, T: Copy, const GROUP: usize, const PACK: usize> BlockTransposedMut<'a, delegate_to_ref!(pub fn as_ptr(&self) -> *const T); delegate_to_ref!(pub fn as_slice(&self) -> &[T]); delegate_to_ref!(#[expect(clippy::missing_safety_doc)] unsafe pub fn block_ptr_unchecked(&self, block: usize) -> *const T); - delegate_to_ref!(pub fn block(&self, block: usize) -> MatrixView<'_, T>); - delegate_to_ref!(pub fn remainder_block(&self) -> Option>); + delegate_to_ref!(pub fn block(&self, block: usize) -> rowmajor::Ref<'_, T>); + delegate_to_ref!(pub fn remainder_block(&self) -> Option>); delegate_to_ref!(pub fn get_element(&self, row: usize, col: usize) -> T); /// Group size (blocking factor `GROUP`). @@ -961,12 +957,12 @@ impl<'a, T: Copy, const GROUP: usize, const PACK: usize> BlockTransposedMut<'a, /// # Panics /// /// Panics if `block >= self.full_blocks()`. - pub fn block_mut(&mut self, block: usize) -> MutMatrixView<'_, T> { + pub fn block_mut(&mut self, block: usize) -> rowmajor::Mut<'_, T> { self.reborrow_mut().block_mut_inner(block) } #[expect(clippy::expect_used)] - fn block_mut_inner(mut self, block: usize) -> MutMatrixView<'a, T> { + fn block_mut_inner(mut self, block: usize) -> rowmajor::Mut<'a, T> { let repr = *self.data.repr(); assert!(block < repr.full_blocks()); let offset = repr.block_offset(block); @@ -979,18 +975,18 @@ impl<'a, T: Copy, const GROUP: usize, const PACK: usize> BlockTransposedMut<'a, stride, ) }; - MutMatrixView::try_from(data, pncols / PACK, GROUP * PACK) + rowmajor::Mut::try_from_data(data, pncols / PACK, GROUP * PACK) .expect("base data should have been sized correctly") } /// Return a mutable view over the remainder block, or `None` if there is no /// remainder. - pub fn remainder_block_mut(&mut self) -> Option> { + pub fn remainder_block_mut(&mut self) -> Option> { self.reborrow_mut().remainder_block_mut_inner() } #[expect(clippy::expect_used)] - fn remainder_block_mut_inner(mut self) -> Option> { + fn remainder_block_mut_inner(mut self) -> Option> { let repr = *self.data.repr(); if repr.remainder() == 0 { None @@ -1006,7 +1002,7 @@ impl<'a, T: Copy, const GROUP: usize, const PACK: usize> BlockTransposedMut<'a, ) }; Some( - MutMatrixView::try_from(data, pncols / PACK, GROUP * PACK) + rowmajor::Mut::try_from_data(data, pncols / PACK, GROUP * PACK) .expect("base data should have been sized correctly"), ) } @@ -1050,8 +1046,8 @@ impl BlockTransposed *const T); delegate_to_ref!(pub fn as_slice(&self) -> &[T]); delegate_to_ref!(#[expect(clippy::missing_safety_doc)] unsafe pub fn block_ptr_unchecked(&self, block: usize) -> *const T); - delegate_to_ref!(pub fn block(&self, block: usize) -> MatrixView<'_, T>); - delegate_to_ref!(pub fn remainder_block(&self) -> Option>); + delegate_to_ref!(pub fn block(&self, block: usize) -> rowmajor::Ref<'_, T>); + delegate_to_ref!(pub fn remainder_block(&self) -> Option>); delegate_to_ref!(pub fn get_element(&self, row: usize, col: usize) -> T); /// Group size (blocking factor `GROUP`). @@ -1084,12 +1080,12 @@ impl BlockTransposed MutMatrixView<'_, T> { + pub fn block_mut(&mut self, block: usize) -> rowmajor::Mut<'_, T> { self.as_view_mut().block_mut_inner(block) } /// See [`BlockTransposedMut::remainder_block_mut`]. - pub fn remainder_block_mut(&mut self) -> Option> { + pub fn remainder_block_mut(&mut self) -> Option> { self.as_view_mut().remainder_block_mut_inner() } @@ -1199,8 +1195,8 @@ impl BlockTransposed) -> Self { + /// Construct a block-transposed matrix by copying data from a [`rowmajor::Ref`]. + pub fn from_matrix_view(v: rowmajor::Ref<'_, T>) -> Self { Self::from_strided(v.into()) } } @@ -1244,7 +1240,10 @@ mod tests { //! parameters to `test_full_api` (`Send`/`Sync`, panic paths, //! non-unit strides, concurrent mutation, etc.). - use diskann_utils::{lazy_format, views::Matrix}; + use diskann_utils::{ + lazy_format, + views::rowmajor::{Matrix, MatrixMut}, + }; use super::*; use crate::utils::div_round_up; @@ -1266,7 +1265,7 @@ mod tests { #[test] fn clone_has_independent_backing_allocation() { - let mut data = Matrix::from_element(5, 3, 0); + let mut data = rowmajor::Owned::from_element(5, 3, 0); data.as_mut_slice() .iter_mut() .enumerate() @@ -1324,7 +1323,7 @@ mod tests { // ── Construction ───────────────────────────────────────── - let mut data = Matrix::from_element(nrows, ncols, T::default()); + let mut data = rowmajor::Owned::from_element(nrows, ncols, T::default()); data.as_mut_slice() .iter_mut() .enumerate() @@ -1878,7 +1877,7 @@ mod tests { ncols: usize, gen_element: fn(usize) -> T, ) { - let mut data = Matrix::from_element(nrows, ncols, T::default()); + let mut data = rowmajor::Owned::from_element(nrows, ncols, T::default()); data.as_mut_slice() .iter_mut() .enumerate() diff --git a/diskann-quantization/src/multi_vector/distance/factory.rs b/diskann-quantization/src/multi_vector/distance/factory.rs index fe87959ccd..17681a272b 100644 --- a/diskann-quantization/src/multi_vector/distance/factory.rs +++ b/diskann-quantization/src/multi_vector/distance/factory.rs @@ -8,6 +8,7 @@ use std::num::NonZeroUsize; +use diskann_utils::views::rowmajor::Matrix; use diskann_vector::distance::InnerProduct; use diskann_vector::{DistanceFunctionMut, PureDistanceFunction}; use diskann_wide::Architecture; diff --git a/diskann-quantization/src/multi_vector/matrix.rs b/diskann-quantization/src/multi_vector/matrix.rs index 3f1914bfea..9d94729241 100644 --- a/diskann-quantization/src/multi_vector/matrix.rs +++ b/diskann-quantization/src/multi_vector/matrix.rs @@ -29,7 +29,7 @@ use std::{alloc::Layout, iter::FusedIterator, marker::PhantomData, ptr::NonNull}; -use diskann_utils::{Reborrow, ReborrowMut, views::MatrixView}; +use diskann_utils::{Reborrow, ReborrowMut, views::rowmajor}; use thiserror::Error; use crate::utils; @@ -751,9 +751,9 @@ impl Mat> { self.as_view().as_slice() } - /// Return a [`MatrixView`] over the backing data. + /// Return a [`rowmajor::Ref`] over the backing data. #[inline] - pub fn as_matrix_view(&self) -> MatrixView<'_, T> { + pub fn as_matrix_view(&self) -> rowmajor::Ref<'_, T> { self.as_view().as_matrix_view() } } @@ -883,13 +883,13 @@ impl<'a, T> MatRef<'a, Standard> { unsafe { std::slice::from_raw_parts(self.ptr.as_ptr().cast::(), len) } } - /// Return a [`MatrixView`] over the backing data. + /// Return a [`rowmajor::Ref`] over the backing data. #[expect(clippy::expect_used)] #[inline] - pub fn as_matrix_view(&self) -> MatrixView<'a, T> { + pub fn as_matrix_view(&self) -> rowmajor::Ref<'a, T> { // `Standard::new` validates that `nrows * ncols` does not overflow, // so `try_from` is infallible here. - MatrixView::try_from(self.as_slice(), self.num_vectors(), self.vector_dim()) + rowmajor::Ref::try_from_data(self.as_slice(), self.num_vectors(), self.vector_dim()) .expect("Standard has valid dimensions") } } @@ -1111,9 +1111,9 @@ impl<'a, T> MatMut<'a, Standard> { self.as_view().as_slice() } - /// Return a [`MatrixView`] over the backing data. + /// Return a [`rowmajor::Ref`] over the backing data. #[inline] - pub fn as_matrix_view(&self) -> MatrixView<'_, T> { + pub fn as_matrix_view(&self) -> rowmajor::Ref<'_, T> { self.as_view().as_matrix_view() } } @@ -1236,6 +1236,7 @@ mod tests { use std::fmt::Display; use diskann_utils::lazy_format; + use diskann_utils::views::rowmajor::Matrix; /// Helper to assert a type is Copy. fn assert_copy(_: &T) {} diff --git a/diskann-quantization/src/product/tables/basic.rs b/diskann-quantization/src/product/tables/basic.rs index c33a4b685f..4816916316 100644 --- a/diskann-quantization/src/product/tables/basic.rs +++ b/diskann-quantization/src/product/tables/basic.rs @@ -5,7 +5,10 @@ use crate::traits::CompressInto; use crate::views::{ChunkOffsetsBase, ChunkOffsetsView}; -use diskann_utils::views::{DenseData, MatrixBase, MatrixView}; +use diskann_utils::views::{ + DenseData, + rowmajor::{self, Matrix}, +}; use diskann_vector::{PureDistanceFunction, distance::SquaredL2}; use thiserror::Error; @@ -26,19 +29,19 @@ use thiserror::Error; #[derive(Debug, Clone)] pub struct BasicTableBase where - T: DenseData, + T: Matrix, U: DenseData, { - pivots: MatrixBase, + pivots: T, offsets: ChunkOffsetsBase, } /// A `BasicTableBase` that owns its contents. -pub type BasicTable = BasicTableBase, Box<[usize]>>; +pub type BasicTable = BasicTableBase, Box<[usize]>>; /// A `BasicTableBase` that references its contents. Construction of such a table will /// not result in a memory allocation. -pub type BasicTableView<'a> = BasicTableBase<&'a [f32], &'a [usize]>; +pub type BasicTableView<'a> = BasicTableBase, &'a [usize]>; #[derive(Error, Debug)] #[non_exhaustive] @@ -54,7 +57,7 @@ pub enum BasicTableError { impl BasicTableBase where - T: DenseData, + T: Matrix, U: DenseData, { /// Construct a new `BasicTableBase` over the pivot table and offsets. @@ -62,10 +65,7 @@ where /// # Error /// /// Returns an error if `pivots.ncols() != offsets.dim()` or if `pivots.nrows() == 0`. - pub fn new( - pivots: MatrixBase, - offsets: ChunkOffsetsBase, - ) -> Result { + pub fn new(pivots: T, offsets: ChunkOffsetsBase) -> Result { let pivot_dim = pivots.ncols(); let offsets_dim = offsets.dim(); @@ -82,7 +82,7 @@ where } /// Return a view over the pivot table. - pub fn view_pivots(&self) -> MatrixView<'_, f32> { + pub fn view_pivots(&self) -> rowmajor::Ref<'_, f32> { self.pivots.as_view() } @@ -122,7 +122,7 @@ pub enum TableCompressionError { impl CompressInto<&[f32], &mut [u8]> for BasicTableBase where - T: DenseData, + T: Matrix, U: DenseData, { type Error = TableCompressionError; @@ -175,7 +175,7 @@ where let range = self.offsets.at(chunk); let slice = &from[range.clone()]; - self.pivots.row_iter().enumerate().for_each(|(index, row)| { + self.pivots.rows().enumerate().for_each(|(index, row)| { let distance: f32 = SquaredL2::evaluate(slice, &row[range.clone()]); if distance < min_distance { min_distance = distance; @@ -200,7 +200,7 @@ where #[cfg(test)] mod tests { - use diskann_utils::{lazy_format, views}; + use diskann_utils::lazy_format; use rand::{ SeedableRng, distr::{Distribution, StandardUniform}, @@ -219,7 +219,7 @@ mod tests { // disagree. #[test] fn error_on_mismatch_dim() { - let pivots = views::Matrix::from_element(3, 5, 0.0); + let pivots = rowmajor::Owned::from_element(3, 5, 0.0); let offsets = crate::views::ChunkOffsets::new(Box::new([0, 1, 6])).unwrap(); let result = BasicTable::new(pivots, offsets); assert!(result.is_err(), "dimensions are not equal"); @@ -232,7 +232,7 @@ mod tests { // Test that the table constructor errors when there are no pivots. #[test] fn error_on_no_pivots() { - let pivots = views::Matrix::from_element(0, 5, 0.0); + let pivots = rowmajor::Owned::from_element(0, 5, 0.0); let offsets = crate::views::ChunkOffsets::new(Box::new([0, 1, 2, 5])).unwrap(); let result = BasicTable::new(pivots, offsets); assert!(result.is_err(), "pivots is empty"); @@ -244,7 +244,7 @@ mod tests { let mut rng = rand::rngs::StdRng::seed_from_u64(0xd96bac968083ec29); for dim in [5, 10, 12] { for total in [1, 2, 3] { - let pivots = views::Matrix::from_fn(total, dim, |_| -> f32 { + let pivots = rowmajor::Owned::from_fn(total, dim, |_| -> f32 { StandardUniform {}.sample(&mut rng) }); let offsets = crate::views::ChunkOffsets::new(Box::new([0, 1, 3, dim])).unwrap(); @@ -299,7 +299,7 @@ mod tests { let (data, expected) = create_dataset(schema, num_centers, num_data, &mut rng); let mut output = vec![0; schema.len()]; - for (input, expected) in std::iter::zip(data.row_iter(), expected.row_iter()) { + for (input, expected) in std::iter::zip(data.rows(), expected.rows()) { table.compress_into(input, &mut output).unwrap(); for (entry, (e, o)) in std::iter::zip(expected.iter(), output.iter()).enumerate() @@ -320,7 +320,7 @@ mod tests { // Set up `ncenters > 256`. { - let pivots = views::Matrix::from_element(257, dim, 0.0); + let pivots = rowmajor::Owned::from_element(257, dim, 0.0); let table = BasicTable::new(pivots, offsets.clone()).unwrap(); let input = vec![f32::default(); dim]; @@ -339,7 +339,7 @@ mod tests { // Setup input dim not equal to expected. { - let pivots = views::Matrix::from_element(10, dim, 0.0); + let pivots = rowmajor::Owned::from_element(10, dim, 0.0); let table = BasicTable::new(pivots, offsets.clone()).unwrap(); let input = vec![f32::default(); dim - 1]; @@ -358,7 +358,7 @@ mod tests { // Setup output dim not equal to expected. { - let pivots = views::Matrix::from_element(10, dim, 0.0); + let pivots = rowmajor::Owned::from_element(10, dim, 0.0); let table = BasicTable::new(pivots, offsets.clone()).unwrap(); let input = vec![f32::default(); dim]; @@ -383,7 +383,7 @@ mod tests { #[test] fn test_table_single_compression_errors() { check_pqtable_single_compression_errors( - &|pivots: views::Matrix, offsets| BasicTable::new(pivots, offsets).unwrap(), + &|pivots: rowmajor::Owned, offsets| BasicTable::new(pivots, offsets).unwrap(), &"BasicTable", ) } diff --git a/diskann-quantization/src/product/tables/test.rs b/diskann-quantization/src/product/tables/test.rs index 6619ffb901..2b7cfeb51e 100644 --- a/diskann-quantization/src/product/tables/test.rs +++ b/diskann-quantization/src/product/tables/test.rs @@ -4,9 +4,7 @@ */ // A collection of test helpers to ensure uniformity across tables. -use diskann_utils::views::Matrix; -#[cfg(not(miri))] -use diskann_utils::views::{MatrixView, MutMatrixView}; +use diskann_utils::views::rowmajor::{self, Matrix, MatrixMut}; #[cfg(not(miri))] use rand::seq::IndexedRandom; use rand::{ @@ -83,8 +81,8 @@ use crate::views::{self, ChunkOffsets, ChunkOffsetsView}; pub(super) fn create_pivot_tables( schema: ChunkOffsets, num_centers: usize, -) -> (Matrix, ChunkOffsets) { - let mut pivots = Matrix::::from_element(num_centers, schema.dim(), 0.0); +) -> (rowmajor::Owned, ChunkOffsets) { + let mut pivots = rowmajor::Owned::::from_element(num_centers, schema.dim(), 0.0); (0..schema.len()).for_each(|chunk| { let range = schema.at(chunk); @@ -118,16 +116,16 @@ pub(super) fn create_pivot_tables( /// /// Returns: /// -/// * The initialized dataset as a Matrix. -/// * The expected center as a Matrix. +/// * The initialized dataset as a rowmajor::Owned. +/// * The expected center as a rowmajor::Owned. pub(super) fn create_dataset( schema: ChunkOffsetsView<'_>, num_centers: usize, num_data: usize, rng: &mut R, -) -> (Matrix, Matrix) { - let mut data = Matrix::::from_element(num_data, schema.dim(), 0.0); - let mut expected = Matrix::::from_element(num_data, schema.len(), 0); +) -> (rowmajor::Owned, rowmajor::Owned) { + let mut data = rowmajor::Owned::::from_element(num_data, schema.dim(), 0.0); + let mut expected = rowmajor::Owned::::from_element(num_data, schema.len(), 0); let dist = Uniform::new(0, num_centers).unwrap(); for row_index in 0..data.nrows() { @@ -160,7 +158,7 @@ pub(super) fn create_dataset( // A cantralized test for error handling in `CompressInto<[f32], [u8]>` pub(super) fn check_pqtable_single_compression_errors( - build: &dyn Fn(Matrix, ChunkOffsets) -> T, + build: &dyn Fn(rowmajor::Owned, ChunkOffsets) -> T, context: &dyn std::fmt::Display, ) where T: for<'a, 'b> CompressInto<&'a [f32], &'b mut [u8]>, @@ -171,7 +169,7 @@ pub(super) fn check_pqtable_single_compression_errors( // Set up `ncenters > 256`. { - let pivots = Matrix::from_element(257, dim, 0.0); + let pivots = rowmajor::Owned::from_element(257, dim, 0.0); let table = build(pivots, offsets.clone()); let input = vec![f32::default(); dim]; @@ -193,7 +191,7 @@ pub(super) fn check_pqtable_single_compression_errors( // Setup input dim not equal to expected. { - let pivots = Matrix::from_element(10, dim, 0.0); + let pivots = rowmajor::Owned::from_element(10, dim, 0.0); let table = build(pivots, offsets.clone()); let input = vec![f32::default(); dim - 1]; @@ -215,7 +213,7 @@ pub(super) fn check_pqtable_single_compression_errors( // Setup output dim not equal to expected. { - let pivots = Matrix::from_element(10, dim, 0.0); + let pivots = rowmajor::Owned::from_element(10, dim, 0.0); let table = build(pivots, offsets.clone()); let input = vec![f32::default(); dim]; @@ -289,16 +287,16 @@ pub(super) fn check_pqtable_single_compression_errors( } //////////////////////////////////////////////////////////////////// -// Testing `CompressInto, MarixView<'_, u8>>` // +// Testing `CompressInto, MarixView<'_, u8>>` // //////////////////////////////////////////////////////////////////// // A cantralized test for error handling in `CompressInto<[f32], [u8]>` #[cfg(not(miri))] pub(super) fn check_pqtable_batch_compression_errors( - build: &dyn Fn(Matrix, ChunkOffsets) -> T, + build: &dyn Fn(rowmajor::Owned, ChunkOffsets) -> T, context: &dyn std::fmt::Display, ) where - T: for<'a> CompressInto, MutMatrixView<'a, u8>>, + T: for<'a> CompressInto, rowmajor::Mut<'a, u8>>, { let dim = 10; let num_chunks = 3; @@ -308,12 +306,12 @@ pub(super) fn check_pqtable_batch_compression_errors( // Set up `ncenters > 256`. { - let pivots = Matrix::from_element(257, dim, 0.0); + let pivots = rowmajor::Owned::from_element(257, dim, 0.0); let table = build(pivots, offsets.clone()); - let input = Matrix::from_element(batchsize, dim, f32::default()); - let mut output = Matrix::from_element(batchsize, num_chunks, u8::MAX); - let result = table.compress_into(input.as_view(), output.as_mut_view()); + let input = rowmajor::Owned::from_element(batchsize, dim, f32::default()); + let mut output = rowmajor::Owned::from_element(batchsize, num_chunks, u8::MAX); + let result = table.compress_into(input.as_view(), output.as_view_mut()); assert!(result.is_err()); assert_eq!( result.unwrap_err().to_string(), @@ -330,12 +328,12 @@ pub(super) fn check_pqtable_batch_compression_errors( // Setup input dim not equal to expected. { - let pivots = Matrix::from_element(10, dim, 0.0); + let pivots = rowmajor::Owned::from_element(10, dim, 0.0); let table = build(pivots, offsets.clone()); - let input = Matrix::from_element(batchsize, dim - 1, f32::default()); - let mut output = Matrix::from_element(batchsize, num_chunks, u8::MAX); - let result = table.compress_into(input.as_view(), output.as_mut_view()); + let input = rowmajor::Owned::from_element(batchsize, dim - 1, f32::default()); + let mut output = rowmajor::Owned::from_element(batchsize, num_chunks, u8::MAX); + let result = table.compress_into(input.as_view(), output.as_view_mut()); assert!(result.is_err()); assert_eq!( result.unwrap_err().to_string(), @@ -352,12 +350,12 @@ pub(super) fn check_pqtable_batch_compression_errors( // Setup output dim not equal to expected. { - let pivots = Matrix::from_element(10, dim, 0.0); + let pivots = rowmajor::Owned::from_element(10, dim, 0.0); let table = build(pivots, offsets.clone()); - let input = Matrix::from_element(batchsize, dim, f32::default()); - let mut output = Matrix::from_element(batchsize, num_chunks - 1, u8::MAX); - let result = table.compress_into(input.as_view(), output.as_mut_view()); + let input = rowmajor::Owned::from_element(batchsize, dim, f32::default()); + let mut output = rowmajor::Owned::from_element(batchsize, num_chunks - 1, u8::MAX); + let result = table.compress_into(input.as_view(), output.as_view_mut()); assert!(result.is_err()); assert_eq!( @@ -379,12 +377,12 @@ pub(super) fn check_pqtable_batch_compression_errors( // Num rows are different. { - let pivots = Matrix::from_element(10, dim, 0.0); + let pivots = rowmajor::Owned::from_element(10, dim, 0.0); let table = build(pivots, offsets.clone()); - let input = Matrix::from_element(batchsize, dim, f32::default()); - let mut output = Matrix::from_element(batchsize - 1, num_chunks, u8::MAX); - let result = table.compress_into(input.as_view(), output.as_mut_view()); + let input = rowmajor::Owned::from_element(batchsize, dim, f32::default()); + let mut output = rowmajor::Owned::from_element(batchsize - 1, num_chunks, u8::MAX); + let result = table.compress_into(input.as_view(), output.as_view_mut()); assert!(result.is_err()); assert_eq!( @@ -416,10 +414,10 @@ pub(super) fn check_pqtable_batch_compression_errors( let table = build(pivots, o); let num_points = 15; - let mut buf = Matrix::::from_element(num_points, offsets.dim(), 0.0); - let mut output = Matrix::::from_element(num_points, offsets.len(), 0); + let mut buf = rowmajor::Owned::::from_element(num_points, offsets.dim(), 0.0); + let mut output = rowmajor::Owned::::from_element(num_points, offsets.len(), 0); - fn clear(mut x: MutMatrixView) { + fn clear(mut x: rowmajor::Mut) { x.as_mut_slice().iter_mut().for_each(|i| *i = T::default()); } @@ -433,11 +431,11 @@ pub(super) fn check_pqtable_batch_compression_errors( let distribution = Uniform::new(range.start, range.end).unwrap(); for row in 0..num_points { - clear(buf.as_mut_view()); + clear(buf.as_view_mut()); let value = *sample.choose(&mut rng).unwrap(); *buf.element_mut(row, distribution.sample(&mut rng)) = value; let err = table - .compress_into(buf.as_view(), output.as_mut_view()) + .compress_into(buf.as_view(), output.as_view_mut()) .expect_err(&format!("expected a value of {}", value)); let message = err.to_string(); diff --git a/diskann-quantization/src/product/tables/transposed/pivots.rs b/diskann-quantization/src/product/tables/transposed/pivots.rs index 88944ae1d5..8909a73f7c 100644 --- a/diskann-quantization/src/product/tables/transposed/pivots.rs +++ b/diskann-quantization/src/product/tables/transposed/pivots.rs @@ -1061,7 +1061,10 @@ where #[cfg(test)] mod tests { - use diskann_utils::{lazy_format, views}; + use diskann_utils::{ + lazy_format, + views::{self, rowmajor::Matrix}, + }; use diskann_vector::{PureDistanceFunction, distance}; use rand::{ SeedableRng, @@ -1507,7 +1510,7 @@ mod tests { #[test] #[should_panic(expected = "row 5 must be less than 5")] fn get_panics_on_row() { - let data = views::Matrix::from_element(5, 10, 0.0); + let data = views::rowmajor::Owned::from_element(5, 10, 0.0); let chunk = Chunk::new(data.as_view().into()).unwrap(); chunk.get(5, 1); } @@ -1515,7 +1518,7 @@ mod tests { #[test] #[should_panic(expected = "col 5 must be less than 5")] fn get_panics_on_col() { - let data = views::Matrix::from_element(10, 5, 0.0); + let data = views::rowmajor::Owned::from_element(10, 5, 0.0); let chunk = Chunk::new(data.as_view().into()).unwrap(); chunk.get(1, 5); } @@ -1534,7 +1537,8 @@ mod tests { fn test_process_into_impl(dim: usize, total: usize, rng: &mut StdRng) { let distribution = Uniform::::new(-10, 10).unwrap(); - let base = views::Matrix::::from_fn(total, dim, |_| distribution.sample(rng) as f32); + let base = + views::rowmajor::Owned::::from_fn(total, dim, |_| distribution.sample(rng) as f32); let chunk = Chunk::new(base.as_view().into()).unwrap(); let mut input = vec![0.0; dim]; @@ -1549,7 +1553,7 @@ mod tests { InnerProduct::process_into(&chunk, &input, &mut output); // Check outputs - std::iter::zip(base.row_iter(), output.iter()).for_each(|(row, got)| { + std::iter::zip(base.rows(), output.iter()).for_each(|(row, got)| { let expected: f32 = distance::InnerProduct::evaluate(row, input.as_slice()); assert_eq!(*got, expected); }); @@ -1558,7 +1562,7 @@ mod tests { SquaredL2::process_into(&chunk, &input, &mut output); // Check outputs - std::iter::zip(base.row_iter(), output.iter()).for_each(|(row, got)| { + std::iter::zip(base.rows(), output.iter()).for_each(|(row, got)| { let expected: f32 = distance::SquaredL2::evaluate(row, input.as_slice()); assert_eq!(*got, expected); }); @@ -1583,7 +1587,7 @@ mod tests { #[test] #[should_panic] fn test_process_into_panics_on_from() { - let data = views::Matrix::::from_element(5, 10, 0.0); + let data = views::rowmajor::Owned::::from_element(5, 10, 0.0); let chunk = Chunk::new(data.as_view().into()).unwrap(); assert_eq!(chunk.dimension(), 10); assert_eq!(chunk.num_centers(), 5); @@ -1597,7 +1601,7 @@ mod tests { #[test] #[should_panic] fn test_process_into_panics_on_into() { - let data = views::Matrix::::from_element(5, 10, 0.0); + let data = views::rowmajor::Owned::::from_element(5, 10, 0.0); let chunk = Chunk::new(data.as_view().into()).unwrap(); assert_eq!(chunk.dimension(), 10); assert_eq!(chunk.num_centers(), 5); diff --git a/diskann-quantization/src/product/tables/transposed/table.rs b/diskann-quantization/src/product/tables/transposed/table.rs index b1f7fe551c..04d67a7d6c 100644 --- a/diskann-quantization/src/product/tables/transposed/table.rs +++ b/diskann-quantization/src/product/tables/transposed/table.rs @@ -13,7 +13,10 @@ use crate::{ }; use diskann_utils::{ strided::Strided, - views::{self, MatrixView, MutMatrixView}, + views::{ + self, + rowmajor::{self, Matrix, MatrixMut}, + }, }; use thiserror::Error; @@ -71,7 +74,7 @@ impl TransposedTable { /// * `pivots.nrows() == 0`: The pivot table cannot be empty. #[expect(clippy::expect_used)] pub fn from_parts( - pivots: views::MatrixView, + pivots: views::rowmajor::Ref, offsets: ChunkOffsets, ) -> Result { let pivot_dim = pivots.ncols(); @@ -154,7 +157,7 @@ impl TransposedTable { #[expect(clippy::expect_used)] pub fn compress_batch( &self, - data: views::MatrixView<'_, T>, + data: views::rowmajor::Ref<'_, T>, mut compression_delegate: F, ) -> Result<(), CompressError> where @@ -187,7 +190,7 @@ impl TransposedTable { let range = self.offsets.at(i); if let Some(chunk_dim) = NonZeroUsize::new(range.len()) { // Construct a view for the packing buffer for this chunk. - let mut packing_view = views::MutMatrixView::try_from( + let mut packing_view = views::rowmajor::Mut::try_from_data( &mut packing_buffer[..SUB_BATCH_SIZE * chunk_dim.get()], SUB_BATCH_SIZE, chunk_dim.get(), @@ -205,7 +208,7 @@ impl TransposedTable { // The control on our loop bounds mean we should never be in a situation // where this would occur, so we'd rather hit the indexing panic early. let mut linear_start = row_start * ncols + range.start; - packing_view.row_iter_mut().for_each(|row| { + packing_view.rows_mut().for_each(|row| { pack(row, &slice[linear_start..linear_start + chunk_dim.get()]); linear_start += dim_nonzero.get(); }); @@ -282,7 +285,7 @@ impl TransposedTable { /// * `query.len() != self.dim()`. /// * `partisl.nrows() != self.nchunks()`. /// * `partisl.ncols() != self.ncenters()`. - pub fn process_into(&self, query: &[f32], mut partials: MutMatrixView<'_, f32>) + pub fn process_into(&self, query: &[f32], mut partials: rowmajor::Mut<'_, f32>) where T: pivots::ProcessInto, { @@ -304,7 +307,7 @@ impl TransposedTable { ); // Loop over each chunk. - std::iter::zip(self.pivots.iter(), partials.row_iter_mut()) + std::iter::zip(self.pivots.iter(), partials.rows_mut()) .enumerate() .for_each(|(i, (pivot, out))| { let range = self.offsets.at(i); @@ -422,7 +425,7 @@ pub enum TableBatchCompressionError { InfinityOrNaN(usize, usize), } -impl CompressInto, MutMatrixView<'_, u8>> for TransposedTable +impl CompressInto, rowmajor::Mut<'_, u8>> for TransposedTable where T: Copy + Into, { @@ -464,8 +467,8 @@ where /// This function is single-threaded. fn compress_into( &self, - from: MatrixView<'_, T>, - mut to: MutMatrixView<'_, u8>, + from: rowmajor::Ref<'_, T>, + mut to: rowmajor::Mut<'_, u8>, ) -> Result<(), Self::Error> { if self.ncenters() > 256 { return Err(Self::Error::CannotCompressToByte(self.ncenters())); @@ -530,7 +533,7 @@ mod test_compression { // disagree. #[test] fn error_on_mismatch_dim() { - let pivots = views::Matrix::from_element(3, 5, 0.0); + let pivots = views::rowmajor::Owned::from_element(3, 5, 0.0); let offsets = ChunkOffsets::new(Box::new([0, 1, 6])).unwrap(); let result = TransposedTable::from_parts(pivots.as_view(), offsets); assert!(result.is_err(), "dimensions are not equal"); @@ -544,7 +547,7 @@ mod test_compression { // disagree. #[test] fn error_on_empty() { - let pivots = views::Matrix::from_element(0, 5, 0.0); + let pivots = views::rowmajor::Owned::from_element(0, 5, 0.0); let offsets = ChunkOffsets::new(Box::new([0, 1, 5])).unwrap(); let result = TransposedTable::from_parts(pivots.as_view(), offsets); assert!(result.is_err(), "dimensions are not equal"); @@ -564,7 +567,7 @@ mod test_compression { for dim in [5, 10, 12] { // Sweep over enough totals to ensure the inner chunks have a non-trivial layout. for total in [1, 2, 3, 7, 8, 9, 10] { - let pivots = views::Matrix::from_fn(total, dim, |_| -> f32 { + let pivots = views::rowmajor::Owned::from_fn(total, dim, |_| -> f32 { StandardUniform {}.sample(&mut rng) }); let offsets = ChunkOffsets::new(Box::new([0, 1, 3, dim])).unwrap(); @@ -682,9 +685,10 @@ mod test_compression { assert_eq!(called.len(), num_data * schema.len()); // Trait Interface. - let mut output = views::Matrix::from_element(num_data, schema.len(), 0); + let mut output = + views::rowmajor::Owned::from_element(num_data, schema.len(), 0); table - .compress_into(data.as_view(), output.as_mut_view()) + .compress_into(data.as_view(), output.as_view_mut()) .unwrap(); assert_eq!(output.nrows(), expected.nrows()); @@ -705,7 +709,7 @@ mod test_compression { // Trait inteface - single step. let mut output = vec![0; schema.len()]; for (i, (row, expected)) in - std::iter::zip(data.row_iter(), expected.row_iter()).enumerate() + std::iter::zip(data.rows(), expected.rows()).enumerate() { table.compress_into(row, output.as_mut_slice()).unwrap(); for (d, (o, e)) in @@ -739,7 +743,7 @@ mod test_compression { // should panic let _ = table.compress_batch( - views::MatrixView::try_from(data.as_slice(), 1, 5).unwrap(), + views::rowmajor::Ref::try_from_data(data.as_slice(), 1, 5).unwrap(), |_, _| -> Result<(), Infallible> { panic!("this shouldn't be called") }, ); } @@ -765,7 +769,7 @@ mod test_compression { let (data, _) = create_dataset(schema, num_centers, num_data, &mut rng); let data_view = - views::MatrixView::try_from(data.as_slice(), num_data, schema.dim()).unwrap(); + views::rowmajor::Ref::try_from_data(data.as_slice(), num_data, schema.dim()).unwrap(); let distribution = rand_distr::StandardUniform {}; for row in 0..data_view.nrows() { @@ -807,7 +811,7 @@ mod test_compression { #[cfg(not(miri))] fn test_table_single_compression_errors() { check_pqtable_single_compression_errors( - &|pivots: views::Matrix, offsets| { + &|pivots: views::rowmajor::Owned, offsets| { TransposedTable::from_parts(pivots.as_view(), offsets).unwrap() }, &"TranposedTable", @@ -818,7 +822,7 @@ mod test_compression { #[cfg(not(miri))] fn test_table_batch_compression_errors() { check_pqtable_batch_compression_errors( - &|pivots: views::Matrix, offsets| { + &|pivots: views::rowmajor::Owned, offsets| { TransposedTable::from_parts(pivots.as_view(), offsets).unwrap() }, &"TranposedTable", @@ -854,19 +858,20 @@ mod test_compression { let offsets = ChunkOffsets::new(offsets.into()).unwrap(); let dim = offsets.dim(); - let pivots = views::Matrix::::from_fn(num_centers, dim, |_| { + let pivots = views::rowmajor::Owned::::from_fn(num_centers, dim, |_| { value_distribution.sample(rng) as f32 }); let table = TransposedTable::from_parts(pivots.as_view(), offsets.clone()).unwrap(); - let mut output = views::Matrix::::from_element(num_chunks, num_centers, 0.0); + let mut output = + views::rowmajor::Owned::::from_element(num_chunks, num_centers, 0.0); let query: Vec<_> = (0..dim) .map(|_| value_distribution.sample(rng) as f32) .collect(); // Inner Product - table.process_into::(&query, output.as_mut_view()); + table.process_into::(&query, output.as_view_mut()); for chunk in 0..num_chunks { let range = offsets.at(chunk); @@ -887,7 +892,7 @@ mod test_compression { } // Squared L2 - table.process_into::(&query, output.as_mut_view()); + table.process_into::(&query, output.as_view_mut()); for chunk in 0..num_chunks { let range = offsets.at(chunk); @@ -932,41 +937,44 @@ mod test_compression { #[should_panic(expected = "query has the wrong number of dimensions")] fn test_process_into_panics_query() { let offsets = ChunkOffsets::new(Box::new([0, 1, 5])).unwrap(); - let data = views::Matrix::::from_element(3, 5, 0.0); + let data = views::rowmajor::Owned::::from_element(3, 5, 0.0); let table = TransposedTable::from_parts(data.as_view(), offsets).unwrap(); assert_eq!(table.dim(), 5); // query has the wrong length. let query = vec![0.0; table.dim() - 1]; - let mut partials = views::Matrix::from_element(table.nchunks(), table.ncenters(), 0.0); - table.process_into::(&query, partials.as_mut_view()); + let mut partials = + views::rowmajor::Owned::from_element(table.nchunks(), table.ncenters(), 0.0); + table.process_into::(&query, partials.as_view_mut()); } #[test] #[should_panic(expected = "output has the wrong number of rows")] fn test_process_into_panics_partials_rows() { let offsets = ChunkOffsets::new(Box::new([0, 1, 5])).unwrap(); - let data = views::Matrix::::from_element(3, 5, 0.0); + let data = views::rowmajor::Owned::::from_element(3, 5, 0.0); let table = TransposedTable::from_parts(data.as_view(), offsets).unwrap(); assert_eq!(table.dim(), 5); let query = vec![0.0; table.dim()]; // partials has the wrong numbers of rows. - let mut partials = views::Matrix::from_element(table.nchunks() - 1, table.ncenters(), 0.0); - table.process_into::(&query, partials.as_mut_view()); + let mut partials = + views::rowmajor::Owned::from_element(table.nchunks() - 1, table.ncenters(), 0.0); + table.process_into::(&query, partials.as_view_mut()); } #[test] #[should_panic(expected = "output has the wrong number of columns")] fn test_process_into_panics_partials_cols() { let offsets = ChunkOffsets::new(Box::new([0, 1, 5])).unwrap(); - let data = views::Matrix::::from_element(3, 5, 0.0); + let data = views::rowmajor::Owned::::from_element(3, 5, 0.0); let table = TransposedTable::from_parts(data.as_view(), offsets).unwrap(); assert_eq!(table.dim(), 5); let query = vec![0.0; table.dim()]; // partials has the wrong numbers of rows. - let mut partials = views::Matrix::from_element(table.nchunks(), table.ncenters() - 1, 0.0); - table.process_into::(&query, partials.as_mut_view()); + let mut partials = + views::rowmajor::Owned::from_element(table.nchunks(), table.ncenters() - 1, 0.0); + table.process_into::(&query, partials.as_view_mut()); } } diff --git a/diskann-quantization/src/product/train.rs b/diskann-quantization/src/product/train.rs index b35e6e81f7..037e2ebddc 100644 --- a/diskann-quantization/src/product/train.rs +++ b/diskann-quantization/src/product/train.rs @@ -5,7 +5,10 @@ use diskann_utils::{ strided::Strided, - views::{self, Matrix}, + views::{ + self, + rowmajor::{self, Matrix, MatrixMut}, + }, }; #[cfg(feature = "rayon")] use rayon::iter::{IntoParallelIterator, ParallelIterator}; @@ -40,15 +43,19 @@ impl LightPQTrainingParameters { pub struct SimplePivots { dim: usize, ncenters: usize, - pivots: Vec>, + pivots: Vec>, } -fn flatten(pivots: &[Matrix], ncenters: usize, dim: usize) -> Matrix { - let mut flattened = Matrix::from_element(ncenters, dim, T::default()); +fn flatten( + pivots: &[rowmajor::Owned], + ncenters: usize, + dim: usize, +) -> rowmajor::Owned { + let mut flattened = rowmajor::Owned::from_element(ncenters, dim, T::default()); let mut col_start = 0; for matrix in pivots { assert_eq!(matrix.nrows(), flattened.nrows()); - for (row_index, row) in matrix.row_iter().enumerate() { + for (row_index, row) in matrix.rows().enumerate() { let dst = &mut flattened.row_mut(row_index)[col_start..col_start + row.len()]; dst.copy_from_slice(row); } @@ -59,7 +66,7 @@ fn flatten(pivots: &[Matrix], ncenters: usize, dim: usize) impl SimplePivots { /// Return the selected pivots for each chunk. - pub fn pivots(&self) -> &[Matrix] { + pub fn pivots(&self) -> &[rowmajor::Owned] { &self.pivots } @@ -77,7 +84,7 @@ pub trait TrainQuantizer { fn train( &self, - data: views::MatrixView, + data: views::rowmajor::Ref, schema: crate::views::ChunkOffsetsView<'_>, parallelism: Parallelism, rng_builder: &R, @@ -108,7 +115,7 @@ impl TrainQuantizer for LightPQTrainingParameters { /// * `NaN` or infinities are observed during the training process. fn train( &self, - data: views::MatrixView, + data: views::rowmajor::Ref, schema: crate::views::ChunkOffsetsView<'_>, parallelism: Parallelism, rng_builder: &R, @@ -123,7 +130,7 @@ impl TrainQuantizer for LightPQTrainingParameters { #[inline(never)] fn train( trainer: &LightPQTrainingParameters, - data: views::MatrixView, + data: views::rowmajor::Ref, schema: crate::views::ChunkOffsetsView<'_>, parallelism: Parallelism, rng_builder: &(dyn BoxedRngBuilder + Sync), @@ -132,7 +139,7 @@ impl TrainQuantizer for LightPQTrainingParameters { // Make sure we're provided sane values for our schema. assert_eq!(data.ncols(), schema.dim()); - let thunk = |i| -> Result, PQTrainingError> { + let thunk = |i| -> Result, PQTrainingError> { let range = schema.at(i); // Check for cancelation. @@ -169,20 +176,22 @@ impl TrainQuantizer for LightPQTrainingParameters { // Allocate scratch data structures. let norms: Vec = view.rows().map(square_norm).collect(); let transpose = BlockTransposed::::from_strided(view); - let mut centers = Matrix::try_from_element(trainer.ncenters, range.len(), 0.0) - .map_err(|err| PQTrainingError { - chunk: i, - of: schema.len(), - dim: range.len(), - kind: PQTrainingErrorKind::InternalError(Box::new(err)), - })?; + let mut centers = + rowmajor::Owned::try_from_element(trainer.ncenters, range.len(), 0.0).map_err( + |err| PQTrainingError { + chunk: i, + of: schema.len(), + dim: range.len(), + kind: PQTrainingErrorKind::InternalError(Box::new(err)), + }, + )?; // Construct the random number generator seeded by the PQ chunk. let mut rng = rng_builder.build_boxed_rng(i); // Initialization kmeans::plusplus::kmeans_plusplus_into_inner( - centers.as_mut_view(), + centers.as_view_mut(), view, transpose.as_view(), &norms, @@ -210,7 +219,7 @@ impl TrainQuantizer for LightPQTrainingParameters { view, &norms, transpose.as_view(), - centers.as_mut_view(), + centers.as_view_mut(), trainer.lloyds_reps, ); Ok(centers) @@ -306,17 +315,18 @@ mod tests { let dim: usize = sub_dims.iter().sum(); // Create the sub matrices. - let matrices: Vec> = std::iter::zip(sub_dims.iter(), prefix_sum.iter()) - .map(|(&this_dim, &offset)| { - let mut m = Matrix::from_element(nrows, this_dim, 0); - for r in 0..nrows { - for c in 0..this_dim { - *m.element_mut(r, c) = dim * r + offset + c; + let matrices: Vec> = + std::iter::zip(sub_dims.iter(), prefix_sum.iter()) + .map(|(&this_dim, &offset)| { + let mut m = rowmajor::Owned::from_element(nrows, this_dim, 0); + for r in 0..nrows { + for c in 0..this_dim { + *m.element_mut(r, c) = dim * r + offset + c; + } } - } - m - }) - .collect(); + m + }) + .collect(); let flattened = flatten(&matrices, nrows, dim); // Check that the output is correct. @@ -332,9 +342,9 @@ mod tests { } struct ClusteredDataset { - data: Matrix, + data: rowmajor::Owned, // The pre-configured center point for the manufactured clusters. - centers: Matrix, + centers: rowmajor::Owned, } impl DatasetBuilder { @@ -364,8 +374,8 @@ mod tests { .map(|chunk| { let dim = schema.at(chunk).len(); - let mut initial = Matrix::from_element(ndata, dim, 0.0); - let mut centers = Matrix::from_element(self.nclusters, 1, 0.0); + let mut initial = rowmajor::Owned::from_element(ndata, dim, 0.0); + let mut centers = rowmajor::Owned::from_element(self.nclusters, 1, 0.0); // The starting offset for clusters. let offset = offsets_distribution.sample(rng); @@ -386,7 +396,7 @@ mod tests { // Shuffle the dataset. indices.shuffle(rng); - let mut piece = Matrix::from_element(ndata, dim, 0.0); + let mut piece = rowmajor::Owned::from_element(ndata, dim, 0.0); for (dst, src) in indices.iter().enumerate() { piece.row_mut(dst).copy_from_slice(initial.row(*src)); } @@ -464,7 +474,7 @@ mod tests { // Start matching pivots to expected centers. let mut seen: Vec = (0..dataset.centers.nrows()).map(|_| false).collect(); - for row in pivot.row_iter() { + for row in pivot.rows() { let mut min_distance = f32::MAX; let mut min_index = 0; for c in 0..dataset.centers.nrows() { @@ -585,7 +595,7 @@ mod tests { // pivots exceeds the number of dataset items. #[test] fn tests_succeeded_with_too_many_pivots() { - let data = Matrix::::from_element(10, 5, 1.0); + let data = rowmajor::Owned::::from_element(10, 5, 1.0); let offsets: Vec = vec![0, 1, 4, 5]; let trainer = LightPQTrainingParameters::new(2 * data.nrows(), 6); @@ -612,7 +622,7 @@ mod tests { "expected pivot 0 to be the non-zero pivot" ); - for (i, row) in flat.row_iter().enumerate() { + for (i, row) in flat.rows().enumerate() { // skip the first row. if i == 0 { continue; @@ -656,7 +666,7 @@ mod tests { assert!(format(&err).contains("infinity")); }; - let mut data = Matrix::::from_element(nrows, ncols, 1.0); + let mut data = rowmajor::Owned::::from_element(nrows, ncols, 1.0); // Positive Infinity *data.element_mut(r, c) = f32::INFINITY; diff --git a/diskann-quantization/src/scalar/mod.rs b/diskann-quantization/src/scalar/mod.rs index db927031d9..88baf10912 100644 --- a/diskann-quantization/src/scalar/mod.rs +++ b/diskann-quantization/src/scalar/mod.rs @@ -62,7 +62,10 @@ //! scalar::{self, train, CompensatedVector, CompensatedIP, CompensatedSquaredL2}, //! num::Positive, //! }; -//! use diskann_utils::{Reborrow, ReborrowMut, views::Matrix}; +//! use diskann_utils::{ +//! Reborrow, ReborrowMut, +//! views::rowmajor::{self, Matrix, MatrixMut}, +//! }; //! use rand::{rngs::StdRng, SeedableRng, distr::Distribution}; //! use rand_distr::StandardNormal; //! use diskann_vector::{PureDistanceFunction, DistanceFunction, distance}; @@ -74,8 +77,8 @@ //! // Construct a set of offsets for each dimension. //! let offset: Vec = (0..dim).map(|_| distribution.sample(&mut rng)).collect(); //! // The output dataset. -//! let mut data = Matrix::::from_element(nvectors, dim, 0.0); -//! for row in data.row_iter_mut() { +//! let mut data = rowmajor::Owned::::from_element(nvectors, dim, 0.0); +//! for row in data.rows_mut() { //! std::iter::zip(row.iter_mut(), offset.iter()).for_each(|(r, i)| { //! let v: f32 = distribution.sample(&mut rng); //! *r = i + v; diff --git a/diskann-quantization/src/scalar/quantizer.rs b/diskann-quantization/src/scalar/quantizer.rs index e131fe8cee..11228d9cea 100644 --- a/diskann-quantization/src/scalar/quantizer.rs +++ b/diskann-quantization/src/scalar/quantizer.rs @@ -38,11 +38,11 @@ use crate::{ /// CompensatedIP, CompensatedSquaredL2, /// } /// }; -/// use diskann_utils::{views::Matrix, Reborrow, ReborrowMut}; +/// use diskann_utils::{views::rowmajor::{self, Matrix, MatrixMut}, Reborrow, ReborrowMut}; /// use diskann_vector::DistanceFunction; /// /// // A small training set consisting of two 5-dimensional vectors. -/// let mut data = Matrix::::from_element(2, 5, 0.0); +/// let mut data = rowmajor::Owned::::from_element(2, 5, 0.0); /// data.row_mut(0).copy_from_slice(&[-1.0, -1.0, -1.0, -1.0, -1.0]); /// data.row_mut(1).copy_from_slice(&[1.0, 1.0, 1.0, 1.0, 1.0]); /// @@ -459,7 +459,10 @@ pub enum SQComparisonError { mod tests { use std::collections::HashSet; - use diskann_utils::{ReborrowMut, views}; + use diskann_utils::{ + ReborrowMut, views, + views::rowmajor::{Matrix, MatrixMut}, + }; use rand::{ SeedableRng, @@ -593,7 +596,7 @@ mod tests { // Create a shuffled matrix of offset values for each dimension. This ensure that // each dimension covers the target dynamic range, but in a different order so // we can rule out cross-coupling of dimensions. - let dim_offsets: views::Matrix = { + let dim_offsets: views::rowmajor::Owned = { let range_min = -min_encodable - 3.0 * scale; let range_max = max_encodable + 3.0 * scale; let mut base: Vec = Vec::new(); @@ -607,7 +610,7 @@ mod tests { // Push one more to have one point above `range_max`. base.push(i); - let mut output = views::Matrix::from_element(base.len(), dim, 0.0); + let mut output = views::rowmajor::Owned::from_element(base.len(), dim, 0.0); (0..dim).for_each(|j| { base.shuffle(rng); for (i, b) in base.iter().enumerate() { diff --git a/diskann-quantization/src/scalar/train.rs b/diskann-quantization/src/scalar/train.rs index becf99fd53..885ae573d8 100644 --- a/diskann-quantization/src/scalar/train.rs +++ b/diskann-quantization/src/scalar/train.rs @@ -8,7 +8,7 @@ use crate::{ num::Positive, utils::{compute_means_and_average_norm, compute_variances}, }; -use diskann_utils::views; +use diskann_utils::views::rowmajor; /// Parameters controlling the generation of the scalar quantization Quantizer. /// @@ -71,7 +71,7 @@ impl ScalarQuantizationParameters { /// # Parallelism /// /// This function is single threaded. - pub fn train(&self, data: views::MatrixView) -> ScalarQuantizer + pub fn train(&self, data: rowmajor::Ref) -> ScalarQuantizer where T: Copy + Into + Into, { @@ -110,6 +110,7 @@ mod tests { use super::*; use crate::test_util::create_test_problem; + use diskann_utils::views::rowmajor::Matrix; fn test_train_impl(nrows: usize, ncols: usize, seed: u64) { // Test Default diff --git a/diskann-quantization/src/spherical/iface.rs b/diskann-quantization/src/spherical/iface.rs index d823805eac..10c9686657 100644 --- a/diskann-quantization/src/spherical/iface.rs +++ b/diskann-quantization/src/spherical/iface.rs @@ -40,10 +40,10 @@ //! spherical::{iface, SupportedMetric, SphericalQuantizer, PreScale}, //! num::PowerOfTwo, //! }; -//! use diskann_utils::views::Matrix; +//! use diskann_utils::views::rowmajor::{self, Matrix}; //! //! // For illustration purposes, the dataset consists of just a single vector. -//! let mut data = Matrix::from_element(1, 4, 1.0); +//! let mut data = rowmajor::Owned::from_element(1, 4, 1.0); //! let quantizer = SphericalQuantizer::train( //! data.as_view(), //! TransformKind::Null, @@ -2077,7 +2077,7 @@ where #[cfg(test)] mod tests { - use diskann_utils::views::{Matrix, MatrixView}; + use diskann_utils::views::rowmajor::{self, Matrix}; use rand::{SeedableRng, rngs::StdRng}; use super::*; @@ -2146,7 +2146,7 @@ mod tests { } #[inline(never)] - fn test_plan(plan: &dyn Quantizer, nbits: usize, dataset: MatrixView) { + fn test_plan(plan: &dyn Quantizer, nbits: usize, dataset: rowmajor::Ref) { // Perform the bit-specific test. if nbits == 1 { test_plan_1_bit(plan); @@ -2366,7 +2366,7 @@ mod tests { } } - fn make_impl(metric: SupportedMetric) -> (Impl, Matrix) + fn make_impl(metric: SupportedMetric) -> (Impl, rowmajor::Owned) where Impl: Constructible, { @@ -2460,7 +2460,7 @@ mod tests { test_plan(&plan, 8, data.as_view()); } - fn test_dataset() -> Matrix { + fn test_dataset() -> rowmajor::Owned { let data = vec![ 0.28657, -0.0318168, @@ -2592,7 +2592,7 @@ mod tests { -0.324718, // row 15 ]; - Matrix::try_from(data.into(), 16, 8).unwrap() + rowmajor::Owned::try_from_data(data.into(), 16, 8).unwrap() } #[cfg(feature = "flatbuffers")] @@ -2609,7 +2609,7 @@ mod tests { quantizer: &dyn Quantizer, deserialized: &dyn Quantizer, nbits: usize, - dataset: MatrixView<'_, f32>, + dataset: rowmajor::Ref<'_, f32>, ) { let scoped_global = ScopedAllocator::global(); @@ -2633,7 +2633,7 @@ mod tests { let mut a = Poly::broadcast(u8::default(), quantizer.bytes(), alloc).unwrap(); let mut b = Poly::broadcast(u8::default(), quantizer.bytes(), alloc).unwrap(); - for row in dataset.row_iter() { + for row in dataset.rows() { quantizer .compress(row, OpaqueMut::new(&mut a), scoped_global) .unwrap(); @@ -2658,14 +2658,14 @@ mod tests { let d_computer = deserialized.distance_computer(GlobalAllocator).unwrap(); let d_computer_ref = deserialized.distance_computer_ref(); - for r0 in dataset.row_iter() { + for r0 in dataset.rows() { quantizer .compress(r0, OpaqueMut::new(&mut a0), scoped_global) .unwrap(); deserialized .compress(r0, OpaqueMut::new(&mut b0), scoped_global) .unwrap(); - for r1 in dataset.row_iter() { + for r1 in dataset.rows() { quantizer .compress(r1, OpaqueMut::new(&mut a1), scoped_global) .unwrap(); @@ -2701,7 +2701,7 @@ mod tests { continue; } - for r in dataset.row_iter() { + for r in dataset.rows() { let q_computer = quantizer .fused_query_computer(r, layout, false, GlobalAllocator, scoped_global) .unwrap(); @@ -2709,7 +2709,7 @@ mod tests { .fused_query_computer(r, layout, false, GlobalAllocator, scoped_global) .unwrap(); - for u in dataset.row_iter() { + for u in dataset.rows() { quantizer .compress(u, OpaqueMut::new(&mut a), scoped_global) .unwrap(); @@ -2731,7 +2731,7 @@ mod tests { fn test_plan_serialization( quantizer: &dyn Quantizer, nbits: usize, - dataset: MatrixView, + dataset: rowmajor::Ref, ) { let global = GlobalAllocator; @@ -3151,11 +3151,14 @@ mod tests { // Helpers // ///////////// - fn compress_dataset(quantizer: &dyn Quantizer, dataset: MatrixView) -> Vec> { + fn compress_dataset( + quantizer: &dyn Quantizer, + dataset: rowmajor::Ref, + ) -> Vec> { let scoped_global = ScopedAllocator::global(); let alloc = AlignedAllocator::new(PowerOfTwo::new(4).unwrap()); dataset - .row_iter() + .rows() .map(|row| { let mut buf = Poly::broadcast(u8::default(), quantizer.bytes(), alloc).unwrap(); quantizer @@ -3168,7 +3171,7 @@ mod tests { fn compute_layout_distances( quantizer: &dyn Quantizer, - dataset: MatrixView, + dataset: rowmajor::Ref, compressed: &[Vec], allow_rescale: bool, ) -> Vec { @@ -3178,7 +3181,7 @@ mod tests { .filter(|&layout| quantizer.is_supported(layout)) .map(|layout| { let distances = dataset - .row_iter() + .rows() .map(|query_row| { let computer = quantizer .fused_query_computer( @@ -3209,7 +3212,7 @@ mod tests { /// per-layout query distances. fn assert_layout_distances( quantizer: &dyn Quantizer, - dataset: MatrixView, + dataset: rowmajor::Ref, compressed: &[Vec], expected: &[LayoutDistances], allow_rescale: bool, @@ -3224,8 +3227,7 @@ mod tests { ); for (qi, (query_row, expected_distances)) in - std::iter::zip(dataset.row_iter(), layout_distances.distances.iter()) - .enumerate() + std::iter::zip(dataset.rows(), layout_distances.distances.iter()).enumerate() { let computer = quantizer .fused_query_computer( @@ -3260,7 +3262,7 @@ mod tests { quantizer: &dyn Quantizer, transform: DataTransform, pre_scale: ScaleConfig, - dataset: MatrixView, + dataset: rowmajor::Ref, ) -> Baseline { let compressed_vectors = compress_dataset(quantizer, dataset); @@ -3379,7 +3381,7 @@ mod tests { fn check_baseline( baseline: &Baseline, - dataset: MatrixView, + dataset: rowmajor::Ref, expected_transform: DataTransform, expected_pre_scale: ScaleConfig, ) { @@ -3497,7 +3499,7 @@ mod tests { metric: SupportedMetric, transform: DataTransform, pre_scale: ScaleConfig, - ) -> (Poly, Matrix) + ) -> (Poly, rowmajor::Owned) where Impl: Constructible + Quantizer, { diff --git a/diskann-quantization/src/spherical/quantizer.rs b/diskann-quantization/src/spherical/quantizer.rs index 5967fc8c06..dcb1072bce 100644 --- a/diskann-quantization/src/spherical/quantizer.rs +++ b/diskann-quantization/src/spherical/quantizer.rs @@ -5,7 +5,10 @@ use std::num::NonZeroUsize; -use diskann_utils::{ReborrowMut, views::MatrixView}; +use diskann_utils::{ + ReborrowMut, + views::rowmajor::{self, Matrix}, +}; use diskann_vector::{ MathematicalValue, Norm, PureDistanceFunction, distance::InnerProduct, norm::FastL2Norm, }; @@ -228,7 +231,7 @@ where /// If argument `pre_scale` is given, then all vectors compressed by this quantizer will /// first be scaled by this value. Note that if given, `pre_scale` **must** be positive. pub fn train( - data: MatrixView, + data: rowmajor::Ref, transform: TransformKind, metric: SupportedMetric, pre_scale: PreScale, @@ -243,7 +246,7 @@ where // cut down on excess monomorphization. #[inline(never)] fn train( - data: MatrixView, + data: rowmajor::Ref, transform: TransformKind, metric: SupportedMetric, pre_scale: PreScale, @@ -1233,7 +1236,7 @@ mod tests { use std::fmt::Display; - use diskann_utils::{ReborrowMut, lazy_format, views::Matrix}; + use diskann_utils::{ReborrowMut, lazy_format, views::rowmajor::MatrixMut}; use diskann_vector::{PureDistanceFunction, norm::FastL2NormSquared}; use diskann_wide::ARCH; use rand::{ @@ -2285,7 +2288,7 @@ mod tests { #[test] fn err_dim_cannot_be_zero() { - let data = Matrix::from_element(10, 0, 0.0f32); + let data = rowmajor::Owned::from_element(10, 0, 0.0f32); let mut rng = StdRng::seed_from_u64(0xe3e9f42ed9f15883); let err = SphericalQuantizer::train( data.as_view(), @@ -2303,7 +2306,7 @@ mod tests { #[test] fn err_norm_must_be_positive() { - let data = Matrix::from_element(10, 10, 0.0f32); + let data = rowmajor::Owned::from_element(10, 10, 0.0f32); let mut rng = StdRng::seed_from_u64(0xe3e9f42ed9f15883); let err = SphericalQuantizer::train( data.as_view(), @@ -2321,7 +2324,7 @@ mod tests { #[test] fn err_norm_cannot_be_infinity() { - let mut data = Matrix::from_element(10, 10, 0.0f32); + let mut data = rowmajor::Owned::from_element(10, 10, 0.0f32); *data.element_mut(2, 5) = f32::INFINITY; let mut rng = StdRng::seed_from_u64(0xe3e9f42ed9f15883); @@ -2341,7 +2344,7 @@ mod tests { #[test] fn err_reciprocal_norm_cannot_be_infinity() { - let mut data = Matrix::from_element(10, 10, 0.0f32); + let mut data = rowmajor::Owned::from_element(10, 10, 0.0f32); *data.element_mut(2, 5) = 2.93863e-39; let mut rng = StdRng::seed_from_u64(0xe3e9f42ed9f15883); @@ -2400,7 +2403,7 @@ mod tests { #[test] fn compression_errors_data() { let mut rng = StdRng::seed_from_u64(0xe3e9f42ed9f15883); - let data = Matrix::::from_fn(16, 12, |_| StandardNormal {}.sample(&mut rng)); + let data = rowmajor::Owned::::from_fn(16, 12, |_| StandardNormal {}.sample(&mut rng)); let quantizer = SphericalQuantizer::train( data.as_view(), diff --git a/diskann-quantization/src/test_util.rs b/diskann-quantization/src/test_util.rs index 8bb9afe307..20bf6fac6e 100644 --- a/diskann-quantization/src/test_util.rs +++ b/diskann-quantization/src/test_util.rs @@ -9,7 +9,7 @@ use std::sync::{ atomic::{AtomicUsize, Ordering}, }; -use diskann_utils::views::Matrix; +use diskann_utils::views::rowmajor::{self, Matrix, MatrixMut}; use rand::{ distr::{Distribution, Uniform}, rngs::StdRng, @@ -105,7 +105,7 @@ pub(crate) fn compute_absolute_error(got: f32, expected: f32) -> f32 { } pub(crate) struct TestProblem { - pub(crate) data: Matrix, + pub(crate) data: rowmajor::Owned, pub(crate) means: Vec, pub(crate) variances: Vec, pub(crate) mean_norm: f64, @@ -169,17 +169,17 @@ pub(crate) fn create_test_problem(nrows: usize, ncols: usize, rng: &mut StdRng) }) .collect(); - let mut data = Matrix::::from_element(nrows, ncols, 0.0); + let mut data = rowmajor::Owned::::from_element(nrows, ncols, 0.0); for col in 0..ncols { offsets.shuffle(rng); - for (row, offset) in std::iter::zip(data.row_iter_mut(), offsets.iter()) { + for (row, offset) in std::iter::zip(data.rows_mut(), offsets.iter()) { row[col] = means[col] + scales[col] * offset; } } // Compute the mean norm directly. let mean_norm = data - .row_iter() + .rows() .map(|row| { row.iter() .map(|&i| { diff --git a/diskann-quantization/src/utils.rs b/diskann-quantization/src/utils.rs index baaadc0214..4f4535acb6 100644 --- a/diskann-quantization/src/utils.rs +++ b/diskann-quantization/src/utils.rs @@ -7,7 +7,7 @@ use std::ptr::NonNull; use thiserror::Error; -use diskann_utils::views::MatrixView; +use diskann_utils::views::rowmajor::{self, Matrix}; /// Specify featres and config flags that will be propagated to `docsrs` config. macro_rules! features { @@ -106,14 +106,14 @@ where pub(crate) struct CannotBeEmpty; /// Compute the mean of each column in `data` as well as the average norm. -pub(crate) fn compute_means_and_average_norm(data: MatrixView) -> (Vec, f64) +pub(crate) fn compute_means_and_average_norm(data: rowmajor::Ref) -> (Vec, f64) where T: Into + Copy, { // Compute the centroid of the dataset as well as the sums of the norms of every // element in the dataset. let mut means: Vec = vec![0.0; data.ncols()]; - let norm_sum = data.row_iter().fold(0.0f64, |accum: f64, row| { + let norm_sum = data.rows().fold(0.0f64, |accum: f64, row| { // Accumulate this row into the means. std::iter::zip(means.iter_mut(), row.iter()).for_each(|(m, r)| { let r: f64 = (*r).into(); @@ -141,7 +141,7 @@ where } /// Compute the mean of each column in `data` as well as the average norm. -pub(crate) fn compute_normalized_means(data: MatrixView) -> Result, CannotBeEmpty> +pub(crate) fn compute_normalized_means(data: rowmajor::Ref) -> Result, CannotBeEmpty> where T: Into + Copy, { @@ -161,7 +161,7 @@ where x * x }; - data.row_iter().for_each(|row| { + data.rows().for_each(|row| { let norm = row.iter().map(square).sum::().sqrt(); let inv_norm = if norm == 0.0 { 1.0 } else { 1.0 / norm }; @@ -177,14 +177,14 @@ where Ok(means) } -pub(crate) fn compute_variances(data: MatrixView, means: &[f64]) -> Vec +pub(crate) fn compute_variances(data: rowmajor::Ref, means: &[f64]) -> Vec where T: Into + Copy, { assert_eq!(data.ncols(), means.len()); let mut variances: Vec = vec![0.0; data.ncols()]; - data.row_iter().for_each(|row| { + data.rows().for_each(|row| { variances .iter_mut() .zip(std::iter::zip(row.iter(), means.iter())) @@ -205,7 +205,7 @@ where #[cfg(test)] mod tests { - use diskann_utils::views::Matrix; + use diskann_utils::views::rowmajor::MatrixMut; use diskann_vector::{Norm, norm::FastL2Norm}; use rand::{SeedableRng, rngs::StdRng}; @@ -271,7 +271,7 @@ mod tests { let test_problem = create_test_problem(nrows, ncols, &mut rng); let mut normalized_data = test_problem.data.clone(); - normalized_data.row_iter_mut().for_each(normalize); + normalized_data.rows_mut().for_each(normalize); // Compute the means and mean norm of the normalized data. // @@ -317,12 +317,12 @@ mod tests { #[test] fn test_normalized_means_corner_cases() { // If the input data has no columns, the returned vector should be empty. - let data = Matrix::from_element(10, 0, 1.0f32); + let data = rowmajor::Owned::from_element(10, 0, 1.0f32); let means = compute_normalized_means(data.as_view()).unwrap(); assert!(means.is_empty()); // If the data has no rows, an error should be returned. - let data = Matrix::from_element(0, 10, 1.0f32); + let data = rowmajor::Owned::from_element(0, 10, 1.0f32); let _: CannotBeEmpty = compute_normalized_means(data.as_view()).unwrap_err(); } diff --git a/diskann-tools/src/bin/compute_streaming_groundtruth.rs b/diskann-tools/src/bin/compute_streaming_groundtruth.rs index 335d1c3a2f..1f6df505de 100644 --- a/diskann-tools/src/bin/compute_streaming_groundtruth.rs +++ b/diskann-tools/src/bin/compute_streaming_groundtruth.rs @@ -31,7 +31,7 @@ use diskann_tools::utils::{ init_subscriber, write_ground_truth, CMDResult, CMDToolError, DataType, }; use diskann_utils::io::read_bin; -use diskann_utils::views::Matrix; +use diskann_utils::views::rowmajor::{self, Matrix}; use diskann_vector::{distance::Metric, DistanceFunction}; use rayon::prelude::*; @@ -44,8 +44,8 @@ trait GroundtruthDistance: Send + Sync { } struct MatrixDistance<'a, V: VectorRepr + Send + Sync> { - dataset: &'a Matrix, - queries: &'a Matrix, + dataset: &'a rowmajor::Owned, + queries: &'a rowmajor::Owned, distance_fn: V::Distance, } diff --git a/diskann-tools/src/bin/generate_minmax.rs b/diskann-tools/src/bin/generate_minmax.rs index 32ba504a71..bbe41bc0b1 100644 --- a/diskann-tools/src/bin/generate_minmax.rs +++ b/diskann-tools/src/bin/generate_minmax.rs @@ -22,6 +22,7 @@ use diskann_quantization::{ CompressInto, }; use diskann_utils::io::Metadata; +use diskann_utils::views::rowmajor::Matrix; use half::f16; use rand::{rngs::StdRng, SeedableRng}; diff --git a/diskann-tools/src/utils/ground_truth.rs b/diskann-tools/src/utils/ground_truth.rs index 2e8de3c48a..d7aae1ddb0 100644 --- a/diskann-tools/src/utils/ground_truth.rs +++ b/diskann-tools/src/utils/ground_truth.rs @@ -20,7 +20,7 @@ use diskann_providers::utils::{ }; use diskann_utils::{ io::{read_bin, Metadata}, - views::Matrix, + views::rowmajor::{self, Matrix}, }; use diskann_vector::{distance::Metric, DistanceFunction}; use itertools::Itertools; @@ -282,7 +282,7 @@ pub fn compute_range_ground_truth_from_datafiles< pub fn compute_range_ground_truth_from_data( distance_function: Metric, dataset_iter: VectorDataIterator, - queries: &Matrix, + queries: &rowmajor::Owned, radius: f32, query_bitmaps: Option>, ) -> CMDResult>>> @@ -295,7 +295,7 @@ where let query_dim = queries.ncols(); let mut ground_truth: Vec>> = vec![Vec::new(); query_num]; - let mut queries_and_result: Vec<_> = queries.row_iter().zip(ground_truth.iter_mut()).collect(); + let mut queries_and_result: Vec<_> = queries.rows().zip(ground_truth.iter_mut()).collect(); let distance_comparer = V::distance(distance_function, Some(query_dim)); @@ -575,7 +575,7 @@ type Npq = Vec>; /// /// * `distance_function` - e.g. L2 /// * `dataset_iter` - The iterator over the dataset vectors and associated data. -/// * `queries` - Query vectors as a row-major `Matrix` of shape `num_queries × query_dim`. +/// * `queries` - Query vectors as a row-major matrix (`rowmajor::Owned`) of shape `num_queries × query_dim`. /// `query_dim` is inferred from `queries.ncols()`. /// * `recall_at` - The number of neighbors to compute for each query. /// * `insert_iter` - Optional iterator containing more dataset vectors. This may be useful if you are testing recall for an index that has points dynamically inserted into it. @@ -584,7 +584,7 @@ type Npq = Vec>; pub fn compute_ground_truth_from_data( distance_function: Metric, dataset_iter: VectorDataIterator, - queries: &Matrix, + queries: &rowmajor::Owned, recall_at: u32, insert_iter: Option>, skip_base: Option, @@ -602,7 +602,7 @@ where .map(|_| NeighborPriorityQueue::new(recall_at as usize)) .collect(); let mut queries_and_neighbor_queue: Vec<_> = - queries.row_iter().zip(neighbor_queues.iter_mut()).collect(); + queries.rows().zip(neighbor_queues.iter_mut()).collect(); let distance_comparer = V::distance(distance_function, Some(query_dim)); @@ -690,8 +690,8 @@ where pub fn compute_multivec_ground_truth_from_data( distance_function: Metric, aggregation_method: MultivecAggregationMethod, - base_vectors: Vec>, - queries: Vec>, + base_vectors: Vec>, + queries: Vec>, query_dim: usize, recall_at: u32, query_bitmaps: Option>, @@ -734,8 +734,8 @@ where let distance = match aggregation_method { MultivecAggregationMethod::AveragePairwise => { let mut total_distance = 0.0; - for query_vec in query_multivec.row_iter() { - for base_vec in base_multivec.row_iter() { + for query_vec in query_multivec.rows() { + for base_vec in base_multivec.rows() { let dist = distance_comparer .evaluate_similarity(query_vec, base_vec); total_distance += dist; @@ -746,8 +746,8 @@ where } MultivecAggregationMethod::MinPairwise => { let mut min_distance = f32::MAX; - for query_vec in query_multivec.row_iter() { - for base_vec in base_multivec.row_iter() { + for query_vec in query_multivec.rows() { + for base_vec in base_multivec.rows() { let dist = distance_comparer .evaluate_similarity(query_vec, base_vec); min_distance = min_distance.min(dist); @@ -757,9 +757,9 @@ where } MultivecAggregationMethod::AvgofMins => { let mut distance = 0_f32; - for query_vec in query_multivec.row_iter() { + for query_vec in query_multivec.rows() { let mut local_min = f32::MAX; - for base_vec in base_multivec.row_iter() { + for base_vec in base_multivec.rows() { let dist = distance_comparer .evaluate_similarity(query_vec, base_vec); local_min = local_min.min(dist); diff --git a/diskann-tools/src/utils/relative_contrast.rs b/diskann-tools/src/utils/relative_contrast.rs index 5ad5f3936f..41aef04aa7 100644 --- a/diskann-tools/src/utils/relative_contrast.rs +++ b/diskann-tools/src/utils/relative_contrast.rs @@ -6,6 +6,7 @@ use diskann::{utils::VectorRepr, ANNError}; use diskann_providers::storage::StorageReadProvider; use diskann_utils::io::read_bin; +use diskann_utils::views::rowmajor::Matrix; use rand::Rng; use crate::utils::{CMDResult, CMDToolError}; @@ -66,9 +67,9 @@ pub fn compute_relative_contrast> = base_data.row_iter().map(|x| x.to_vec()).collect(); - let query: Vec> = query_data.row_iter().map(|x| x.to_vec()).collect(); - let gt: Vec> = gt_data.row_iter().map(|x| x.to_vec()).collect(); + let base: Vec> = base_data.rows().map(|x| x.to_vec()).collect(); + let query: Vec> = query_data.rows().map(|x| x.to_vec()).collect(); + let gt: Vec> = gt_data.rows().map(|x| x.to_vec()).collect(); let mut mean_rc = 0.0; diff --git a/diskann-utils/README.md b/diskann-utils/README.md index 96e243e5c9..251cbfee81 100644 --- a/diskann-utils/README.md +++ b/diskann-utils/README.md @@ -4,7 +4,7 @@ Traits, structs, and algorithms to be shared between the `diskann` specific crat auxiliary crates like `quantization`. The goal is to allow crates like `diskann` to use traits like `Reborrow` and structs -like `MatrixView` without relying on `quantization`. +like `rowmajor::Matrix` without relying on `quantization`. Dependencies of this crate should be kept to a minimum as like `vector` and `wide`, it sits at the very bottom of the DiskANN dependency stack. diff --git a/diskann-utils/src/internal.rs b/diskann-utils/src/internal.rs index b4a30ef5f5..cac5ec9404 100644 --- a/diskann-utils/src/internal.rs +++ b/diskann-utils/src/internal.rs @@ -3,7 +3,34 @@ * Licensed under the MIT license. */ -pub(crate) fn slice_to_nonnull(s: &[T]) -> std::ptr::NonNull { +use std::ptr::NonNull; + +pub(crate) fn slice_to_nonnull(s: &[T]) -> NonNull { // SAFETY: slices are guaranteed to have non-null base pointers. unsafe { std::ptr::NonNull::new_unchecked(s.as_ptr().cast_mut()) } } + +pub(crate) fn mut_slice_to_nonnull(s: &mut [T]) -> NonNull { + // SAFETY: slices are guaranteed to have non-null base pointers. + unsafe { std::ptr::NonNull::new_unchecked(s.as_mut_ptr()) } +} + +pub(crate) fn box_to_nonnull(b: Box<[T]>) -> NonNull { + let ptr = Box::into_raw(b).cast::(); + // SAFETY: boxes are guaranteed to have non-null base pointers. + unsafe { NonNull::new_unchecked(ptr) } +} + +/// # Safety +/// +/// `p` must have come from [`box_to_nonnull`] where the argument boxed slice had a length +/// of exactly `len`. +/// +/// Additionally, this must pair uniquely with said call to [`box_to_nonnull`] to avoid +/// a double or otherwise invalid free. +pub(crate) unsafe fn nonnull_to_box(p: NonNull, len: usize) -> Box<[T]> { + let slice = std::ptr::slice_from_raw_parts_mut(p.as_ptr(), len); + + // SAFETY: Inherited from caller. + unsafe { Box::from_raw(slice) } +} diff --git a/diskann-utils/src/io.rs b/diskann-utils/src/io.rs index 08b649a6f9..9b67198bc5 100644 --- a/diskann-utils/src/io.rs +++ b/diskann-utils/src/io.rs @@ -16,12 +16,12 @@ use std::io::{Read, Seek, Write}; use diskann_wide::{LoHi, SplitJoin}; use thiserror::Error; -use crate::views::{Layout, Matrix, MatrixView}; +use crate::views::rowmajor::{self, Layout, Matrix, MatrixMut}; /// Read a matrix of `T` from the DiskANN binary format (see [module docs](self)). /// /// Validates that the reader contains enough data before allocating. -pub fn read_bin(reader: &mut (impl Read + Seek)) -> Result, ReadBinError> +pub fn read_bin(reader: &mut (impl Read + Seek)) -> Result, ReadBinError> where T: bytemuck::Pod, { @@ -57,7 +57,8 @@ where }); } - let mut data = Matrix::from_element_with_layout(layout, ::zeroed()); + let mut data = + rowmajor::Owned::from_element_with_layout(layout, ::zeroed()); reader.read_exact(bytemuck::must_cast_slice_mut::(data.as_mut_slice()))?; Ok(data) } @@ -65,7 +66,10 @@ where /// Write a matrix of `T` in the DiskANN binary format (see [module docs](self)). /// /// Returns the total number of bytes written. -pub fn write_bin(data: MatrixView<'_, T>, writer: &mut impl Write) -> Result +pub fn write_bin( + data: rowmajor::Ref<'_, T>, + writer: &mut impl Write, +) -> Result where T: bytemuck::Pod, { @@ -219,7 +223,7 @@ mod tests { #[test] fn round_trip_f32() { let mut counter = 1.0f32; - let matrix = Matrix::::from_fn(3, 4, |_| { + let matrix = rowmajor::Owned::::from_fn(3, 4, |_| { let v = counter; counter += 1.0; v diff --git a/diskann-utils/src/sampling/latin_hypercube.rs b/diskann-utils/src/sampling/latin_hypercube.rs index 18e847f9aa..630e961603 100644 --- a/diskann-utils/src/sampling/latin_hypercube.rs +++ b/diskann-utils/src/sampling/latin_hypercube.rs @@ -3,37 +3,37 @@ * Licensed under the MIT license. */ -use crate::views::{Matrix, MatrixView}; +use crate::views::rowmajor::{self, Matrix, MatrixMut}; use rand::{rngs::StdRng, Rng, SeedableRng}; /// Return multiple rows sampled using Latin Hypercube Sampling in `data` that aproximetely uniformly distributed. /// This makes the assumtion that the data is uniformly distributed. pub trait SampleLatinHyperCube: Sized + Copy + Default { fn sample_latin_hypercube( - data: MatrixView, + data: rowmajor::Ref<'_, Self>, num_samples: usize, seed: Option, - ) -> Matrix; + ) -> rowmajor::Owned; } impl SampleLatinHyperCube for T { fn sample_latin_hypercube( - data: MatrixView, + data: rowmajor::Ref<'_, Self>, num_samples: usize, seed: Option, - ) -> Matrix { + ) -> rowmajor::Owned { let nrows = data.nrows(); let ncols = data.ncols(); if ncols == 0 || nrows == 0 { - return Matrix::from_element(num_samples, ncols, T::default()); + return rowmajor::Owned::from_element(num_samples, ncols, T::default()); } let seed = seed.unwrap_or(0xaf2f5fa0b5161acf); let mut rng = StdRng::seed_from_u64(seed); - let mut result: Matrix = Matrix::from_element(num_samples, ncols, T::default()); + let mut result = rowmajor::Owned::from_element(num_samples, ncols, T::default()); // sample a random partitions down the diagonal - for (s, res) in result.row_iter_mut().enumerate() { + for (s, res) in result.rows_mut().enumerate() { for (idx, val) in res.iter_mut().enumerate() { let step = nrows / num_samples; let value = data @@ -63,11 +63,12 @@ impl SampleLatinHyperCube for T { // Tests // /////////// +#[cfg(not(miri))] #[cfg(test)] mod tests { use std::fmt::Display; - use crate::{assert_contains, views::Matrix}; + use crate::assert_contains; use diskann_vector::conversion::CastFromSlice; use half::f16; @@ -79,7 +80,7 @@ mod tests { use super::*; - fn example_dataset() -> Matrix { + fn example_dataset() -> rowmajor::Owned { let data: Vec = vec![ // row 0 0.203688, @@ -143,10 +144,10 @@ mod tests { 0.329328, ]; - Matrix::::try_from(data.into(), 10, 5).unwrap() + rowmajor::Owned::::try_from_data(data.into(), 10, 5).unwrap() } - fn example_dataset_u8() -> Matrix { + fn example_dataset_u8() -> rowmajor::Owned { let data: Vec = vec![ 52, 215, 218, 204, 192, // row 0 79, 55, 16, 89, 255, // row 1 @@ -156,11 +157,11 @@ mod tests { 145, 111, 142, 122, 181, // row 5 -- this is the medoid ]; - Matrix::::try_from(data.into(), 6, 5).unwrap() + rowmajor::Owned::::try_from_data(data.into(), 6, 5).unwrap() } // This is a test for the i8 function. Each entry is between -128 and 127. - fn example_dataset_i8() -> Matrix { + fn example_dataset_i8() -> rowmajor::Owned { let data: Vec = vec![ -76, 87, 90, 76, 64, // row 0 -49, -73, -112, -39, 127, // row 1 @@ -170,26 +171,26 @@ mod tests { 17, -17, 14, -6, 53, // row 5 -- this is the medoid ]; - Matrix::::try_from(data.into(), 6, 5).unwrap() + rowmajor::Owned::::try_from_data(data.into(), 6, 5).unwrap() } - fn test_for_type(data: Matrix) + fn test_for_type(data: rowmajor::Owned) where T: SampleLatinHyperCube + PartialEq + std::fmt::Debug + Display, StandardUniform: Distribution, { // No Rows - let x = Matrix::::from_element(0, 10, T::default()); + let x = rowmajor::Owned::::from_element(0, 10, T::default()); assert_eq!( T::sample_latin_hypercube(x.as_view(), 1, None), - Matrix::::from_element(1, x.ncols(), T::default()) + rowmajor::Owned::::from_element(1, x.ncols(), T::default()) ); // No Cols0 - let x = Matrix::::from_element(1, 0, T::default()); + let x = rowmajor::Owned::::from_element(1, 0, T::default()); assert_eq!( T::sample_latin_hypercube(x.as_view(), 1, None), - Matrix::::from_element(1, x.ncols(), T::default()) + rowmajor::Owned::::from_element(1, x.ncols(), T::default()) ); let mut rng: StdRng = StdRng::seed_from_u64(0xaf2f5fa0b5161acf); @@ -197,16 +198,17 @@ mod tests { // One row let dist = StandardUniform; for dim in 1..20 { - let x = Matrix::::from_fn(1, dim, |_| dist.sample(&mut rng)); + let x = rowmajor::Owned::::from_fn(1, dim, |_| dist.sample(&mut rng)); assert_eq!( T::sample_latin_hypercube(x.as_view(), 1, None), - Matrix::::try_from(x.row(0).to_vec().into_boxed_slice(), 1, dim).unwrap() + rowmajor::Owned::::try_from_data(x.row(0).to_vec().into_boxed_slice(), 1, dim) + .unwrap() ); } // Example dataset let starts = T::sample_latin_hypercube(data.as_view(), 2, None); - for s in starts.row_iter() { + for s in starts.rows() { for (col, &val) in s.iter().enumerate() { let col_vals: Vec = (0..data.nrows()) .map(|row| { @@ -236,7 +238,8 @@ mod tests { #[test] fn test_f16() { let data = example_dataset(); - let mut data_f16 = Matrix::::from_element(data.nrows(), data.ncols(), f16::default()); + let mut data_f16 = + rowmajor::Owned::::from_element(data.nrows(), data.ncols(), f16::default()); data_f16.as_mut_slice().cast_from_slice(data.as_slice()); test_for_type(data_f16); } diff --git a/diskann-utils/src/sampling/medoid.rs b/diskann-utils/src/sampling/medoid.rs index 543828d90f..bf0ed45d2a 100644 --- a/diskann-utils/src/sampling/medoid.rs +++ b/diskann-utils/src/sampling/medoid.rs @@ -3,23 +3,23 @@ * Licensed under the MIT license. */ -use crate::views::MatrixView; +use crate::views::rowmajor::{self, Matrix}; use diskann_vector::{conversion::CastFromSlice, distance::SquaredL2, PureDistanceFunction}; use half::f16; /// Return the row in `data` that is closest to the medoid of all rows. pub trait ComputeMedoid: Sized { - fn compute_medoid(data: MatrixView) -> Vec; + fn compute_medoid(data: rowmajor::Ref<'_, Self>) -> Vec; } impl ComputeMedoid for f32 { - fn compute_medoid(data: MatrixView) -> Vec { + fn compute_medoid(data: rowmajor::Ref<'_, Self>) -> Vec { if data.ncols() == 0 { return vec![]; } let mut sum = vec![0.0f64; data.ncols()]; - data.row_iter().for_each(|r| { + data.rows().for_each(|r| { std::iter::zip(sum.iter_mut(), r.iter()).for_each(|(o, i)| { let i: f64 = (*i).into(); *o += i; @@ -33,7 +33,7 @@ impl ComputeMedoid for f32 { let mut min_dist: f32 = f32::MAX; let mut medoid = None; - data.row_iter().for_each(|r| { + data.rows().for_each(|r| { let d = SquaredL2::evaluate(m.as_slice(), r); if d < min_dist { min_dist = d; @@ -48,14 +48,14 @@ impl ComputeMedoid for f32 { } impl ComputeMedoid for f16 { - fn compute_medoid(data: MatrixView) -> Vec { + fn compute_medoid(data: rowmajor::Ref<'_, Self>) -> Vec { if data.ncols() == 0 { return vec![]; } let mut sum = vec![0.0f64; data.ncols()]; let mut buffer = vec![0.0f32; data.ncols()]; - data.row_iter().for_each(|r| { + data.rows().for_each(|r| { buffer.cast_from_slice(r); std::iter::zip(sum.iter_mut(), buffer.iter()).for_each(|(o, i)| { let i: f64 = (*i).into(); @@ -69,7 +69,7 @@ impl ComputeMedoid for f16 { let mut min_dist: f32 = f32::MAX; let mut medoid = None; - data.row_iter().for_each(|r| { + data.rows().for_each(|r| { let d = SquaredL2::evaluate(buffer.as_slice(), r); if d < min_dist { min_dist = d; @@ -84,13 +84,13 @@ impl ComputeMedoid for f16 { } impl ComputeMedoid for u8 { - fn compute_medoid(data: MatrixView) -> Vec { + fn compute_medoid(data: rowmajor::Ref<'_, Self>) -> Vec { if data.ncols() == 0 { return vec![]; } let mut sum = vec![0.0f64; data.ncols()]; - data.row_iter().for_each(|r| { + data.rows().for_each(|r| { std::iter::zip(sum.iter_mut(), r.iter()).for_each(|(o, i)| { let i: f64 = (*i).into(); *o += i; @@ -105,7 +105,7 @@ impl ComputeMedoid for u8 { let mut min_dist: f32 = f32::MAX; let mut medoid = None; let mut as_float = vec![0.0f32; data.ncols()]; - data.row_iter().for_each(|r| { + data.rows().for_each(|r| { std::iter::zip(as_float.iter_mut(), r.iter()) .for_each(|(dst, src)| *dst = (*src).into()); let d = SquaredL2::evaluate(m.as_slice(), &*as_float); @@ -120,13 +120,13 @@ impl ComputeMedoid for u8 { } impl ComputeMedoid for i8 { - fn compute_medoid(data: MatrixView) -> Vec { + fn compute_medoid(data: rowmajor::Ref<'_, Self>) -> Vec { if data.ncols() == 0 { return vec![]; } let mut sum = vec![0.0f64; data.ncols()]; - data.row_iter().for_each(|r| { + data.rows().for_each(|r| { std::iter::zip(sum.iter_mut(), r.iter()).for_each(|(o, i)| { let i: f64 = (*i).into(); *o += i; @@ -141,7 +141,7 @@ impl ComputeMedoid for i8 { let mut min_dist: f32 = f32::MAX; let mut medoid = None; let mut as_float = vec![0.0f32; data.ncols()]; - data.row_iter().for_each(|r| { + data.rows().for_each(|r| { std::iter::zip(as_float.iter_mut(), r.iter()) .for_each(|(dst, src)| *dst = (*src).into()); let d = SquaredL2::evaluate(m.as_slice(), &*as_float); @@ -159,18 +159,21 @@ impl ComputeMedoid for i8 { // Tests // /////////// +#[cfg(not(miri))] #[cfg(test)] mod tests { - use crate::views::Matrix; + use super::*; + + use diskann_wide::cast_f32_to_f16; use rand::{ distr::{Distribution, StandardUniform}, rngs::StdRng, SeedableRng, }; - use super::*; + use crate::views::rowmajor::MatrixMut; - fn example_dataset() -> (Matrix, Vec) { + fn example_dataset() -> (rowmajor::Owned, Vec) { let data: Vec = vec![ // row 0 0.203688, @@ -234,7 +237,7 @@ mod tests { 0.329328, ]; - let data = Matrix::::try_from(data.into(), 10, 5).unwrap(); + let data = rowmajor::Owned::::try_from_data(data.into(), 10, 5).unwrap(); let expected: Vec = data.row(5).into(); (data, expected) } @@ -242,11 +245,11 @@ mod tests { #[test] fn test_f32() { // No Rows - let x = Matrix::::from_element(0, 10, 0.0f32); + let x = rowmajor::Owned::::from_element(0, 10, 0.0f32); assert_eq!(f32::compute_medoid(x.as_view()), vec![0.0; x.ncols()]); // No Cols - let x = Matrix::::from_element(10, 0, 0.0f32); + let x = rowmajor::Owned::::from_element(10, 0, 0.0f32); assert_eq!(f32::compute_medoid(x.as_view()), Vec::::new()); let mut rng = StdRng::seed_from_u64(0xaf2f5fa0b5161acf); @@ -254,7 +257,7 @@ mod tests { // One row let dist = StandardUniform; for dim in 1..20 { - let x = Matrix::::from_fn(1, dim, |_| dist.sample(&mut rng)); + let x = rowmajor::Owned::::from_fn(1, dim, |_| dist.sample(&mut rng)); assert_eq!(&*f32::compute_medoid(x.as_view()), x.row(0)); } @@ -267,14 +270,14 @@ mod tests { #[test] fn test_f16() { // No Rows - let x = Matrix::::from_element(0, 10, f16::default()); + let x = rowmajor::Owned::::from_element(0, 10, f16::default()); assert_eq!( f16::compute_medoid(x.as_view()), vec![f16::default(); x.ncols()] ); // No Cols - let x = Matrix::::from_element(10, 0, f16::default()); + let x = rowmajor::Owned::::from_element(10, 0, f16::default()); assert_eq!(f16::compute_medoid(x.as_view()), Vec::::new()); let mut rng = StdRng::seed_from_u64(0x88e2f7096fc9b90e); @@ -282,13 +285,15 @@ mod tests { // One row let dist = StandardUniform; for dim in 1..20 { - let x = Matrix::::from_fn(1, dim, |_| f16::from_f32(dist.sample(&mut rng))); + let x = + rowmajor::Owned::::from_fn(1, dim, |_| cast_f32_to_f16(dist.sample(&mut rng))); assert_eq!(&*f16::compute_medoid(x.as_view()), x.row(0)); } // Example dataset let (data, expected) = example_dataset(); - let mut data_f16 = Matrix::::from_element(data.nrows(), data.ncols(), f16::default()); + let mut data_f16 = + rowmajor::Owned::::from_element(data.nrows(), data.ncols(), f16::default()); data_f16.as_mut_slice().cast_from_slice(data.as_slice()); let mut expected_f16 = vec![f16::default(); expected.len()]; @@ -298,7 +303,7 @@ mod tests { assert_eq!(m, expected_f16); } - fn example_dataset_u8() -> (Matrix, Vec) { + fn example_dataset_u8() -> (rowmajor::Owned, Vec) { let data: Vec = vec![ 52, 215, 218, 204, 192, // row 0 79, 55, 16, 89, 255, // row 1 @@ -308,7 +313,7 @@ mod tests { 145, 111, 142, 122, 181, // row 5 -- this is the medoid ]; - let data = Matrix::::try_from(data.into(), 6, 5).unwrap(); + let data = rowmajor::Owned::::try_from_data(data.into(), 6, 5).unwrap(); let expected: Vec = data.row(5).into(); (data, expected) } @@ -316,18 +321,18 @@ mod tests { #[test] fn test_u8() { // No Rows - let x = Matrix::::from_element(0, 10, 0u8); + let x = rowmajor::Owned::::from_element(0, 10, 0u8); assert_eq!(u8::compute_medoid(x.as_view()), vec![0u8; x.ncols()]); // No Cols - let x = Matrix::::from_element(10, 0, 0u8); + let x = rowmajor::Owned::::from_element(10, 0, 0u8); assert_eq!(u8::compute_medoid(x.as_view()), Vec::::new()); let mut rng = StdRng::seed_from_u64(0x8f2f5fa0b5161acf); // One row let dist = StandardUniform; for dim in 1..20 { - let x = Matrix::::from_fn(1, dim, |_| dist.sample(&mut rng)); + let x = rowmajor::Owned::::from_fn(1, dim, |_| dist.sample(&mut rng)); assert_eq!(&*u8::compute_medoid(x.as_view()), x.row(0)); } @@ -338,7 +343,7 @@ mod tests { } // This is a test for the i8 medoid function. Each entry is between -128 and 127. - fn example_dataset_i8() -> (Matrix, Vec) { + fn example_dataset_i8() -> (rowmajor::Owned, Vec) { let data: Vec = vec![ -76, 87, 90, 76, 64, // row 0 -49, -73, -112, -39, 127, // row 1 @@ -348,7 +353,7 @@ mod tests { 17, -17, 14, -6, 53, // row 5 -- this is the medoid ]; - let data = Matrix::::try_from(data.into(), 6, 5).unwrap(); + let data = rowmajor::Owned::::try_from_data(data.into(), 6, 5).unwrap(); let expected: Vec = data.row(5).into(); (data, expected) } @@ -356,11 +361,11 @@ mod tests { #[test] fn test_i8() { // No Rows - let x = Matrix::::from_element(0, 10, 0i8); + let x = rowmajor::Owned::::from_element(0, 10, 0i8); assert_eq!(i8::compute_medoid(x.as_view()), vec![0i8; x.ncols()]); // No Cols - let x = Matrix::::from_element(10, 0, 0i8); + let x = rowmajor::Owned::::from_element(10, 0, 0i8); assert_eq!(i8::compute_medoid(x.as_view()), Vec::::new()); let mut rng = StdRng::seed_from_u64(0x8f2f5fa0b5161acf); @@ -368,7 +373,7 @@ mod tests { // One row let dist = StandardUniform; for dim in 1..20 { - let x = Matrix::::from_fn(1, dim, |_| dist.sample(&mut rng)); + let x = rowmajor::Owned::::from_fn(1, dim, |_| dist.sample(&mut rng)); assert_eq!(&*i8::compute_medoid(x.as_view()), x.row(0)); } diff --git a/diskann-utils/src/sampling/random.rs b/diskann-utils/src/sampling/random.rs index 6bb643862b..74c13a76de 100644 --- a/diskann-utils/src/sampling/random.rs +++ b/diskann-utils/src/sampling/random.rs @@ -84,6 +84,7 @@ where vec.into_iter().map(f).collect() } +#[cfg(not(miri))] #[cfg(test)] mod tests { use super::*; diff --git a/diskann-utils/src/strided.rs b/diskann-utils/src/strided.rs index 053ebbabd7..90084efe59 100644 --- a/diskann-utils/src/strided.rs +++ b/diskann-utils/src/strided.rs @@ -8,7 +8,7 @@ use thiserror::Error; use crate::{ internal, - views::{self}, + views::rowmajor::{self, Matrix}, Reborrow, }; @@ -80,8 +80,8 @@ impl Clone for Layout { impl Copy for Layout {} -impl From> for Layout { - fn from(layout: views::Layout) -> Self { +impl From> for Layout { + fn from(layout: rowmajor::Layout) -> Self { Self { nrows: layout.nrows(), ncols: layout.ncols(), @@ -184,7 +184,7 @@ impl fmt::Display for LayoutErrorInner { /// A row-major strided matrix. /// -/// This is a generalization of the `MatrixBase` class as it does not mandate a dense +/// This is a generalization of the [`Matrix`] trait as it does not mandate a dense /// layout in memory. /// /// ```text @@ -413,13 +413,13 @@ pub enum TryFromError { InvalidLength { got: usize, expected: usize }, } -impl<'a, T> From> for Strided<'a, T> { - fn from(matrix: views::MatrixView<'a, T>) -> Self { +impl<'a, T> From> for Strided<'a, T> { + fn from(matrix: rowmajor::Ref<'a, T>) -> Self { let layout = Layout::from(matrix.layout()); - // SAFETY: `MatrixView` guarantees that the length of the base slice for `matrix` + // SAFETY: `rowmajor::Ref` guarantees that the length of the base slice for `matrix` // is exactly `layout.linear_length()`. - unsafe { Self::from_data_unchecked(matrix.into_inner(), layout) } + unsafe { Self::from_data_unchecked(matrix.into_slice(), layout) } } } @@ -509,6 +509,8 @@ impl std::iter::FusedIterator for Rows<'_, T> {} mod tests { use super::*; + use crate::views::rowmajor::MatrixMut; + #[test] fn test_linear_length() { // If the number of rows is zero - the output should always be zero. @@ -602,7 +604,7 @@ mod tests { #[test] fn test_try_from_data_errors() { - let m = views::Matrix::::from_element(10, 10, 0); + let m = rowmajor::Owned::::from_element(10, 10, 0); let nrows = m.nrows(); let ncols = m.ncols(); @@ -755,7 +757,7 @@ mod tests { } // Test that the contents of `dut` match those in the dense 2d matrix. - fn test_indexing(dut: Strided<'_, usize>, expected: views::MatrixView<'_, usize>) { + fn test_indexing(dut: Strided<'_, usize>, expected: rowmajor::Ref<'_, usize>) { assert_eq!(dut.nrows(), expected.nrows()); assert_eq!(dut.ncols(), expected.ncols()); @@ -802,7 +804,7 @@ mod tests { } // Compare via row iterators. - assert!(dut.rows().eq(expected.row_iter())); + assert!(dut.rows().eq(expected.rows())); } // Create a base Matrix with the following pattern: @@ -812,9 +814,9 @@ mod tests { // 2*ncols 2*ncols+1 2*ncols+2 ... 3*ncols-1 // ... // ``` - fn create_test_matrix(nrows: usize, ncols: usize) -> views::Matrix { + fn create_test_matrix(nrows: usize, ncols: usize) -> rowmajor::Owned { let mut i = 0; - views::Matrix::from_fn(nrows, ncols, |_| { + rowmajor::Owned::from_fn(nrows, ncols, |_| { let v = i; i += 1; v @@ -846,7 +848,7 @@ mod tests { assert_eq!(v.as_ptr(), ptr, "base pointer was not preserved"); // Create the expected matrix. - let mut expected = views::Matrix::from_element(5, 2, 0); + let mut expected = rowmajor::Owned::from_element(5, 2, 0); for row in 0..expected.nrows() { for col in 0..expected.ncols() { *expected.element_mut(row, col) = *m.element(row, col); @@ -856,7 +858,7 @@ mod tests { // Create a strided view over the last two columns. let v = Strided::try_from_data(&(m.as_slice()[1..]), m.nrows(), 2, m.ncols()).unwrap(); - let mut expected = views::Matrix::from_element(5, 2, 0); + let mut expected = rowmajor::Owned::from_element(5, 2, 0); for row in 0..expected.nrows() { for col in 0..expected.ncols() { *expected.element_mut(row, col) = *m.element(row, col + 1); @@ -899,7 +901,7 @@ mod tests { #[test] fn test_try_shrink_from() { // Exact is okay. - let m = views::Matrix::::from_element(10, 10, 0); + let m = rowmajor::Owned::::from_element(10, 10, 0); let nrows = m.nrows(); let ncols = m.ncols(); let s = Strided::try_from_data(m.as_slice(), nrows, ncols, ncols).unwrap(); @@ -918,7 +920,7 @@ mod tests { fn test_invalid_stride_is_an_error_not_a_panic() { // Constructing a `Strided` with an invalid layout (`cstride < ncols`) returns an // `Err` rather than panicking - only unwrapping the result panics. - let m = views::Matrix::::from_element(4, 4, 0); + let m = rowmajor::Owned::::from_element(4, 4, 0); let err = Strided::try_from_data(m.as_slice(), 2, 2, 1).unwrap_err(); assert!(matches!(err, TryFromError::LayoutError(_))); } diff --git a/diskann-utils/src/views.rs b/diskann-utils/src/views.rs deleted file mode 100644 index 2097868ca3..0000000000 --- a/diskann-utils/src/views.rs +++ /dev/null @@ -1,2856 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -use std::{marker::PhantomData, num::NonZeroUsize}; - -#[cfg(feature = "rayon")] -use rayon::prelude::{IndexedParallelIterator, ParallelIterator, ParallelSlice, ParallelSliceMut}; -use thiserror::Error; - -/// Various view types (types such as [`MatrixView`] that add semantic meaning to blobs -/// of data) need both immutable and mutable variants. -/// -/// This trait can be implemented by wrappers for immutable and mutable slice references, -/// allowing for a common code path for immutable and mutable view types. -/// -/// The main goal is to provide a way of retrieving an underlying dense slice, which can -/// then be used as the building block for higher level abstractions. -/// -/// # Safety -/// -/// This trait is unsafe because it requires `as_slice` to be idempotent (and unsafe code -/// relies on this). -/// -/// In other words: `as_slice` must **always** return the same slice with the same length. -pub unsafe trait DenseData { - type Elem; - - /// Return the underlying data as a slice. - fn as_slice(&self) -> &[Self::Elem]; -} - -/// A mutable companion to `DenseData`. -/// -/// This trait allows mutable methods on view types to be selectively enabled when data -/// underlying the type is mutable. -/// -/// # Safety -/// -/// This trait is unsafe because it requires `as_slice` to be idempotent (and unsafe code -/// relies on this). -/// -/// In other words: `as_slice` must **always** return the same slice with the same length. -/// -/// Additionally, the returned slice must span the exact same memory as `as_slice`. -pub unsafe trait MutDenseData: DenseData { - fn as_mut_slice(&mut self) -> &mut [Self::Elem]; -} - -// SAFETY: This fulfills the idempotency requirement. -unsafe impl DenseData for &[T] { - type Elem = T; - fn as_slice(&self) -> &[Self::Elem] { - self - } -} - -// SAFETY: This fulfills the idempotency requirement. -unsafe impl DenseData for &mut [T] { - type Elem = T; - fn as_slice(&self) -> &[Self::Elem] { - self - } -} - -// SAFETY: This fulfills the idempotency requirement and returns a slice spanning the same -// range as `as_slice`. -unsafe impl MutDenseData for &mut [T] { - fn as_mut_slice(&mut self) -> &mut [Self::Elem] { - self - } -} - -// SAFETY: This fulfills the idempotency requirement. -unsafe impl DenseData for Box<[T]> { - type Elem = T; - fn as_slice(&self) -> &[Self::Elem] { - self - } -} - -// SAFETY: This fulfills the idempotency requirement and returns a slice spanning the same -// memory as `as_slice`. -unsafe impl MutDenseData for Box<[T]> { - fn as_mut_slice(&mut self) -> &mut [Self::Elem] { - self - } -} - -/////////////////// -// Matrix Layout // -/////////////////// - -/// A validated layout for [`MatrixBase`]. -/// -/// This type guarantees the following invariants: -/// -/// * `self.nrows() * self.ncols()` does not exceed `usize::MAX`. -/// * `self.nrows() * self.ncols() * std::mem::size_of::()` does not exceed `isize::MAX`. -pub struct Layout { - nrows: usize, - ncols: usize, - _type: PhantomData T>, -} - -impl Layout { - /// Construct a new [`Layout`], validating the following: - /// - /// * `nrows * ncols` does not exceed `usize::MAX`. - /// * `nrows * ncols * std::mem::size_of::()` does not exceed `isize::MAX` (the maximum - /// addressable byte span). - pub const fn new(nrows: usize, ncols: usize) -> Result { - match LayoutError::check::(nrows, ncols) { - Ok(()) => Ok(Self { - nrows, - ncols, - _type: PhantomData, - }), - Err(err) => Err(err), - } - } - - /// Construct a layout without validating its dimensions. - /// - /// # Safety - /// - /// `LayoutError::check::(nrows, ncols)` must succeed. - unsafe fn new_unchecked(nrows: usize, ncols: usize) -> Self { - debug_assert!(LayoutError::check::(nrows, ncols).is_ok()); - Self { - nrows, - ncols, - _type: PhantomData, - } - } - - /// Return the product `self.nrows() * self.ncols()`. - pub fn num_elements(&self) -> usize { - self.nrows() * self.ncols() - } - - /// Return the number of rows. - pub fn nrows(&self) -> usize { - self.nrows - } - - /// Return the number of columns. - pub fn ncols(&self) -> usize { - self.ncols - } - - /// Rebind the element type to `U`. - /// - /// # Errors - /// - /// Returns an error if the rebound layout's byte size would exceed `isize::MAX`. - pub fn rebind(&self) -> Result, LayoutError> { - if std::mem::size_of::() <= std::mem::size_of::() { - // This branch is mainly to communicate to the compiler situations where an - // erroring branch can be avoided. - // - // SAFETY: `self` already has a validated layout. Since we know - // `self.nrows() * self.ncols()` cannot overflow, the only danger is allocation - // overflow. If we are staying or decreasing size, no need to revalidate. - Ok(unsafe { Layout::new_unchecked(self.nrows(), self.ncols()) }) - } else { - Layout::new(self.nrows(), self.ncols()) - } - } - - /// Swap the rows and columns. - pub fn transpose(&self) -> Layout { - // SAFETY: We've already validated the relationship between `self.nrows` and - // `self.ncols`. Since multiplication is commutative, swapping rows and cols does - // not invalidate the relationship. - unsafe { Layout::new_unchecked(self.ncols, self.nrows) } - } -} - -impl Clone for Layout { - fn clone(&self) -> Self { - *self - } -} - -impl Copy for Layout {} - -impl std::fmt::Debug for Layout { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("Layout") - .field("nrows", &self.nrows) - .field("ncols", &self.ncols) - .field("elsize", &std::mem::size_of::()) - .finish() - } -} - -impl PartialEq for Layout { - fn eq(&self, other: &Self) -> bool { - self.nrows == other.nrows && self.ncols == other.ncols - } -} - -impl Eq for Layout {} - -/// Errors in the invariants guaranteed by [`Layout`]. -#[derive(Debug, Clone, Copy)] -pub struct LayoutError { - nrows: usize, - ncols: usize, - elsize: Option, -} - -impl LayoutError { - pub(crate) const fn check(nrows: usize, ncols: usize) -> Result<(), Self> { - // Guard the element count itself so that `num_elements()` can never overflow. - let elsize = std::mem::size_of::(); - let num_elements = match nrows.checked_mul(ncols) { - Some(num_elements) => num_elements, - None => { - return Err(Self { - nrows, - ncols, - elsize: None, - }) - } - }; - - if let Some(len) = num_elements.checked_mul(std::mem::size_of::()) { - if len <= isize::MAX as usize { - return Ok(()); - } - } - - Err(Self { - nrows, - ncols, - elsize: NonZeroUsize::new(elsize), - }) - } -} - -impl std::fmt::Display for LayoutError { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self.elsize { - Some(elsize) => { - write!( - f, - "a matrix of size {}x{} with elements of size {} exceeds `isize::MAX` bytes", - self.nrows, self.ncols, elsize - ) - } - None => { - write!( - f, - "a matrix of size {}x{} has a length exceeding `usize::MAX`", - self.nrows, self.ncols, - ) - } - } - } -} - -impl std::error::Error for LayoutError {} - -//////////// -// Matrix // -//////////// - -/// A view over a dense chunk of memory, interpreting that memory as a 2-dimensional matrix -/// laid out in row-major order. -/// -/// When this type views immutable memory, it is `Copy`. -/// -/// # Temporary Note -/// -/// This data structure is in the process of a representation migration. It currently -/// needs two type parameters: -/// -/// * `T`: The type of the container holding the data in the matrix. -/// * `E`: The element type of the matrix. -/// -/// The latter is needed to establish proper covariance with respect to the element type. -#[derive(Debug, Clone, Copy, PartialEq)] -pub struct MatrixBase::Elem> -where - T: DenseData, -{ - data: T, - layout: Layout, -} - -/// An initializer argument for the closure provided to [`Matrix::from_fn`], -/// [`Matrix::try_from_fn`], and [`Matrix::from_fn_with_layout`] to remove ambiguity of -/// the row and column being initialized. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct RowCol { - pub row: usize, - pub col: usize, -} - -impl MatrixBase> { - // NOTE: For constructors, keep `from_fn` and `from_element` first. - // - // Rust suggests methods in their declaration order, so this keeps the most common - // methods as top suggestions. - - /// Construct a new matrix using `init`. - /// - /// Elements are initialized in memory order. - /// - /// ``` - /// use diskann_utils::views::Matrix; - /// - /// let mat = Matrix::from_fn(2, 3, |rc| 3 * rc.row + rc.col); - /// - /// assert_eq!(mat.row(0), &[0, 1, 2]); - /// assert_eq!(mat.row(1), &[3, 4, 5]); - /// ``` - /// - /// # Panics - /// - /// Panics if `nrows * ncols` overflows `usize::MAX`, or if the allocation size exceeds - /// `isize::MAX`. - #[track_caller] - pub fn from_fn(nrows: usize, ncols: usize, init: F) -> Self - where - F: FnMut(RowCol) -> T, - { - match Self::try_from_fn(nrows, ncols, init) { - Ok(matrix) => matrix, - Err(error) => panic!("Matrix::from_fn failed with: {error}"), - } - } - - /// Construct a new matrix using `init`. - /// - /// Elements are initialized in memory order. - /// - /// ``` - /// use diskann_utils::views::Matrix; - /// - /// let mat = Matrix::try_from_fn(2, 3, |rc| 3 * rc.row + rc.col).unwrap(); - /// - /// assert_eq!(mat.row(0), &[0, 1, 2]); - /// assert_eq!(mat.row(1), &[3, 4, 5]); - /// ``` - /// - /// # Errors - /// - /// Returns an error if `nrows * ncols` overflows `usize::MAX`, or if the allocation size - /// exceeds `isize::MAX`. - pub fn try_from_fn(nrows: usize, ncols: usize, init: F) -> Result - where - F: FnMut(RowCol) -> T, - { - let layout = Layout::new(nrows, ncols)?; - Ok(Self::from_fn_with_layout(layout, init)) - } - - /// Construct a new matrix by cloning `element`. - /// - /// Elements are initialized in memory order. - /// - /// ``` - /// use diskann_utils::views::Matrix; - /// - /// let mat = Matrix::from_element(2, 3, 0u32); - /// - /// assert_eq!(mat.row(0), &[0, 0, 0]); - /// assert_eq!(mat.row(1), &[0, 0, 0]); - /// ``` - /// - /// # Panics - /// - /// Panics if `nrows * ncols` overflows `usize::MAX`, or if the allocation size exceeds - /// `isize::MAX`. - #[track_caller] - pub fn from_element(nrows: usize, ncols: usize, element: T) -> Self - where - T: Clone, - { - match Self::try_from_element(nrows, ncols, element) { - Ok(matrix) => matrix, - Err(error) => panic!("Matrix::from_element failed with: {error}"), - } - } - - /// Construct a new matrix by cloning `element`. - /// - /// Elements are initialized in memory order. - /// - /// ``` - /// use diskann_utils::views::Matrix; - /// - /// let mat = Matrix::try_from_element(2, 3, 0u32).unwrap(); - /// - /// assert_eq!(mat.row(0), &[0, 0, 0]); - /// assert_eq!(mat.row(1), &[0, 0, 0]); - /// ``` - /// - /// # Errors - /// - /// Returns an error if `nrows * ncols` overflows `usize::MAX`, or if the allocation size - /// exceeds `isize::MAX`. - pub fn try_from_element(nrows: usize, ncols: usize, element: T) -> Result - where - T: Clone, - { - let layout = Layout::new(nrows, ncols)?; - Ok(Self::from_element_with_layout(layout, element)) - } - - // Less common constructors. - - /// Construct a new matrix using `init`. - /// - /// Elements are initialized in memory order. - pub fn from_fn_with_layout(layout: Layout, mut init: F) -> Self - where - F: FnMut(RowCol) -> T, - { - let mut row = 0; - let mut col = 0; - - let data: Box<[T]> = (0..layout.num_elements()) - .map(|_| { - let v = (init)(RowCol { row, col }); - col += 1; - if col == layout.ncols() { - col = 0; - row += 1; - } - v - }) - .collect(); - - Self { data, layout } - } - - /// Construct a new matrix by cloning `element`. - /// - /// Elements are initialized in memory order. - pub fn from_element_with_layout(layout: Layout, element: T) -> Self - where - T: Clone, - { - let data: Box<[T]> = std::iter::repeat_n(element, layout.num_elements()).collect(); - Self { data, layout } - } -} - -impl MatrixBase -where - T: DenseData, -{ - /// Try to construct a `MatrixBase` over the provided base. If the size of the base - /// is incorrect, return a `TryFromError` containing the base. - /// - /// The length of the base must be equal to `nrows * ncols`. - /// - /// # Errors - /// - /// Returns an error containing `data` if the dimensions do not describe a valid layout - /// or if `data.as_slice().len()` does not equal `nrows * ncols`. - pub fn try_from(data: T, nrows: usize, ncols: usize) -> Result> { - let layout = match Layout::::new(nrows, ncols) { - Ok(layout) => layout, - Err(err) => return Err(TryFromError::layout(data, err)), - }; - - let len = data.as_slice().len(); - if len != layout.num_elements() { - Err(TryFromError::mismatch( - data, - layout.nrows(), - layout.ncols(), - len, - )) - } else { - Ok(Self { data, layout }) - } - } - - /// Construct a matrix without validating that the data length matches the layout. - /// - /// # Safety - /// - /// `data.as_slice().len()` must equal `layout.num_elements()`. - unsafe fn from_data_unchecked(data: T, layout: Layout) -> Self { - debug_assert_eq!(data.as_slice().len(), layout.num_elements()); - Self { data, layout } - } - - /// Return the [`Layout`] for this matrix. - pub fn layout(&self) -> Layout { - self.layout - } - - /// Return the number of columns in the matrix. - pub fn ncols(&self) -> usize { - self.layout().ncols() - } - - /// Return the number of rows in the matrix. - pub fn nrows(&self) -> usize { - self.layout().nrows() - } - - /// Create a new [`Matrix`] by applying the closure `f` to each element. - /// - /// The returned matrix has the same shape as `self`. - /// - /// # Panics - /// - /// Panics if the resulting matrix's byte size would exceed `isize::MAX`. - #[track_caller] - pub fn map(&self, f: F) -> Matrix - where - F: FnMut(&T::Elem) -> R, - { - match self.try_map(f) { - Ok(matrix) => matrix, - Err(error) => panic!("Matrix::map failed with: {error}"), - } - } - - /// Create a new [`Matrix`] by applying the closure `f` to each element. - /// - /// The returned matrix has the same shape as `self`. - /// - /// # Errors - /// - /// Returns an error if the resulting matrix's byte size would exceed `isize::MAX`. - pub fn try_map(&self, f: F) -> Result, LayoutError> - where - F: FnMut(&T::Elem) -> R, - { - let layout: Layout = self.layout().rebind::()?; - let data: Box<[_]> = self.as_slice().iter().map(f).collect(); - - // SAFETY: By construction, `data.len() == layout.num_elements()`. - Ok(unsafe { Matrix::from_data_unchecked(data, layout) }) - } - - /// Return the underlying data as a slice. - pub fn as_slice(&self) -> &[T::Elem] { - self.data.as_slice() - } - - /// Return the underlying data as a mutable slice. - pub fn as_mut_slice(&mut self) -> &mut [T::Elem] - where - T: MutDenseData, - { - self.data.as_mut_slice() - } - - /// Return row `row` as a slice. - /// - /// # Panics - /// - /// Panics if `row >= self.nrows()`. - pub fn row(&self, row: usize) -> &[T::Elem] { - assert!( - row < self.nrows(), - "tried to access row {row} of a matrix with {} rows", - self.nrows() - ); - - // SAFETY: `row` is in-bounds. - unsafe { self.get_row_unchecked(row) } - } - - /// Construct a new `MatrixBase` over the raw data. - /// - /// The returned `MatrixBase` will only have a single row with contents equal to `data`. - pub fn row_vector(data: T) -> Self { - let ncols = data.as_slice().len(); - - // SAFETY: The `Layout` construction is valid because `ncols` comes from the length - // of a slice, so we know the size of that slice cannot exceed `isize::MAX` bytes. - // - // By construction, `data.len() == layout.num_elements()`. - unsafe { Self::from_data_unchecked(data, Layout::new_unchecked(1, ncols)) } - } - - /// Construct a new `MatrixBase` over the raw data. - /// - /// The returned `MatrixBase` will only have a single column with contents equal to `data`. - pub fn column_vector(data: T) -> Self { - let nrows = data.as_slice().len(); - - // SAFETY: The `Layout` construction is valid because `nrows` comes from the length - // of a slice, so we know the size of that slice cannot exceed `isize::MAX` bytes. - // - // By construction, `data.len() == layout.num_elements()`. - unsafe { Self::from_data_unchecked(data, Layout::new_unchecked(nrows, 1)) } - } - - /// Return row `row` if `row < self.nrows()`. Otherwise, return `None`. - pub fn get_row(&self, row: usize) -> Option<&[T::Elem]> { - if row < self.nrows() { - // SAFETY: `row` is in-bounds. - Some(unsafe { self.get_row_unchecked(row) }) - } else { - None - } - } - - /// Returns the requested row without boundschecking. - /// - /// # Safety - /// - /// The following conditions must hold to avoid undefined behavior: - /// * `row < self.nrows()`. - pub unsafe fn get_row_unchecked(&self, row: usize) -> &[T::Elem] { - debug_assert!(row < self.nrows()); - let ncols = self.ncols(); - let start = row * ncols; - - debug_assert!(start + ncols <= self.as_slice().len()); - // SAFETY: The idempotency requirement of `as_slice` and our audited constructors - // mean that `self.as_slice()` has a length of `self.nrows * self.ncols`. - // - // Therefore, this access is in-bounds. - unsafe { self.as_slice().get_unchecked(start..start + ncols) } - } - - /// Return row `row` as a mutable slice. - /// - /// # Panics - /// - /// Panics if `row >= self.nrows()`. - pub fn row_mut(&mut self, row: usize) -> &mut [T::Elem] - where - T: MutDenseData, - { - assert!( - row < self.nrows(), - "tried to access row {row} of a matrix with {} rows", - self.nrows() - ); - - // SAFETY: `row` is in-bounds. - unsafe { self.get_row_unchecked_mut(row) } - } - - /// Returns the requested row without boundschecking. - /// - /// # Safety - /// - /// The following conditions must hold to avoid undefined behavior: - /// * `row < self.nrows()`. - pub unsafe fn get_row_unchecked_mut(&mut self, row: usize) -> &mut [T::Elem] - where - T: MutDenseData, - { - debug_assert!(row < self.nrows()); - let ncols = self.ncols(); - let start = row * ncols; - - debug_assert!(start + ncols <= self.as_slice().len()); - // SAFETY: The idempotency requirement of `as_mut_slice` and our audited constructors - // mean that `self.as_mut_slice()` has a length of `self.nrows * self.ncols`. - // - // Therefore, this access is in-bounds. - unsafe { - self.data - .as_mut_slice() - .get_unchecked_mut(start..start + ncols) - } - } - - /// Return a iterator over all rows in the matrix. - /// - /// Rows are yielded sequentially beginning with row 0. - pub fn row_iter(&self) -> impl ExactSizeIterator { - self.data.as_slice().chunks_exact(self.ncols()) - } - - /// Return a mutable iterator over all rows in the matrix. - /// - /// Rows are yielded sequentially beginning with row 0. - pub fn row_iter_mut(&mut self) -> impl ExactSizeIterator - where - T: MutDenseData, - { - let ncols = self.ncols(); - self.data.as_mut_slice().chunks_exact_mut(ncols) - } - - /// Return an iterator that divides the matrix into sub-matrices with (up to) - /// `batchsize` rows with `self.ncols()` columns. - /// - /// It is possible for yielded sub-matrices to have fewer than `batchsize` rows if the - /// number of rows in the parent matrix is not evenly divisible by `batchsize`. - /// - /// # Panics - /// - /// Panics if `batchsize = 0`. - pub fn window_iter(&self, batchsize: usize) -> impl Iterator> - where - T::Elem: Sync, - { - assert!(batchsize != 0, "window_iter batchsize cannot be zero"); - let ncols = self.ncols(); - - self.data - .as_slice() - .chunks(ncols * batchsize) - .map(move |data| { - let blobsize = data.len(); - let nrows = blobsize / ncols; - assert_eq!(blobsize % ncols, 0); - - // SAFETY: `self` contains a valid `Layout`. Since `nrows <= self.nrows()`, - // the resulting `Layout` (it is no bigger than self's layout). - // - // Further, we've verified that `data.len() == nrows * ncols`. - unsafe { - MatrixView::from_data_unchecked(data, Layout::new_unchecked(nrows, ncols)) - } - }) - } - - /// Return a parallel iterator that divides the matrix into sub-matrices with (up to) - /// `batchsize` rows with `self.ncols()` columns. - /// - /// This allows workers in parallel algorithms to work on dense subsets of the whole - /// matrix for better locality. - /// - /// It is possible for yielded sub-matrices to have fewer than `batchsize` rows if the - /// number of rows in the parent matrix is not evenly divisible by `batchsize`. - /// - /// # Panics - /// - /// Panics if `batchsize = 0`. - #[cfg(feature = "rayon")] - pub fn par_window_iter( - &self, - batchsize: usize, - ) -> impl IndexedParallelIterator> - where - T::Elem: Sync, - { - assert!(batchsize != 0, "par_window_iter batchsize cannot be zero"); - let ncols = self.ncols(); - self.data - .as_slice() - .par_chunks(ncols * batchsize) - .map(move |data| { - let blobsize = data.len(); - let nrows = blobsize / ncols; - assert_eq!(blobsize % ncols, 0); - - // SAFETY: `self` contains a valid `Layout`. Since `nrows <= self.nrows()`, - // the resulting `Layout` (it is no bigger than self's layout). - // - // Further, we've verified that `data.len() == nrows * ncols`. - unsafe { - MatrixView::from_data_unchecked(data, Layout::new_unchecked(nrows, ncols)) - } - }) - } - - /// Return a parallel iterator that divides the matrix into mutable sub-matrices with - /// (up to) `batchsize` rows with `self.ncols()` columns. - /// - /// This allows workers in parallel algorithms to work on dense subsets of the whole - /// matrix for better locality. - /// - /// It is possible for yielded sub-matrices to have fewer than `batchsize` rows if the - /// number of rows in the parent matrix is not evenly divisible by `batchsize`. - /// - /// # Panics - /// - /// Panics if `batchsize = 0`. - #[cfg(feature = "rayon")] - pub fn par_window_iter_mut( - &mut self, - batchsize: usize, - ) -> impl IndexedParallelIterator> - where - T: MutDenseData, - T::Elem: Send, - { - assert!( - batchsize != 0, - "par_window_iter_mut batchsize cannot be zero" - ); - let ncols = self.ncols(); - self.data - .as_mut_slice() - .par_chunks_mut(ncols * batchsize) - .map(move |data| { - let blobsize = data.len(); - let nrows = blobsize / ncols; - assert_eq!(blobsize % ncols, 0); - - // SAFETY: `self` contains a valid `Layout`. Since `nrows <= self.nrows()`, - // the resulting `Layout` (it is no bigger than self's layout). - // - // Further, we've verified that `data.len() == nrows * ncols`. - unsafe { - MutMatrixView::from_data_unchecked(data, Layout::new_unchecked(nrows, ncols)) - } - }) - } - - /// Return a parallel iterator over the rows of the matrix. - #[cfg(feature = "rayon")] - pub fn par_row_iter(&self) -> impl IndexedParallelIterator - where - T::Elem: Sync, - { - self.as_slice().par_chunks_exact(self.ncols()) - } - - /// Return a parallel iterator over the rows of the matrix. - #[cfg(feature = "rayon")] - pub fn par_row_iter_mut(&mut self) -> impl IndexedParallelIterator - where - T: MutDenseData, - T::Elem: Send, - { - let ncols = self.ncols(); - self.as_mut_slice().par_chunks_exact_mut(ncols) - } - - /// Consume the matrix, returning the inner representation. - /// - /// This loses the information about the number of rows and columns. - pub fn into_inner(self) -> T { - self.data - } - - /// Return a view over the matrix. - pub fn as_view(&self) -> MatrixView<'_, T::Elem> { - // SAFETY: This propagates `self`'s already valid layout and underlying slice. - unsafe { MatrixView::from_data_unchecked(self.as_slice(), self.layout()) } - } - - /// Return a mutable view over the matrix. - pub fn as_mut_view(&mut self) -> MutMatrixView<'_, T::Elem> - where - T: MutDenseData, - { - let layout = self.layout(); - // SAFETY: This propagates `self`'s already valid layout and underlying slice. - unsafe { MutMatrixView::from_data_unchecked(self.as_mut_slice(), layout) } - } - - /// Return a view over the specified rows of the matrix. - /// - /// If the specified range is out of bounds, return `None`. - /// - /// ```rust - /// use diskann_utils::views::Matrix; - /// - /// let mut mat = Matrix::from_element(4, 3, 0usize); - /// - /// // Fill the matrix with some data. - /// mat.row_iter_mut().enumerate().for_each(|(i, row)| row.fill(i)); - /// - /// // Creating a subview into an offset portion of the matrix. - /// let subview = mat.subview(1..3).unwrap(); - /// assert_eq!(subview.nrows(), 2); - /// assert_eq!(subview.row(0), &[1, 1, 1]); - /// assert_eq!(subview.row(1), &[2, 2, 2]); - /// - /// // A trying to access out-of-bounds returns `None` - /// assert!(mat.subview(3..5).is_none()); - /// ``` - pub fn subview(&self, rows: std::ops::Range) -> Option> { - if rows.start > rows.end || rows.end > self.nrows() { - return None; - } - - let ncols = self.ncols(); - let lower = rows.start * ncols; - let upper = rows.end * ncols; - - if let Some(data) = self.as_slice().get(lower..upper) { - // SAFETY: The successful checked index into `self.as_slice()` attests that - // `rows.len()` is no greater than `self.nrows()`. So the `Layout` is valid. - // - // By construction, `data.len() == rows * self.ncols()`. - Some(unsafe { - MatrixView::from_data_unchecked( - data, - Layout::new_unchecked(rows.len(), self.ncols()), - ) - }) - } else { - None - } - } - - /// Return a pointer to the base of the matrix. - pub fn as_ptr(&self) -> *const T::Elem { - self.as_slice().as_ptr() - } - - /// Return a pointer to the base of the matrix. - pub fn as_mut_ptr(&mut self) -> *mut T::Elem - where - T: MutDenseData, - { - self.as_mut_slice().as_mut_ptr() - } - - /// Return a reference to the element at the specified `row` and `col`. - /// - /// # Safety - /// - /// The following conditions must hold to avoid undefined behavior: - /// - /// * `row < self.nrows()`. - /// * `col < self.ncols()`. - pub unsafe fn element_unchecked(&self, row: usize, col: usize) -> &T::Elem { - debug_assert!(row < self.nrows()); - debug_assert!(col < self.ncols()); - self.as_slice().get_unchecked(row * self.ncols() + col) - } - - /// Return a reference to the element at the specified `row` and `col`. - /// - /// If either index is out-of-bounds, return `None`. - pub fn get_element(&self, row: usize, col: usize) -> Option<&T::Elem> { - if row >= self.nrows() || col >= self.ncols() { - None - } else { - // SAFETY: We just verified that `row` and `col` are in-bounds. - Some(unsafe { self.element_unchecked(row, col) }) - } - } - - /// Return a reference to the element at the specified `row` and `col`. - /// - /// # Panics - /// - /// Panics if `row >= self.nrows()` or `col >= self.ncols()`. - pub fn element(&self, row: usize, col: usize) -> &T::Elem { - assert!( - row < self.nrows(), - "row {row} is out of bounds (max: {})", - self.nrows() - ); - assert!( - col < self.ncols(), - "col {col} is out of bounds (max: {})", - self.ncols() - ); - - // SAFETY: We have checked that `row` and `col` are in-bounds. - unsafe { self.element_unchecked(row, col) } - } - - /// Return a reference to the element at the specified `row` and `col`. - /// - /// # Safety - /// - /// The following conditions must hold to avoid undefined behavior: - /// - /// * `row < self.nrows()`. - /// * `col < self.ncols()`. - pub unsafe fn element_unchecked_mut(&mut self, row: usize, col: usize) -> &mut T::Elem - where - T: MutDenseData, - { - let ncols = self.ncols(); - debug_assert!(row < self.nrows()); - debug_assert!(col < self.ncols()); - self.as_mut_slice().get_unchecked_mut(row * ncols + col) - } - - /// Return a mutable reference to the element at the specified `row` and `col`. - /// - /// Returns `None` if `row >= self.nrows()` or `col >= self.ncols()`. - pub fn get_element_mut(&mut self, row: usize, col: usize) -> Option<&mut T::Elem> - where - T: MutDenseData, - { - if row >= self.nrows() || col >= self.ncols() { - None - } else { - // SAFETY: We have checked that `row` and `col` are in-bounds. - Some(unsafe { self.element_unchecked_mut(row, col) }) - } - } - - /// Return a mutable reference to the element at the specified `row` and `col`. - /// - /// # Panics - /// - /// Panics if `row >= self.nrows()` or `col >= self.ncols()`. - pub fn element_mut(&mut self, row: usize, col: usize) -> &mut T::Elem - where - T: MutDenseData, - { - assert!( - row < self.nrows(), - "row {row} is out of bounds (max: {})", - self.nrows() - ); - assert!( - col < self.ncols(), - "col {col} is out of bounds (max: {})", - self.ncols() - ); - - // SAFETY: We have checked that `row` and `col` are in-bounds. - unsafe { self.element_unchecked_mut(row, col) } - } - - pub fn to_owned(&self) -> Matrix - where - T::Elem: Clone, - { - // SAFETY: This propagates `self`'s already valid layout and underlying slice. - unsafe { Matrix::from_data_unchecked(self.data.as_slice().into(), self.layout()) } - } - - /// Transpose the elements in `self`. - pub fn transpose(&self) -> Matrix - where - T::Elem: Clone, - { - Matrix::from_fn_with_layout(self.layout.transpose(), |RowCol { row, col }| { - // SAFETY: By construction, `col < self.nrows()` and `row < self.ncols()`. - unsafe { self.element_unchecked(col, row).clone() } - }) - } -} - -/// Represents an owning, 2-dimensional view of a contiguous block of memory, -/// interpreted as a matrix in row-major order. -pub type Matrix = MatrixBase, T>; - -/// Represents a non-owning, 2-dimensional view of a contiguous block of memory, -/// interpreted as a matrix in row-major order. -/// -/// This type is useful for functions that need to read matrix data without taking ownership. -/// By accepting a `MatrixView`, such functions can operate on both owned matrices (by converting them -/// to a `MatrixView`) and existing non-owning views. -pub type MatrixView<'a, T> = MatrixBase<&'a [T], T>; - -/// Represents a mutable non-owning, 2-dimensional view of a contiguous block of memory, -/// interpreted as a matrix in row-major order. -/// -/// This type is useful for functions that need to modify matrix data without taking ownership. -/// By accepting a `MutMatrixView`, such functions can operate on both owned matrices (by converting them -/// to a `MutMatrixView`) and existing non-owning mutable views. -pub type MutMatrixView<'a, T> = MatrixBase<&'a mut [T], T>; - -/// Allow matrix views to be converted directly to slices. -impl<'a, T> From> for &'a [T] { - fn from(view: MatrixView<'a, T>) -> Self { - view.data - } -} - -/// Allow mutable matrix views to be converted directly to slices. -impl<'a, T> From> for &'a [T] { - fn from(view: MutMatrixView<'a, T>) -> Self { - view.data - } -} - -/// Errors from [`MatrixBase::try_from`]. -pub struct TryFromError { - data: T, - inner: TryFromErrorInner, -} - -impl TryFromError { - /// Consume the error and return the base data. - pub fn into_inner(self) -> T { - self.data - } - - /// Return a variation of `Self` that is guaranteed to be `'static` by removing the - /// data that was passed to the original constructor. - pub fn as_static(&self) -> TryFromErrorLight { - TryFromErrorLight(self.inner) - } - - //--------------// - // Constructors // - //--------------// - - fn layout(data: T, error: LayoutError) -> Self { - Self { - data, - inner: TryFromErrorInner::Layout(error), - } - } - - fn mismatch(data: T, nrows: usize, ncols: usize, len: usize) -> Self { - Self { - data, - inner: TryFromErrorInner::Mismatch { nrows, ncols, len }, - } - } -} - -impl std::fmt::Debug for TryFromError { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("TryFromError") - .field("data", &"") - .field("inner", &self.inner) - .finish() - } -} - -impl std::fmt::Display for TryFromError { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - self.inner.fmt(f) - } -} - -impl std::error::Error for TryFromError { - fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { - match &self.inner { - TryFromErrorInner::Layout(error) => Some(error), - TryFromErrorInner::Mismatch { .. } => None, - } - } -} - -#[derive(Debug, Error)] -#[error(transparent)] -pub struct TryFromErrorLight(TryFromErrorInner); - -#[derive(Debug, Error, Clone, Copy)] -enum TryFromErrorInner { - #[error(transparent)] - Layout(LayoutError), - #[error( - "tried to construct a {}x{} matrix over a span of length {}", - nrows, - ncols, - len - )] - Mismatch { - nrows: usize, - ncols: usize, - len: usize, - }, -} - -/////////// -// Tests // -/////////// - -#[cfg(test)] -mod tests { - use super::*; - use crate::{assert_contains, lazy_format}; - - /// This function is only callable with copyable types. - /// - /// This lets us test for types we expect to be `Copy`. - fn is_copyable(_x: T) -> bool { - true - } - - /// This function attests that `MatrixView` is covariant in the view lifetime. - fn _matrix_view_is_covariant<'a, 'b>(m: MatrixView<'a, f32>) -> MatrixView<'b, f32> - where - 'a: 'b, - { - m - } - - fn _matrix_view_is_covariant_in_t<'a, 'b, 'm>( - m: MatrixView<'m, &'a f32>, - ) -> MatrixView<'m, &'b f32> - where - 'a: 'b, - { - m - } - - fn _matrix_is_covariant_in_t<'a, 'b, 'm>(m: &'m Matrix<&'a f32>) -> &'m Matrix<&'b f32> - where - 'a: 'b, - { - m - } - - /// Test the that provided representation yields a slice with the expected base pointer - /// and length. - fn test_dense_data_repr( - ptr: *const T, - len: usize, - repr: Repr, - context: &dyn std::fmt::Display, - ) where - T: Copy, - Repr: DenseData, - { - let retrieved = repr.as_slice(); - assert_eq!(retrieved.len(), len, "{}", context); - assert_eq!(retrieved.as_ptr(), ptr, "{}", context); - } - - /// Set the underlying data for the provided representation to the following: - /// - /// [base, base + increment, base + increment + increment, ...] - fn set_mut_dense_data_repr(repr: &mut Repr, base: T, increment: T) - where - T: Copy + std::ops::Add, - Repr: DenseData + MutDenseData, - { - let slice = repr.as_mut_slice(); - for i in 0..slice.len() { - if i == 0 { - slice[i] = base; - } else { - slice[i] = slice[i - 1] + increment; - } - } - } - - #[test] - fn slice_implements_dense_data_repr() { - for len in 0..10 { - let context = lazy_format!("len = {}", len); - let data: Vec = vec![0.0; len]; - let slice = data.as_slice(); - test_dense_data_repr(slice.as_ptr(), slice.len(), slice, &context); - } - } - - #[test] - fn mut_slice_mplements_dense_data_repr() { - for len in 0..10 { - let context = lazy_format!("len = {}", len); - let mut data: Vec = vec![0.0; len]; - let slice = data.as_mut_slice(); - - let ptr = slice.as_ptr(); - let len = slice.len(); - test_dense_data_repr(ptr, len, slice, &context); - } - } - - #[test] - fn mut_slice_implements_mut_dense_data_repr() { - for len in 0..10 { - let context = lazy_format!("len = {}", len); - let mut data: Vec = vec![0.0; len]; - let mut slice = data.as_mut_slice(); - - let base = 2.0; - let increment = 1.0; - set_mut_dense_data_repr(&mut slice, base, increment); - - for (i, &v) in slice.iter().enumerate() { - let context = lazy_format!("entry {}, {}", i, context); - assert_eq!(v, base + increment * (i as f32), "{}", context); - } - } - } - - //--------// - // Layout // - //--------// - - #[test] - fn test_layout() { - // Happy path - for rows in 0..5 { - for cols in 0..5 { - let layout = Layout::::new(rows, cols).unwrap(); - assert_eq!(layout.nrows(), rows); - assert_eq!(layout.ncols(), cols); - assert_eq!(layout.num_elements(), rows * cols); - - let transpose = layout.transpose(); - assert_eq!(transpose.nrows(), cols); - assert_eq!(transpose.ncols(), rows); - assert_eq!(transpose.num_elements(), rows * cols); - - let rebind = layout.rebind::().unwrap(); - assert_eq!(rebind.nrows(), rows); - assert_eq!(rebind.ncols(), cols); - assert_eq!(rebind.num_elements(), rows * cols); - - is_copyable(layout); - } - } - - #[expect(unused, reason = "we need this so the size is non-zero")] - struct NotDebugOrEq(u32); - - assert_eq!( - Layout::::new(10, 20).unwrap(), - Layout::::new(10, 20).unwrap(), - ); - - assert_eq!( - Layout::::new(20, 0).unwrap(), - Layout::::new(20, 0).unwrap(), - ); - - assert_ne!( - Layout::::new(10, 20).unwrap(), - Layout::::new(20, 0).unwrap(), - ); - - let fmt = format!("{:?}", Layout::::new(5, 6).unwrap()); - assert_eq!(fmt, "Layout { nrows: 5, ncols: 6, elsize: 4 }"); - - // Overflowing the element count returns an error. - let error = Layout::::new(usize::MAX, 2).unwrap_err(); - assert_eq!( - error.to_string(), - format!( - "a matrix of size {}x2 has a length exceeding `usize::MAX`", - usize::MAX - ) - ); - - // The largest possible byte span is valid without allocating it. - let layout = Layout::::new(isize::MAX as usize, 1).unwrap(); - assert_eq!(layout.num_elements(), isize::MAX as usize); - - let transpose = layout.transpose(); - assert_eq!(transpose.nrows(), 1); - assert_eq!(transpose.ncols(), isize::MAX as usize); - assert_eq!(transpose.num_elements(), layout.num_elements()); - - // One byte beyond the maximum span returns an error. - let error = Layout::::new(isize::MAX as usize + 1, 1).unwrap_err(); - assert_eq!( - error.to_string(), - format!( - "a matrix of size {}x1 with elements of size 1 exceeds `isize::MAX` bytes", - isize::MAX as usize + 1 - ) - ); - - // Rebinding to a larger element type revalidates the byte span. - let rebound = Layout::::new(3, 4).unwrap().rebind::().unwrap(); - assert_eq!(rebound.nrows(), 3); - assert_eq!(rebound.ncols(), 4); - assert_eq!(rebound.num_elements(), 12); - - let error = layout.rebind::().unwrap_err(); - assert_eq!( - error.to_string(), - format!( - "a matrix of size {}x1 with elements of size 2 exceeds `isize::MAX` bytes", - isize::MAX - ) - ); - } - - ///////////////// - // Matrix View // - ///////////////// - - #[test] - fn fallible_matrix_constructors() { - let err = Matrix::try_from_element(usize::MAX, usize::MAX, 0u32).unwrap_err(); - let msg = err.to_string(); - assert_contains!(msg, "exceeding `usize::MAX`"); - - let err = Matrix::try_from_element(isize::MAX as usize, 1, 0u32).unwrap_err(); - let msg = err.to_string(); - assert_contains!(msg, "exceeds `isize::MAX` bytes"); - - // Panicking - let err = std::panic::catch_unwind(|| { - Matrix::from_element(usize::MAX, usize::MAX, 0u32); - }) - .unwrap_err() - .downcast::() - .unwrap(); - - let msg = err.to_string(); - assert_contains!(msg, "exceeding `usize::MAX`"); - - let err = std::panic::catch_unwind(|| { - Matrix::from_element(isize::MAX as usize, 1, 0u32); - }) - .unwrap_err() - .downcast::() - .unwrap(); - let msg = err.to_string(); - assert_contains!(msg, "exceeds `isize::MAX` bytes"); - - // Construction fails without invoking the generator. - let err = Matrix::try_from_fn(usize::MAX, usize::MAX, |_| panic!("boom")).unwrap_err(); - let msg = err.to_string(); - assert_contains!(msg, "exceeding `usize::MAX`"); - } - - fn make_test_matrix() -> Vec { - // Construct a matrix with 4 rows of length 3. - // The expected layout is as follows: - // - // 0, 1, 2, - // 1, 2, 3, - // 2, 3, 4, - // 3, 4, 5 - // - vec![0, 1, 2, 1, 2, 3, 2, 3, 4, 3, 4, 5] - } - - #[cfg(feature = "rayon")] - fn test_basic_indexing_parallel(m: MatrixView<'_, usize>) { - // Par window iters. - let batchsize = 2; - m.par_window_iter(batchsize) - .enumerate() - .for_each(|(i, submatrix)| { - assert_eq!(submatrix.nrows(), batchsize); - assert_eq!(submatrix.ncols(), m.ncols()); - - // Make sure we are in the correct window of the original matrix. - let base = i * batchsize; - assert_eq!(*submatrix.element(0, 0), base); - assert_eq!(*submatrix.element(0, 1), base + 1); - assert_eq!(*submatrix.element(0, 2), base + 2); - - assert_eq!(*submatrix.element(1, 0), base + 1); - assert_eq!(*submatrix.element(1, 1), base + 2); - assert_eq!(*submatrix.element(1, 2), base + 3); - }); - - // Try again, but with a batch size of 3 to ensure that we correctly handle cases - // where the last block is under-sized. - let batchsize = 3; - m.par_window_iter(batchsize) - .enumerate() - .for_each(|(i, submatrix)| { - if i == 0 { - assert_eq!(submatrix.nrows(), batchsize); - assert_eq!(submatrix.ncols(), m.ncols()); - - // Check indexing - assert_eq!(*submatrix.element(0, 0), 0); - assert_eq!(*submatrix.element(0, 1), 1); - assert_eq!(*submatrix.element(0, 2), 2); - - assert_eq!(*submatrix.element(1, 0), 1); - assert_eq!(*submatrix.element(1, 1), 2); - assert_eq!(*submatrix.element(1, 2), 3); - - assert_eq!(*submatrix.element(2, 0), 2); - assert_eq!(*submatrix.element(2, 1), 3); - assert_eq!(*submatrix.element(2, 2), 4); - } else { - assert_eq!(submatrix.nrows(), 1); - assert_eq!(submatrix.ncols(), m.ncols()); - - // Check indexing - assert_eq!(*submatrix.element(0, 0), 3); - assert_eq!(*submatrix.element(0, 1), 4); - assert_eq!(*submatrix.element(0, 2), 5); - } - }); - - // par-row-iter - let seen_rows: Box<[usize]> = m - .par_row_iter() - .enumerate() - .map(|(i, row)| { - let expected: Box<[usize]> = (0..m.ncols()).map(|j| j + i).collect(); - assert_eq!(row, &*expected); - i - }) - .collect(); - - let expected: Box<[usize]> = (0..m.nrows()).collect(); - assert_eq!(seen_rows, expected); - } - - fn test_basic_indexing(m: &MatrixBase) - where - T: DenseData + Sync, - { - assert_eq!(m.nrows(), 4); - assert_eq!(m.ncols(), 3); - - // Basic indexing - assert_eq!(*m.element(0, 0), 0); - assert_eq!(*m.element(0, 1), 1); - assert_eq!(*m.element(0, 2), 2); - - assert_eq!(*m.element(1, 0), 1); - assert_eq!(*m.element(1, 1), 2); - assert_eq!(*m.element(1, 2), 3); - - assert_eq!(*m.element(2, 0), 2); - assert_eq!(*m.element(2, 1), 3); - assert_eq!(*m.element(2, 2), 4); - - assert_eq!(*m.element(3, 0), 3); - assert_eq!(*m.element(3, 1), 4); - assert_eq!(*m.element(3, 2), 5); - - assert_eq!(*m.get_element(0, 0).unwrap(), 0); - assert_eq!(*m.get_element(0, 1).unwrap(), 1); - assert_eq!(*m.get_element(0, 2).unwrap(), 2); - - assert_eq!(*m.get_element(1, 0).unwrap(), 1); - assert_eq!(*m.get_element(1, 1).unwrap(), 2); - assert_eq!(*m.get_element(1, 2).unwrap(), 3); - - assert_eq!(*m.get_element(2, 0).unwrap(), 2); - assert_eq!(*m.get_element(2, 1).unwrap(), 3); - assert_eq!(*m.get_element(2, 2).unwrap(), 4); - - assert_eq!(*m.get_element(3, 0).unwrap(), 3); - assert_eq!(*m.get_element(3, 1).unwrap(), 4); - assert_eq!(*m.get_element(3, 2).unwrap(), 5); - - // Row indexing. - assert_eq!(m.row(0), &[0, 1, 2]); - assert_eq!(m.row(1), &[1, 2, 3]); - assert_eq!(m.row(2), &[2, 3, 4]); - assert_eq!(m.row(3), &[3, 4, 5]); - - let rows: Vec> = m.row_iter().map(|x| x.to_vec()).collect(); - assert_eq!(m.row(0), &rows[0]); - assert_eq!(m.row(1), &rows[1]); - assert_eq!(m.row(2), &rows[2]); - assert_eq!(m.row(3), &rows[3]); - - // Window Iters. - let batchsize = 2; - m.window_iter(batchsize) - .enumerate() - .for_each(|(i, submatrix)| { - assert_eq!(submatrix.nrows(), batchsize); - assert_eq!(submatrix.ncols(), m.ncols()); - - // Make sure we are in the correct window of the original matrix. - let base = i * batchsize; - assert_eq!(*submatrix.element(0, 0), base); - assert_eq!(*submatrix.element(0, 1), base + 1); - assert_eq!(*submatrix.element(0, 2), base + 2); - - assert_eq!(*submatrix.element(1, 0), base + 1); - assert_eq!(*submatrix.element(1, 1), base + 2); - assert_eq!(*submatrix.element(1, 2), base + 3); - }); - - // Try again, but with a batch size of 3 to ensure that we correctly handle cases - // where the last block is under-sized. - let batchsize = 3; - m.window_iter(batchsize) - .enumerate() - .for_each(|(i, submatrix)| { - if i == 0 { - assert_eq!(submatrix.nrows(), batchsize); - assert_eq!(submatrix.ncols(), m.ncols()); - - // Check indexing - assert_eq!(*submatrix.element(0, 0), 0); - assert_eq!(*submatrix.element(0, 1), 1); - assert_eq!(*submatrix.element(0, 2), 2); - - assert_eq!(*submatrix.element(1, 0), 1); - assert_eq!(*submatrix.element(1, 1), 2); - assert_eq!(*submatrix.element(1, 2), 3); - - assert_eq!(*submatrix.element(2, 0), 2); - assert_eq!(*submatrix.element(2, 1), 3); - assert_eq!(*submatrix.element(2, 2), 4); - } else { - assert_eq!(submatrix.nrows(), 1); - assert_eq!(submatrix.ncols(), m.ncols()); - - // Check indexing - assert_eq!(*submatrix.element(0, 0), 3); - assert_eq!(*submatrix.element(0, 1), 4); - assert_eq!(*submatrix.element(0, 2), 5); - } - }); - - #[cfg(all(not(miri), feature = "rayon"))] - test_basic_indexing_parallel(m.as_view()); - } - - #[test] - fn matrix_happy_path() { - let data = make_test_matrix(); - let m = Matrix::try_from(data.into(), 4, 3).unwrap(); - test_basic_indexing(&m); - - // Get the base pointer of the matrix and make sure view-conversion preserves this - // value. - let ptr = m.as_ptr(); - let view = m.as_view(); - assert!(is_copyable(view)); - assert_eq!(view.as_ptr(), ptr); - assert_eq!(view.nrows(), m.nrows()); - assert_eq!(view.ncols(), m.ncols()); - test_basic_indexing(&view); - } - - #[test] - fn matrix_try_from_construction_error() { - let data = make_test_matrix(); - let ptr = data.as_ptr(); - let len = data.len(); - - let m = Matrix::try_from(data.into(), 5, 4); - assert!(m.is_err()); - let err = m.unwrap_err(); - assert_eq!( - err.to_string(), - "tried to construct a 5x4 matrix over a span of length 12" - ); - - // Make sure that we can retrieve the original allocation from the interior. - let data = err.into_inner(); - assert_eq!(data.as_ptr(), ptr); - assert_eq!(data.len(), len); - - let m = MatrixView::try_from(&data, 5, 4); - assert!(m.is_err()); - assert_eq!( - m.unwrap_err().to_string(), - "tried to construct a 5x4 matrix over a span of length 12" - ); - } - - #[test] - fn matrix_mut_view() { - let mut m = Matrix::::from_element(4, 3, 0); - assert_eq!(m.nrows(), 4); - assert_eq!(m.ncols(), 3); - assert!(m.as_slice().iter().all(|&i| i == 0)); - let ptr = m.as_ptr(); - let mut_ptr = m.as_mut_ptr(); - assert_eq!(ptr, mut_ptr); - - let mut view = m.as_mut_view(); - assert_eq!(view.nrows(), 4); - assert_eq!(view.ncols(), 3); - assert_eq!(view.as_ptr(), ptr); - assert_eq!(view.as_mut_ptr(), mut_ptr); - - // Construct the test matrix manually. - for i in 0..view.nrows() { - for j in 0..view.ncols() { - *view.element_mut(i, j) = i + j; - } - } - - // Drop the view and test the original matrix. - test_basic_indexing(&m); - - let inner = m.into_inner(); - assert_eq!(inner.as_ptr(), ptr); - assert_eq!(inner.len(), 4 * 3); - } - - #[test] - fn matrix_view_zero_sizes() { - let data: Vec = vec![]; - // Zero rows, but non-zero columns. - let m = MatrixView::try_from(data.as_slice(), 0, 10).unwrap(); - assert_eq!(m.nrows(), 0); - assert_eq!(m.ncols(), 10); - - // Non-zero rows, but zero columns. - let m = MatrixView::try_from(data.as_slice(), 3, 0).unwrap(); - assert_eq!(m.nrows(), 3); - assert_eq!(m.ncols(), 0); - let empty: &[usize] = &[]; - assert_eq!(m.row(0), empty); - assert_eq!(m.row(1), empty); - assert_eq!(m.row(2), empty); - - // Zero rows and columns. - let m = MatrixView::try_from(data.as_slice(), 0, 0).unwrap(); - assert_eq!(m.nrows(), 0); - assert_eq!(m.ncols(), 0); - } - - #[test] - fn matrix_view_construction_elementwise() { - let mut m = Matrix::::from_element(4, 3, 0); - - // Construct the test matrix manually. - for i in 0..m.nrows() { - for j in 0..m.ncols() { - *m.element_mut(i, j) = i + j; - } - } - test_basic_indexing(&m); - } - - #[test] - fn matrix_construction_by_row() { - let mut m = Matrix::::from_element(4, 3, 0); - assert!(m.as_slice().iter().all(|i| *i == 0)); - - let ncols = m.ncols(); - for i in 0..m.nrows() { - let row = m.row_mut(i); - assert_eq!(row.len(), ncols); - row[0] = i; - row[1] = i + 1; - row[2] = i + 2; - } - test_basic_indexing(&m); - } - - #[test] - fn matrix_construction_by_rowiter() { - let mut m = Matrix::::from_element(4, 3, 0); - assert!(m.as_slice().iter().all(|i| *i == 0)); - - let ncols = m.ncols(); - m.row_iter_mut().enumerate().for_each(|(i, row)| { - assert_eq!(row.len(), ncols); - row[0] = i; - row[1] = i + 1; - row[2] = i + 2; - }); - test_basic_indexing(&m); - } - - #[cfg(all(not(miri), feature = "rayon"))] - #[test] - fn matrix_construction_by_par_windows() { - let mut m = Matrix::::from_element(4, 3, 0); - assert!(m.as_slice().iter().all(|i| *i == 0)); - - let ncols = m.ncols(); - for batchsize in 1..=4 { - m.par_window_iter_mut(batchsize) - .enumerate() - .for_each(|(i, mut submatrix)| { - let base = i * batchsize; - submatrix.row_iter_mut().enumerate().for_each(|(j, row)| { - assert_eq!(row.len(), ncols); - row[0] = base + j; - row[1] = base + j + 1; - row[2] = base + j + 2; - }); - }); - test_basic_indexing(&m); - } - } - - #[test] - fn matrix_construction_happens_in_memory_order() { - let mut i = 0; - let ncols = 3; - let initializer = |_| { - let value = (i % ncols) + (i / ncols); - i += 1; - value - }; - - let m = Matrix::from_fn(4, 3, initializer); - test_basic_indexing(&m); - } - - // Panics - #[test] - #[should_panic(expected = "tried to access row 3 of a matrix with 3 rows")] - fn test_get_row_panics() { - let m = Matrix::::from_element(3, 7, 0); - m.row(3); - } - - #[test] - #[should_panic(expected = "tried to access row 3 of a matrix with 3 rows")] - fn test_get_row_mut_panics() { - let mut m = Matrix::::from_element(3, 7, 0); - m.row_mut(3); - } - - #[test] - #[should_panic(expected = "row 3 is out of bounds (max: 3)")] - fn test_element_panics_row() { - let m = Matrix::::from_element(3, 7, 0); - assert!(m.get_element(3, 2).is_none()); - let _ = m.element(3, 2); - } - - #[test] - #[should_panic(expected = "col 7 is out of bounds (max: 7)")] - fn test_element_panics_col() { - let m = Matrix::::from_element(3, 7, 0); - assert!(m.get_element(2, 7).is_none()); - let _ = m.element(2, 7); - } - - #[test] - #[should_panic(expected = "row 3 is out of bounds (max: 3)")] - fn test_element_mut_panics_row() { - let mut m = Matrix::::from_element(3, 7, 0); - assert!(m.get_element_mut(3, 2).is_none()); - *m.element_mut(3, 2) = 1; - } - - #[test] - #[should_panic(expected = "col 7 is out of bounds (max: 7)")] - fn test_element_mut_panics_col() { - let mut m = Matrix::::from_element(3, 7, 0); - assert!(m.get_element_mut(2, 7).is_none()); - *m.element_mut(2, 7) = 1; - } - - #[test] - #[cfg(feature = "rayon")] - #[should_panic(expected = "par_window_iter batchsize cannot be zero")] - fn test_par_window_iter_panics() { - let m = Matrix::::from_element(4, 4, 0); - let _ = m.par_window_iter(0); - } - - #[test] - #[cfg(feature = "rayon")] - #[should_panic(expected = "par_window_iter_mut batchsize cannot be zero")] - fn test_par_window_iter_mut_panics() { - let mut m = Matrix::::from_element(4, 4, 0); - let _ = m.par_window_iter_mut(0); - } - - // Additional tests for better coverage - - #[test] - fn test_box_slice_dense_data_impls() { - // Test Box<[T]> implementations - let data: Box<[f32]> = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0].into(); - let ptr = data.as_ptr(); - let len = data.len(); - - // Test DenseData impl for Box<[T]> - test_dense_data_repr(ptr, len, data, &lazy_format!("Box<[T]> DenseData")); - - // Test MutDenseData impl for Box<[T]> - let mut data: Box<[f32]> = vec![0.0; 6].into(); - set_mut_dense_data_repr(&mut data, 1.0, 2.0); - for (i, &v) in data.iter().enumerate() { - assert_eq!( - v, - 1.0 + 2.0 * (i as f32), - "Box<[T]> MutDenseData at index {}", - i - ); - } - } - - #[test] - fn test_try_from_error_light() { - // Incorrect slice - let data = vec![1, 2, 3]; - let err = MatrixView::try_from(data.as_slice(), 2, 3).unwrap_err(); - - // Test `as_static` method - let err_static = err.as_static(); - let msg = err_static.to_string(); - assert_contains!( - msg, - "tried to construct a 2x3 matrix over a span of length 3", - ); - // Test `into_inner` method - let recovered_data = err.into_inner(); - assert_eq!(recovered_data, data.as_slice()); - - // Invalid length. - let err = MatrixView::try_from(data.as_slice(), 2, usize::MAX).unwrap_err(); - let msg = err.to_string(); - assert_contains!(msg, "usize::MAX"); - - assert_eq!(data.as_slice(), err.into_inner()); - } - - #[test] - fn test_map_errors() { - #[derive(Debug, Clone, Copy)] - struct Zst; - - // Create a large ZST slice without taking forever on debug builds. - let b = Box::<[Zst]>::new_uninit_slice((isize::MAX as usize) + 1); - - // SAFETY: `b` has zero-sized elements, so all elements are initialized. - let b = unsafe { b.assume_init() }; - - let m = Matrix::column_vector(b); - let err = m.try_map(|_: &Zst| 0u8).unwrap_err(); - let msg = err.to_string(); - assert!(msg.contains("isize::MAX"), "{msg}"); - - // Panicking variant. - let err = std::panic::catch_unwind(|| m.map(|_: &Zst| 0u8)) - .unwrap_err() - .downcast::() - .unwrap(); - let msg = err.to_string(); - assert!(msg.contains("isize::MAX"), "{msg}"); - } - - #[test] - fn test_get_row_optional() { - let data = make_test_matrix(); - let m = MatrixView::try_from(data.as_slice(), 4, 3).unwrap(); - - // Test successful get_row - assert_eq!(m.get_row(0), Some(&[0, 1, 2][..])); - assert_eq!(m.get_row(1), Some(&[1, 2, 3][..])); - assert_eq!(m.get_row(3), Some(&[3, 4, 5][..])); - - // Test out-of-bounds get_row - assert_eq!(m.get_row(4), None); - assert_eq!(m.get_row(100), None); - } - - #[test] - fn test_unsafe_get_unchecked_methods() { - let data = make_test_matrix(); - let mut m = Matrix::try_from(data.into(), 4, 3).unwrap(); - - // Safety: derives from known size of matrix and access element ids - unsafe { - assert_eq!(*m.element_unchecked(0, 0), 0); - assert_eq!(*m.element_unchecked(1, 2), 3); - assert_eq!(*m.element_unchecked(3, 1), 4); - } - - // Safety: derives from known size of matrix and access element ids - unsafe { - *m.element_unchecked_mut(0, 0) = 100; - *m.element_unchecked_mut(1, 2) = 200; - } - - assert_eq!(*m.element(0, 0), 100); - assert_eq!(*m.element(1, 2), 200); - - // Safety: derives from known size of matrix and access element ids - unsafe { - let row0 = m.get_row_unchecked(0); - assert_eq!(row0[0], 100); - assert_eq!(row0[1], 1); - assert_eq!(row0[2], 2); - } - - // Safety: derives from known size of matrix and access element ids - unsafe { - let row1 = m.get_row_unchecked_mut(1); - row1[0] = 300; - } - - assert_eq!(*m.element(1, 0), 300); - } - - #[test] - fn test_to_owned() { - let data = make_test_matrix(); - let view = MatrixView::try_from(data.as_slice(), 4, 3).unwrap(); - - // Test to_owned creates a proper clone - let owned = view.to_owned(); - assert_eq!(owned.nrows(), view.nrows()); - assert_eq!(owned.ncols(), view.ncols()); - assert_eq!(owned.as_slice(), view.as_slice()); - - // Verify it's actually owned (different memory location) - assert_ne!(owned.as_ptr(), view.as_ptr()); - - // Test the owned matrix works properly - test_basic_indexing(&owned); - } - - #[test] - fn test_matrix_from_conversions() { - let data = make_test_matrix(); - let m = Matrix::try_from(data.into(), 4, 3).unwrap(); - - // Test MatrixView to slice conversion - let view = m.as_view(); - let slice: &[usize] = view.into(); - assert_eq!(slice.len(), 12); - assert_eq!(slice[0], 0); - assert_eq!(slice[11], 5); - - // Test MutMatrixView to slice conversion - let data2 = make_test_matrix(); - let mut m2 = Matrix::try_from(data2.into(), 4, 3).unwrap(); - let mut_view = m2.as_mut_view(); - let slice2: &[usize] = mut_view.into(); - assert_eq!(slice2.len(), 12); - assert_eq!(slice2[0], 0); - assert_eq!(slice2[11], 5); - } - - #[test] - fn test_matrix_construction_edge_cases() { - // Test 1x1 matrix - let m = Matrix::from_element(1, 1, 42); - assert_eq!(m.nrows(), 1); - assert_eq!(m.ncols(), 1); - assert_eq!(*m.element(0, 0), 42); - assert_eq!(*m.get_element(0, 0).unwrap(), 42); - - // Test single row matrix - let m = Matrix::from_element(1, 5, 7); - assert_eq!(m.nrows(), 1); - assert_eq!(m.ncols(), 5); - assert!(m.as_slice().iter().all(|&x| x == 7)); - - // Test single column matrix - let m = Matrix::from_element(5, 1, 9); - assert_eq!(m.nrows(), 5); - assert_eq!(m.ncols(), 1); - assert!(m.as_slice().iter().all(|&x| x == 9)); - } - - #[test] - fn test_matrix_view_edge_cases_with_data() { - // Test matrix with actual data for edge cases - let data = vec![10, 20]; - - // 2x1 matrix - let m = MatrixView::try_from(data.as_slice(), 2, 1).unwrap(); - assert_eq!(m.nrows(), 2); - assert_eq!(m.ncols(), 1); - assert_eq!(*m.element(0, 0), 10); - assert_eq!(*m.element(1, 0), 20); - assert_eq!(*m.get_element(0, 0).unwrap(), 10); - assert_eq!(*m.get_element(1, 0).unwrap(), 20); - assert_eq!(m.row(0), &[10]); - assert_eq!(m.row(1), &[20]); - - // 1x2 matrix - let m = MatrixView::try_from(data.as_slice(), 1, 2).unwrap(); - assert_eq!(m.nrows(), 1); - assert_eq!(m.ncols(), 2); - assert_eq!(*m.element(0, 0), 10); - assert_eq!(*m.element(0, 1), 20); - assert_eq!(*m.get_element(0, 0).unwrap(), 10); - assert_eq!(*m.get_element(0, 1).unwrap(), 20); - assert_eq!(m.row(0), &[10, 20]); - } - - #[test] - fn test_row_vector() { - let data = vec![1, 2, 3]; - let m = MatrixView::row_vector(data.as_slice()); - assert_eq!(m.nrows(), 1); - assert_eq!(m.ncols(), 3); - assert_eq!(m.as_slice(), &[1, 2, 3]); - assert_eq!(m.row(0), &[1, 2, 3]); - - // Empty - let empty: &[i32] = &[]; - let m = MatrixView::row_vector(empty); - assert_eq!(m.nrows(), 1); - assert_eq!(m.ncols(), 0); - - // Owned - let m = Matrix::row_vector(vec![10u64, 20].into_boxed_slice()); - assert_eq!(m.nrows(), 1); - assert_eq!(m.ncols(), 2); - assert_eq!(*m.element(0, 0), 10); - assert_eq!(*m.element(0, 1), 20); - } - - #[test] - fn test_column_vector() { - let data = vec![1, 2, 3]; - let m = MatrixView::column_vector(data.as_slice()); - assert_eq!(m.nrows(), 3); - assert_eq!(m.ncols(), 1); - assert_eq!(m.as_slice(), &[1, 2, 3]); - assert_eq!(*m.element(0, 0), 1); - assert_eq!(*m.element(1, 0), 2); - assert_eq!(*m.element(2, 0), 3); - assert_eq!(m.row(0), &[1]); - assert_eq!(m.row(1), &[2]); - assert_eq!(m.row(2), &[3]); - - // Empty - let empty: &[i32] = &[]; - let m = MatrixView::column_vector(empty); - assert_eq!(m.nrows(), 0); - assert_eq!(m.ncols(), 1); - - // Owned - let m = Matrix::column_vector(vec![10u64, 20].into_boxed_slice()); - assert_eq!(m.nrows(), 2); - assert_eq!(m.ncols(), 1); - assert_eq!(*m.element(0, 0), 10); - assert_eq!(*m.element(1, 0), 20); - } - - #[test] - fn test_map() { - let m = Matrix::try_from(vec![1u32, 2, 3, 4].into(), 2, 2).unwrap(); - let doubled = m.map(|&x| x * 2); - assert_eq!(doubled.as_slice(), &[2, 4, 6, 8]); - assert_eq!(doubled.nrows(), 2); - assert_eq!(doubled.ncols(), 2); - - // Type-changing map - let as_f64 = m.map(|&x| x as f64); - assert_eq!(as_f64.as_slice(), &[1.0, 2.0, 3.0, 4.0]); - } - - #[test] - fn test_get_element() { - let mut m = Matrix::try_from(vec![1, 2, 3, 4, 5, 6].into(), 2, 3).unwrap(); - assert_eq!(m.get_element(0, 0), Some(&1)); - assert_eq!(m.get_element(1, 2), Some(&6)); - assert_eq!(m.get_element(2, 0), None); - assert_eq!(m.get_element(0, 3), None); - - *m.get_element_mut(1, 2).unwrap() = 7; - assert_eq!(m.get_element(1, 2), Some(&7)); - assert_eq!(m.get_element_mut(2, 0), None); - assert_eq!(m.get_element_mut(0, 3), None); - } - - #[test] - fn test_subview() { - let data = make_test_matrix(); - let m = Matrix::try_from(data.into(), 4, 3).unwrap(); - - // Create a subview of the first two rows - { - let subview = m.subview(0..4).unwrap(); - assert_eq!(subview.nrows(), 4); - assert_eq!(subview.ncols(), 3); - - assert_eq!(subview.row(0), &[0, 1, 2]); - assert_eq!(subview.row(1), &[1, 2, 3]); - assert_eq!(subview.row(2), &[2, 3, 4]); - assert_eq!(subview.row(3), &[3, 4, 5]); - assert!(subview.get_row(4).is_none()); - } - - // Sub view over a subset that touches the end. - { - let subview = m.subview(1..4).unwrap(); - assert_eq!(subview.nrows(), 3); - assert_eq!(subview.ncols(), 3); - - assert_eq!(subview.row(0), &[1, 2, 3]); - assert_eq!(subview.row(1), &[2, 3, 4]); - assert_eq!(subview.row(2), &[3, 4, 5]); - assert!(subview.get_row(3).is_none()); - } - - // Sub view over a subset that is in the middle - { - let subview = m.subview(1..3).unwrap(); - assert_eq!(subview.nrows(), 2); - assert_eq!(subview.ncols(), 3); - - assert_eq!(subview.row(0), &[1, 2, 3]); - assert_eq!(subview.row(1), &[2, 3, 4]); - assert!(subview.get_row(2).is_none()); - } - - // Empty sub-view. - { - let subview = m.subview(2..2).unwrap(); - assert_eq!(subview.nrows(), 0); - assert_eq!(subview.ncols(), 3); - } - - // Empty subview in bounds - { - let subview = m.subview(0..0).unwrap(); - assert_eq!(subview.nrows(), 0); - assert_eq!(subview.ncols(), 3); - - let subview = m.subview(4..4).unwrap(); - assert_eq!(subview.nrows(), 0); - assert_eq!(subview.ncols(), 3); - } - - // Empty out-of-bounds subview - assert!(m.subview(5..5).is_none()); - - // View too-large - assert!(m.subview(0..6).is_none()); - assert!(m.subview(2..10).is_none()); - - // View disjoint. - assert!(m.subview(10..100).is_none()); - - // Negative bounds - #[expect( - clippy::reversed_empty_ranges, - reason = "we want to make sure it doesn't work" - )] - let empty = 3..2; - assert!(m.subview(empty).is_none()); - - #[expect( - clippy::reversed_empty_ranges, - reason = "we want to make sure it doesn't work" - )] - let empty = 3..1; - assert!(m.subview(empty).is_none()); - - // Bounds that overflow. - assert!(m.subview(usize::MAX - 1..usize::MAX).is_none()); - assert!(m.subview(0..usize::MAX).is_none()); - } - - #[expect( - clippy::reversed_empty_ranges, - reason = "we want to make sure it doesn't work" - )] - #[test] - fn test_subview_zero_cols() { - let m = Matrix::from_element(10, 0, 0u32); - - // Out-of-bounds indexing - assert!(m.subview(100..200).is_none()); - assert!(m.subview(200..100).is_none()); - - assert!(m.subview(10..11).is_none()); - assert!(m.subview(11..10).is_none()); - - assert!(m.subview(0..11).is_none()); - assert!(m.subview(11..0).is_none()); - - assert!(m.subview(10..0).is_none()); - assert!(m.subview(5..4).is_none()); - - // In-bounds. - let v = m.subview(5..10).unwrap(); - assert_eq!(v.nrows(), 5); - assert_eq!(v.ncols(), 0); - - let v = m.subview(0..0).unwrap(); - assert_eq!(v.nrows(), 0); - assert_eq!(v.ncols(), 0); - - let v = m.subview(10..10).unwrap(); - assert_eq!(v.nrows(), 0); - assert_eq!(v.ncols(), 0); - - let v = m.subview(0..10).unwrap(); - assert_eq!(v.nrows(), 10); - assert_eq!(v.ncols(), 0); - } - - #[test] - #[cfg(all(not(miri), feature = "rayon"))] - fn test_parallel_methods_edge_cases() { - let data = make_test_matrix(); - let m = Matrix::try_from(data.into(), 4, 3).unwrap(); - - // Test par_window_iter with batchsize larger than matrix - let windows: Vec<_> = m.par_window_iter(10).collect(); - assert_eq!(windows.len(), 1); - assert_eq!(windows[0].nrows(), 4); - assert_eq!(windows[0].ncols(), 3); - - // Test par_row_iter - let rows: Vec<_> = m.par_row_iter().collect(); - assert_eq!(rows.len(), 4); - assert_eq!(rows[0], &[0, 1, 2]); - assert_eq!(rows[3], &[3, 4, 5]); - - // Test par_window_iter_mut and par_row_iter_mut - let mut m2 = Matrix::from_element(4, 3, 0); - - // Use par_row_iter_mut to set values - m2.par_row_iter_mut().enumerate().for_each(|(i, row)| { - for (j, elem) in row.iter_mut().enumerate() { - *elem = i + j; - } - }); - test_basic_indexing(&m2); - - // Test par_window_iter_mut with larger batchsize - let mut m3 = Matrix::from_element(4, 3, 0); - m3.par_window_iter_mut(10) - .enumerate() - .for_each(|(_, mut window)| { - window.row_iter_mut().enumerate().for_each(|(i, row)| { - for (j, elem) in row.iter_mut().enumerate() { - *elem = i + j; - } - }); - }); - test_basic_indexing(&m3); - } - - #[test] - fn test_matrix_pointers() { - let mut m = Matrix::from_element(3, 4, 42); - - // Test as_ptr and as_mut_ptr return the same address - let const_ptr = m.as_ptr(); - let mut_ptr = m.as_mut_ptr(); - assert_eq!(const_ptr, mut_ptr as *const _); - - // Test that view pointers match original - let view = m.as_view(); - assert_eq!(view.as_ptr(), const_ptr); - - let mut mut_view = m.as_mut_view(); - assert_eq!(mut_view.as_ptr(), const_ptr); - assert_eq!(mut_view.as_mut_ptr(), mut_ptr); - } - - #[test] - fn test_matrix_iteration_empty_cases() { - // Test construction of empty matrices (we don't iterate over 0x0 matrices - // since chunks_exact requires non-zero chunk size) - let empty_data: Vec = vec![]; - - // Matrix with 0 rows but non-zero cols can be constructed - let _empty_matrix = MatrixView::try_from(empty_data.as_slice(), 0, 5).unwrap(); - - // Test with actual single row to verify iterator works normally - let data = vec![1, 2, 3]; - let single_row = MatrixView::try_from(data.as_slice(), 1, 3).unwrap(); - let rows: Vec<_> = single_row.row_iter().collect(); - assert_eq!(rows.len(), 1); - assert_eq!(rows[0], &[1, 2, 3]); - - // Test iteration over matrix with multiple rows but single column - let data = vec![1, 2, 3]; - let single_col = MatrixView::try_from(data.as_slice(), 3, 1).unwrap(); - let rows: Vec<_> = single_col.row_iter().collect(); - assert_eq!(rows.len(), 3); - assert_eq!(rows[0], &[1]); - assert_eq!(rows[1], &[2]); - assert_eq!(rows[2], &[3]); - } - - #[test] - fn test_matrix_init_generator_various_types() { - // Test with different types and generators - use std::sync::atomic::{AtomicUsize, Ordering}; - - let counter = AtomicUsize::new(0); - let m = Matrix::from_fn(2, 3, |_| counter.fetch_add(1, Ordering::SeqCst)); - - // Should be filled in memory order - assert_eq!(*m.element(0, 0), 0); - assert_eq!(*m.element(0, 1), 1); - assert_eq!(*m.element(0, 2), 2); - assert_eq!(*m.element(1, 0), 3); - assert_eq!(*m.element(1, 1), 4); - assert_eq!(*m.element(1, 2), 5); - } - - #[test] - fn test_transpose() { - { - let v = Matrix::from_element(0, 0, 0); - let t = v.transpose(); - assert_eq!(t.nrows(), 0); - assert_eq!(t.ncols(), 0); - } - - { - let v = Matrix::from_element(0, 10, 0); - let t = v.transpose(); - assert_eq!(t.nrows(), 10); - assert_eq!(t.ncols(), 0); - } - - { - let v = Matrix::from_element(10, 0, 0); - let t = v.transpose(); - assert_eq!(t.nrows(), 0); - assert_eq!(t.ncols(), 10); - } - - { - let v = Matrix::::try_from(Box::new([1, 2, 3, 4, 5, 6]), 2, 3).unwrap(); - let t = v.transpose(); - - assert_eq!(t.row(0), &[1, 4]); - assert_eq!(t.row(1), &[2, 5]); - assert_eq!(t.row(2), &[3, 6]); - } - } - - #[test] - fn test_debug_error_formatting() { - // Test Debug implementation for TryFromError - let data = vec![1, 2, 3]; - let err = Matrix::try_from(data.into(), 2, 3).unwrap_err(); - let debug_str = format!("{:?}", err); - assert_contains!(debug_str, "TryFromError"); - - // Ensure Debug doesn't require T: Debug by using a non-Debug type - #[derive(Clone)] - struct NonDebug(#[expect(dead_code)] i32); - - let non_debug_data: Box<[NonDebug]> = vec![NonDebug(1), NonDebug(2)].into(); - let non_debug_err = match Matrix::try_from(non_debug_data, 1, 3) { - Ok(_) => panic!("should not have succeeded!"), - Err(err) => err, - }; - let debug_str = format!("{:?}", non_debug_err); - assert_contains!(debug_str, "TryFromError"); - } - - // Comprehensive tests for rayon-specific functionality - - #[test] - #[cfg(feature = "rayon")] - fn test_par_window_iter_comprehensive() { - use rayon::prelude::*; - - // Create a larger test matrix for more comprehensive testing - let data: Vec = (0..24).collect(); // 6x4 matrix - let m = MatrixView::try_from(data.as_slice(), 6, 4).unwrap(); - - // Test various batch sizes - for batchsize in 1..=8 { - let context = lazy_format!("batchsize = {}", batchsize); - let windows: Vec<_> = m.par_window_iter(batchsize).collect(); - - // Calculate expected number of windows - let expected_windows = (m.nrows()).div_ceil(batchsize); - assert_eq!(windows.len(), expected_windows, "{}", context); - - // Verify each window's properties - let mut total_rows_seen = 0; - for (window_idx, window) in windows.iter().enumerate() { - let expected_rows = if window_idx == windows.len() - 1 { - // Last window may have fewer rows - m.nrows() - (windows.len() - 1) * batchsize - } else { - batchsize - }; - - assert_eq!( - window.nrows(), - expected_rows, - "window {} - {}", - window_idx, - context - ); - assert_eq!( - window.ncols(), - m.ncols(), - "window {} - {}", - window_idx, - context - ); - - // Verify data integrity - for (row_idx, row) in window.row_iter().enumerate() { - let global_row = window_idx * batchsize + row_idx; - let expected: Vec = - (0..m.ncols()).map(|j| global_row * m.ncols() + j).collect(); - assert_eq!( - row, - expected.as_slice(), - "window {}, row {} - {}", - window_idx, - row_idx, - context - ); - } - - total_rows_seen += window.nrows(); - } - - assert_eq!(total_rows_seen, m.nrows(), "{}", context); - } - - // Test with batchsize equal to matrix rows - let windows: Vec<_> = m.par_window_iter(m.nrows()).collect(); - assert_eq!(windows.len(), 1); - assert_eq!(windows[0].nrows(), m.nrows()); - assert_eq!(windows[0].ncols(), m.ncols()); - - // Test with batchsize larger than matrix rows - let windows: Vec<_> = m.par_window_iter(m.nrows() * 2).collect(); - assert_eq!(windows.len(), 1); - assert_eq!(windows[0].nrows(), m.nrows()); - assert_eq!(windows[0].ncols(), m.ncols()); - } - - #[test] - #[cfg(feature = "rayon")] - fn test_par_window_iter_mut_comprehensive() { - use rayon::prelude::*; - - // Test various matrix sizes and batch sizes - for nrows in [1, 2, 3, 5, 8, 10] { - for ncols in [1, 3, 4] { - for batchsize in [1, 2, 3, 7] { - let context = lazy_format!("{}x{}, batchsize={}", nrows, ncols, batchsize); - - let mut m = Matrix::from_element(nrows, ncols, 0usize); - - // Use par_window_iter_mut to fill matrix - m.par_window_iter_mut(batchsize).enumerate().for_each( - |(window_idx, mut window)| { - let base_row = window_idx * batchsize; - window - .row_iter_mut() - .enumerate() - .for_each(|(row_offset, row)| { - let global_row = base_row + row_offset; - for (col, elem) in row.iter_mut().enumerate() { - *elem = global_row * ncols + col; - } - }); - }, - ); - - // Verify the matrix was filled correctly - for row in 0..nrows { - for col in 0..ncols { - let expected = row * ncols + col; - assert_eq!( - *m.element(row, col), - expected, - "pos ({}, {}) - {}", - row, - col, - context - ); - } - } - } - } - } - } - - #[test] - #[cfg(feature = "rayon")] - fn test_par_row_iter_comprehensive() { - use rayon::prelude::*; - - // Create test matrix with predictable pattern - let nrows = 7; - let ncols = 5; - let data: Vec = (0..(nrows * ncols) as i32).collect(); - let m = MatrixView::try_from(data.as_slice(), nrows, ncols).unwrap(); - - // Test that par_row_iter preserves order and data - let collected_rows: Vec> = m.par_row_iter().map(|row| row.to_vec()).collect(); - - assert_eq!(collected_rows.len(), nrows); - - for (row_idx, row) in collected_rows.iter().enumerate() { - assert_eq!(row.len(), ncols); - let expected: Vec = ((row_idx * ncols)..((row_idx + 1) * ncols)) - .map(|x| x as i32) - .collect(); - assert_eq!(row, &expected, "row {} mismatch", row_idx); - } - - // Test parallel enumeration - let enumerated_rows: Vec<(usize, Vec)> = m - .par_row_iter() - .enumerate() - .map(|(idx, row)| (idx, row.to_vec())) - .collect(); - - // Sort by index to ensure we got all indices - let mut sorted_rows = enumerated_rows; - sorted_rows.sort_by_key(|(idx, _)| *idx); - - assert_eq!(sorted_rows.len(), nrows); - for (expected_idx, (actual_idx, row)) in sorted_rows.iter().enumerate() { - assert_eq!(*actual_idx, expected_idx); - assert_eq!(row.len(), ncols); - } - - // Test parallel reduction operations - let sum: i32 = m.par_row_iter().map(|row| row.iter().sum::()).sum(); - - let expected_sum: i32 = data.iter().sum(); - assert_eq!(sum, expected_sum); - - // Test parallel find operations - let target_row = 3; - let found_row = m - .par_row_iter() - .enumerate() - .find_any(|(idx, _)| *idx == target_row) - .map(|(_, row)| row.to_vec()); - - assert!(found_row.is_some()); - let expected_row: Vec = ((target_row * ncols)..((target_row + 1) * ncols)) - .map(|x| x as i32) - .collect(); - assert_eq!(found_row.unwrap(), expected_row); - } - - #[test] - #[cfg(feature = "rayon")] - fn test_par_row_iter_mut_comprehensive() { - use rayon::prelude::*; - use std::sync::atomic::{AtomicUsize, Ordering}; - - let nrows = 6; - let ncols = 4; - let mut m = Matrix::from_element(nrows, ncols, 0u32); - - // Test parallel modification - m.par_row_iter_mut().enumerate().for_each(|(row_idx, row)| { - for (col_idx, elem) in row.iter_mut().enumerate() { - *elem = (row_idx * ncols + col_idx) as u32; - } - }); - - // Verify modifications were applied correctly - for row in 0..nrows { - for col in 0..ncols { - let expected = (row * ncols + col) as u32; - assert_eq!(*m.element(row, col), expected, "pos ({}, {})", row, col); - } - } - - // Test parallel accumulation with atomic counter - let counter = AtomicUsize::new(0); - m.par_row_iter_mut().for_each(|row| { - counter.fetch_add(1, Ordering::Relaxed); - // Multiply each element by 2 - for elem in row { - *elem *= 2; - } - }); - - assert_eq!(counter.load(Ordering::Relaxed), nrows); - - // Verify all elements were doubled - for row in 0..nrows { - for col in 0..ncols { - let expected = ((row * ncols + col) * 2) as u32; - assert_eq!( - *m.element(row, col), - expected, - "doubled pos ({}, {})", - row, - col - ); - } - } - } - - #[test] - #[cfg(feature = "rayon")] - fn test_parallel_iterators_with_single_dimensions() { - use rayon::prelude::*; - - // Test single row matrix - let data = vec![1, 2, 3, 4, 5]; - let single_row = MatrixView::try_from(data.as_slice(), 1, 5).unwrap(); - - let windows: Vec<_> = single_row.par_window_iter(1).collect(); - assert_eq!(windows.len(), 1); - assert_eq!(windows[0].nrows(), 1); - assert_eq!(windows[0].ncols(), 5); - - let rows: Vec<_> = single_row.par_row_iter().collect(); - assert_eq!(rows.len(), 1); - assert_eq!(rows[0], &[1, 2, 3, 4, 5]); - - // Test single column matrix - let data = vec![1, 2, 3, 4, 5]; - let single_col = MatrixView::try_from(data.as_slice(), 5, 1).unwrap(); - - let windows: Vec<_> = single_col.par_window_iter(2).collect(); - assert_eq!(windows.len(), 3); // ceil(5/2) = 3 - assert_eq!(windows[0].nrows(), 2); - assert_eq!(windows[1].nrows(), 2); - assert_eq!(windows[2].nrows(), 1); // Last window has remainder - - let rows: Vec<_> = single_col.par_row_iter().collect(); - assert_eq!(rows.len(), 5); - for (i, row) in rows.iter().enumerate() { - assert_eq!(row, &[i + 1]); - } - - // Test 1x1 matrix - let data = vec![42]; - let tiny = MatrixView::try_from(data.as_slice(), 1, 1).unwrap(); - - let windows: Vec<_> = tiny.par_window_iter(1).collect(); - assert_eq!(windows.len(), 1); - assert_eq!(*windows[0].element(0, 0), 42); - - let rows: Vec<_> = tiny.par_row_iter().collect(); - assert_eq!(rows.len(), 1); - assert_eq!(rows[0], &[42]); - } - - #[test] - #[cfg(feature = "rayon")] - fn test_parallel_window_properties() { - use rayon::prelude::*; - - // Test that windows maintain proper matrix properties - let data: Vec = (0..30).collect(); - let m = MatrixView::try_from(data.as_slice(), 6, 5).unwrap(); - - // Test window indexing works correctly - m.par_window_iter(2) - .enumerate() - .for_each(|(window_idx, window)| { - for row_idx in 0..window.nrows() { - for col_idx in 0..window.ncols() { - let global_row = window_idx * 2 + row_idx; - let expected = global_row * 5 + col_idx; - assert_eq!( - *window.element(row_idx, col_idx), - expected, - "window {}, pos ({}, {})", - window_idx, - row_idx, - col_idx - ); - } - } - }); - - // Test window as_slice consistency - m.par_window_iter(3) - .enumerate() - .for_each(|(window_idx, window)| { - let slice = window.as_slice(); - assert_eq!(slice.len(), window.nrows() * window.ncols()); - - for (slice_idx, &value) in slice.iter().enumerate() { - let row = slice_idx / window.ncols(); - let col = slice_idx % window.ncols(); - assert_eq!( - value, - *window.element(row, col), - "window {}, slice_idx {}", - window_idx, - slice_idx - ); - } - }); - - // Test window row iteration - m.par_window_iter(2).for_each(|window| { - let rows_via_iter: Vec<_> = window.row_iter().collect(); - assert_eq!(rows_via_iter.len(), window.nrows()); - - for (row_idx, row) in rows_via_iter.iter().enumerate() { - assert_eq!(row.len(), window.ncols()); - for (col_idx, &value) in row.iter().enumerate() { - assert_eq!(value, *window.element(row_idx, col_idx)); - } - } - }); - } - - #[test] - #[cfg(feature = "rayon")] - fn test_parallel_performance_characteristics() { - use rayon::prelude::*; - use std::sync::atomic::{AtomicUsize, Ordering}; - - // Create a larger matrix to test parallelism benefits - let nrows = 100; - let ncols = 10; - let mut m = Matrix::from_element(nrows, ncols, 0usize); - - // Test that parallel operations can be chained - let work_counter = AtomicUsize::new(0); - - m.par_window_iter_mut(10) - .enumerate() - .for_each(|(window_idx, mut window)| { - work_counter.fetch_add(1, Ordering::Relaxed); - - // Nested parallel operation within window - window - .row_iter_mut() - .enumerate() - .for_each(|(row_offset, row)| { - let global_row = window_idx * 10 + row_offset; - for (col, elem) in row.iter_mut().enumerate() { - *elem = global_row * ncols + col; - } - }); - }); - - // Should have processed 10 windows (100 rows / 10 batch size) - assert_eq!(work_counter.load(Ordering::Relaxed), 10); - - // Verify correctness - for row in 0..nrows { - for col in 0..ncols { - assert_eq!(*m.element(row, col), row * ncols + col); - } - } - - // Test parallel reduction across windows - let total_sum: usize = m - .par_window_iter(15) - .map(|window| { - window - .row_iter() - .map(|row| row.iter().sum::()) - .sum::() - }) - .sum(); - - let expected_sum: usize = (0..(nrows * ncols)).sum(); - assert_eq!(total_sum, expected_sum); - } - - #[test] - #[cfg(feature = "rayon")] - fn test_rayon_trait_bounds_validation() { - use rayon::prelude::*; - - // Test that the Sync/Send bounds work correctly - let data: Vec = (0..20).collect(); - let m = MatrixView::try_from(data.as_slice(), 4, 5).unwrap(); - - // This should compile because u64 is Sync - let _: Vec<_> = m.par_window_iter(2).collect(); - let _: Vec<_> = m.par_row_iter().collect(); - - // Test with mutable matrix - let mut m = Matrix::from_element(4, 5, 0u64); - - // This should compile because u64 is Send - m.par_window_iter_mut(2).for_each(|mut window| { - window.row_iter_mut().for_each(|row| { - for elem in row { - *elem = 42; - } - }); - }); - - m.par_row_iter_mut().for_each(|row| { - for elem in row { - *elem += 1; - } - }); - - // Verify all elements are 43 - assert!(m.as_slice().iter().all(|&x| x == 43)); - } -} diff --git a/diskann-utils/src/views/mod.rs b/diskann-utils/src/views/mod.rs new file mode 100644 index 0000000000..59cb7f4ebd --- /dev/null +++ b/diskann-utils/src/views/mod.rs @@ -0,0 +1,191 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +pub mod rowmajor; + +/// This trait can be implemented by wrappers for immutable and mutable slice references, +/// allowing for a common code path for immutable and mutable view types. +/// +/// The main goal is to provide a way of retrieving an underlying dense slice, which can +/// then be used as the building block for higher level abstractions. +/// +/// # Safety +/// +/// This trait is unsafe because it requires `as_slice` to be idempotent (and unsafe code +/// relies on this). +/// +/// In other words: `as_slice` must **always** return the same slice with the same length. +pub unsafe trait DenseData { + type Elem; + + /// Return the underlying data as a slice. + fn as_slice(&self) -> &[Self::Elem]; +} + +/// A mutable companion to [`DenseData`]. +/// +/// This trait allows mutable methods on view types to be selectively enabled when data +/// underlying the type is mutable. +/// +/// # Safety +/// +/// This trait is unsafe because it requires `as_slice` to be idempotent (and unsafe code +/// relies on this). +/// +/// In other words: `as_slice` must **always** return the same slice with the same length. +/// +/// Additionally, the returned slice must span the exact same memory as `as_slice`. +pub unsafe trait MutDenseData: DenseData { + fn as_mut_slice(&mut self) -> &mut [Self::Elem]; +} + +// SAFETY: This fulfills the idempotency requirement. +unsafe impl DenseData for &[T] { + type Elem = T; + fn as_slice(&self) -> &[Self::Elem] { + self + } +} + +// SAFETY: This fulfills the idempotency requirement. +unsafe impl DenseData for &mut [T] { + type Elem = T; + fn as_slice(&self) -> &[Self::Elem] { + self + } +} + +// SAFETY: This fulfills the idempotency requirement and returns a slice spanning the same +// range as `as_slice`. +unsafe impl MutDenseData for &mut [T] { + fn as_mut_slice(&mut self) -> &mut [Self::Elem] { + self + } +} + +// SAFETY: This fulfills the idempotency requirement. +unsafe impl DenseData for Box<[T]> { + type Elem = T; + fn as_slice(&self) -> &[Self::Elem] { + self + } +} + +// SAFETY: This fulfills the idempotency requirement and returns a slice spanning the same +// memory as `as_slice`. +unsafe impl MutDenseData for Box<[T]> { + fn as_mut_slice(&mut self) -> &mut [Self::Elem] { + self + } +} + +/////////// +// Tests // +/////////// + +#[cfg(test)] +mod tests { + use super::*; + + use crate::lazy_format; + + /// Test the that provided representation yields a slice with the expected base pointer + /// and length. + fn test_dense_data_repr( + ptr: *const T, + len: usize, + repr: Repr, + context: &dyn std::fmt::Display, + ) where + T: Copy, + Repr: DenseData, + { + let retrieved = repr.as_slice(); + assert_eq!(retrieved.len(), len, "{}", context); + assert_eq!(retrieved.as_ptr(), ptr, "{}", context); + } + + /// Set the underlying data for the provided representation to the following: + /// + /// [base, base + increment, base + increment + increment, ...] + fn set_mut_dense_data_repr(repr: &mut Repr, base: T, increment: T) + where + T: Copy + std::ops::Add, + Repr: DenseData + MutDenseData, + { + let slice = repr.as_mut_slice(); + for i in 0..slice.len() { + if i == 0 { + slice[i] = base; + } else { + slice[i] = slice[i - 1] + increment; + } + } + } + + #[test] + fn slice_implements_dense_data_repr() { + for len in 0..10 { + let context = lazy_format!("len = {}", len); + let data: Vec = vec![0.0; len]; + let slice = data.as_slice(); + test_dense_data_repr(slice.as_ptr(), slice.len(), slice, &context); + } + } + + #[test] + fn mut_slice_implements_dense_data_repr() { + for len in 0..10 { + let context = lazy_format!("len = {}", len); + let mut data: Vec = vec![0.0; len]; + let slice = data.as_mut_slice(); + + let ptr = slice.as_ptr(); + let len = slice.len(); + test_dense_data_repr(ptr, len, slice, &context); + } + } + + #[test] + fn mut_slice_implements_mut_dense_data_repr() { + for len in 0..10 { + let context = lazy_format!("len = {}", len); + let mut data: Vec = vec![0.0; len]; + let mut slice = data.as_mut_slice(); + + let base = 2.0; + let increment = 1.0; + set_mut_dense_data_repr(&mut slice, base, increment); + + for (i, &v) in slice.iter().enumerate() { + let context = lazy_format!("entry {}, {}", i, context); + assert_eq!(v, base + increment * (i as f32), "{}", context); + } + } + } + + #[test] + fn test_box_slice_dense_data_impls() { + // Test Box<[T]> implementations + let data: Box<[f32]> = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0].into(); + let ptr = data.as_ptr(); + let len = data.len(); + + // Test DenseData impl for Box<[T]> + test_dense_data_repr(ptr, len, data, &lazy_format!("Box<[T]> DenseData")); + + // Test MutDenseData impl for Box<[T]> + let mut data: Box<[f32]> = vec![0.0; 6].into(); + set_mut_dense_data_repr(&mut data, 1.0, 2.0); + for (i, &v) in data.iter().enumerate() { + assert_eq!( + v, + 1.0 + 2.0 * (i as f32), + "Box<[T]> MutDenseData at index {}", + i + ); + } + } +} diff --git a/diskann-utils/src/views/rowmajor.rs b/diskann-utils/src/views/rowmajor.rs new file mode 100644 index 0000000000..ffd569394d --- /dev/null +++ b/diskann-utils/src/views/rowmajor.rs @@ -0,0 +1,2561 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use std::{marker::PhantomData, mem::ManuallyDrop, num::NonZeroUsize, ptr::NonNull}; + +#[cfg(feature = "rayon")] +use rayon::prelude::{ + IndexedParallelIterator, IntoParallelIterator, ParallelIterator, ParallelSliceMut, +}; +use thiserror::Error; + +pub mod iter; + +use crate::{internal, Reborrow, ReborrowMut}; + +//////////// +// Matrix // +//////////// + +/// Tools for treating contiguous regions of memory as a dense, row-major matrix. +/// +/// See also: [`MatrixMut`], [`Owned`], [`Ref`], [`Mut`]. +/// +/// # Safety +/// +/// For any shared reference `matrix: &Self`, the following must hold for the entire lifetime +/// of that reference: +/// +/// * Every call to `matrix.as_nonnull()` returns the same pointer. +/// +/// * Every call to `matrix.layout()` returns the same layout. +/// +/// * The pointer and layout describe the same valid initialized span throughout that lifetime. +/// Given: +/// +/// ```text +/// let data = matrix.as_nonnull(); +/// let layout = matrix.layout(); +/// ``` +/// +/// constructing the following slice is valid: +/// +/// ```text +/// unsafe { std::slice::from_raw_parts(data.as_ptr().cast_const(), layout.num_elements()) }; +/// ``` +/// +/// In particular: +/// +/// - `data` must be properly aligned, including when `layout.num_elements() == 0`. +/// +/// - If both `layout.num_elements()` and `size_of::()` are nonzero, the +/// described span must lie within a single allocation. +/// +/// - The span must contain `layout.num_elements()` consecutive, properly initialized +/// values of type `Self::Element`. +/// +/// - The referenced memory must not be mutated for the lifetime of `matrix`, except +/// through an `UnsafeCell`. +pub unsafe trait Matrix { + /// The type of the element stored in the matrix. + type Element; + + /// Return the base pointer for the matrix. + fn as_nonnull(&self) -> NonNull; + + /// Return the [`Layout`] for the matrix. + fn layout(&self) -> Layout; + + //----------// + // Provided // + //----------// + + /// Return the number of rows in the matrix. + fn nrows(&self) -> usize { + self.layout().nrows() + } + + /// Return the number of columns in the matrix. + fn ncols(&self) -> usize { + self.layout().ncols() + } + + /// Returns the requested row without bounds checking. + /// + /// # Safety + /// + /// The following conditions must hold to avoid undefined behavior: + /// + /// * `row < self.nrows()`. + unsafe fn row_unchecked(&self, row: usize) -> &[Self::Element] { + let layout = self.layout(); + debug_assert!(row < layout.nrows()); + + // SAFETY: The caller guarantees that `row` is in-bounds. The validated layout + // therefore places this row entirely within the initialized span required by the + // `Matrix` invariant. The returned lifetime is tied to the borrow of `self`. + unsafe { + std::slice::from_raw_parts(self.as_ptr().add(layout.ncols() * row), layout.ncols()) + } + } + + /// Return a pointer to the base of the matrix. + fn as_ptr(&self) -> *const Self::Element { + self.as_nonnull().as_ptr().cast_const() + } + + /// Return the underlying data as a slice. + fn as_slice(&self) -> &[Self::Element] { + // SAFETY: The `Matrix` invariant requires this pointer and length to describe a + // properly aligned, initialized span that remains immutable for this borrow. + unsafe { std::slice::from_raw_parts(self.as_ptr(), self.layout().num_elements()) } + } + + /// Return row `row` as a slice. + /// + /// # Panic + /// + /// Panics if `row >= self.nrows()`. + fn row(&self, row: usize) -> &[Self::Element] { + assert!( + row < self.nrows(), + "tried to access row {row} of a matrix with {} rows", + self.nrows() + ); + + // SAFETY: `row` is in-bounds. + unsafe { self.row_unchecked(row) } + } + + /// Return row `row` if `row < self.nrows()`. Otherwise, return `None`. + fn get_row(&self, row: usize) -> Option<&[Self::Element]> { + if row < self.nrows() { + // SAFETY: `row` is in-bounds. + Some(unsafe { self.row_unchecked(row) }) + } else { + None + } + } + + /// Return a iterator over all rows in the matrix. + /// + /// Rows are yielded sequentially beginning with row 0. + fn rows(&self) -> iter::Rows<'_, Self::Element> { + iter::Rows::new(self.as_view()) + } + + /// Returns a reference to an element without bounds checking. + /// + /// # Safety + /// + /// The following conditions must hold to avoid undefined behavior: + /// * `row < self.nrows()`. + /// * `col < self.ncols()`. + unsafe fn element_unchecked(&self, row: usize, col: usize) -> &Self::Element { + let layout = self.layout(); + debug_assert!(row < layout.nrows()); + debug_assert!(col < layout.ncols()); + + // SAFETY: The caller guarantees that both indices are in-bounds, so the validated + // layout places this element within the initialized span required by the `Matrix` + // invariant. The returned lifetime is tied to the borrow of `self`. + unsafe { &*self.as_ptr().add(row * layout.ncols() + col) } + } + + /// Return the value at the specified `row` and `col`. + /// + /// If either index is out-of-bounds, return `None`. + fn get_element(&self, row: usize, col: usize) -> Option<&Self::Element> { + if row >= self.nrows() || col >= self.ncols() { + None + } else { + // SAFETY: We just verified that `row` and `col` are in-bounds. + Some(unsafe { self.element_unchecked(row, col) }) + } + } + + /// Return the value at the specified `row` and `col`. + /// + /// # Panics + /// + /// Panics if either `row` or `col` is out-of-bounds. + fn element(&self, row: usize, col: usize) -> &Self::Element { + assert!( + row < self.nrows(), + "row {row} is out of bounds (max: {})", + self.nrows() + ); + assert!( + col < self.ncols(), + "col {col} is out of bounds (max: {})", + self.ncols() + ); + + // SAFETY: We just verified that `row` and `col` are in-bounds. + unsafe { self.element_unchecked(row, col) } + } + + /// Return a view over the matrix. + fn as_view(&self) -> Ref<'_, Self::Element> { + Ref { + ptr: self.as_nonnull(), + layout: self.layout(), + _lifetime: PhantomData, + } + } + + /// Return a view over the rows in `rows`, or `None` when the range is invalid. + fn subview(&self, rows: std::ops::Range) -> Option> { + if rows.start > rows.end || rows.end > self.nrows() { + return None; + } + + let ncols = self.ncols(); + // SAFETY: `rows.start <= self.nrows()`, so the validated parent layout makes the + // offset representable and places it within or one past the matrix span. Pointer + // arithmetic within that span preserves non-nullness. + let ptr = + unsafe { NonNull::new_unchecked(self.as_ptr().add(rows.start * ncols).cast_mut()) }; + // SAFETY: The selected rows are a subset of the validated parent layout, with the + // same column count. + let layout = unsafe { Layout::new_unchecked(rows.end - rows.start, ncols) }; + Some(Ref { + ptr, + layout, + _lifetime: PhantomData, + }) + } + + /// Return an iterator that divides the matrix into sub-matrices with (up to) + /// `batchsize` rows with `self.ncols()` columns. + /// + /// It is possible for yielded sub-matrices to have fewer than `batchsize` rows if the + /// number of rows in the parent matrix is not evenly divisible by `batchsize`. + fn window_iter(&self, batchsize: NonZeroUsize) -> iter::Windows<'_, Self::Element> { + iter::Windows::new(self.as_view(), batchsize) + } + + /// Return an [`Owned`] with the same shape as `self` and cloned contents. + fn to_rowmajor_owned(&self) -> Owned + where + Self::Element: Clone, + { + // SAFETY: The `Matrix` invariant makes `self.as_slice().len()` equal to + // `self.layout().num_elements()`. + unsafe { Owned::from_data_unchecked(self.as_slice().into(), self.layout()) } + } + + /// Create a new [`Matrix`] by applying the closure `f` to each element. + /// + /// The returned matrix has the same shape as `self`. + fn try_map(&self, f: F) -> Result, LayoutError> + where + F: FnMut(&Self::Element) -> R, + { + let layout = self.layout().rebind::()?; + let data: Box<[_]> = self.as_slice().iter().map(f).collect(); + + // SAFETY: Mapping preserves the element count, and `layout` is the source layout + // validated for the destination element type. + Ok(unsafe { Owned::from_data_unchecked(data, layout) }) + } + + /// Create a new [`Matrix`] by applying the closure `f` to each element. + /// + /// The returned matrix has the same shape as `self`. + /// + /// # Panics + /// + /// Panics if allocating space for [`Owned`] would overflow `isize::MAX`. + #[track_caller] + fn map(&self, f: F) -> Owned + where + F: FnMut(&Self::Element) -> R, + { + match self.try_map(f) { + Ok(owned) => owned, + Err(error) => panic!("`Matrix::map` failed: {error}"), + } + } + + /// Transpose the elements in `self`. + fn transpose(&self) -> Owned + where + Self::Element: Clone, + { + Owned::from_fn_with_layout(self.layout().transpose(), |RowCol { row, col }| { + // SAFETY: By construction, `col < self.nrows()` and `row < self.ncols()`. + unsafe { self.element_unchecked(col, row).clone() } + }) + } + + //-------// + // Rayon // + //-------// + + /// Return a parallel iterator over the rows of the matrix. + #[cfg(feature = "rayon")] + fn par_rows(&self) -> impl IndexedParallelIterator + where + Self::Element: Sync, + { + let r = self.as_view(); + + (0..r.nrows()).into_par_iter().map(move |row| { + // SAFETY: `row` comes from `0..r.nrows()`. + unsafe { r.into_row_unchecked(row) } + }) + } + + /// Return a parallel iterator that divides the matrix into sub-matrices with (up to) + /// `batchsize` rows with `self.ncols()` columns. + /// + /// This allows workers in parallel algorithms to work on dense subsets of the whole + /// matrix for better locality. + /// + /// It is possible for yielded sub-matrices to have fewer than `batchsize` rows if the + /// number of rows in the parent matrix is not evenly divisible by `batchsize`. + /// + /// # Panics + /// + /// Panics if `batchsize = 0`. + #[cfg(feature = "rayon")] + fn par_window_iter( + &self, + batchsize: usize, + ) -> impl IndexedParallelIterator> + where + Self::Element: Sync, + { + assert!(batchsize != 0, "par_window_iter batchsize cannot be zero"); + + let r = self.as_view(); + (0..r.nrows()) + .into_par_iter() + .step_by(batchsize) + .map(move |start| { + let end = start.saturating_add(batchsize).min(r.nrows()); + + // SAFETY: `start` comes from `0..nrows` and `end` is clamped to `nrows`. + unsafe { r.into_subview_unchecked(start..end) } + }) + } +} + +/////////////// +// MatrixMut // +/////////////// + +/// Tools for treating contiguous regions of mutable memory as a dense, row-major matrix. +/// +/// See also: [`Owned`], [`Ref`], [`Mut`]. +/// +/// # Safety +/// +/// In addition to the requirements of [`Matrix`], implementations must ensure that for +/// any **mutable** borrow `matrix: &mut Self`, it is valid to construct: +/// +/// ```text +/// let data = matrix.as_nonnull_mut(); +/// let layout = matrix.layout(); +/// +/// unsafe { +/// std::slice::from_raw_parts_mut(data.as_ptr(), layout.num_elements()) +/// } +/// ``` +/// +/// The resulting slice must have exclusive access to the described elements for its entire +/// lifetime. No other reference may be used to access those elements during that lifetime. +/// +/// Implementations may change their pointer or layout through exclusive access, such as when +/// resizing or reallocating. They must not do so while any reference derived from the previous +/// pointer and layout remains live. +/// +/// For a given state of `matrix`, the pointers from `matrix.as_nonnull_mut()` and +/// `matrix.as_nonnull()` must have the same address. However, the pointer obtained from +/// `matrix.as_nonnull_mut()` must have provenance permitting mutable access to the span +/// described by `matrix.layout()`. +pub unsafe trait MatrixMut: Matrix { + /// Return the base pointer for the matrix. + /// + /// Calling this method must not change the matrix's pointer, layout, or contents. + /// + /// The returned pointer must have provenance permitting mutable access to the span + /// described by [`Matrix::layout`] and for a given matrix state, must have the same + /// address as [`Matrix::as_nonnull`]. + fn as_nonnull_mut(&mut self) -> NonNull; + + //----------// + // Provided // + //----------// + + /// Returns the requested row without bounds checking. + /// + /// # Safety + /// + /// The following conditions must hold to avoid undefined behavior: + /// + /// * `row < self.nrows()`. + unsafe fn row_unchecked_mut(&mut self, row: usize) -> &mut [Self::Element] { + let layout = self.layout(); + + debug_assert!(row < layout.nrows()); + + // SAFETY: The caller guarantees that `row` is in-bounds. The validated layout + // therefore places this row within the initialized span, and the `MatrixMut` + // invariant grants exclusive access for the returned borrow. + unsafe { + std::slice::from_raw_parts_mut( + self.as_mut_ptr().add(layout.ncols() * row), + layout.ncols(), + ) + } + } + + /// Return a pointer to the base of the matrix. + fn as_mut_ptr(&mut self) -> *mut Self::Element { + self.as_nonnull_mut().as_ptr() + } + + /// Return the underlying data as a mutable slice. + fn as_mut_slice(&mut self) -> &mut [Self::Element] { + // SAFETY: The `MatrixMut` invariant requires this pointer and length to describe + // the initialized matrix span and grants exclusive access for this mutable borrow. + unsafe { std::slice::from_raw_parts_mut(self.as_mut_ptr(), self.layout().num_elements()) } + } + + /// Return row `row` as a mutable slice. + /// + /// # Panics + /// + /// Panics if `row >= self.nrows()`. + fn row_mut(&mut self, row: usize) -> &mut [Self::Element] { + assert!( + row < self.nrows(), + "tried to access row {row} of a matrix with {} rows", + self.nrows() + ); + + // SAFETY: `row` is in-bounds. + unsafe { self.row_unchecked_mut(row) } + } + + /// Return row `row` if `row < self.nrows()`. Otherwise, return `None`. + fn get_row_mut(&mut self, row: usize) -> Option<&mut [Self::Element]> { + if row < self.nrows() { + // SAFETY: `row` is in-bounds. + Some(unsafe { self.row_unchecked_mut(row) }) + } else { + None + } + } + + /// Return a mutable iterator over all rows in the matrix. + /// + /// Rows are yielded sequentially beginning with row 0. + fn rows_mut(&mut self) -> iter::RowsMut<'_, Self::Element> { + iter::RowsMut::new(self.as_view_mut()) + } + + /// Returns a mutable reference to an element without bounds checking. + /// + /// # Safety + /// + /// The following conditions must hold to avoid undefined behavior: + /// * `row < self.nrows()`. + /// * `col < self.ncols()`. + unsafe fn element_unchecked_mut(&mut self, row: usize, col: usize) -> &mut Self::Element { + let layout = self.layout(); + debug_assert!(row < layout.nrows()); + debug_assert!(col < layout.ncols()); + + // SAFETY: The caller guarantees that both indices are in-bounds, so the validated + // layout places this element within the initialized matrix span. The `MatrixMut` + // invariant grants exclusive access for the returned borrow. + unsafe { &mut *self.as_mut_ptr().add(row * layout.ncols() + col) } + } + + /// Return the value at the specified `row` and `col`. + /// + /// If either index is out-of-bounds, return `None`. + fn get_element_mut(&mut self, row: usize, col: usize) -> Option<&mut Self::Element> { + if row >= self.nrows() || col >= self.ncols() { + None + } else { + // SAFETY: We just verified that `row` and `col` are in-bounds. + Some(unsafe { self.element_unchecked_mut(row, col) }) + } + } + + /// Return the value at the specified `row` and `col`. + /// + /// # Panics + /// + /// Panics if either `row` or `col` is out-of-bounds. + fn element_mut(&mut self, row: usize, col: usize) -> &mut Self::Element { + assert!( + row < self.nrows(), + "row {row} is out of bounds (max: {})", + self.nrows() + ); + assert!( + col < self.ncols(), + "col {col} is out of bounds (max: {})", + self.ncols() + ); + + // SAFETY: We just verified that `row` and `col` are in-bounds. + unsafe { self.element_unchecked_mut(row, col) } + } + + /// Return a view over the matrix. + fn as_view_mut(&mut self) -> Mut<'_, Self::Element> { + Mut { + ptr: self.as_nonnull_mut(), + layout: self.layout(), + _lifetime: PhantomData, + } + } + + //-------// + // Rayon // + //-------// + + /// Return a parallel iterator over the rows of the matrix. + /// + /// # Panics + /// + /// Panics if `self.ncols() == 0 && self.nrows() != 0`. + #[cfg(feature = "rayon")] + fn par_rows_mut(&mut self) -> impl IndexedParallelIterator + where + Self::Element: Send, + { + let ncols = self.ncols(); + assert!( + ncols != 0 || self.nrows() == 0, + "`MatrixMut::par_rows_mut` does not support matrices with rows and zero columns" + ); + self.as_mut_slice().par_chunks_exact_mut(ncols.max(1)) + } + + /// Return a parallel iterator that divides the matrix into mutable sub-matrices with + /// (up to) `batchsize` rows with `self.ncols()` columns. + /// + /// This allows workers in parallel algorithms to work on dense subsets of the whole + /// matrix for better locality. + /// + /// It is possible for yielded sub-matrices to have fewer than `batchsize` rows if the + /// number of rows in the parent matrix is not evenly divisible by `batchsize`. + /// + /// # Panics + /// + /// Panics if `batchsize = 0` or `self.ncols() == 0 && self.nrows() != 0`. + #[cfg(feature = "rayon")] + fn par_window_iter_mut( + &mut self, + batchsize: usize, + ) -> impl IndexedParallelIterator> + where + Self::Element: Send, + { + assert!( + batchsize != 0, + "par_window_iter_mut batchsize cannot be zero" + ); + + let ncols = self.ncols(); + assert!( + ncols != 0 || self.nrows() == 0, + "`MatrixMut::par_window_iter_mut` does not support matrices with rows and zero columns" + ); + + // Ensure that `batchsize * ncols` does not overflow. + let batchsize = batchsize.min(self.nrows()); + self.as_mut_slice() + .par_chunks_mut((ncols * batchsize).max(1)) + .map(move |data| { + let blobsize = data.len(); + let nrows = blobsize / ncols; + assert_eq!(blobsize % ncols, 0); + + // SAFETY: + // + // * `Layout::new_unchecked` is safe because `ncols` is the parent column + // count and `nrows <= self.nrows()`, so this layout cannot exceed the + // validated parent layout. + // + // * `Mut::from_data_unchecked` is safe because by construction, + // `data.len() == ncols * nrows`. + unsafe { Mut::from_data_unchecked(data, Layout::new_unchecked(nrows, ncols)) } + }) + } +} + +/////////////////// +// Matrix Layout // +/////////////////// + +/// A validated layout for [`Matrix`] or [`MatrixMut`]. +/// +/// This type guarantees the following invariants: +/// +/// * `self.nrows() * self.ncols()` does not exceed `usize::MAX`. +/// * `self.nrows() * self.ncols() * std::mem::size_of::()` does not exceed `isize::MAX`. +pub struct Layout { + nrows: usize, + ncols: usize, + _type: PhantomData T>, +} + +impl Layout { + /// Construct a new [`Layout`], validating the following: + /// + /// * `nrows * ncols` does not exceed `usize::MAX`. + /// * `nrows * ncols * std::mem::size_of::()` does not exceed `isize::MAX` (the maximum + /// addressable byte span). + pub const fn new(nrows: usize, ncols: usize) -> Result { + match LayoutError::check::(nrows, ncols) { + Ok(()) => Ok(Self { + nrows, + ncols, + _type: PhantomData, + }), + Err(err) => Err(err), + } + } + + /// Construct a layout without validating its dimensions. + /// + /// # Safety + /// + /// `LayoutError::check::(nrows, ncols)` must succeed. + unsafe fn new_unchecked(nrows: usize, ncols: usize) -> Self { + debug_assert!(LayoutError::check::(nrows, ncols).is_ok()); + Self { + nrows, + ncols, + _type: PhantomData, + } + } + + /// Return the product `self.nrows() * self.ncols()`. + pub fn num_elements(&self) -> usize { + self.nrows() * self.ncols() + } + + /// Return the number of rows. + pub fn nrows(&self) -> usize { + self.nrows + } + + /// Return the number of columns. + pub fn ncols(&self) -> usize { + self.ncols + } + + /// Rebind the element type to `U`. + /// + /// # Errors + /// + /// Returns an error if the rebound layout's byte size would exceed `isize::MAX`. + pub fn rebind(&self) -> Result, LayoutError> { + if std::mem::size_of::() <= std::mem::size_of::() { + // This branch is mainly to communicate to the compiler situations where an + // erroring branch can be avoided. + // + // SAFETY: `self` is validated, and rebinding to an equally sized or smaller + // element type cannot increase its byte span. + Ok(unsafe { Layout::new_unchecked(self.nrows(), self.ncols()) }) + } else { + Layout::new(self.nrows(), self.ncols()) + } + } + + /// Swap the rows and columns. + pub fn transpose(&self) -> Layout { + // SAFETY: Swapping the dimensions preserves both the element count and byte span + // of this validated layout. + unsafe { Layout::new_unchecked(self.ncols, self.nrows) } + } +} + +impl Clone for Layout { + fn clone(&self) -> Self { + *self + } +} + +impl Copy for Layout {} + +impl std::fmt::Debug for Layout { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Layout") + .field("nrows", &self.nrows) + .field("ncols", &self.ncols) + .field("elsize", &std::mem::size_of::()) + .finish() + } +} + +impl PartialEq for Layout { + fn eq(&self, other: &Self) -> bool { + self.nrows == other.nrows && self.ncols == other.ncols + } +} + +impl Eq for Layout {} + +/// Errors in the invariants guaranteed by [`Layout`]. +#[derive(Debug, Clone, Copy)] +pub struct LayoutError { + nrows: usize, + ncols: usize, + elsize: Option, +} + +impl LayoutError { + pub(crate) const fn check(nrows: usize, ncols: usize) -> Result<(), Self> { + // Guard the element count itself so that `num_elements()` can never overflow. + let elsize = std::mem::size_of::(); + let num_elements = match nrows.checked_mul(ncols) { + Some(num_elements) => num_elements, + None => { + return Err(Self { + nrows, + ncols, + elsize: None, + }) + } + }; + + if let Some(len) = num_elements.checked_mul(std::mem::size_of::()) { + if len <= isize::MAX as usize { + return Ok(()); + } + } + + Err(Self { + nrows, + ncols, + elsize: NonZeroUsize::new(elsize), + }) + } +} + +impl std::fmt::Display for LayoutError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self.elsize { + Some(elsize) => { + write!( + f, + "a matrix of size {}x{} with elements of size {} exceeds `isize::MAX` bytes", + self.nrows, self.ncols, elsize + ) + } + None => { + write!( + f, + "a matrix of size {}x{} has a length exceeding `usize::MAX`", + self.nrows, self.ncols, + ) + } + } + } +} + +impl std::error::Error for LayoutError {} + +//---------------// +// Helper Macros // +//---------------// + +macro_rules! constructors { + ($element:ident, $data:ty) => { + /// Try to construct directly from `data`. + /// + /// Returns an error if [`Layout::new`] fails for `nrows` and `ncols` or `data.len()` + /// is not equal to `nrows * ncols`. + pub fn try_from_data( + data: $data, + nrows: usize, + ncols: usize, + ) -> Result> { + let layout = match Layout::<$element>::new(nrows, ncols) { + Ok(layout) => layout, + Err(err) => return Err(TryFromError::layout(data, err)), + }; + + let len = data.len(); + if len == layout.num_elements() { + // SAFETY: We've checked that `data.len() == layout.num_elements()`. + Ok(unsafe { Self::from_data_unchecked(data, layout) }) + } else { + Err(TryFromError::mismatch( + data, + layout.nrows(), + layout.ncols(), + len, + )) + } + } + + /// Construct a row vector directly from `data`. + pub fn row_vector(data: $data) -> Self { + // SAFETY: An existing slice's byte span cannot exceed `isize::MAX`, and + // `1 * data.len()` cannot overflow. + let layout = unsafe { Layout::new_unchecked(1, data.len()) }; + + // SAFETY: By construction, `data.len() == layout.num_elements()`. + unsafe { Self::from_data_unchecked(data, layout) } + } + + /// Construct a column vector directly from `data`. + pub fn column_vector(data: $data) -> Self { + // SAFETY: An existing slice's byte span cannot exceed `isize::MAX`, and + // `data.len() * 1` cannot overflow. + let layout = unsafe { Layout::new_unchecked(data.len(), 1) }; + + // SAFETY: By construction, `data.len() == layout.num_elements()`. + unsafe { Self::from_data_unchecked(data, layout) } + } + }; +} + +/////////// +// Owned // +/////////// + +/// An initializer argument for the closure provided to [`Owned::from_fn`] and +/// [`Owned::try_from_fn`] to remove ambiguity of the row and column being initialiazed. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct RowCol { + pub row: usize, + pub col: usize, +} + +/// A [`Matrix`]/[`MatrixMut`] that owns its data. +#[derive(Debug)] +pub struct Owned { + ptr: NonNull, + layout: Layout, +} + +impl Owned { + constructors!(T, Box<[T]>); + + // NOTE: For constructors, keep `from_fn` and `from_element` first. + // + // Rust suggests methods in their declaration order, so this keeps the most common + // methods as top suggestions. + + /// Construct a new matrix using `init`. + /// + /// Elements are initialized in memory order. + /// + /// ``` + /// use diskann_utils::views::rowmajor::{self, Matrix}; + /// + /// let mat = rowmajor::Owned::from_fn(2, 3, |rc| 3 * rc.row + rc.col); + /// + /// assert_eq!(mat.row(0), &[0, 1, 2]); + /// assert_eq!(mat.row(1), &[3, 4, 5]); + /// ``` + /// + /// # Panics + /// + /// Panics if `nrows * ncols` overflows `usize::MAX`, or if the allocation size exceeds + /// `isize::MAX`. + #[track_caller] + pub fn from_fn(nrows: usize, ncols: usize, init: F) -> Self + where + F: FnMut(RowCol) -> T, + { + match Self::try_from_fn(nrows, ncols, init) { + Ok(matrix) => matrix, + Err(error) => panic!("Owned::from_fn failed with: {error}"), + } + } + + /// Construct a new matrix using `init`. + /// + /// Elements are initialized in memory order. + /// + /// ``` + /// use diskann_utils::views::rowmajor::{self, Matrix}; + /// + /// let mat = rowmajor::Owned::try_from_fn(2, 3, |rc| 3 * rc.row + rc.col).unwrap(); + /// + /// assert_eq!(mat.row(0), &[0, 1, 2]); + /// assert_eq!(mat.row(1), &[3, 4, 5]); + /// ``` + /// + /// # Errors + /// + /// Returns an error if `nrows * ncols` overflows `usize::MAX`, or if the allocation size + /// exceeds `isize::MAX`. + pub fn try_from_fn(nrows: usize, ncols: usize, init: F) -> Result + where + F: FnMut(RowCol) -> T, + { + let layout = Layout::new(nrows, ncols)?; + Ok(Self::from_fn_with_layout(layout, init)) + } + + /// Construct a new matrix by cloning `element`. + /// + /// Elements are initialized in memory order. + /// + /// ``` + /// use diskann_utils::views::rowmajor::{self, Matrix}; + /// + /// let mat = rowmajor::Owned::from_element(2, 3, 0u32); + /// + /// assert_eq!(mat.row(0), &[0, 0, 0]); + /// assert_eq!(mat.row(1), &[0, 0, 0]); + /// ``` + /// + /// # Panics + /// + /// Panics if `nrows * ncols` overflows `usize::MAX`, or if the allocation size exceeds + /// `isize::MAX`. + #[track_caller] + pub fn from_element(nrows: usize, ncols: usize, element: T) -> Self + where + T: Clone, + { + match Self::try_from_element(nrows, ncols, element) { + Ok(matrix) => matrix, + Err(error) => panic!("Owned::from_element failed with: {error}"), + } + } + + /// Construct a new matrix by cloning `element`. + /// + /// Elements are initialized in memory order. + /// + /// ``` + /// use diskann_utils::views::rowmajor::{self, Matrix}; + /// + /// let mat = rowmajor::Owned::try_from_element(2, 3, 0u32).unwrap(); + /// + /// assert_eq!(mat.row(0), &[0, 0, 0]); + /// assert_eq!(mat.row(1), &[0, 0, 0]); + /// ``` + /// + /// # Errors + /// + /// Returns an error if `nrows * ncols` overflows `usize::MAX`, or if the allocation size + /// exceeds `isize::MAX`. + pub fn try_from_element(nrows: usize, ncols: usize, element: T) -> Result + where + T: Clone, + { + let layout = Layout::new(nrows, ncols)?; + Ok(Self::from_element_with_layout(layout, element)) + } + + // Less common constructors. + + /// Construct a new matrix using `init`. + /// + /// Elements are initialized in memory order. + pub fn from_fn_with_layout(layout: Layout, mut init: F) -> Self + where + F: FnMut(RowCol) -> T, + { + let mut row = 0; + let mut col = 0; + + let data: Box<[T]> = (0..layout.num_elements()) + .map(|_| { + let v = (init)(RowCol { row, col }); + col += 1; + if col == layout.ncols() { + col = 0; + row += 1; + } + v + }) + .collect(); + + // SAFETY: We constructed `data` to have length exactly `layout.num_elements()`. + unsafe { Self::from_data_unchecked(data, layout) } + } + + /// Construct a new matrix by cloning `element`. + /// + /// Elements are initialized in memory order. + pub fn from_element_with_layout(layout: Layout, element: T) -> Self + where + T: Clone, + { + let data: Box<[T]> = std::iter::repeat_n(element, layout.num_elements()).collect(); + + // SAFETY: We constructed `data` to have length exactly `layout.num_elements()`. + unsafe { Self::from_data_unchecked(data, layout) } + } + + /// # Safety + /// + /// `b.len()` must equal `layout.num_elements()`. + unsafe fn from_data_unchecked(b: Box<[T]>, layout: Layout) -> Self { + debug_assert_eq!(b.len(), layout.num_elements()); + Self { + ptr: internal::box_to_nonnull(b), + layout, + } + } + + /// Consume `self`, returning the unmodified contents as a boxed slice. + /// + /// ``` + /// use diskann_utils::views::rowmajor::{Matrix, Owned}; + /// + /// let mat = Owned::from_fn(2, 3, |rc| rc.col); + /// assert_eq!(mat.row(0), &[0, 1, 2]); + /// assert_eq!(mat.row(1), &[0, 1, 2]); + /// + /// let b: Box<[usize]> = mat.into_inner(); + /// assert_eq!(&*b, &[0, 1, 2, 0, 1, 2]); + /// ``` + pub fn into_inner(self) -> Box<[T]> { + let me = ManuallyDrop::new(self); + + // SAFETY: `me.ptr` came from exactly one `Box<[T]>` of this length. Suppressing + // `Owned::drop` transfers that allocation back to the reconstructed box. + unsafe { internal::nonnull_to_box(me.ptr, me.layout.num_elements()) } + } +} + +// SAFETY: `Owned` has the ownership semantics of `Box<[T]>`, which is `Send` when `T` is. +unsafe impl Send for Owned where T: Send {} +// SAFETY: Shared access to `Owned` exposes only shared access to `T`. +unsafe impl Sync for Owned where T: Sync {} + +impl Drop for Owned { + fn drop(&mut self) { + // SAFETY: `self.ptr` came from exactly one `Box<[T]>` of this length, and `drop` + // is the unique place that reconstructs it. + let _ = unsafe { internal::nonnull_to_box(self.ptr, self.layout.num_elements()) }; + } +} + +impl Clone for Owned +where + T: Clone, +{ + fn clone(&self) -> Self { + // SAFETY: The constructed boxed slice has length exactly `self.layout.num_elements()`. + unsafe { Owned::from_data_unchecked(self.as_slice().into(), self.layout()) } + } +} + +// SAFETY: `Owned` keeps a stable pointer and layout to its initialized boxed slice. +unsafe impl Matrix for Owned { + type Element = T; + + fn as_nonnull(&self) -> NonNull { + self.ptr + } + + fn layout(&self) -> Layout { + self.layout + } +} + +// SAFETY: A mutable borrow of `Owned` has exclusive access to its boxed slice. +unsafe impl MatrixMut for Owned { + fn as_nonnull_mut(&mut self) -> NonNull { + self.ptr + } +} + +impl PartialEq for Owned +where + T: PartialEq, +{ + fn eq(&self, other: &Self) -> bool { + Matrix::as_view(self).eq(&Matrix::as_view(other)) + } +} + +impl<'a, T> Reborrow<'a> for Owned { + type Target = Ref<'a, T>; + fn reborrow(&'a self) -> Self::Target { + Matrix::as_view(self) + } +} + +impl<'a, T> ReborrowMut<'a> for Owned { + type Target = Mut<'a, T>; + fn reborrow_mut(&'a mut self) -> Self::Target { + MatrixMut::as_view_mut(self) + } +} + +//-----// +// Ref // +//-----// + +/// A [`Matrix`] implementation that references its data. +#[derive(Debug)] +pub struct Ref<'a, T> { + ptr: NonNull, + layout: Layout, + _lifetime: PhantomData<&'a [T]>, +} + +// SAFETY: `Ref<'_, T>` has the ownership semantics of `&[T]`, which is `Send` when `T` is +// `Sync`. +unsafe impl Send for Ref<'_, T> where T: Sync {} +// SAFETY: `Ref<'_, T>` exposes only shared access to `T`. +unsafe impl Sync for Ref<'_, T> where T: Sync {} + +impl<'a, T> Ref<'a, T> { + constructors!(T, &'a [T]); + + /// # Safety + /// + /// `b.len()` must equal `layout.num_elements()`. + unsafe fn from_data_unchecked(b: &'a [T], layout: Layout) -> Self { + debug_assert_eq!(b.len(), layout.num_elements()); + Self { + ptr: internal::slice_to_nonnull(b), + layout, + _lifetime: PhantomData, + } + } + + /// Return the contents of `self` as a slice. + /// + /// Unlike [`Matrix::as_slice`], the returned slices inherits the lifetime of the [`Ref`]. + pub fn into_slice(self) -> &'a [T] { + // SAFETY: `Ref` represents a valid `&'a [T]` of exactly this length. + unsafe { std::slice::from_raw_parts(self.as_ptr(), self.layout().num_elements()) } + } + + /// Get the indicated `row` with the lifetime `'a`. + /// + /// # Safety + /// + /// `row < self.nrows()`. + #[cfg(feature = "rayon")] + unsafe fn into_row_unchecked(self, row: usize) -> &'a [T] { + let layout = self.layout(); + debug_assert!(row < layout.nrows()); + + // SAFETY: The caller guarantees that `row` is in-bounds, so the validated layout + // places this row within the `&'a [T]` represented by `self`. + unsafe { + std::slice::from_raw_parts(self.as_ptr().add(layout.ncols() * row), layout.ncols()) + } + } + + /// Return a `Ref` containing the indicated rows with the lifetime `'a`. + /// + /// # Safety + /// + /// `rows.start <= rows.end` and `rows.end <= self.nrows()`. + #[cfg(feature = "rayon")] + unsafe fn into_subview_unchecked(self, rows: std::ops::Range) -> Ref<'a, T> { + debug_assert!(rows.start <= rows.end); + debug_assert!(rows.end <= self.nrows()); + + let ncols = self.ncols(); + Self { + // SAFETY: The validated range starts within or one past the matrix span, and + // the parent layout guarantees that the offset is representable. + ptr: unsafe { self.ptr.add(rows.start * ncols) }, + // SAFETY: This layout is no larger than `self`'s layout. + layout: unsafe { Layout::new_unchecked(rows.end - rows.start, ncols) }, + _lifetime: PhantomData, + } + } +} + +impl Clone for Ref<'_, T> { + fn clone(&self) -> Self { + *self + } +} + +impl Copy for Ref<'_, T> {} + +// SAFETY: `Ref` keeps a stable pointer and layout to the initialized slice it borrows. +unsafe impl Matrix for Ref<'_, T> { + type Element = T; + + fn as_nonnull(&self) -> NonNull { + self.ptr + } + + fn layout(&self) -> Layout { + self.layout + } +} + +impl<'a, T> Reborrow<'a> for Ref<'_, T> { + type Target = Ref<'a, T>; + fn reborrow(&'a self) -> Self::Target { + Matrix::as_view(self) + } +} + +impl PartialEq for Ref<'_, T> +where + T: PartialEq, +{ + fn eq(&self, other: &Self) -> bool { + self.layout() == other.layout() && self.as_slice() == other.as_slice() + } +} + +//-----// +// Mut // +//-----// + +/// A [`Matrix`]/[`MatrixMut`] implementation that mutably references its data. +#[derive(Debug)] +pub struct Mut<'a, T> { + ptr: NonNull, + layout: Layout, + _lifetime: PhantomData<&'a mut [T]>, +} + +// SAFETY: `Mut<'_, T>` has the ownership semantics of `&mut [T]`, which is `Send` when +// `T` is `Send`. +unsafe impl Send for Mut<'_, T> where T: Send {} +// SAFETY: Shared access to `Mut<'_, T>` exposes only shared access to `T`. +unsafe impl Sync for Mut<'_, T> where T: Sync {} + +impl<'a, T> Mut<'a, T> { + constructors!(T, &'a mut [T]); + + /// # Safety + /// + /// `b.len()` must equal `layout.num_elements()`. + unsafe fn from_data_unchecked(b: &'a mut [T], layout: Layout) -> Self { + debug_assert_eq!(b.len(), layout.num_elements()); + Self { + ptr: internal::mut_slice_to_nonnull(b), + layout, + _lifetime: PhantomData, + } + } + + /// Consume `self` and return the underlying data as a mutable slice. + pub fn into_mut_slice(self) -> &'a mut [T] { + // SAFETY: `self.ptr` and `self.layout` together describe a valid `&'a mut [T]` of + // length `self.layout.num_elements()`, per the invariants of `Mut`. + unsafe { std::slice::from_raw_parts_mut(self.ptr.as_ptr(), self.layout.num_elements()) } + } +} + +// SAFETY: `Mut` keeps a stable pointer and layout to the initialized slice it borrows. +unsafe impl Matrix for Mut<'_, T> { + type Element = T; + + fn as_nonnull(&self) -> NonNull { + self.ptr + } + + fn layout(&self) -> Layout { + self.layout + } +} + +// SAFETY: A mutable borrow of `Mut` has exclusive access to its borrowed slice. +unsafe impl MatrixMut for Mut<'_, T> { + fn as_nonnull_mut(&mut self) -> NonNull { + self.ptr + } +} + +impl PartialEq for Mut<'_, T> +where + T: PartialEq, +{ + fn eq(&self, other: &Self) -> bool { + Matrix::as_view(self).eq(&Matrix::as_view(other)) + } +} + +impl<'a, T> Reborrow<'a> for Mut<'_, T> { + type Target = Ref<'a, T>; + fn reborrow(&'a self) -> Self::Target { + Matrix::as_view(self) + } +} + +impl<'a, T> ReborrowMut<'a> for Mut<'_, T> { + type Target = Mut<'a, T>; + fn reborrow_mut(&'a mut self) -> Self::Target { + MatrixMut::as_view_mut(self) + } +} + +//--------// +// Errors // +//--------// + +/// Errors from [`Owned::try_from_data`], [`Ref::try_from_data`], and [`Mut::try_from_data`]. +pub struct TryFromError { + data: T, + inner: TryFromErrorInner, +} + +impl TryFromError { + /// Consume the error and return the base data. + pub fn into_inner(self) -> T { + self.data + } + + /// Return a variation of `Self` that is guaranteed to be `'static` by removing the + /// data that was passed to the original constructor. + pub fn as_static(&self) -> TryFromErrorLight { + TryFromErrorLight(self.inner) + } + + //--------------// + // Constructors // + //--------------// + + fn layout(data: T, error: LayoutError) -> Self { + Self { + data, + inner: TryFromErrorInner::Layout(error), + } + } + + fn mismatch(data: T, nrows: usize, ncols: usize, len: usize) -> Self { + Self { + data, + inner: TryFromErrorInner::Mismatch { nrows, ncols, len }, + } + } +} + +impl std::fmt::Debug for TryFromError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("TryFromError") + .field("data", &"") + .field("inner", &self.inner) + .finish() + } +} + +impl std::fmt::Display for TryFromError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + self.inner.fmt(f) + } +} + +impl std::error::Error for TryFromError {} + +/// A guaranteed `'static` version of [`TryFromError`]. +#[derive(Debug, Error)] +#[error(transparent)] +pub struct TryFromErrorLight(TryFromErrorInner); + +#[derive(Debug, Error, Clone, Copy)] +enum TryFromErrorInner { + #[error(transparent)] + Layout(LayoutError), + #[error( + "tried to construct a {}x{} matrix over a span of length {}", + nrows, + ncols, + len + )] + Mismatch { + nrows: usize, + ncols: usize, + len: usize, + }, +} + +/////////// +// Tests // +/////////// + +#[cfg(test)] +mod tests { + use super::*; + use crate::{assert_contains, lazy_format}; + + /// This function is only callable with copyable types. + /// + /// This lets us test for types we expect to be `Copy`. + fn is_copyable(_x: T) -> bool { + true + } + + /// This function attests that `Ref` is covariant in the view lifetime. + fn _matrix_view_is_covariant<'a, 'b>(m: Ref<'a, f32>) -> Ref<'b, f32> + where + 'a: 'b, + { + m + } + + fn _matrix_view_is_covariant_in_t<'a, 'b, 'm>(m: Ref<'m, &'a f32>) -> Ref<'m, &'b f32> + where + 'a: 'b, + { + m + } + + fn _matrix_is_covariant_in_t<'a, 'b, 'm>(m: &'m Owned<&'a f32>) -> &'m Owned<&'b f32> + where + 'a: 'b, + { + m + } + + //--------// + // Layout // + //--------// + + #[test] + fn test_layout() { + // Happy path + for rows in 0..5 { + for cols in 0..5 { + let layout = Layout::::new(rows, cols).unwrap(); + assert_eq!(layout.nrows(), rows); + assert_eq!(layout.ncols(), cols); + assert_eq!(layout.num_elements(), rows * cols); + + let transpose = layout.transpose(); + assert_eq!(transpose.nrows(), cols); + assert_eq!(transpose.ncols(), rows); + assert_eq!(transpose.num_elements(), rows * cols); + + let rebind = layout.rebind::().unwrap(); + assert_eq!(rebind.nrows(), rows); + assert_eq!(rebind.ncols(), cols); + assert_eq!(rebind.num_elements(), rows * cols); + + is_copyable(layout); + } + } + + #[expect(unused, reason = "we need this so the size is non-zero")] + struct NotDebugOrEq(u32); + + assert_eq!( + Layout::::new(10, 20).unwrap(), + Layout::::new(10, 20).unwrap(), + ); + + assert_eq!( + Layout::::new(20, 0).unwrap(), + Layout::::new(20, 0).unwrap(), + ); + + assert_ne!( + Layout::::new(10, 20).unwrap(), + Layout::::new(20, 0).unwrap(), + ); + + let fmt = format!("{:?}", Layout::::new(5, 6).unwrap()); + assert_eq!(fmt, "Layout { nrows: 5, ncols: 6, elsize: 4 }"); + + // Overflowing the element count returns an error. + let error = Layout::::new(usize::MAX, 2).unwrap_err(); + assert_eq!( + error.to_string(), + format!( + "a matrix of size {}x2 has a length exceeding `usize::MAX`", + usize::MAX + ) + ); + + // The largest possible byte span is valid without allocating it. + let layout = Layout::::new(isize::MAX as usize, 1).unwrap(); + assert_eq!(layout.num_elements(), isize::MAX as usize); + + let transpose = layout.transpose(); + assert_eq!(transpose.nrows(), 1); + assert_eq!(transpose.ncols(), isize::MAX as usize); + assert_eq!(transpose.num_elements(), layout.num_elements()); + + // One byte beyond the maximum span returns an error. + let error = Layout::::new(isize::MAX as usize + 1, 1).unwrap_err(); + assert_eq!( + error.to_string(), + format!( + "a matrix of size {}x1 with elements of size 1 exceeds `isize::MAX` bytes", + isize::MAX as usize + 1 + ) + ); + + // Rebinding to a larger element type revalidates the byte span. + let rebound = Layout::::new(3, 4).unwrap().rebind::().unwrap(); + assert_eq!(rebound.nrows(), 3); + assert_eq!(rebound.ncols(), 4); + assert_eq!(rebound.num_elements(), 12); + + let error = layout.rebind::().unwrap_err(); + assert_eq!( + error.to_string(), + format!( + "a matrix of size {}x1 with elements of size 2 exceeds `isize::MAX` bytes", + isize::MAX + ) + ); + } + + //////////// + // Matrix // + //////////// + + #[test] + fn test_sizes() { + let expected = 3 * std::mem::size_of::(); + assert_eq!(std::mem::size_of::>(), expected); + assert_eq!(std::mem::size_of::>>(), expected); + + assert_eq!(std::mem::size_of::>(), expected); + assert_eq!(std::mem::size_of::>>(), expected); + + assert_eq!(std::mem::size_of::>(), expected); + assert_eq!(std::mem::size_of::>>(), expected); + } + + #[test] + fn fallible_matrix_constructors() { + let err = Owned::try_from_element(usize::MAX, usize::MAX, 0u32).unwrap_err(); + let msg = err.to_string(); + assert_contains!(msg, "exceeding `usize::MAX`"); + + let err = Owned::try_from_element(isize::MAX as usize, 1, 0u32).unwrap_err(); + let msg = err.to_string(); + assert_contains!(msg, "exceeds `isize::MAX` bytes"); + + // Panicking + let err = std::panic::catch_unwind(|| { + Owned::from_element(usize::MAX, usize::MAX, 0u32); + }) + .unwrap_err() + .downcast::() + .unwrap(); + + let msg = err.to_string(); + assert_contains!(msg, "exceeding `usize::MAX`"); + + let err = std::panic::catch_unwind(|| { + Owned::from_element(isize::MAX as usize, 1, 0u32); + }) + .unwrap_err() + .downcast::() + .unwrap(); + let msg = err.to_string(); + assert_contains!(msg, "exceeds `isize::MAX` bytes"); + + // Construction fails without invoking the generator. + let err = Owned::try_from_fn(usize::MAX, usize::MAX, |_| panic!("boom")).unwrap_err(); + let msg = err.to_string(); + assert_contains!(msg, "exceeding `usize::MAX`"); + + let err = std::panic::catch_unwind(|| { + Owned::from_fn(usize::MAX, usize::MAX, |_| { + panic!("initializer must not run") + }); + }) + .unwrap_err() + .downcast::() + .unwrap(); + let msg = err.to_string(); + assert_contains!(msg, "Owned::from_fn failed"); + assert_contains!(msg, "exceeding `usize::MAX`"); + } + + fn make_test_matrix() -> Vec { + // Construct a matrix with 4 rows of length 3. + // The expected layout is as follows: + // + // 0, 1, 2, + // 1, 2, 3, + // 2, 3, 4, + // 3, 4, 5 + // + vec![0, 1, 2, 1, 2, 3, 2, 3, 4, 3, 4, 5] + } + + fn striped_matrix(nrows: usize, ncols: usize) -> Owned { + Owned::from_fn(nrows, ncols, |rc| rc.row * 100 + rc.col) + } + + fn assert_exact_fused(mut iter: I, expected: usize) + where + I: ExactSizeIterator + std::iter::FusedIterator, + { + assert_eq!(iter.len(), expected); + assert_eq!(iter.size_hint(), (expected, Some(expected))); + + for remaining in (0..expected).rev() { + assert!(iter.next().is_some()); + assert_eq!(iter.len(), remaining); + assert_eq!(iter.size_hint(), (remaining, Some(remaining))); + } + + assert!(iter.next().is_none()); + assert!(iter.next().is_none()); + assert_eq!(iter.len(), 0); + assert_eq!(iter.size_hint(), (0, Some(0))); + } + + /// Assert that the collection of rows matches the rows in `m`. + fn assert_rows_match_scalar(m: Ref<'_, usize>, rows: Vec<&[usize]>) { + let context = lazy_format!("nrows = {}, ncols = {}", m.nrows(), m.ncols()); + + assert_eq!(rows.len(), m.nrows(), "{context}"); + + for (row_index, row) in rows.into_iter().enumerate() { + assert_eq!(row.len(), m.ncols(), "row = {row_index} -- {context}"); + for (col_index, value) in row.iter().enumerate() { + assert_eq!( + value, + m.element(row_index, col_index), + "row = {row_index}, col = {col_index} -- {context}" + ); + } + } + } + + fn assert_windows_match_scalar( + m: Ref<'_, usize>, + batchsize: usize, + windows: Vec>, + ) { + let context = lazy_format!( + "nrows = {}, ncols = {}, batchsize = {batchsize}", + m.nrows(), + m.ncols() + ); + + assert_eq!(windows.len(), m.nrows().div_ceil(batchsize), "{context}"); + + for (window_index, window) in windows.into_iter().enumerate() { + let window_context = lazy_format!("window = {window_index} -- {context}"); + let first_row = window_index * batchsize; + let expected_rows = batchsize.min(m.nrows() - first_row); + assert_eq!(window.nrows(), expected_rows, "{window_context}"); + assert_eq!(window.ncols(), m.ncols(), "{window_context}"); + + for row_index in 0..window.nrows() { + for col_index in 0..window.ncols() { + assert_eq!( + window.element(row_index, col_index), + m.element(first_row + row_index, col_index), + "row = {row_index}, col = {col_index} -- {window_context}" + ); + } + } + } + } + + #[cfg(all(not(miri), feature = "rayon"))] + fn assert_parallel_rows_match_scalar(m: Ref<'_, usize>) { + let rows: Vec<_> = m.par_rows().collect(); + assert_rows_match_scalar(m, rows); + } + + #[cfg(all(not(miri), feature = "rayon"))] + fn assert_parallel_windows_match_scalar(m: Ref<'_, usize>, batchsize: usize) { + let windows: Vec<_> = m.par_window_iter(batchsize).collect(); + assert_windows_match_scalar(m, batchsize, windows); + } + + /// This expects `m` to look like the result s of [`make_test_matrix`]. + fn test_basic_indexing(m: &T) + where + T: Matrix + Sync, + { + assert_eq!(m.nrows(), 4); + assert_eq!(m.ncols(), 3); + + // Basic indexing + assert_eq!(*m.element(0, 0), 0); + assert_eq!(*m.element(0, 1), 1); + assert_eq!(*m.element(0, 2), 2); + + assert_eq!(*m.element(1, 0), 1); + assert_eq!(*m.element(1, 1), 2); + assert_eq!(*m.element(1, 2), 3); + + assert_eq!(*m.element(2, 0), 2); + assert_eq!(*m.element(2, 1), 3); + assert_eq!(*m.element(2, 2), 4); + + assert_eq!(*m.element(3, 0), 3); + assert_eq!(*m.element(3, 1), 4); + assert_eq!(*m.element(3, 2), 5); + + assert_eq!(*m.get_element(0, 0).unwrap(), 0); + assert_eq!(*m.get_element(0, 1).unwrap(), 1); + assert_eq!(*m.get_element(0, 2).unwrap(), 2); + + assert_eq!(*m.get_element(1, 0).unwrap(), 1); + assert_eq!(*m.get_element(1, 1).unwrap(), 2); + assert_eq!(*m.get_element(1, 2).unwrap(), 3); + + assert_eq!(*m.get_element(2, 0).unwrap(), 2); + assert_eq!(*m.get_element(2, 1).unwrap(), 3); + assert_eq!(*m.get_element(2, 2).unwrap(), 4); + + assert_eq!(*m.get_element(3, 0).unwrap(), 3); + assert_eq!(*m.get_element(3, 1).unwrap(), 4); + assert_eq!(*m.get_element(3, 2).unwrap(), 5); + + // Row indexing. + assert_eq!(m.row(0), &[0, 1, 2]); + assert_eq!(m.row(1), &[1, 2, 3]); + assert_eq!(m.row(2), &[2, 3, 4]); + assert_eq!(m.row(3), &[3, 4, 5]); + + let rows: Vec> = m.rows().map(|x| x.to_vec()).collect(); + assert_eq!(m.row(0), &rows[0]); + assert_eq!(m.row(1), &rows[1]); + assert_eq!(m.row(2), &rows[2]); + assert_eq!(m.row(3), &rows[3]); + + // Window Iters. + let batchsize = 2; + m.window_iter(NonZeroUsize::new(batchsize).unwrap()) + .enumerate() + .for_each(|(i, submatrix)| { + assert_eq!(submatrix.nrows(), batchsize); + assert_eq!(submatrix.ncols(), m.ncols()); + + // Make sure we are in the correct window of the original matrix. + let base = i * batchsize; + assert_eq!(*submatrix.element(0, 0), base); + assert_eq!(*submatrix.element(0, 1), base + 1); + assert_eq!(*submatrix.element(0, 2), base + 2); + + assert_eq!(*submatrix.element(1, 0), base + 1); + assert_eq!(*submatrix.element(1, 1), base + 2); + assert_eq!(*submatrix.element(1, 2), base + 3); + }); + + // Try again, but with a batch size of 3 to ensure that we correctly handle cases + // where the last block is under-sized. + let batchsize = 3; + m.window_iter(NonZeroUsize::new(batchsize).unwrap()) + .enumerate() + .for_each(|(i, submatrix)| { + if i == 0 { + assert_eq!(submatrix.nrows(), batchsize); + assert_eq!(submatrix.ncols(), m.ncols()); + + // Check indexing + assert_eq!(*submatrix.element(0, 0), 0); + assert_eq!(*submatrix.element(0, 1), 1); + assert_eq!(*submatrix.element(0, 2), 2); + + assert_eq!(*submatrix.element(1, 0), 1); + assert_eq!(*submatrix.element(1, 1), 2); + assert_eq!(*submatrix.element(1, 2), 3); + + assert_eq!(*submatrix.element(2, 0), 2); + assert_eq!(*submatrix.element(2, 1), 3); + assert_eq!(*submatrix.element(2, 2), 4); + } else { + assert_eq!(submatrix.nrows(), 1); + assert_eq!(submatrix.ncols(), m.ncols()); + + // Check indexing + assert_eq!(*submatrix.element(0, 0), 3); + assert_eq!(*submatrix.element(0, 1), 4); + assert_eq!(*submatrix.element(0, 2), 5); + } + }); + } + + #[test] + fn matrix_happy_path() { + let data = make_test_matrix(); + let m = Owned::try_from_data(data.into(), 4, 3).unwrap(); + test_basic_indexing(&m); + + // Get the base pointer of the matrix and make sure view-conversion preserves this + // value. + let ptr = m.as_ptr(); + let view = m.as_view(); + assert!(is_copyable(view)); + assert_eq!(view.as_ptr(), ptr); + assert_eq!(view.nrows(), m.nrows()); + assert_eq!(view.ncols(), m.ncols()); + test_basic_indexing(&view); + } + + #[test] + fn matrix_try_from_construction_error() { + let data = make_test_matrix(); + let ptr = data.as_ptr(); + let len = data.len(); + + let m = Owned::try_from_data(data.into(), 5, 4); + assert!(m.is_err()); + let err = m.unwrap_err(); + assert_eq!( + err.to_string(), + "tried to construct a 5x4 matrix over a span of length 12" + ); + + // Make sure that we can retrieve the original allocation from the interior. + let data = err.into_inner(); + assert_eq!(data.as_ptr(), ptr); + assert_eq!(data.len(), len); + + let m = Ref::try_from_data(&data, 5, 4); + assert!(m.is_err()); + assert_eq!( + m.unwrap_err().to_string(), + "tried to construct a 5x4 matrix over a span of length 12" + ); + } + + #[test] + fn mutable_matrix_direct_construction() { + let mut data = make_test_matrix(); + + { + let mut m = Mut::try_from_data(data.as_mut_slice(), 4, 3).unwrap(); + *m.element_mut(1, 2) = 30; + } + assert_eq!(data[5], 30); + + let err = Mut::try_from_data(data.as_mut_slice(), 5, 4).unwrap_err(); + assert_eq!( + err.to_string(), + "tried to construct a 5x4 matrix over a span of length 12" + ); + let recovered = err.into_inner(); + recovered[0] = 10; + assert_eq!(data[0], 10); + } + + #[test] + fn matrix_mut_view() { + let mut m = Owned::::from_element(4, 3, 0); + assert_eq!(m.nrows(), 4); + assert_eq!(m.ncols(), 3); + assert!(m.as_slice().iter().all(|&i| i == 0)); + let ptr = m.as_ptr(); + let mut_ptr = m.as_mut_ptr(); + assert_eq!(ptr, mut_ptr); + + let mut view = m.as_view_mut(); + assert_eq!(view.nrows(), 4); + assert_eq!(view.ncols(), 3); + assert_eq!(view.as_ptr(), ptr); + assert_eq!(view.as_mut_ptr(), mut_ptr); + + // Construct the test matrix manually. + for i in 0..view.nrows() { + for j in 0..view.ncols() { + *view.element_mut(i, j) = i + j; + } + } + + // Drop the view and test the original matrix. + test_basic_indexing(&m); + + // Check `PartialEq`. + let mut m_clone = m.clone(); + assert_eq!(m.as_view_mut(), m_clone.as_view_mut()); + + let inner = m.into_inner(); + assert_eq!(inner.as_ptr(), ptr); + assert_eq!(inner.len(), 4 * 3); + } + + #[test] + fn matrix_view_zero_sizes() { + let data: Vec = vec![]; + // Zero rows, but non-zero columns. + let m = Ref::try_from_data(data.as_slice(), 0, 10).unwrap(); + assert_eq!(m.nrows(), 0); + assert_eq!(m.ncols(), 10); + + // Non-zero rows, but zero columns. + let m = Ref::try_from_data(data.as_slice(), 3, 0).unwrap(); + assert_eq!(m.nrows(), 3); + assert_eq!(m.ncols(), 0); + let empty: &[usize] = &[]; + assert_eq!(m.row(0), empty); + assert_eq!(m.row(1), empty); + assert_eq!(m.row(2), empty); + + // Zero rows and columns. + let m = Ref::try_from_data(data.as_slice(), 0, 0).unwrap(); + assert_eq!(m.nrows(), 0); + assert_eq!(m.ncols(), 0); + } + + #[test] + fn matrix_construction_by_row() { + let mut m = Owned::::from_element(4, 3, 0); + assert!(m.as_slice().iter().all(|i| *i == 0)); + + let ncols = m.ncols(); + for i in 0..m.nrows() { + let row = m.row_mut(i); + assert_eq!(row.len(), ncols); + row[0] = i; + row[1] = i + 1; + row[2] = i + 2; + } + test_basic_indexing(&m); + } + + // Panics + #[test] + #[should_panic(expected = "tried to access row 3 of a matrix with 3 rows")] + fn test_get_row_panics() { + let m = Owned::::from_element(3, 7, 0); + m.row(3); + } + + #[test] + #[should_panic(expected = "tried to access row 3 of a matrix with 3 rows")] + fn test_get_row_mut_panics() { + let mut m = Owned::::from_element(3, 7, 0); + m.row_mut(3); + } + + #[test] + #[should_panic(expected = "row 3 is out of bounds (max: 3)")] + fn test_element_panics_row() { + let m = Owned::::from_element(3, 7, 0); + assert!(m.get_element(3, 2).is_none()); + let _ = m.element(3, 2); + } + + #[test] + #[should_panic(expected = "col 7 is out of bounds (max: 7)")] + fn test_element_panics_col() { + let m = Owned::::from_element(3, 7, 0); + assert!(m.get_element(2, 7).is_none()); + let _ = m.element(2, 7); + } + + #[test] + #[should_panic(expected = "row 3 is out of bounds (max: 3)")] + fn test_element_mut_panics_row() { + let mut m = Owned::::from_element(3, 7, 0); + assert!(m.get_element_mut(3, 2).is_none()); + *m.element_mut(3, 2) = 1; + } + + #[test] + #[should_panic(expected = "col 7 is out of bounds (max: 7)")] + fn test_element_mut_panics_col() { + let mut m = Owned::::from_element(3, 7, 0); + assert!(m.get_element_mut(2, 7).is_none()); + *m.element_mut(2, 7) = 1; + } + + #[test] + #[cfg(feature = "rayon")] + #[should_panic(expected = "par_window_iter batchsize cannot be zero")] + fn test_par_window_iter_panics() { + let m = Owned::::from_element(4, 4, 0); + let _ = m.par_window_iter(0); + } + + #[test] + #[cfg(feature = "rayon")] + #[should_panic(expected = "par_window_iter_mut batchsize cannot be zero")] + fn test_par_window_iter_mut_panics() { + let mut m = Owned::::from_element(4, 4, 0); + let _ = m.par_window_iter_mut(0); + } + + // Additional tests for better coverage + + #[test] + fn test_try_from_error_light() { + // Incorrect slice + let data = vec![1, 2, 3]; + let err = Ref::try_from_data(data.as_slice(), 2, 3).unwrap_err(); + + // Test `as_static` method + let err_static = err.as_static(); + let msg = err_static.to_string(); + assert_contains!( + msg, + "tried to construct a 2x3 matrix over a span of length 3", + ); + // Test `into_inner` method + let recovered_data = err.into_inner(); + assert_eq!(recovered_data, data.as_slice()); + + // Invalid length. + let err = Ref::try_from_data(data.as_slice(), 2, usize::MAX).unwrap_err(); + let msg = err.to_string(); + assert_contains!(msg, "usize::MAX"); + + assert_eq!(data.as_slice(), err.into_inner()); + } + + #[test] + fn test_map_errors() { + #[derive(Debug, Clone, Copy)] + struct Zst; + + // Create a large ZST slice without taking forever on debug builds. + let b = Box::<[Zst]>::new_uninit_slice((isize::MAX as usize) + 1); + + // SAFETY: `b` has zero-sized elements, so all elements are initialized. + let b = unsafe { b.assume_init() }; + + let m = Owned::column_vector(b); + let err = m.try_map(|_: &Zst| 0u8).unwrap_err(); + let msg = err.to_string(); + assert!(msg.contains("isize::MAX"), "{msg}"); + + // Panicking variant. + let err = std::panic::catch_unwind(|| m.map(|_: &Zst| 0u8)) + .unwrap_err() + .downcast::() + .unwrap(); + let msg = err.to_string(); + assert!(msg.contains("isize::MAX"), "{msg}"); + } + + #[test] + fn test_get_row_optional() { + let data = make_test_matrix(); + let mut m = Owned::try_from_data(data.into(), 4, 3).unwrap(); + + assert_eq!(m.get_row(0), Some(&[0, 1, 2][..])); + assert_eq!(m.get_row(1), Some(&[1, 2, 3][..])); + assert_eq!(m.get_row(3), Some(&[3, 4, 5][..])); + assert_eq!(m.get_row(4), None); + assert_eq!(m.get_row(100), None); + + let row = m.get_row_mut(1).unwrap(); + assert_eq!(row, &[1, 2, 3]); + row[0] = 10; + assert_eq!(m.row(1), &[10, 2, 3]); + assert!(m.get_row_mut(4).is_none()); + assert!(m.get_row_mut(100).is_none()); + } + + #[test] + fn test_unsafe_get_unchecked_methods() { + let data = make_test_matrix(); + let mut m = Owned::try_from_data(data.into(), 4, 3).unwrap(); + + // Safety: derives from known size of matrix and access element ids + unsafe { + assert_eq!(*m.element_unchecked(0, 0), 0); + assert_eq!(*m.element_unchecked(1, 2), 3); + assert_eq!(*m.element_unchecked(3, 1), 4); + } + + // Safety: derives from known size of matrix and access element ids + unsafe { + *m.element_unchecked_mut(0, 0) = 100; + *m.element_unchecked_mut(1, 2) = 200; + } + + assert_eq!(*m.element(0, 0), 100); + assert_eq!(*m.element(1, 2), 200); + + // Safety: derives from known size of matrix and access element ids + unsafe { + let row0 = m.row_unchecked(0); + assert_eq!(row0[0], 100); + assert_eq!(row0[1], 1); + assert_eq!(row0[2], 2); + } + + // Safety: derives from known size of matrix and access element ids + unsafe { + let row1 = m.row_unchecked_mut(1); + row1[0] = 300; + } + + assert_eq!(*m.element(1, 0), 300); + } + + #[test] + fn test_to_owned() { + let data = make_test_matrix(); + let view = Ref::try_from_data(data.as_slice(), 4, 3).unwrap(); + + // Test to_owned creates a proper clone + let owned: Owned<_> = view.to_rowmajor_owned(); + assert_eq!(owned.nrows(), view.nrows()); + assert_eq!(owned.ncols(), view.ncols()); + assert_eq!(owned.as_slice(), view.as_slice()); + + // Verify it's actually owned (different memory location) + assert_ne!(owned.as_ptr(), view.as_ptr()); + + // Test the owned matrix works properly + test_basic_indexing(&owned); + } + + #[test] + fn test_matrix_from_conversions() { + let data = make_test_matrix(); + let m = Owned::try_from_data(data.into(), 4, 3).unwrap(); + + // Test Ref to slice conversion + let view = m.as_view(); + let slice: &[usize] = view.into_slice(); + assert_eq!(slice.len(), 12); + assert_eq!(slice[0], 0); + assert_eq!(slice[11], 5); + + // Test Mut to mutable slice conversion. + let data2 = make_test_matrix(); + let mut m2 = Owned::try_from_data(data2.into(), 4, 3).unwrap(); + let mut_view = m2.as_view_mut(); + let slice2: &mut [usize] = mut_view.into_mut_slice(); + assert_eq!(slice2.len(), 12); + assert_eq!(slice2[0], 0); + assert_eq!(slice2[11], 5); + slice2[11] = 6; + assert_eq!(*m2.element(3, 2), 6); + } + + #[test] + fn test_row_vector() { + let data = vec![1, 2, 3]; + let m = Ref::row_vector(data.as_slice()); + assert_eq!(m.nrows(), 1); + assert_eq!(m.ncols(), 3); + assert_eq!(m.as_slice(), &[1, 2, 3]); + assert_eq!(m.row(0), &[1, 2, 3]); + + // Empty + let empty: &[i32] = &[]; + let m = Ref::row_vector(empty); + assert_eq!(m.nrows(), 1); + assert_eq!(m.ncols(), 0); + + // Owned + let m = Owned::row_vector(vec![10u64, 20].into_boxed_slice()); + assert_eq!(m.nrows(), 1); + assert_eq!(m.ncols(), 2); + assert_eq!(*m.element(0, 0), 10); + assert_eq!(*m.element(0, 1), 20); + } + + #[test] + fn test_column_vector() { + let data = vec![1, 2, 3]; + let m = Ref::column_vector(data.as_slice()); + assert_eq!(m.nrows(), 3); + assert_eq!(m.ncols(), 1); + assert_eq!(m.as_slice(), &[1, 2, 3]); + assert_eq!(*m.element(0, 0), 1); + assert_eq!(*m.element(1, 0), 2); + assert_eq!(*m.element(2, 0), 3); + assert_eq!(m.row(0), &[1]); + assert_eq!(m.row(1), &[2]); + assert_eq!(m.row(2), &[3]); + + // Empty + let empty: &[i32] = &[]; + let m = Ref::column_vector(empty); + assert_eq!(m.nrows(), 0); + assert_eq!(m.ncols(), 1); + + // Owned + let m = Owned::column_vector(vec![10u64, 20].into_boxed_slice()); + assert_eq!(m.nrows(), 2); + assert_eq!(m.ncols(), 1); + assert_eq!(*m.element(0, 0), 10); + assert_eq!(*m.element(1, 0), 20); + } + + #[test] + fn test_map() { + let m = Owned::try_from_data(vec![1u32, 2, 3, 4].into(), 2, 2).unwrap(); + let doubled = m.map(|&x| x * 2); + assert_eq!(doubled.as_slice(), &[2, 4, 6, 8]); + assert_eq!(doubled.nrows(), 2); + assert_eq!(doubled.ncols(), 2); + + // Type-changing map + let as_f64 = m.map(|&x| x as f64); + assert_eq!(as_f64.as_slice(), &[1.0, 2.0, 3.0, 4.0]); + } + + #[test] + fn test_get_element() { + let mut m = Owned::try_from_data(vec![1, 2, 3, 4, 5, 6].into(), 2, 3).unwrap(); + assert_eq!(m.get_element(0, 0), Some(&1)); + assert_eq!(m.get_element(1, 2), Some(&6)); + assert_eq!(m.get_element(2, 0), None); + assert_eq!(m.get_element(0, 3), None); + + *m.get_element_mut(1, 2).unwrap() = 7; + assert_eq!(m.get_element(1, 2), Some(&7)); + assert_eq!(m.get_element_mut(2, 0), None); + assert_eq!(m.get_element_mut(0, 3), None); + } + + #[test] + fn test_subview() { + let data = make_test_matrix(); + let m = Owned::try_from_data(data.into(), 4, 3).unwrap(); + + // Create a subview of the first two rows + { + let subview = m.subview(0..4).unwrap(); + assert_eq!(subview.nrows(), 4); + assert_eq!(subview.ncols(), 3); + + assert_eq!(subview.row(0), &[0, 1, 2]); + assert_eq!(subview.row(1), &[1, 2, 3]); + assert_eq!(subview.row(2), &[2, 3, 4]); + assert_eq!(subview.row(3), &[3, 4, 5]); + assert!(subview.get_row(4).is_none()); + } + + // Sub view over a subset that touches the end. + { + let subview = m.subview(1..4).unwrap(); + assert_eq!(subview.nrows(), 3); + assert_eq!(subview.ncols(), 3); + + assert_eq!(subview.row(0), &[1, 2, 3]); + assert_eq!(subview.row(1), &[2, 3, 4]); + assert_eq!(subview.row(2), &[3, 4, 5]); + assert!(subview.get_row(3).is_none()); + } + + // Empty subview in the middle. + { + let subview = m.subview(2..2).unwrap(); + assert_eq!(subview.nrows(), 0); + assert_eq!(subview.ncols(), 3); + } + + // Empty subviews at both boundaries. + { + let subview = m.subview(0..0).unwrap(); + assert_eq!(subview.nrows(), 0); + assert_eq!(subview.ncols(), 3); + + let subview = m.subview(4..4).unwrap(); + assert_eq!(subview.nrows(), 0); + assert_eq!(subview.ncols(), 3); + } + + // Empty out-of-bounds subview + assert!(m.subview(5..5).is_none()); + + // End is out of bounds. + assert!(m.subview(2..10).is_none()); + + // Reversed bounds. + #[expect( + clippy::reversed_empty_ranges, + reason = "we want to make sure it doesn't work" + )] + let empty = 3..2; + assert!(m.subview(empty).is_none()); + + // Extreme out-of-bounds values do not overflow internal calculations. + assert!(m.subview(usize::MAX - 1..usize::MAX).is_none()); + } + + #[expect( + clippy::reversed_empty_ranges, + reason = "we want to make sure it doesn't work" + )] + #[test] + fn test_subview_zero_cols() { + let m = Owned::from_element(10, 0, 0u32); + + // A fully disjoint range is rejected. + assert!(m.subview(100..200).is_none()); + + // A range extending one row beyond the matrix is rejected. + assert!(m.subview(10..11).is_none()); + + // Reversed bounds are rejected. + assert!(m.subview(5..4).is_none()); + + // An in-bounds suffix preserves its logical shape. + let v = m.subview(5..10).unwrap(); + assert_eq!(v.nrows(), 5); + assert_eq!(v.ncols(), 0); + + // An empty range at the end is valid. + let v = m.subview(10..10).unwrap(); + assert_eq!(v.nrows(), 0); + assert_eq!(v.ncols(), 0); + + // The full range preserves all logical rows despite having no elements. + let v = m.subview(0..10).unwrap(); + assert_eq!(v.nrows(), 10); + assert_eq!(v.ncols(), 0); + } + + #[test] + #[cfg(all(not(miri), feature = "rayon"))] + fn parallel_immutable_iterators_match_scalar_indexing() { + for (nrows, ncols) in [(0, 0), (0, 4), (3, 0), (1, 1), (1, 4), (4, 1), (5, 3)] { + let m = striped_matrix(nrows, ncols); + let view = m.as_view(); + + assert_parallel_rows_match_scalar(view); + for batchsize in [1, 2, 3, usize::MAX] { + assert_parallel_windows_match_scalar(view, batchsize); + } + } + } + + #[test] + #[cfg(all(not(miri), feature = "rayon"))] + fn parallel_mutable_iterators_match_scalar_indexing() { + use rayon::prelude::*; + + for (nrows, ncols) in [(0, 0), (0, 4), (1, 1), (1, 4), (4, 1), (5, 3)] { + let context = lazy_format!("nrows = {nrows}, ncols = {ncols}"); + + let original = striped_matrix(nrows, ncols); + let mut rows = original.clone(); + + let row_views: Vec<_> = rows.par_rows_mut().collect(); + assert_eq!(row_views.len(), nrows, "{context}"); + for (row_index, row) in row_views.into_iter().enumerate() { + assert_eq!(row.len(), ncols, "row = {row_index} -- {context}"); + for value in row { + *value = value.wrapping_add(row_index); + } + } + + for row in 0..nrows { + for col in 0..ncols { + assert_eq!( + *rows.element(row, col), + original.element(row, col).wrapping_add(row), + "row = {}, col = {} -- {}", + row, + col, + context, + ); + } + } + + for batchsize in [1, 2, 3, usize::MAX] { + let context = lazy_format!("{context}, batchsize = {batchsize}"); + + let mut windows = original.clone(); + let window_views: Vec<_> = windows.par_window_iter_mut(batchsize).collect(); + assert_eq!(window_views.len(), nrows.div_ceil(batchsize), "{context}"); + for (window_index, mut window) in window_views.into_iter().enumerate() { + let first_row = window_index * batchsize; + assert_eq!( + window.nrows(), + batchsize.min(nrows - first_row), + "window = {window_index} -- {context}" + ); + assert_eq!( + window.ncols(), + ncols, + "window = {window_index} -- {context}" + ); + for value in window.as_mut_slice() { + *value = value.wrapping_add(window_index); + } + } + + for row in 0..nrows { + for col in 0..ncols { + assert_eq!( + *windows.element(row, col), + original.element(row, col).wrapping_add(row / batchsize), + "row = {}, col = {} -- {}", + row, + col, + context, + ); + } + } + } + } + } + + #[test] + #[cfg(feature = "rayon")] + #[should_panic( + expected = "`MatrixMut::par_rows_mut` does not support matrices with rows and zero columns" + )] + fn par_rows_mut_rejects_nonempty_zero_column_matrix() { + let mut m = striped_matrix(3, 0); + let _ = m.par_rows_mut(); + } + + #[test] + #[cfg(feature = "rayon")] + #[should_panic( + expected = "`MatrixMut::par_window_iter_mut` does not support matrices with rows and zero columns" + )] + fn par_window_iter_mut_rejects_nonempty_zero_column_matrix() { + let mut m = striped_matrix(3, 0); + let _ = m.par_window_iter_mut(2); + } + + #[test] + fn matrix_iterators_match_scalar_indexing() { + for (nrows, ncols) in [(0, 0), (0, 4), (3, 0), (1, 1), (1, 4), (4, 1), (5, 3)] { + let original = striped_matrix(nrows, ncols); + let view = original.as_view(); + + let rows = view.rows().collect(); + assert_rows_match_scalar(view, rows); + for batchsize in [1, 2, 3, usize::MAX] { + let windows = view + .window_iter(NonZeroUsize::new(batchsize).unwrap()) + .collect(); + assert_windows_match_scalar(view, batchsize, windows); + } + + let context = lazy_format!("nrows = {nrows}, ncols = {ncols}"); + let mut mutable = original.clone(); + let rows: Vec<_> = mutable.rows_mut().collect(); + assert_eq!(rows.len(), nrows, "{context}"); + + for (row_index, row) in rows.into_iter().enumerate() { + assert_eq!(row.len(), ncols, "row = {row_index} -- {context}"); + for (col_index, value) in row.iter_mut().enumerate() { + assert_eq!( + *value, + *original.element(row_index, col_index), + "row = {row_index}, col = {col_index} -- {context}" + ); + *value = value.wrapping_add(1); + } + } + + for row in 0..nrows { + for col in 0..ncols { + assert_eq!( + *mutable.element(row, col), + original.element(row, col).wrapping_add(1), + "row = {row}, col = {col} -- {context}" + ); + } + } + } + } + + #[test] + fn matrix_iterators_track_exact_remaining_lengths() { + for (nrows, ncols) in [(0, 0), (0, 4), (3, 0), (5, 3)] { + let mut m = striped_matrix(nrows, ncols); + + assert_exact_fused(m.rows(), nrows); + assert_exact_fused(m.rows_mut(), nrows); + + for batchsize in [1, 2, usize::MAX] { + assert_exact_fused( + m.window_iter(NonZeroUsize::new(batchsize).unwrap()), + nrows.div_ceil(batchsize), + ); + } + } + } + + #[test] + fn matrix_transformations_preserve_empty_shapes() { + for (nrows, ncols) in [(0, 0), (0, 4), (3, 0)] { + let m = striped_matrix(nrows, ncols); + + let owned = m.to_rowmajor_owned(); + assert_eq!(owned.nrows(), nrows); + assert_eq!(owned.ncols(), ncols); + + let mapped = m.map(|_| -> u8 { unreachable!("empty matrix has no elements") }); + assert_eq!(mapped.nrows(), nrows); + assert_eq!(mapped.ncols(), ncols); + } + } + + #[test] + fn owned_from_fn_initializes_in_memory_order() { + let mut value = 0; + let m = Owned::from_fn(2, 3, |_| { + let result = value; + value += 1; + result + }); + + assert_eq!(*m.element(0, 0), 0); + assert_eq!(*m.element(0, 1), 1); + assert_eq!(*m.element(0, 2), 2); + assert_eq!(*m.element(1, 0), 3); + assert_eq!(*m.element(1, 1), 4); + assert_eq!(*m.element(1, 2), 5); + } + + #[test] + fn test_transpose() { + { + let v = Owned::from_element(0, 0, 0); + let t = v.transpose(); + assert_eq!(t.nrows(), 0); + assert_eq!(t.ncols(), 0); + } + + { + let v = Owned::from_element(0, 10, 0); + let t = v.transpose(); + assert_eq!(t.nrows(), 10); + assert_eq!(t.ncols(), 0); + } + + { + let v = Owned::from_element(10, 0, 0); + let t = v.transpose(); + assert_eq!(t.nrows(), 0); + assert_eq!(t.ncols(), 10); + } + + { + let v = Owned::::try_from_data(Box::new([1, 2, 3, 4, 5, 6]), 2, 3).unwrap(); + let t = v.transpose(); + + assert_eq!(t.row(0), &[1, 4]); + assert_eq!(t.row(1), &[2, 5]); + assert_eq!(t.row(2), &[3, 6]); + } + } + + #[test] + fn test_debug_error_formatting() { + // Test Debug implementation for TryFromError + let data = vec![1, 2, 3]; + let err = Owned::try_from_data(data.into(), 2, 3).unwrap_err(); + let debug_str = format!("{:?}", err); + assert_contains!(debug_str, "TryFromError"); + + // Ensure Debug doesn't require T: Debug by using a non-Debug type + #[derive(Clone)] + struct NonDebug(#[expect(dead_code)] i32); + + let non_debug_data: Box<[NonDebug]> = vec![NonDebug(1), NonDebug(2)].into(); + let non_debug_err = match Owned::try_from_data(non_debug_data, 1, 3) { + Ok(_) => panic!("should not have succeeded!"), + Err(err) => err, + }; + let debug_str = format!("{:?}", non_debug_err); + assert_contains!(debug_str, "TryFromError"); + } +} diff --git a/diskann-utils/src/views/rowmajor/iter.rs b/diskann-utils/src/views/rowmajor/iter.rs new file mode 100644 index 0000000000..91a765e551 --- /dev/null +++ b/diskann-utils/src/views/rowmajor/iter.rs @@ -0,0 +1,189 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use std::{marker::PhantomData, num::NonZeroUsize, ptr::NonNull}; + +use crate::views::rowmajor::{Layout, Matrix, Mut, Ref}; + +//------// +// Rows // +//------// + +/// An iterator over rows in a matrix. See: [`Matrix::rows`]. +#[derive(Debug)] +pub struct Rows<'a, T> { + ptr: NonNull, + remaining: usize, + ncols: usize, + _lifetime: PhantomData<&'a [T]>, +} + +impl<'a, T> Rows<'a, T> { + pub(super) fn new(m: Ref<'a, T>) -> Self { + let layout = m.layout(); + Self { + ptr: m.as_nonnull(), + remaining: layout.nrows(), + ncols: layout.ncols(), + _lifetime: PhantomData, + } + } +} + +// SAFETY: `Rows<'_, T>` owns a shared slice borrow, so sending it requires `T: Sync`. +unsafe impl Send for Rows<'_, T> where T: Sync {} +// SAFETY: Shared access to `Rows<'_, T>` exposes only shared access to `T`. +unsafe impl Sync for Rows<'_, T> where T: Sync {} + +impl<'a, T> Iterator for Rows<'a, T> { + type Item = &'a [T]; + fn next(&mut self) -> Option<&'a [T]> { + self.remaining.checked_sub(1).map(|remaining| { + // SAFETY: Construction from a valid `Ref` guarantees that each remaining row + // contains `ncols` initialized elements beginning at `self.ptr`. + let item = + unsafe { std::slice::from_raw_parts(self.ptr.as_ptr().cast_const(), self.ncols) }; + self.remaining = remaining; + + // SAFETY: Advancing by one row remains within or one past the original matrix + // span. The validated parent layout guarantees that the offset is representable. + self.ptr = unsafe { self.ptr.add(self.ncols) }; + item + }) + } + + fn size_hint(&self) -> (usize, Option) { + (self.remaining, Some(self.remaining)) + } +} + +impl ExactSizeIterator for Rows<'_, T> {} +impl std::iter::FusedIterator for Rows<'_, T> {} + +//---------// +// RowsMut // +//---------// + +/// An iterator over mutable rows in a matrix. See: [`crate::views::rowmajor::MatrixMut::rows_mut`]. +#[derive(Debug)] +pub struct RowsMut<'a, T> { + ptr: NonNull, + remaining: usize, + ncols: usize, + _lifetime: PhantomData<&'a mut [T]>, +} + +impl<'a, T> RowsMut<'a, T> { + pub(super) fn new(m: Mut<'a, T>) -> Self { + let layout = m.layout(); + Self { + ptr: m.as_nonnull(), + remaining: layout.nrows(), + ncols: layout.ncols(), + _lifetime: PhantomData, + } + } +} + +// SAFETY: `RowsMut<'_, T>` owns an exclusive slice borrow, so sending it requires `T: Send`. +unsafe impl Send for RowsMut<'_, T> where T: Send {} +// SAFETY: Shared access to `RowsMut<'_, T>` exposes only shared access to `T`. +unsafe impl Sync for RowsMut<'_, T> where T: Sync {} + +impl<'a, T> Iterator for RowsMut<'a, T> { + type Item = &'a mut [T]; + fn next(&mut self) -> Option<&'a mut [T]> { + self.remaining.checked_sub(1).map(|remaining| { + // SAFETY: Construction from a valid `Mut` guarantees that each remaining row + // contains `ncols` initialized elements beginning at `self.ptr`. Advancing the + // pointer after every yield makes nonempty returned rows disjoint; zero-length + // rows do not access memory and may share an address. + let item = unsafe { std::slice::from_raw_parts_mut(self.ptr.as_ptr(), self.ncols) }; + self.remaining = remaining; + + // SAFETY: Advancing by one row remains within or one past the original matrix + // span. The validated parent layout guarantees that the offset is representable. + self.ptr = unsafe { self.ptr.add(self.ncols) }; + item + }) + } + + fn size_hint(&self) -> (usize, Option) { + (self.remaining, Some(self.remaining)) + } +} + +impl ExactSizeIterator for RowsMut<'_, T> {} +impl std::iter::FusedIterator for RowsMut<'_, T> {} + +//---------// +// Windows // +//---------// + +/// An iterator over rows in a matrix. See: [`Matrix::window_iter`]. +#[derive(Debug)] +pub struct Windows<'a, T> { + ptr: NonNull, + remaining: usize, + batchsize: NonZeroUsize, + ncols: usize, + _lifetime: PhantomData<&'a [T]>, +} + +impl<'a, T> Windows<'a, T> { + pub(super) fn new(m: Ref<'a, T>, batchsize: NonZeroUsize) -> Self { + let layout = m.layout(); + Self { + ptr: m.as_nonnull(), + remaining: layout.nrows(), + batchsize, + ncols: layout.ncols(), + _lifetime: PhantomData, + } + } +} + +// SAFETY: `Windows<'_, T>` owns a shared slice borrow, so sending it requires `T: Sync`. +unsafe impl Send for Windows<'_, T> where T: Sync {} +// SAFETY: Shared access to `Windows<'_, T>` exposes only shared access to `T`. +unsafe impl Sync for Windows<'_, T> where T: Sync {} + +impl<'a, T> Iterator for Windows<'a, T> { + type Item = Ref<'a, T>; + fn next(&mut self) -> Option> { + if self.remaining == 0 { + None + } else { + let next_remaining = self.remaining.saturating_sub(self.batchsize.get()); + let nrows = self.remaining - next_remaining; + + // SAFETY: `self.ptr` starts the remaining suffix of a valid `Ref`, and `nrows` + // does not exceed that suffix. Keeping the parent's column count therefore + // produces a valid subview and a layout no larger than the parent layout. + let window = unsafe { + Ref { + ptr: self.ptr, + layout: Layout::new_unchecked(nrows, self.ncols), + _lifetime: PhantomData, + } + }; + + // SAFETY: Advancing by the yielded window remains within or one past the + // original matrix span. The validated parent layout guarantees that the + // multiplication and pointer offset are representable. + self.ptr = unsafe { self.ptr.add(nrows * self.ncols) }; + self.remaining = next_remaining; + Some(window) + } + } + + fn size_hint(&self) -> (usize, Option) { + let remaining = self.remaining.div_ceil(self.batchsize.get()); + (remaining, Some(remaining)) + } +} + +impl ExactSizeIterator for Windows<'_, T> {} +impl std::iter::FusedIterator for Windows<'_, T> {} diff --git a/diskann/src/error/ann_error.rs b/diskann/src/error/ann_error.rs index d810f8e8a6..a16a83c7f8 100644 --- a/diskann/src/error/ann_error.rs +++ b/diskann/src/error/ann_error.rs @@ -262,8 +262,8 @@ convert_error!(TryFromIntError); convert_error!(TryFromSliceError); convert_error!(diskann_utils::io::ReadBinError); convert_error!(diskann_utils::io::SaveBinError); -convert_error!(diskann_utils::views::LayoutError); -convert_error!(diskann_utils::views::TryFromErrorLight); +convert_error!(diskann_utils::views::rowmajor::LayoutError); +convert_error!(diskann_utils::views::rowmajor::TryFromErrorLight); // Convert from `mpsc::SendError` to `ANNError` impl From> for ANNError @@ -287,12 +287,9 @@ where } } -impl From> for ANNError -where - T: diskann_utils::views::DenseData, -{ +impl From> for ANNError { #[track_caller] - fn from(err: diskann_utils::views::TryFromError) -> Self { + fn from(err: diskann_utils::views::rowmajor::TryFromError) -> Self { Self::from(err.as_static()) } } diff --git a/diskann/src/flat/test/harness.rs b/diskann/src/flat/test/harness.rs index 097760a0e9..b2b8035efa 100644 --- a/diskann/src/flat/test/harness.rs +++ b/diskann/src/flat/test/harness.rs @@ -10,6 +10,7 @@ use std::{convert::Infallible, num::NonZeroUsize}; +use diskann_utils::views::rowmajor::Matrix; use diskann_vector::{PreprocessedDistanceFunction, distance::Metric}; use crate::{ @@ -210,7 +211,7 @@ pub(crate) fn brute_force_topk( let mut neighbors: Vec> = provider .items() - .row_iter() + .rows() .enumerate() .map(|(id, element)| Neighbor::new(id as u32, computer.evaluate_similarity(element))) .collect(); diff --git a/diskann/src/flat/test/provider.rs b/diskann/src/flat/test/provider.rs index 3a8f6b2900..5212cbf5d2 100644 --- a/diskann/src/flat/test/provider.rs +++ b/diskann/src/flat/test/provider.rs @@ -11,7 +11,10 @@ use std::{ future::Future, }; -use diskann_utils::{future::SendFuture, views::Matrix}; +use diskann_utils::{ + future::SendFuture, + views::rowmajor::{self, Matrix}, +}; use diskann_vector::{PreprocessedDistanceFunction, distance::Metric}; use thiserror::Error; @@ -43,7 +46,7 @@ convert_error!(ProviderError); /// In-memory test provider for flat search. #[derive(Debug)] pub struct Provider { - items: Matrix, + items: rowmajor::Owned, get_element: Counter, } @@ -53,7 +56,7 @@ impl Provider { /// # Errors /// /// Returns an error if the matrix is empty or has zero-width columns. - pub fn new(items: Matrix) -> Result { + pub fn new(items: rowmajor::Owned) -> Result { if items.nrows() == 0 { return Err(ProviderError::Empty); } @@ -97,7 +100,7 @@ impl Provider { } /// Expose the items for brute force. - pub fn items(&self) -> &Matrix { + pub fn items(&self) -> &rowmajor::Owned { &self.items } } @@ -341,7 +344,7 @@ impl DistancesUnordered for Visitor<'_> { F: Send + FnMut(Self::Id, f32), { async move { - for (i, vector) in self.provider.items.row_iter().enumerate() { + for (i, vector) in self.provider.items.rows().enumerate() { let id = i as u32; if let Some(ids) = &self.transient_ids && ids.contains(&id) diff --git a/diskann/src/graph/glue.rs b/diskann/src/graph/glue.rs index e143394aef..e052260566 100644 --- a/diskann/src/graph/glue.rs +++ b/diskann/src/graph/glue.rs @@ -97,7 +97,7 @@ use std::{future::Future, sync::Arc}; -use diskann_utils::{Reborrow, future::SendFuture}; +use diskann_utils::{Reborrow, future::SendFuture, views::rowmajor::Matrix}; use diskann_vector::DistanceFunction; use futures_util::FutureExt; @@ -1031,7 +1031,7 @@ where /// [`Overlay`](crate::graph::workingset::map::Overlay) for a working set seed compatible /// with [`Batch`]. /// -/// The primary implementation of this trait is [`Matrix`](diskann_utils::views::Matrix). +/// The primary implementation of this trait is [`Owned`](diskann_utils::views::rowmajor::Owned). pub trait Batch: Send + Sync + 'static { /// The element type of the batch. type Element<'a>: Copy; @@ -1048,7 +1048,7 @@ pub trait Batch: Send + Sync + 'static { } } -impl Batch for diskann_utils::views::Matrix { +impl Batch for diskann_utils::views::rowmajor::Owned { type Element<'a> = &'a [T]; fn len(&self) -> usize { diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index ff9bec1409..1a6db06b85 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -10,7 +10,7 @@ //! offers it to the nearest sets of both points. use crate::ANNResult; -use diskann_utils::views::{MatrixView, MutMatrixView}; +use diskann_utils::views::rowmajor::{self, Matrix, MatrixMut}; use super::{ check_output_rows, distance_scratch, @@ -42,8 +42,8 @@ pub(super) struct LeafKernelWorkspace { /// distance matrix size overflows `usize`. pub(super) fn select_leaf_neighbors( arch: A, - points: MatrixView<'_, f32>, - output: MutMatrixView<'_, Candidate>, + points: rowmajor::Ref<'_, f32>, + output: rowmajor::Mut<'_, Candidate>, workspace: &mut LeafKernelWorkspace, ) -> ANNResult<()> where @@ -54,7 +54,7 @@ where check_output_rows(point_count, output.nrows())?; let mut distances = distance_scratch(&mut workspace.distance_scratch, point_count, point_count)?; - M::compute_distances(points, distances.as_mut_view())?; + M::compute_distances(points, distances.as_view_mut())?; select_top_k_symmetric( arch, distances.as_view(), @@ -96,7 +96,7 @@ mod tests { ]; let values = test_support::packed_points(&coordinates, 2, scalar_metric == Metric::CosineNormalized); - let points = MatrixView::try_from(values.as_slice(), 5, 2).unwrap(); + let points = rowmajor::Ref::try_from_data(values.as_slice(), 5, 2).unwrap(); // Widths above four exceed the other points in the leaf. for neighbors in [0, 1, 2, 3, 4, 6] { let mut output = vec![Candidate::new(0, -100.0); 5 * neighbors]; @@ -104,7 +104,7 @@ mod tests { select_leaf_neighbors::<_, M>( ARCH, points, - MutMatrixView::try_from(output.as_mut_slice(), 5, neighbors).unwrap(), + rowmajor::Mut::try_from_data(output.as_mut_slice(), 5, neighbors).unwrap(), &mut LeafKernelWorkspace::default(), ) .unwrap_or_else(|error| panic!("neighbors={neighbors}: {error}")); @@ -171,7 +171,8 @@ mod tests { test_support::normalize(&mut values, dimensions); } let points = - MatrixView::try_from(values.as_slice(), point_count, dimensions).unwrap(); + rowmajor::Ref::try_from_data(values.as_slice(), point_count, dimensions) + .unwrap(); let tolerance = test_support::dense_tolerance(metric, dimensions); // The oracle sorts every other point by its scalar distance. let oracle: Vec> = (0..point_count) @@ -202,8 +203,12 @@ mod tests { select_leaf_neighbors::( arch, points, - MutMatrixView::try_from(output.as_mut_slice(), point_count, neighbors) - .unwrap(), + rowmajor::Mut::try_from_data( + output.as_mut_slice(), + point_count, + neighbors, + ) + .unwrap(), &mut LeafKernelWorkspace::default(), ) .unwrap_or_else(|error| panic!("{context}: {error}")); @@ -258,11 +263,11 @@ mod tests { // Grow, shrink, change K, and finish with a singleton wider than its leaf. for (count, neighbors) in [(4, 1), (2, 1), (5, 3), (1, 2)] { output.resize(count * neighbors, Candidate::new(4, -100.0)); - let points = MatrixView::try_from(&values[..count], count, 1).unwrap(); + let points = rowmajor::Ref::try_from_data(&values[..count], count, 1).unwrap(); select_leaf_neighbors::<_, L2>( ARCH, points, - MutMatrixView::try_from(output.as_mut_slice(), count, neighbors).unwrap(), + rowmajor::Mut::try_from_data(output.as_mut_slice(), count, neighbors).unwrap(), &mut workspace, ) .unwrap(); @@ -292,8 +297,8 @@ mod tests { let error = select_leaf_neighbors::<_, L2>( ARCH, - MatrixView::try_from(&values[..], 3, 1).unwrap(), - MutMatrixView::try_from(&mut output[..], 2, 1).unwrap(), + rowmajor::Ref::try_from_data(&values[..], 3, 1).unwrap(), + rowmajor::Mut::try_from_data(&mut output[..], 2, 1).unwrap(), &mut LeafKernelWorkspace::default(), ) .unwrap_err(); @@ -318,8 +323,8 @@ mod tests { select_top_k_symmetric( ARCH, - MatrixView::try_from(&distances[..], 4, 4).unwrap(), - MutMatrixView::try_from(&mut output[..], 4, 2).unwrap(), + rowmajor::Ref::try_from_data(&distances[..], 4, 4).unwrap(), + rowmajor::Mut::try_from_data(&mut output[..], 4, 2).unwrap(), &mut kth_distances, ); diff --git a/diskann/src/graph/pipnn/leaf_metric.rs b/diskann/src/graph/pipnn/leaf_metric.rs index 6d1fce29fc..24ebacc3f3 100644 --- a/diskann/src/graph/pipnn/leaf_metric.rs +++ b/diskann/src/graph/pipnn/leaf_metric.rs @@ -10,7 +10,7 @@ //! triangle, which holds each pair below the diagonal. use crate::{ANNError, ANNResult}; -use diskann_utils::views::{MatrixView, MutMatrixView}; +use diskann_utils::views::rowmajor::{self, Matrix, MatrixMut}; use diskann_vector::{ Norm, norm::{FastL2Norm, FastL2NormSquared}, @@ -31,23 +31,23 @@ pub(super) trait LeafMetric: Send + Sync + 'static { /// /// `storage` has one row and one column per point. fn compute_distances( - points: MatrixView<'_, f32>, - storage: MutMatrixView<'_, f32>, + points: rowmajor::Ref<'_, f32>, + storage: rowmajor::Mut<'_, f32>, ) -> ANNResult<()>; } impl LeafMetric for L2 { fn compute_distances( - points: MatrixView<'_, f32>, - mut storage: MutMatrixView<'_, f32>, + points: rowmajor::Ref<'_, f32>, + mut storage: rowmajor::Mut<'_, f32>, ) -> ANNResult<()> { // The expanded L2 formula is `||x||² + ||y||² - 2(x·y)`. let squared_norms: Vec = points - .row_iter() + .rows() .map(|point| FastL2NormSquared.evaluate(point)) .collect(); // Initialize the norm terms before GEMM adds the dot-product term. - for (source, row) in storage.row_iter_mut().enumerate() { + for (source, row) in storage.rows_mut().enumerate() { let source_norm = squared_norms[source]; for (distance, &target_norm) in row[..=source].iter_mut().zip(&squared_norms) { *distance = source_norm + target_norm; @@ -66,8 +66,8 @@ impl LeafMetric for L2 { impl LeafMetric for Cosine { fn compute_distances( - points: MatrixView<'_, f32>, - mut storage: MutMatrixView<'_, f32>, + points: rowmajor::Ref<'_, f32>, + mut storage: rowmajor::Mut<'_, f32>, ) -> ANNResult<()> { diskann_linalg::sgemm_aat_lower( points.nrows(), @@ -78,11 +78,11 @@ impl LeafMetric for Cosine { ) .map_err(ANNError::new)?; let norms: Vec = points - .row_iter() + .rows() .map(|point| FastL2Norm.evaluate(point)) .collect(); // Convert each lower-triangle dot to the bounded cosine distance. - for (source, row) in storage.row_iter_mut().enumerate() { + for (source, row) in storage.rows_mut().enumerate() { let source_norm = norms[source]; for (distance, &target_norm) in row[..=source].iter_mut().zip(&norms) { *distance = cosine_distance(*distance, source_norm, target_norm); @@ -94,8 +94,8 @@ impl LeafMetric for Cosine { impl LeafMetric for InnerProduct { fn compute_distances( - points: MatrixView<'_, f32>, - mut storage: MutMatrixView<'_, f32>, + points: rowmajor::Ref<'_, f32>, + mut storage: rowmajor::Mut<'_, f32>, ) -> ANNResult<()> { diskann_linalg::sgemm_aat_lower( points.nrows(), @@ -110,14 +110,14 @@ impl LeafMetric for InnerProduct { impl LeafMetric for CosineNormalized { fn compute_distances( - points: MatrixView<'_, f32>, - mut storage: MutMatrixView<'_, f32>, + points: rowmajor::Ref<'_, f32>, + mut storage: rowmajor::Mut<'_, f32>, ) -> ANNResult<()> { - InnerProduct::compute_distances(points, storage.as_mut_view())?; + InnerProduct::compute_distances(points, storage.as_view_mut())?; // Keep the constant of `1 - dot`. Near neighbors then have distances near // zero, where floating-point spacing is finest. Later stages quantize these // distances, so the constant keeps more near neighbors distinguishable. - for (source, row) in storage.row_iter_mut().enumerate() { + for (source, row) in storage.rows_mut().enumerate() { row[..=source] .iter_mut() .for_each(|distance| *distance += 1.0); @@ -142,13 +142,13 @@ mod tests { values[128] = 0.0; values[129..].fill(0.0); values[257] = 1.0; - let points = MatrixView::try_from(values.as_slice(), 2, 129).unwrap(); + let points = rowmajor::Ref::try_from_data(values.as_slice(), 2, 129).unwrap(); let mut output = [f32::NAN; 4]; let expected = 16_777_344.0; // 4096^2 + 128. L2::compute_distances( points, - MutMatrixView::try_from(&mut output[..], 2, 2).unwrap(), + rowmajor::Mut::try_from_data(&mut output[..], 2, 2).unwrap(), ) .unwrap(); @@ -184,12 +184,13 @@ mod tests { test_support::normalize(&mut values, dimensions); } let points = - MatrixView::try_from(values.as_slice(), point_count, dimensions).unwrap(); + rowmajor::Ref::try_from_data(values.as_slice(), point_count, dimensions) + .unwrap(); let mut output = vec![f32::NAN; point_count * point_count]; M::compute_distances( points, - MutMatrixView::try_from(output.as_mut_slice(), point_count, point_count) + rowmajor::Mut::try_from_data(output.as_mut_slice(), point_count, point_count) .unwrap(), ) .unwrap_or_else(|error| { @@ -244,12 +245,14 @@ mod tests { if scalar_metric == Metric::CosineNormalized { test_support::normalize(&mut values, dimensions); } - let points = MatrixView::try_from(values.as_slice(), point_count, dimensions).unwrap(); + let points = + rowmajor::Ref::try_from_data(values.as_slice(), point_count, dimensions).unwrap(); let mut output = vec![f32::NAN; point_count * point_count]; M::compute_distances( points, - MutMatrixView::try_from(output.as_mut_slice(), point_count, point_count).unwrap(), + rowmajor::Mut::try_from_data(output.as_mut_slice(), point_count, point_count) + .unwrap(), ) .unwrap_or_else(|error| panic!("shape={shape:?}: {error}")); @@ -292,12 +295,12 @@ mod tests { #[case] values: &[f32], #[case] expected: [f32; 6], ) { - let points = MatrixView::try_from(values, 3, 2).unwrap(); + let points = rowmajor::Ref::try_from_data(values, 3, 2).unwrap(); let mut output = [42.0; 9]; M::compute_distances( points, - MutMatrixView::try_from(&mut output[..], 3, 3).unwrap(), + rowmajor::Mut::try_from_data(&mut output[..], 3, 3).unwrap(), ) .unwrap(); @@ -314,12 +317,12 @@ mod tests { #[case::squared_norm_underflows(f32::MIN_POSITIVE)] fn cosine_gives_unit_distance_to_points_with_small_norms(#[case] coordinate: f32) { let values = [coordinate, 0.0, 0.0, 2.0]; - let points = MatrixView::try_from(&values[..], 2, 2).unwrap(); + let points = rowmajor::Ref::try_from_data(&values[..], 2, 2).unwrap(); let mut output = [42.0; 4]; Cosine::compute_distances( points, - MutMatrixView::try_from(&mut output[..], 2, 2).unwrap(), + rowmajor::Mut::try_from_data(&mut output[..], 2, 2).unwrap(), ) .unwrap(); @@ -339,8 +342,8 @@ mod tests { let mut output = [42.0; 9]; M::compute_distances( - MatrixView::try_from(&values[..], 3, 2).unwrap(), - MutMatrixView::try_from(&mut output[..], 3, 3).unwrap(), + rowmajor::Ref::try_from_data(&values[..], 3, 2).unwrap(), + rowmajor::Mut::try_from_data(&mut output[..], 3, 3).unwrap(), ) .unwrap(); diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 882788dc79..123743ef63 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -31,7 +31,7 @@ mod simd; mod topk; use crate::{ANNError, ANNResult}; -use diskann_utils::views::MutMatrixView; +use diskann_utils::views::rowmajor; /// Squared Euclidean distance. pub(super) struct L2; @@ -79,7 +79,7 @@ fn distance_scratch( storage: &mut Vec, rows: usize, columns: usize, -) -> ANNResult> { +) -> ANNResult> { let len = rows.checked_mul(columns).ok_or_else(|| { ANNError::message(format!( "distance matrix size overflows for {rows} x {columns}" @@ -88,7 +88,11 @@ fn distance_scratch( if storage.len() < len { storage.resize(len, 0.0); } - Ok(MutMatrixView::try_from(&mut storage[..len], rows, columns)?) + Ok(rowmajor::Mut::try_from_data( + &mut storage[..len], + rows, + columns, + )?) } #[cfg(test)] diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index 0af3df6c28..4b31b472d3 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -12,7 +12,7 @@ //! send each point. use crate::ANNResult; -use diskann_utils::views::{MatrixView, MutMatrixView}; +use diskann_utils::views::rowmajor::{self, Matrix, MatrixMut}; use super::{ check_output_rows, distance_scratch, @@ -46,9 +46,9 @@ pub(super) struct PartitionKernelWorkspace { /// different dimensions. pub(super) fn assign_leaders( arch: A, - points: MatrixView<'_, f32>, + points: rowmajor::Ref<'_, f32>, leaders: &M::Leaders<'_>, - output: MutMatrixView<'_, u32>, + output: rowmajor::Mut<'_, u32>, workspace: &mut PartitionKernelWorkspace, ) -> ANNResult<()> where @@ -62,7 +62,7 @@ where point_count, M::leader_count(leaders), )?; - M::compute_distances(points, leaders, distances.as_mut_view())?; + M::compute_distances(points, leaders, distances.as_view_mut())?; select_top_k_ids( arch, distances.as_view(), @@ -102,17 +102,18 @@ mod tests { 2, unit_norm, ); - let leaders = - M::create_leaders(MatrixView::try_from(leader_values.as_slice(), 4, 2).unwrap()); + let leaders = M::create_leaders( + rowmajor::Ref::try_from_data(leader_values.as_slice(), 4, 2).unwrap(), + ); // A width of six exceeds the four leaders. for assignments in [1, 2, 3, 4, 6] { let mut output = vec![0; 3 * assignments]; assign_leaders::<_, M>( ARCH, - MatrixView::try_from(point_values.as_slice(), 3, 2).unwrap(), + rowmajor::Ref::try_from_data(point_values.as_slice(), 3, 2).unwrap(), &leaders, - MutMatrixView::try_from(output.as_mut_slice(), 3, assignments).unwrap(), + rowmajor::Mut::try_from_data(output.as_mut_slice(), 3, assignments).unwrap(), &mut PartitionKernelWorkspace::default(), ) .unwrap_or_else(|error| panic!("assignments={assignments}: {error}")); @@ -162,12 +163,18 @@ mod tests { test_support::normalize(&mut point_values, dimensions); test_support::normalize(&mut leader_values, dimensions); } - let points = - MatrixView::try_from(point_values.as_slice(), point_count, dimensions) - .unwrap(); - let leader_matrix = - MatrixView::try_from(leader_values.as_slice(), leader_count, dimensions) - .unwrap(); + let points = rowmajor::Ref::try_from_data( + point_values.as_slice(), + point_count, + dimensions, + ) + .unwrap(); + let leader_matrix = rowmajor::Ref::try_from_data( + leader_values.as_slice(), + leader_count, + dimensions, + ) + .unwrap(); let leaders = M::create_leaders(leader_matrix); let tolerance = test_support::dense_tolerance(metric, dimensions); // The oracle sorts every leader by its scalar distance to the point. @@ -199,7 +206,7 @@ mod tests { arch, points, &leaders, - MutMatrixView::try_from( + rowmajor::Mut::try_from_data( output.as_mut_slice(), point_count, assignments, @@ -246,7 +253,8 @@ mod tests { #[test] fn workspace_reuse_does_not_mix_results_from_different_stripes() { let leader_values = [0.0, 5.0, 12.0]; - let leaders = L2::create_leaders(MatrixView::try_from(&leader_values[..], 3, 1).unwrap()); + let leaders = + L2::create_leaders(rowmajor::Ref::try_from_data(&leader_values[..], 3, 1).unwrap()); let point_values = [1.0, 7.0, 11.0]; let expected_ids = [[0, 1, 2], [1, 2, 0], [2, 1, 0]]; let mut output = Vec::new(); @@ -256,9 +264,10 @@ mod tests { output.resize(point_count * assignments, 0); assign_leaders::<_, L2>( ARCH, - MatrixView::try_from(&point_values[..point_count], point_count, 1).unwrap(), + rowmajor::Ref::try_from_data(&point_values[..point_count], point_count, 1).unwrap(), &leaders, - MutMatrixView::try_from(output.as_mut_slice(), point_count, assignments).unwrap(), + rowmajor::Mut::try_from_data(output.as_mut_slice(), point_count, assignments) + .unwrap(), &mut workspace, ) .unwrap(); @@ -293,8 +302,8 @@ mod tests { select_top_k_ids( ARCH, - MatrixView::try_from(&distances[..], 3, 4).unwrap(), - MutMatrixView::try_from(&mut output[..], 3, 3).unwrap(), + rowmajor::Ref::try_from_data(&distances[..], 3, 4).unwrap(), + rowmajor::Mut::try_from_data(&mut output[..], 3, 3).unwrap(), &mut candidates, ); @@ -310,14 +319,15 @@ mod tests { fn output_without_one_row_per_point_is_rejected() { let point_values = [1.0, 6.0]; let leader_values = [0.0, 5.0, 12.0]; - let leaders = L2::create_leaders(MatrixView::try_from(&leader_values[..], 3, 1).unwrap()); + let leaders = + L2::create_leaders(rowmajor::Ref::try_from_data(&leader_values[..], 3, 1).unwrap()); let mut output = [2]; let error = assign_leaders::<_, L2>( ARCH, - MatrixView::try_from(&point_values[..], 2, 1).unwrap(), + rowmajor::Ref::try_from_data(&point_values[..], 2, 1).unwrap(), &leaders, - MutMatrixView::try_from(&mut output[..], 1, 1).unwrap(), + rowmajor::Mut::try_from_data(&mut output[..], 1, 1).unwrap(), &mut PartitionKernelWorkspace::default(), ) .unwrap_err(); diff --git a/diskann/src/graph/pipnn/partition_metric.rs b/diskann/src/graph/pipnn/partition_metric.rs index 29078bde87..a2471fe65c 100644 --- a/diskann/src/graph/pipnn/partition_metric.rs +++ b/diskann/src/graph/pipnn/partition_metric.rs @@ -13,7 +13,7 @@ use crate::{ANNError, ANNResult}; use diskann_linalg::Transpose; -use diskann_utils::views::{MatrixView, MutMatrixView}; +use diskann_utils::views::rowmajor::{self, Matrix, MatrixMut}; use diskann_vector::{ Norm, norm::{FastL2Norm, FastL2NormSquared}, @@ -27,7 +27,7 @@ use super::{Cosine, CosineNormalized, InnerProduct, L2, cosine_distance}; /// when the leader set is created, and all point stripes share them. Inner product /// and normalized cosine store no norms. pub(super) struct PartitionLeaders<'a, Norms> { - values: MatrixView<'a, f32>, + values: rowmajor::Ref<'a, f32>, norms: Norms, } @@ -43,7 +43,7 @@ pub(super) trait PartitionMetric: Send + Sync + 'static { /// /// Partitioning always samples at least one leader. The kernels do not support /// an empty leader set. - fn create_leaders<'a>(values: MatrixView<'a, f32>) -> Self::Leaders<'a>; + fn create_leaders<'a>(values: rowmajor::Ref<'a, f32>) -> Self::Leaders<'a>; /// Return the number of leaders. fn leader_count(leaders: &Self::Leaders<'_>) -> usize; @@ -56,17 +56,17 @@ pub(super) trait PartitionMetric: Send + Sync + 'static { /// [`assign_leaders`](super::partition_kernel::assign_leaders) creates it with /// this shape. A zero distance can have either sign. fn compute_distances( - points: MatrixView<'_, f32>, + points: rowmajor::Ref<'_, f32>, leaders: &Self::Leaders<'_>, - storage: MutMatrixView<'_, f32>, + storage: rowmajor::Mut<'_, f32>, ) -> ANNResult<()>; } /// Compute the L2 norm of each row. Points and leaders both use this function, so /// their norms round the same way. -fn cosine_norms(vectors: MatrixView<'_, f32>) -> Vec { +fn cosine_norms(vectors: rowmajor::Ref<'_, f32>) -> Vec { vectors - .row_iter() + .rows() .map(|vector| FastL2Norm.evaluate(vector)) .collect() } @@ -74,11 +74,11 @@ fn cosine_norms(vectors: MatrixView<'_, f32>) -> Vec { impl PartitionMetric for L2 { type Leaders<'a> = PartitionLeaders<'a, Vec>; - fn create_leaders<'a>(values: MatrixView<'a, f32>) -> Self::Leaders<'a> { + fn create_leaders<'a>(values: rowmajor::Ref<'a, f32>) -> Self::Leaders<'a> { PartitionLeaders { values, norms: values - .row_iter() + .rows() .map(|leader| FastL2NormSquared.evaluate(leader)) .collect(), } @@ -89,14 +89,14 @@ impl PartitionMetric for L2 { } fn compute_distances( - points: MatrixView<'_, f32>, + points: rowmajor::Ref<'_, f32>, leaders: &Self::Leaders<'_>, - mut storage: MutMatrixView<'_, f32>, + mut storage: rowmajor::Mut<'_, f32>, ) -> ANNResult<()> { // The ranking distance is `||l||² - 2(p·l)`: the squared L2 distance without // `||p||²`, which is equal for every leader of the point. Start each row with // the leader norms, and let GEMM add the dot-product term. - for row in storage.row_iter_mut() { + for row in storage.rows_mut() { row.copy_from_slice(&leaders.norms); } diskann_linalg::sgemm( @@ -119,7 +119,7 @@ impl PartitionMetric for L2 { impl PartitionMetric for Cosine { type Leaders<'a> = PartitionLeaders<'a, Vec>; - fn create_leaders<'a>(values: MatrixView<'a, f32>) -> Self::Leaders<'a> { + fn create_leaders<'a>(values: rowmajor::Ref<'a, f32>) -> Self::Leaders<'a> { PartitionLeaders { values, norms: cosine_norms(values), @@ -131,9 +131,9 @@ impl PartitionMetric for Cosine { } fn compute_distances( - points: MatrixView<'_, f32>, + points: rowmajor::Ref<'_, f32>, leaders: &Self::Leaders<'_>, - mut storage: MutMatrixView<'_, f32>, + mut storage: rowmajor::Mut<'_, f32>, ) -> ANNResult<()> { diskann_linalg::sgemm( Transpose::None, @@ -152,7 +152,7 @@ impl PartitionMetric for Cosine { let leader_norms = &leaders.norms; // Convert each dot to cosine distance. The leader norms come from the leader // set, so each stripe computes only its point norms. - for (row, &point_norm) in storage.row_iter_mut().zip(point_norms.iter()) { + for (row, &point_norm) in storage.rows_mut().zip(point_norms.iter()) { for (distance, &leader_norm) in row.iter_mut().zip(leader_norms.iter()) { *distance = cosine_distance(*distance, point_norm, leader_norm); } @@ -164,7 +164,7 @@ impl PartitionMetric for Cosine { impl PartitionMetric for InnerProduct { type Leaders<'a> = PartitionLeaders<'a, ()>; - fn create_leaders<'a>(values: MatrixView<'a, f32>) -> Self::Leaders<'a> { + fn create_leaders<'a>(values: rowmajor::Ref<'a, f32>) -> Self::Leaders<'a> { PartitionLeaders { values, norms: () } } @@ -173,9 +173,9 @@ impl PartitionMetric for InnerProduct { } fn compute_distances( - points: MatrixView<'_, f32>, + points: rowmajor::Ref<'_, f32>, leaders: &Self::Leaders<'_>, - mut storage: MutMatrixView<'_, f32>, + mut storage: rowmajor::Mut<'_, f32>, ) -> ANNResult<()> { diskann_linalg::sgemm( Transpose::None, @@ -197,7 +197,7 @@ impl PartitionMetric for InnerProduct { impl PartitionMetric for CosineNormalized { type Leaders<'a> = ::Leaders<'a>; - fn create_leaders<'a>(values: MatrixView<'a, f32>) -> Self::Leaders<'a> { + fn create_leaders<'a>(values: rowmajor::Ref<'a, f32>) -> Self::Leaders<'a> { InnerProduct::create_leaders(values) } @@ -206,9 +206,9 @@ impl PartitionMetric for CosineNormalized { } fn compute_distances( - points: MatrixView<'_, f32>, + points: rowmajor::Ref<'_, f32>, leaders: &Self::Leaders<'_>, - storage: MutMatrixView<'_, f32>, + storage: rowmajor::Mut<'_, f32>, ) -> ANNResult<()> { // The constant in `1 - dot` does not change nearest-first order. InnerProduct::compute_distances(points, leaders, storage) @@ -233,13 +233,14 @@ mod tests { values[129..].fill(0.0); values[129] = 4096.0; values[130] = 4.0; - let leaders = L2::create_leaders(MatrixView::try_from(values.as_slice(), 2, 129).unwrap()); + let leaders = + L2::create_leaders(rowmajor::Ref::try_from_data(values.as_slice(), 2, 129).unwrap()); let mut output = [f32::NAN; 2]; L2::compute_distances( - MatrixView::try_from(&points[..], 1, 129).unwrap(), + rowmajor::Ref::try_from_data(&points[..], 1, 129).unwrap(), &leaders, - MutMatrixView::try_from(&mut output[..], 1, 2).unwrap(), + rowmajor::Mut::try_from_data(&mut output[..], 1, 2).unwrap(), ) .unwrap(); @@ -284,19 +285,25 @@ mod tests { test_support::normalize(&mut point_values, dimensions); test_support::normalize(&mut leader_values, dimensions); } - let points = - MatrixView::try_from(point_values.as_slice(), point_count, dimensions) - .unwrap(); - let leader_matrix = - MatrixView::try_from(leader_values.as_slice(), leader_count, dimensions) - .unwrap(); + let points = rowmajor::Ref::try_from_data( + point_values.as_slice(), + point_count, + dimensions, + ) + .unwrap(); + let leader_matrix = rowmajor::Ref::try_from_data( + leader_values.as_slice(), + leader_count, + dimensions, + ) + .unwrap(); let leaders = M::create_leaders(leader_matrix); let mut output = vec![f32::NAN; point_count * leader_count]; M::compute_distances( points, &leaders, - MutMatrixView::try_from(output.as_mut_slice(), point_count, leader_count) + rowmajor::Mut::try_from_data(output.as_mut_slice(), point_count, leader_count) .unwrap(), ) .unwrap_or_else(|error| { @@ -371,16 +378,19 @@ mod tests { test_support::normalize(&mut leader_values, dimensions); } let points = - MatrixView::try_from(point_values.as_slice(), point_count, dimensions).unwrap(); + rowmajor::Ref::try_from_data(point_values.as_slice(), point_count, dimensions) + .unwrap(); let leader_matrix = - MatrixView::try_from(leader_values.as_slice(), leader_count, dimensions).unwrap(); + rowmajor::Ref::try_from_data(leader_values.as_slice(), leader_count, dimensions) + .unwrap(); let leaders = M::create_leaders(leader_matrix); let mut output = vec![f32::NAN; point_count * leader_count]; M::compute_distances( points, &leaders, - MutMatrixView::try_from(output.as_mut_slice(), point_count, leader_count).unwrap(), + rowmajor::Mut::try_from_data(output.as_mut_slice(), point_count, leader_count) + .unwrap(), ) .unwrap_or_else(|error| panic!("shape={shape:?}: {error}")); @@ -432,14 +442,14 @@ mod tests { #[case] leader_values: &[f32], #[case] expected: [f32; 6], ) { - let leaders = M::create_leaders(MatrixView::try_from(leader_values, 3, 2).unwrap()); + let leaders = M::create_leaders(rowmajor::Ref::try_from_data(leader_values, 3, 2).unwrap()); let mut output = [-100.0; 3]; for (point, expected) in point_values.chunks_exact(2).zip(expected.chunks_exact(3)) { M::compute_distances( - MatrixView::try_from(point, 1, 2).unwrap(), + rowmajor::Ref::try_from_data(point, 1, 2).unwrap(), &leaders, - MutMatrixView::try_from(&mut output[..], 1, 3).unwrap(), + rowmajor::Mut::try_from_data(&mut output[..], 1, 3).unwrap(), ) .unwrap(); @@ -452,13 +462,13 @@ mod tests { let point_values = [0.0, 0.0, 0.0, 2.0]; let leader_values = [3.0, 0.0, 0.0, 0.0, 0.0, -4.0]; let leaders = - Cosine::create_leaders(MatrixView::try_from(&leader_values[..], 3, 2).unwrap()); + Cosine::create_leaders(rowmajor::Ref::try_from_data(&leader_values[..], 3, 2).unwrap()); let mut output = [42.0; 6]; Cosine::compute_distances( - MatrixView::try_from(&point_values[..], 2, 2).unwrap(), + rowmajor::Ref::try_from_data(&point_values[..], 2, 2).unwrap(), &leaders, - MutMatrixView::try_from(&mut output[..], 2, 3).unwrap(), + rowmajor::Mut::try_from_data(&mut output[..], 2, 3).unwrap(), ) .unwrap(); @@ -475,13 +485,14 @@ mod tests { ) { let point_values = [1.0, 2.0, 3.0]; let leader_values = [1.0, 2.0, 3.0, 4.0]; - let leaders = M::create_leaders(MatrixView::try_from(&leader_values[..], 2, 2).unwrap()); + let leaders = + M::create_leaders(rowmajor::Ref::try_from_data(&leader_values[..], 2, 2).unwrap()); let mut output = [17.0; 2]; let error = M::compute_distances( - MatrixView::try_from(&point_values[..], 1, 3).unwrap(), + rowmajor::Ref::try_from_data(&point_values[..], 1, 3).unwrap(), &leaders, - MutMatrixView::try_from(&mut output[..], 1, 2).unwrap(), + rowmajor::Mut::try_from_data(&mut output[..], 1, 2).unwrap(), ) .unwrap_err(); diff --git a/diskann/src/graph/pipnn/topk.rs b/diskann/src/graph/pipnn/topk.rs index 665f12cddc..dbc471ece6 100644 --- a/diskann/src/graph/pipnn/topk.rs +++ b/diskann/src/graph/pipnn/topk.rs @@ -25,7 +25,7 @@ //! The SIMD loops run inside `run2` or `run3` of architecture `A`, which compiles //! them with the target features of `A`. -use diskann_utils::views::{MatrixView, MutMatrixView}; +use diskann_utils::views::rowmajor::{self, Matrix, MatrixMut}; use diskann_wide::{SIMDPartialOrd, SIMDVector}; use super::simd::{LANES, Simd}; @@ -75,8 +75,8 @@ impl Default for Candidate { /// of k without a fixed-size nearest set. pub(super) fn select_top_k_ids( arch: A, - distances: MatrixView<'_, f32>, - output: MutMatrixView<'_, u32>, + distances: rowmajor::Ref<'_, f32>, + output: rowmajor::Mut<'_, u32>, scratch: &mut Vec, ) { debug_assert_eq!( @@ -114,8 +114,8 @@ pub(super) fn select_top_k_ids( /// after the scan. pub(super) fn select_top_k_symmetric( arch: A, - distances: MatrixView<'_, f32>, - output: MutMatrixView<'_, Candidate>, + distances: rowmajor::Ref<'_, f32>, + output: rowmajor::Mut<'_, Candidate>, kth_distances: &mut Vec, ) { let points = output.nrows(); @@ -124,7 +124,7 @@ pub(super) fn select_top_k_symmetric( "a symmetric scan needs one distance row and column per point" ); let k = output.ncols(); - let slots = output.into_inner(); + let slots = output.into_mut_slice(); // Fixed-size nearest sets serve the common `leaf_k` values 1, 2, and 3. On AVX2 // at k = 3, slices took 1.1x to 2.5x the ranking time of fixed-size sets. match k { @@ -143,17 +143,15 @@ pub(super) fn select_top_k_symmetric( /// the matching output row. All rows reuse the storage of `nearest`. fn select_ids_with( arch: A, - distances: MatrixView<'_, f32>, - mut output: MutMatrixView<'_, u32>, + distances: rowmajor::Ref<'_, f32>, + mut output: rowmajor::Mut<'_, u32>, nearest: &mut Nearest, ) where A: Simd, Nearest: AsMut<[Candidate]> + ?Sized, { - // `row_iter` panics on a matrix without columns. `row` returns an empty row, - // so every slot of that row becomes `UNASSIGNED`. - for (row, ids) in output.row_iter_mut().enumerate() { - select_nearest(arch, nearest, distances.row(row)); + for (ids, distances) in std::iter::zip(output.rows_mut(), distances.rows()) { + select_nearest(arch, nearest, distances); for (id, candidate) in ids.iter_mut().zip(nearest.as_mut().iter()) { *id = candidate.local_idx; } @@ -193,7 +191,7 @@ where /// `kth_distances`. fn scan_pairs( arch: A, - distances: MatrixView<'_, f32>, + distances: rowmajor::Ref<'_, f32>, neighborhoods: &mut [Nearest], kth_distances: &mut Vec, ) where @@ -398,8 +396,8 @@ mod tests { select_top_k_symmetric( ARCH, - MatrixView::try_from(distances.as_slice(), points, points).unwrap(), - MutMatrixView::try_from(output.as_flattened_mut(), points, 2).unwrap(), + rowmajor::Ref::try_from_data(distances.as_slice(), points, points).unwrap(), + rowmajor::Mut::try_from_data(output.as_flattened_mut(), points, 2).unwrap(), &mut kth_distances, ); @@ -417,8 +415,8 @@ mod tests { select_top_k_symmetric( ARCH, - MatrixView::try_from(&distances[..], 2, 2).unwrap(), - MutMatrixView::try_from(&mut output[..], 3, 1).unwrap(), + rowmajor::Ref::try_from_data(&distances[..], 2, 2).unwrap(), + rowmajor::Mut::try_from_data(&mut output[..], 3, 1).unwrap(), &mut Vec::new(), ); } @@ -514,8 +512,8 @@ mod tests { select_top_k_ids( arch, - MatrixView::try_from(distances.as_slice(), 1, count).unwrap(), - MutMatrixView::try_from(output.as_mut_slice(), 1, k).unwrap(), + rowmajor::Ref::try_from_data(distances.as_slice(), 1, count).unwrap(), + rowmajor::Mut::try_from_data(output.as_mut_slice(), 1, k).unwrap(), &mut Vec::new(), ); @@ -669,22 +667,26 @@ mod tests { } }) .collect(); - let distances = - MatrixView::try_from(distances.as_slice(), point_count, point_count) - .unwrap(); + let distances = rowmajor::Ref::try_from_data( + distances.as_slice(), + point_count, + point_count, + ) + .unwrap(); // k = 1, 2, and 3 use fixed-size nearest sets. The others use slices. for k in [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 17] { let mut output = vec![Candidate::EMPTY; point_count * k]; let mut kth_distances = Vec::new(); let mut neighborhoods = - MutMatrixView::try_from(output.as_mut_slice(), point_count, k).unwrap(); + rowmajor::Mut::try_from_data(output.as_mut_slice(), point_count, k) + .unwrap(); // When: offer every non-self pair once, through the production // width dispatch. select_top_k_symmetric( arch, distances, - neighborhoods.as_mut_view(), + neighborhoods.as_view_mut(), &mut kth_distances, ); @@ -882,14 +884,14 @@ mod tests { select_top_k_symmetric( ARCH, - MatrixView::try_from(&distances[..], 3, 3).unwrap(), - MutMatrixView::try_from(&mut [][..], 3, 0).unwrap(), + rowmajor::Ref::try_from_data(&distances[..], 3, 3).unwrap(), + rowmajor::Mut::try_from_data(&mut [][..], 3, 0).unwrap(), &mut kth_distances, ); select_top_k_ids( ARCH, - MatrixView::try_from(&distances[..], 3, 3).unwrap(), - MutMatrixView::try_from(&mut [][..], 3, 0).unwrap(), + rowmajor::Ref::try_from_data(&distances[..], 3, 3).unwrap(), + rowmajor::Mut::try_from_data(&mut [][..], 3, 0).unwrap(), &mut Vec::new(), ); @@ -903,8 +905,8 @@ mod tests { select_top_k_ids( ARCH, - MatrixView::try_from(&[][..], 2, 0).unwrap(), - MutMatrixView::try_from(&mut output[..], 2, 2).unwrap(), + rowmajor::Ref::try_from_data(&[][..], 2, 0).unwrap(), + rowmajor::Mut::try_from_data(&mut output[..], 2, 2).unwrap(), &mut Vec::new(), ); diff --git a/diskann/src/graph/start_point.rs b/diskann/src/graph/start_point.rs index fd95977615..b619229cd8 100644 --- a/diskann/src/graph/start_point.rs +++ b/diskann/src/graph/start_point.rs @@ -9,7 +9,7 @@ use thiserror::Error; use diskann_utils::{ sampling::WithApproximateNorm, - views::{Matrix, MatrixView}, + views::rowmajor::{self, Matrix, MatrixMut}, }; /// 'StartPointStrategy' is an enum that represents the different strategies to select @@ -69,7 +69,7 @@ pub enum StartPointError { MatrixRowError { row_id: usize }, #[error(transparent)] - MatrixLayout(#[from] diskann_utils::views::LayoutError), + MatrixLayout(#[from] diskann_utils::views::rowmajor::LayoutError), } impl StartPointStrategy { @@ -83,7 +83,10 @@ impl StartPointStrategy { } } - pub fn compute(&self, train_data: MatrixView<'_, T>) -> Result, StartPointError> + pub fn compute( + &self, + train_data: rowmajor::Ref<'_, T>, + ) -> Result, StartPointError> where T: Copy + SampleableForStart, { @@ -101,15 +104,18 @@ impl StartPointStrategy { let indices = rand::seq::index::sample(&mut rng, train_data.nrows(), nsamples.get()); - let mut points = - Matrix::try_from_element(nsamples.get(), train_data.ncols(), T::default())?; - std::iter::zip(points.row_iter_mut(), indices).for_each(|(dst, src)| { + let mut points = rowmajor::Owned::try_from_element( + nsamples.get(), + train_data.ncols(), + T::default(), + )?; + std::iter::zip(points.rows_mut(), indices).for_each(|(dst, src)| { dst.copy_from_slice(train_data.row(src)); }); Ok(points) } - StartPointStrategy::Medoid => Ok(Matrix::row_vector( + StartPointStrategy::Medoid => Ok(rowmajor::Owned::row_vector( T::compute_medoid(train_data.as_view()).into(), )), StartPointStrategy::RandomVectors { @@ -119,8 +125,9 @@ impl StartPointStrategy { } => { let mut rng = StdRng::seed_from_u64(*seed); let dim = train_data.ncols(); - let mut points = Matrix::try_from_element(nsamples.get(), dim, T::default())?; - points.row_iter_mut().for_each(|row| { + let mut points = + rowmajor::Owned::try_from_element(nsamples.get(), dim, T::default())?; + points.rows_mut().for_each(|row| { row.copy_from_slice(&WithApproximateNorm::with_approximate_norm( dim, *norm, &mut rng, )) @@ -134,7 +141,7 @@ impl StartPointStrategy { Some(*seed), )), StartPointStrategy::FirstVector => match train_data.get_row(0) { - Some(row) => Ok(Matrix::row_vector(row.into())), + Some(row) => Ok(rowmajor::Owned::row_vector(row.into())), None => Err(StartPointError::NotEnoughTrainingData { requested: 1, found: 0, @@ -266,7 +273,7 @@ mod tests { let data = vec![ 1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, ]; - let matrix = Matrix::try_from(data.into(), 4, 3).unwrap(); + let matrix = rowmajor::Owned::try_from_data(data.into(), 4, 3).unwrap(); let strategy = StartPointStrategy::FirstVector; let start_points = strategy.compute(matrix.as_view()).unwrap(); assert_eq!(start_points.nrows(), 1); @@ -279,7 +286,7 @@ mod tests { let data = vec![ 1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, ]; - let matrix = Matrix::try_from(data.into(), 4, 3).unwrap(); + let matrix = rowmajor::Owned::try_from_data(data.into(), 4, 3).unwrap(); let strategy = StartPointStrategy::Medoid; let start_points = strategy.compute(matrix.as_view()).unwrap(); assert_eq!(start_points.nrows(), 1); @@ -293,7 +300,7 @@ mod tests { let data = vec![ 1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, ]; - let matrix = Matrix::try_from(data.into(), 4, 3).unwrap(); + let matrix = rowmajor::Owned::try_from_data(data.into(), 4, 3).unwrap(); let strategy = StartPointStrategy::RandomVectors { norm: 5.0, nsamples: NonZeroUsize::new(2).unwrap(), @@ -314,7 +321,7 @@ mod tests { let data = vec![ 1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, ]; - let matrix = Matrix::try_from(data.into(), 4, 3).unwrap(); + let matrix = rowmajor::Owned::try_from_data(data.into(), 4, 3).unwrap(); let strategy = StartPointStrategy::RandomSamples { nsamples: NonZeroUsize::new(2).unwrap(), seed: 42, @@ -324,7 +331,7 @@ mod tests { assert_eq!(start_points.ncols(), 3); for i in 0..2 { let row = start_points.get_row(i).unwrap(); - assert!(matrix.row_iter().any(|r| r == row)); + assert!(matrix.rows().any(|r| r == row)); } } @@ -333,7 +340,7 @@ mod tests { let data = vec![ 1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, ]; - let matrix = Matrix::try_from(data.into(), 4, 3).unwrap(); + let matrix = rowmajor::Owned::try_from_data(data.into(), 4, 3).unwrap(); let strategy = StartPointStrategy::LatinHyperCube { nsamples: NonZeroUsize::new(2).unwrap(), seed: 42, diff --git a/diskann/src/graph/test/cases/grid_insert.rs b/diskann/src/graph/test/cases/grid_insert.rs index ed26f970a6..462ce3deb1 100644 --- a/diskann/src/graph/test/cases/grid_insert.rs +++ b/diskann/src/graph/test/cases/grid_insert.rs @@ -18,7 +18,7 @@ use std::{num::NonZeroUsize, sync::Arc}; -use diskann_utils::views::MatrixView; +use diskann_utils::views::rowmajor::{self, Matrix}; use diskann_vector::distance::Metric; use crate::{ @@ -120,7 +120,7 @@ fn build_index( /// to `true` (A) builds the same graph and (B) results in fewer `get_vector` calls. fn run_build( index: &Arc>, - data: MatrixView<'_, f32>, + data: rowmajor::Ref<'_, f32>, batchsize: Option, working_set_reuse: bool, runtime: &tokio::runtime::Runtime, @@ -130,7 +130,7 @@ fn run_build( match batchsize { None => { - for (id, vector) in data.row_iter().enumerate() { + for (id, vector) in data.rows().enumerate() { runtime .block_on(index.insert(&strategy, &context, &(id as u32), vector)) .unwrap(); @@ -140,7 +140,7 @@ fn run_build( let mut start = 0; while start < data.nrows() { let stop = (start + batchsize.get()).min(data.nrows()); - let batch = Arc::new(data.subview(start..stop).unwrap().to_owned()); + let batch = Arc::new(data.subview(start..stop).unwrap().to_rowmajor_owned()); runtime .block_on(index.multi_insert::( strategy.clone(), diff --git a/diskann/src/graph/test/cases/helpers.rs b/diskann/src/graph/test/cases/helpers.rs index 5733d6ae0f..9c68a7ac11 100644 --- a/diskann/src/graph/test/cases/helpers.rs +++ b/diskann/src/graph/test/cases/helpers.rs @@ -7,6 +7,7 @@ use std::{iter, sync::Arc}; +use diskann_utils::views::rowmajor::Matrix; use diskann_vector::distance::Metric; use crate::{ @@ -62,7 +63,7 @@ pub(super) fn setup_2d_square( .unwrap(); let points = vectors - .row_iter() + .rows() .zip(adjacency_lists.into_iter().take(num_points)) .enumerate() .map(|(id, (row, adj))| (id as u32, row.to_vec(), adj)); diff --git a/diskann/src/graph/test/provider.rs b/diskann/src/graph/test/provider.rs index 7504c4e58f..d6613bcb7d 100644 --- a/diskann/src/graph/test/provider.rs +++ b/diskann/src/graph/test/provider.rs @@ -13,7 +13,7 @@ use std::{ }; use dashmap::{DashMap, mapref::entry::Entry}; -use diskann_utils::views::Matrix; +use diskann_utils::views::rowmajor; use diskann_vector::{PreprocessedDistanceFunction, distance::Metric}; use thiserror::Error; @@ -1338,7 +1338,7 @@ impl<'a> glue::InsertStrategy<'a, Provider, &'a [f32]> for Strategy { } } -impl glue::MultiInsertStrategy> for Strategy { +impl glue::MultiInsertStrategy> for Strategy { type Seed = workingset::map::Builder>; type FinishError = Infallible; type PruneStrategy = Self; @@ -1352,7 +1352,7 @@ impl glue::MultiInsertStrategy> for Strategy { &self, _provider: &Provider, _ctx: &Context, - batch: &Arc>, + batch: &Arc>, ids: Itr, ) -> impl std::future::Future> + Send where diff --git a/diskann/src/graph/test/synthetic.rs b/diskann/src/graph/test/synthetic.rs index 8d6ba097d5..5884a146fd 100644 --- a/diskann/src/graph/test/synthetic.rs +++ b/diskann/src/graph/test/synthetic.rs @@ -3,7 +3,7 @@ * Licensed under the MIT license. */ -use diskann_utils::views::Matrix; +use diskann_utils::views::rowmajor::{self, Matrix}; use crate::graph::AdjacencyList; @@ -29,7 +29,7 @@ impl Grid { /// Return the generated grid with `f32` elements. /// /// See [`Self::data_as`] for documentation on the order of generation. - pub fn data(self, size: usize) -> Matrix { + pub fn data(self, size: usize) -> rowmajor::Owned { Self::data_as(self, size, |i: usize| i as f32) } @@ -47,6 +47,7 @@ impl Grid { /// /// ```rust /// use diskann::graph::test::synthetic::Grid; + /// use diskann_utils::views::rowmajor::Matrix; /// /// fn identity(x: usize) -> usize { /// x @@ -99,15 +100,15 @@ impl Grid { /// /// // etc. /// ``` - pub fn data_as(self, size: usize, mut f: F) -> Matrix + pub fn data_as(self, size: usize, mut f: F) -> rowmajor::Owned where F: FnMut(usize) -> R, { match self { - Self::One => Matrix::from_fn(size, 1, |rc| f(rc.row)), + Self::One => rowmajor::Owned::from_fn(size, 1, |rc| f(rc.row)), Self::Two => { let mut v = [0; 2]; - Matrix::from_fn(size.pow(self.dim().into()), 2, |rc| { + rowmajor::Owned::from_fn(size.pow(self.dim().into()), 2, |rc| { let value = f(v[rc.col]); if rc.col == 1 { increment(&mut v, size); @@ -121,7 +122,7 @@ impl Grid { // // Is it overkill? Yes. Is it fun? Also yes! let mut v = [0; 3]; - Matrix::from_fn(size.pow(self.dim().into()), 3, |rc| { + rowmajor::Owned::from_fn(size.pow(self.dim().into()), 3, |rc| { let value = f(v[rc.col]); if rc.col == 2 { increment(&mut v, size); @@ -131,7 +132,7 @@ impl Grid { } Self::Four => { let mut v = [0; 4]; - Matrix::from_fn(size.pow(self.dim().into()), 4, |rc| { + rowmajor::Owned::from_fn(size.pow(self.dim().into()), 4, |rc| { let value = f(v[rc.col]); if rc.col == 3 { increment(&mut v, size); @@ -356,7 +357,7 @@ pub(super) struct Setup { start_id: u32, start_neighbors: AdjacencyList, - data: Matrix, + data: rowmajor::Owned, neighbors: Vec>, } @@ -376,7 +377,7 @@ impl Setup { pub(super) fn setup(&self) -> impl Iterator, AdjacencyList)> { let mut i = 0u32; self.data - .row_iter() + .rows() .zip(self.neighbors.iter()) .map(move |(data, neighbors)| { let id = i; diff --git a/diskann/src/graph/workingset/map.rs b/diskann/src/graph/workingset/map.rs index 35733eb550..83809ec90d 100644 --- a/diskann/src/graph/workingset/map.rs +++ b/diskann/src/graph/workingset/map.rs @@ -101,11 +101,11 @@ //! //! ``` //! use std::sync::Arc; -//! use diskann_utils::views::Matrix; +//! use diskann_utils::views::rowmajor; //! //! use diskann::graph::workingset::{Map, View, map::{Overlay, Ref, Builder, Capacity}}; //! -//! let batch = Matrix::::row_vector(Box::new([1, 2, 3])); +//! let batch = rowmajor::Owned::::row_vector(Box::new([1, 2, 3])); //! //! // Construct an "overlay" with just one id. //! // @@ -646,14 +646,14 @@ where /// /// ``` /// use std::sync::Arc; -/// use diskann_utils::views::Matrix; +/// use diskann_utils::views::rowmajor; /// /// use diskann::graph::workingset::map::{Overlay, Ref}; /// -/// let data = Matrix::::column_vector(Box::new([10, 20, 30])); +/// let data = rowmajor::Owned::::column_vector(Box::new([10, 20, 30])); /// let ids = [5, 3, 2]; /// -/// // The `Ref<[u32]>` projection yields slices, which is compatible with a `Matrix`. +/// // The `Ref<[u32]>` projection yields slices, which is compatible with a `rowmajor::Owned`. /// let overlay = Overlay::>::from_batch(Arc::new(data), ids); /// /// assert_eq!(overlay.get(&5).unwrap(), &[10]); @@ -923,7 +923,7 @@ mod tests { use std::sync::Arc; - use diskann_utils::views::Matrix; + use diskann_utils::views::rowmajor::{self, Matrix}; use crate::graph::workingset::View as WorkingSetView; @@ -1015,20 +1015,21 @@ mod tests { /// 3.0 4.0 /// 5.0 6.0 /// ``` - fn test_matrix() -> Matrix { - Matrix::try_from(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0].into_boxed_slice(), 3, 2).unwrap() + fn test_matrix() -> rowmajor::Owned { + rowmajor::Owned::try_from_data(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0].into_boxed_slice(), 3, 2) + .unwrap() } type TestOverlay = Overlay>; - fn test_overlay() -> (Arc>, TestOverlay) { + fn test_overlay() -> (Arc>, TestOverlay) { let batch = Arc::new(test_matrix()); let ids = [10u32, 20, 30]; let overlay = Overlay::from_batch(batch.clone(), ids); (batch, overlay) } - fn test_overlay_projected() -> (Arc>, Overlay) { + fn test_overlay_projected() -> (Arc>, Overlay) { let batch = Arc::new(test_matrix()); let ids = [10u32, 20, 30]; let overlay = Overlay::from_batch(batch.clone(), ids); @@ -1816,7 +1817,7 @@ mod tests { #[test] fn overlay_from_batch_empty() { - let batch = Arc::new(Matrix::try_from(Box::new([]), 0, 2).unwrap()); + let batch = Arc::new(rowmajor::Owned::try_from_data(Box::new([]), 0, 2).unwrap()); let overlay = Overlay::>::from_batch(batch, std::iter::empty()); assert!(overlay.get(&0).is_none()); assert!(!overlay.contains_key(&0)); @@ -1824,7 +1825,7 @@ mod tests { #[test] fn overlay_from_batch_single_element() { - let batch = Arc::new(Matrix::try_from(Box::new([1.0, 2.0]), 1, 2).unwrap()); + let batch = Arc::new(rowmajor::Owned::try_from_data(Box::new([1.0, 2.0]), 1, 2).unwrap()); let overlay = Overlay::>::from_batch(batch, [42u32]); assert_eq!(overlay.get(&42).unwrap(), &[1.0, 2.0]); } @@ -1983,7 +1984,7 @@ mod tests { #[test] fn fill_skips_seeded_entries() { // Seed with a batch containing a different value for ID 0. - let batch = Arc::new(Matrix::try_from(Box::new([99.0, 88.0]), 2, 1).unwrap()); + let batch = Arc::new(rowmajor::Owned::try_from_data(Box::new([99.0, 88.0]), 2, 1).unwrap()); let overlay = Overlay::>::from_batch(batch, [0u32, 1]); let mut map = seeded_map(overlay, Capacity::Unbounded);