From 7881bebabe2c0e3a17cb5bae3c5a007e1539c984 Mon Sep 17 00:00:00 2001 From: "Alex Razumov (from Dev Box)" Date: Mon, 5 Oct 2026 15:04:29 -0700 Subject: [PATCH 1/7] fix(utils): support zero-column mutable iterators Preserve logical rows and windows for N x 0 matrices during parallel mutable iteration instead of panicking. --- diskann-utils/src/views/rowmajor.rs | 105 ++++++++++++----------- diskann-utils/src/views/rowmajor/iter.rs | 49 +++++++++++ 2 files changed, 105 insertions(+), 49 deletions(-) diff --git a/diskann-utils/src/views/rowmajor.rs b/diskann-utils/src/views/rowmajor.rs index ffd569394d..8d7773291b 100644 --- a/diskann-utils/src/views/rowmajor.rs +++ b/diskann-utils/src/views/rowmajor.rs @@ -5,6 +5,8 @@ use std::{marker::PhantomData, mem::ManuallyDrop, num::NonZeroUsize, ptr::NonNull}; +#[cfg(feature = "rayon")] +use rayon::iter::Either; #[cfg(feature = "rayon")] use rayon::prelude::{ IndexedParallelIterator, IntoParallelIterator, ParallelIterator, ParallelSliceMut, @@ -524,21 +526,19 @@ pub unsafe trait MatrixMut: Matrix { //-------// /// 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 nrows = self.nrows(); 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)) + if ncols == 0 { + let matrix = iter::ZeroColumnMut::new(self); + Either::Left((0..nrows).into_par_iter().map(move |_| matrix.row())) + } else { + Either::Right(self.as_mut_slice().par_chunks_exact_mut(ncols)) + } } /// Return a parallel iterator that divides the matrix into mutable sub-matrices with @@ -552,7 +552,7 @@ pub unsafe trait MatrixMut: Matrix { /// /// # Panics /// - /// Panics if `batchsize = 0` or `self.ncols() == 0 && self.nrows() != 0`. + /// Panics if `batchsize = 0`. #[cfg(feature = "rayon")] fn par_window_iter_mut( &mut self, @@ -566,31 +566,42 @@ pub unsafe trait MatrixMut: Matrix { "par_window_iter_mut batchsize cannot be zero" ); + let nrows = self.nrows(); let ncols = self.ncols(); - assert!( - ncols != 0 || self.nrows() == 0, - "`MatrixMut::par_window_iter_mut` does not support matrices with rows and zero columns" - ); + if ncols == 0 { + let matrix = iter::ZeroColumnMut::new(self); + return Either::Left( + (0..nrows) + .into_par_iter() + .step_by(batchsize) + .map(move |start| { + let window_nrows = batchsize.min(nrows - start); + matrix.window(window_nrows) + }), + ); + } // 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)) } - }) + let batchsize = batchsize.min(nrows); + Either::Right( + 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)) } + }), + ) } } @@ -2325,7 +2336,7 @@ mod tests { 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)] { + for (nrows, ncols) in [(0, 0), (0, 4), (3, 0), (1, 1), (1, 4), (4, 1), (5, 3)] { let context = lazy_format!("nrows = {nrows}, ncols = {ncols}"); let original = striped_matrix(nrows, ncols); @@ -2394,22 +2405,18 @@ mod tests { #[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); + fn zero_column_mut_views_can_coexist() { + let mut matrix = striped_matrix(3, 0); + let ptr = matrix.as_ptr(); + let zero_column = iter::ZeroColumnMut::new(&mut matrix); + + let rows = [zero_column.row(), zero_column.row(), zero_column.row()]; + let windows = [zero_column.window(2), zero_column.window(1)]; + + assert!(rows.iter().all(|row| row.is_empty() && row.as_ptr() == ptr)); + assert!(windows.iter().all(|window| { + window.ncols() == 0 && window.as_slice().is_empty() && window.as_ptr() == ptr + })); } #[test] diff --git a/diskann-utils/src/views/rowmajor/iter.rs b/diskann-utils/src/views/rowmajor/iter.rs index 91a765e551..63473598a2 100644 --- a/diskann-utils/src/views/rowmajor/iter.rs +++ b/diskann-utils/src/views/rowmajor/iter.rs @@ -5,6 +5,8 @@ use std::{marker::PhantomData, num::NonZeroUsize, ptr::NonNull}; +#[cfg(feature = "rayon")] +use crate::views::rowmajor::MatrixMut; use crate::views::rowmajor::{Layout, Matrix, Mut, Ref}; //------// @@ -187,3 +189,50 @@ impl<'a, T> Iterator for Windows<'a, T> { impl ExactSizeIterator for Windows<'_, T> {} impl std::iter::FusedIterator for Windows<'_, T> {} + +//---------------// +// ZeroColumnMut // +//---------------// + +#[cfg(feature = "rayon")] +pub(super) struct ZeroColumnMut<'a, T> { + ptr: NonNull, + _lifetime: PhantomData<&'a mut [T]>, +} + +// SAFETY: `ZeroColumnMut` has the ownership semantics of an empty `&mut [T]`. +#[cfg(feature = "rayon")] +unsafe impl Send for ZeroColumnMut<'_, T> {} +// SAFETY: Shared access can only create zero-length mutable views, so it cannot expose +// overlapping access to any element. +#[cfg(feature = "rayon")] +unsafe impl Sync for ZeroColumnMut<'_, T> {} + +#[cfg(feature = "rayon")] +impl<'a, T> ZeroColumnMut<'a, T> { + pub(super) fn new(matrix: &'a mut M) -> Self + where + M: MatrixMut + ?Sized, + { + debug_assert_eq!(matrix.ncols(), 0); + Self { + ptr: matrix.as_nonnull_mut(), + _lifetime: PhantomData, + } + } + + pub(super) fn row(&self) -> &'a mut [T] { + // SAFETY: The matrix pointer is valid for its zero-element layout. The returned + // slice preserves that pointer and cannot access or overlap any element. + unsafe { std::slice::from_raw_parts_mut(self.ptr.as_ptr(), 0) } + } + + pub(super) fn window(&self, nrows: usize) -> Mut<'a, T> { + Mut { + ptr: self.ptr, + // SAFETY: Every zero-column layout has zero elements. + layout: unsafe { Layout::new_unchecked(nrows, 0) }, + _lifetime: PhantomData, + } + } +} From 1f96dda5e1ed7ff9ab5acaab1ae5b0790f9e4a7a Mon Sep 17 00:00:00 2001 From: "Alex Razumov (from Dev Box)" Date: Mon, 5 Oct 2026 15:40:29 -0700 Subject: [PATCH 2/7] docs(utils): clarify zero-column iterator semantics --- diskann-utils/src/views/rowmajor.rs | 5 +++++ diskann-utils/src/views/rowmajor/iter.rs | 4 ++++ 2 files changed, 9 insertions(+) diff --git a/diskann-utils/src/views/rowmajor.rs b/diskann-utils/src/views/rowmajor.rs index 8d7773291b..941cc7072e 100644 --- a/diskann-utils/src/views/rowmajor.rs +++ b/diskann-utils/src/views/rowmajor.rs @@ -526,6 +526,8 @@ pub unsafe trait MatrixMut: Matrix { //-------// /// Return a parallel iterator over the rows of the matrix. + /// + /// A matrix with zero columns yields one empty mutable slice per row. #[cfg(feature = "rayon")] fn par_rows_mut(&mut self) -> impl IndexedParallelIterator where @@ -550,6 +552,9 @@ pub unsafe trait MatrixMut: Matrix { /// 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`. /// + /// A matrix with zero columns yields zero-column sub-matrices containing up to + /// `batchsize` rows. + /// /// # Panics /// /// Panics if `batchsize = 0`. diff --git a/diskann-utils/src/views/rowmajor/iter.rs b/diskann-utils/src/views/rowmajor/iter.rs index 63473598a2..a53e6e87bd 100644 --- a/diskann-utils/src/views/rowmajor/iter.rs +++ b/diskann-utils/src/views/rowmajor/iter.rs @@ -195,6 +195,10 @@ impl std::iter::FusedIterator for Windows<'_, T> {} //---------------// #[cfg(feature = "rayon")] +/// Carries an exclusive borrow of a zero-column matrix across Rayon workers. +/// +/// Multiple views may share the matrix pointer because every view contains zero elements +/// and therefore exposes no overlapping element access. pub(super) struct ZeroColumnMut<'a, T> { ptr: NonNull, _lifetime: PhantomData<&'a mut [T]>, From b8d1c07ed21c9d51e288f867b14f867789447522 Mon Sep 17 00:00:00 2001 From: "Alex Razumov (from Dev Box)" Date: Tue, 6 Oct 2026 11:41:46 -0700 Subject: [PATCH 3/7] test(utils): benchmark parallel mutable row iteration Compare the original chunk-based par_rows_mut implementation with the current Either-backed implementation for 5x100 and 1,000,000x100 matrices. Run with: cargo bench -p diskann-utils --bench bench_main Ten-run stability summary (Criterion midpoint estimates): Implementation Median MAD CV Range 5x100 old 100.695 us 1.520 us 15.57% 54.311-104.710 us 5x100 current 100.595 us 1.715 us 3.49% 94.418-105.270 us 1Mx100 old 7.988 ms 0.088 ms 1.98% 7.680-8.189 ms 1Mx100 current 7.929 ms 0.078 ms 1.43% 7.747-8.108 ms The old 5x100 CV is inflated by one anomalous 54.311 us run. The ratio of medians is -0.10% for 5x100 and -0.74% for 1Mx100, showing no measurable regression. --- Cargo.lock | 1 + diskann-utils/Cargo.toml | 5 ++ diskann-utils/benches/bench_main.rs | 79 +++++++++++++++++++++++++++++ 3 files changed, 85 insertions(+) create mode 100644 diskann-utils/benches/bench_main.rs diff --git a/Cargo.lock b/Cargo.lock index bc9903c01c..686b732ce8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -788,6 +788,7 @@ version = "0.59.0" dependencies = [ "bytemuck", "cfg-if", + "criterion", "diskann-vector", "diskann-wide", "half", diff --git a/diskann-utils/Cargo.toml b/diskann-utils/Cargo.toml index d0dbb9b873..9e68549c65 100644 --- a/diskann-utils/Cargo.toml +++ b/diskann-utils/Cargo.toml @@ -30,9 +30,14 @@ workspace = true [dev-dependencies] cfg-if.workspace = true +criterion.workspace = true rand.workspace = true rstest.workspace = true +[[bench]] +name = "bench_main" +harness = false + [features] default = ["rayon"] diff --git a/diskann-utils/benches/bench_main.rs b/diskann-utils/benches/bench_main.rs new file mode 100644 index 0000000000..f99d4bbea0 --- /dev/null +++ b/diskann-utils/benches/bench_main.rs @@ -0,0 +1,79 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use std::{hint::black_box, time::Duration}; + +use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput}; +use diskann_utils::views::rowmajor::{Matrix, MatrixMut, Owned}; +use rayon::prelude::{IndexedParallelIterator, ParallelIterator, ParallelSliceMut}; + +const NCOLS: usize = 100; + +fn par_rows_mut_old(matrix: &mut M) -> impl IndexedParallelIterator +where + M: MatrixMut, + M::Element: Send, +{ + let ncols = matrix.ncols(); + assert!( + ncols != 0 || matrix.nrows() == 0, + "`MatrixMut::par_rows_mut` does not support matrices with rows and zero columns" + ); + matrix.as_mut_slice().par_chunks_exact_mut(ncols.max(1)) +} + +fn update_rows<'a, I>(rows: I) +where + I: ParallelIterator, +{ + rows.for_each(|row| { + let (first, rest) = row.split_first_mut().unwrap(); + *first = first.wrapping_add(1); + + let mut previous = *first; + for value in rest { + *value = value.wrapping_add(previous).wrapping_add(1); + previous = *value; + } + }); +} + +fn benchmark_shape(c: &mut Criterion, nrows: usize) { + let mut group = c.benchmark_group(format!("par_rows_mut/{nrows}x{NCOLS}")); + group.throughput(Throughput::Elements((nrows * NCOLS) as u64)); + + group.bench_function(BenchmarkId::new("old", nrows), |b| { + let mut matrix = Owned::from_element(nrows, NCOLS, 0u32); + b.iter(|| { + update_rows(par_rows_mut_old(black_box(&mut matrix))); + black_box(matrix.as_slice()); + }); + }); + + group.bench_function(BenchmarkId::new("current", nrows), |b| { + let mut matrix = Owned::from_element(nrows, NCOLS, 0u32); + b.iter(|| { + update_rows(black_box(&mut matrix).par_rows_mut()); + black_box(matrix.as_slice()); + }); + }); + + group.finish(); +} + +fn benchmark_par_rows_mut(c: &mut Criterion) { + benchmark_shape(c, 5); + benchmark_shape(c, 1_000_000); +} + +criterion_group! { + name = benches; + config = Criterion::default() + .sample_size(10) + .warm_up_time(Duration::from_secs(2)) + .measurement_time(Duration::from_secs(5)); + targets = benchmark_par_rows_mut +} +criterion_main!(benches); From 055e0a5134cd0de2da48c61f50348442d25e3598 Mon Sep 17 00:00:00 2001 From: "Alex Razumov (from Dev Box)" Date: Tue, 6 Oct 2026 12:30:45 -0700 Subject: [PATCH 4/7] Removed "before" bench --- diskann-utils/benches/bench_main.rs | 55 +++++++++-------------------- 1 file changed, 17 insertions(+), 38 deletions(-) diff --git a/diskann-utils/benches/bench_main.rs b/diskann-utils/benches/bench_main.rs index f99d4bbea0..4fd6ee57bf 100644 --- a/diskann-utils/benches/bench_main.rs +++ b/diskann-utils/benches/bench_main.rs @@ -7,23 +7,10 @@ use std::{hint::black_box, time::Duration}; use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput}; use diskann_utils::views::rowmajor::{Matrix, MatrixMut, Owned}; -use rayon::prelude::{IndexedParallelIterator, ParallelIterator, ParallelSliceMut}; +use rayon::prelude::ParallelIterator; const NCOLS: usize = 100; -fn par_rows_mut_old(matrix: &mut M) -> impl IndexedParallelIterator -where - M: MatrixMut, - M::Element: Send, -{ - let ncols = matrix.ncols(); - assert!( - ncols != 0 || matrix.nrows() == 0, - "`MatrixMut::par_rows_mut` does not support matrices with rows and zero columns" - ); - matrix.as_mut_slice().par_chunks_exact_mut(ncols.max(1)) -} - fn update_rows<'a, I>(rows: I) where I: ParallelIterator, @@ -40,34 +27,26 @@ where }); } -fn benchmark_shape(c: &mut Criterion, nrows: usize) { - let mut group = c.benchmark_group(format!("par_rows_mut/{nrows}x{NCOLS}")); - group.throughput(Throughput::Elements((nrows * NCOLS) as u64)); - - group.bench_function(BenchmarkId::new("old", nrows), |b| { - let mut matrix = Owned::from_element(nrows, NCOLS, 0u32); - b.iter(|| { - update_rows(par_rows_mut_old(black_box(&mut matrix))); - black_box(matrix.as_slice()); - }); - }); - - group.bench_function(BenchmarkId::new("current", nrows), |b| { - let mut matrix = Owned::from_element(nrows, NCOLS, 0u32); - b.iter(|| { - update_rows(black_box(&mut matrix).par_rows_mut()); - black_box(matrix.as_slice()); - }); - }); +fn benchmark_par_rows_mut(c: &mut Criterion) { + let mut group = c.benchmark_group("par_rows_mut"); + + for nrows in [5, 1_000_000] { + group.throughput(Throughput::Elements((nrows * NCOLS) as u64)); + group.bench_function( + BenchmarkId::from_parameter(format!("{nrows}x{NCOLS}")), + |b| { + let mut matrix = Owned::from_element(nrows, NCOLS, 0u32); + b.iter(|| { + update_rows(black_box(&mut matrix).par_rows_mut()); + black_box(matrix.as_slice()); + }); + }, + ); + } group.finish(); } -fn benchmark_par_rows_mut(c: &mut Criterion) { - benchmark_shape(c, 5); - benchmark_shape(c, 1_000_000); -} - criterion_group! { name = benches; config = Criterion::default() From 9d819f4938db514fcd7558c95c7bc81ba7492861 Mon Sep 17 00:00:00 2001 From: "Alex Razumov (from Dev Box)" Date: Tue, 6 Oct 2026 13:40:07 -0700 Subject: [PATCH 5/7] required-features = ["rayon"] --- diskann-utils/Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/diskann-utils/Cargo.toml b/diskann-utils/Cargo.toml index 9e68549c65..f76b971cb0 100644 --- a/diskann-utils/Cargo.toml +++ b/diskann-utils/Cargo.toml @@ -37,7 +37,7 @@ rstest.workspace = true [[bench]] name = "bench_main" harness = false - +required-features = ["rayon"] [features] default = ["rayon"] From 95d3864c3a1b358bc69415ab21e4d70a99014638 Mon Sep 17 00:00:00 2001 From: "Alex Razumov (from Dev Box)" Date: Wed, 7 Oct 2026 13:12:52 -0700 Subject: [PATCH 6/7] revert(utils): restore par_rows_mut baseline benchmark Restore the side-by-side comparison removed by 055e0a51 while naming the original implementation baseline and retaining current for the production implementation. --- diskann-utils/benches/bench_main.rs | 57 ++++++++++++++++++++--------- 1 file changed, 40 insertions(+), 17 deletions(-) diff --git a/diskann-utils/benches/bench_main.rs b/diskann-utils/benches/bench_main.rs index 4fd6ee57bf..a33b636b9e 100644 --- a/diskann-utils/benches/bench_main.rs +++ b/diskann-utils/benches/bench_main.rs @@ -7,10 +7,25 @@ use std::{hint::black_box, time::Duration}; use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput}; use diskann_utils::views::rowmajor::{Matrix, MatrixMut, Owned}; -use rayon::prelude::ParallelIterator; +use rayon::prelude::{IndexedParallelIterator, ParallelIterator, ParallelSliceMut}; const NCOLS: usize = 100; +fn par_rows_mut_baseline( + matrix: &mut M, +) -> impl IndexedParallelIterator +where + M: MatrixMut, + M::Element: Send, +{ + let ncols = matrix.ncols(); + assert!( + ncols != 0 || matrix.nrows() == 0, + "`MatrixMut::par_rows_mut` does not support matrices with rows and zero columns" + ); + matrix.as_mut_slice().par_chunks_exact_mut(ncols.max(1)) +} + fn update_rows<'a, I>(rows: I) where I: ParallelIterator, @@ -27,26 +42,34 @@ where }); } -fn benchmark_par_rows_mut(c: &mut Criterion) { - let mut group = c.benchmark_group("par_rows_mut"); - - for nrows in [5, 1_000_000] { - group.throughput(Throughput::Elements((nrows * NCOLS) as u64)); - group.bench_function( - BenchmarkId::from_parameter(format!("{nrows}x{NCOLS}")), - |b| { - let mut matrix = Owned::from_element(nrows, NCOLS, 0u32); - b.iter(|| { - update_rows(black_box(&mut matrix).par_rows_mut()); - black_box(matrix.as_slice()); - }); - }, - ); - } +fn benchmark_shape(c: &mut Criterion, nrows: usize) { + let mut group = c.benchmark_group(format!("par_rows_mut/{nrows}x{NCOLS}")); + group.throughput(Throughput::Elements((nrows * NCOLS) as u64)); + + group.bench_function(BenchmarkId::new("baseline", nrows), |b| { + let mut matrix = Owned::from_element(nrows, NCOLS, 0u32); + b.iter(|| { + update_rows(par_rows_mut_baseline(black_box(&mut matrix))); + black_box(matrix.as_slice()); + }); + }); + + group.bench_function(BenchmarkId::new("current", nrows), |b| { + let mut matrix = Owned::from_element(nrows, NCOLS, 0u32); + b.iter(|| { + update_rows(black_box(&mut matrix).par_rows_mut()); + black_box(matrix.as_slice()); + }); + }); group.finish(); } +fn benchmark_par_rows_mut(c: &mut Criterion) { + benchmark_shape(c, 5); + benchmark_shape(c, 1_000_000); +} + criterion_group! { name = benches; config = Criterion::default() From 4f25b43748cc49412364c338088050aaa1e10498 Mon Sep 17 00:00:00 2001 From: "Alex Razumov (from Dev Box)" Date: Wed, 7 Oct 2026 14:20:49 -0700 Subject: [PATCH 7/7] refactor(utils): unify mutable parallel iterators Use index-based row and window ranges with a shared ParMut helper instead of branching between zero-column views and Rayon slice chunks. Keep disjoint-view safety contracts beside the implementation and cover zero-column and nonempty coexistence under Miri. --- diskann-utils/src/views/rowmajor.rs | 82 +++---------- diskann-utils/src/views/rowmajor/iter.rs | 144 +++++++++++++++++++---- 2 files changed, 138 insertions(+), 88 deletions(-) diff --git a/diskann-utils/src/views/rowmajor.rs b/diskann-utils/src/views/rowmajor.rs index 941cc7072e..25be0b880b 100644 --- a/diskann-utils/src/views/rowmajor.rs +++ b/diskann-utils/src/views/rowmajor.rs @@ -6,11 +6,7 @@ use std::{marker::PhantomData, mem::ManuallyDrop, num::NonZeroUsize, ptr::NonNull}; #[cfg(feature = "rayon")] -use rayon::iter::Either; -#[cfg(feature = "rayon")] -use rayon::prelude::{ - IndexedParallelIterator, IntoParallelIterator, ParallelIterator, ParallelSliceMut, -}; +use rayon::prelude::{IndexedParallelIterator, IntoParallelIterator, ParallelIterator}; use thiserror::Error; pub mod iter; @@ -533,14 +529,12 @@ pub unsafe trait MatrixMut: Matrix { where Self::Element: Send, { - let nrows = self.nrows(); - let ncols = self.ncols(); - if ncols == 0 { - let matrix = iter::ZeroColumnMut::new(self); - Either::Left((0..nrows).into_par_iter().map(move |_| matrix.row())) - } else { - Either::Right(self.as_mut_slice().par_chunks_exact_mut(ncols)) - } + let matrix = iter::ParMut::new(self); + (0..matrix.nrows()).into_par_iter().map(move |row| { + // SAFETY: The range produces each in-bounds row exactly once. Distinct rows + // have disjoint element ranges; zero-column rows contain no elements. + unsafe { matrix.row_disjoint_unchecked(row) } + }) } /// Return a parallel iterator that divides the matrix into mutable sub-matrices with @@ -571,42 +565,18 @@ pub unsafe trait MatrixMut: Matrix { "par_window_iter_mut batchsize cannot be zero" ); - let nrows = self.nrows(); - let ncols = self.ncols(); - if ncols == 0 { - let matrix = iter::ZeroColumnMut::new(self); - return Either::Left( - (0..nrows) - .into_par_iter() - .step_by(batchsize) - .map(move |start| { - let window_nrows = batchsize.min(nrows - start); - matrix.window(window_nrows) - }), - ); - } + let matrix = iter::ParMut::new(self); + let nrows = matrix.nrows(); + (0..nrows) + .into_par_iter() + .step_by(batchsize) + .map(move |start| { + let end = start.saturating_add(batchsize).min(nrows); - // Ensure that `batchsize * ncols` does not overflow. - let batchsize = batchsize.min(nrows); - Either::Right( - 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)) } - }), - ) + // SAFETY: `start` comes from an in-bounds range and `end` is clamped to + // `nrows`. Stepping by `batchsize` makes the yielded ranges disjoint. + unsafe { matrix.window_disjoint_unchecked(start..end) } + }) } } @@ -2408,22 +2378,6 @@ mod tests { } } - #[test] - #[cfg(feature = "rayon")] - fn zero_column_mut_views_can_coexist() { - let mut matrix = striped_matrix(3, 0); - let ptr = matrix.as_ptr(); - let zero_column = iter::ZeroColumnMut::new(&mut matrix); - - let rows = [zero_column.row(), zero_column.row(), zero_column.row()]; - let windows = [zero_column.window(2), zero_column.window(1)]; - - assert!(rows.iter().all(|row| row.is_empty() && row.as_ptr() == ptr)); - assert!(windows.iter().all(|window| { - window.ncols() == 0 && window.as_slice().is_empty() && window.as_ptr() == ptr - })); - } - #[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)] { diff --git a/diskann-utils/src/views/rowmajor/iter.rs b/diskann-utils/src/views/rowmajor/iter.rs index a53e6e87bd..a7a5f7a148 100644 --- a/diskann-utils/src/views/rowmajor/iter.rs +++ b/diskann-utils/src/views/rowmajor/iter.rs @@ -190,53 +190,149 @@ impl<'a, T> Iterator for Windows<'a, T> { impl ExactSizeIterator for Windows<'_, T> {} impl std::iter::FusedIterator for Windows<'_, T> {} -//---------------// -// ZeroColumnMut // -//---------------// +//--------// +// ParMut // +//--------// #[cfg(feature = "rayon")] -/// Carries an exclusive borrow of a zero-column matrix across Rayon workers. -/// -/// Multiple views may share the matrix pointer because every view contains zero elements -/// and therefore exposes no overlapping element access. -pub(super) struct ZeroColumnMut<'a, T> { +/// Carries an exclusive matrix borrow across Rayon workers. +pub(super) struct ParMut<'a, T> { ptr: NonNull, + layout: Layout, _lifetime: PhantomData<&'a mut [T]>, } -// SAFETY: `ZeroColumnMut` has the ownership semantics of an empty `&mut [T]`. +// SAFETY: `ParMut` owns an exclusive slice borrow, so sending it requires `T: Send`. #[cfg(feature = "rayon")] -unsafe impl Send for ZeroColumnMut<'_, T> {} -// SAFETY: Shared access can only create zero-length mutable views, so it cannot expose -// overlapping access to any element. +unsafe impl Send for ParMut<'_, T> {} +// SAFETY: The only methods that produce mutable views are unsafe and require callers to +// ensure disjointness. Sending those views between workers requires `T: Send`. #[cfg(feature = "rayon")] -unsafe impl Sync for ZeroColumnMut<'_, T> {} +unsafe impl Sync for ParMut<'_, T> {} #[cfg(feature = "rayon")] -impl<'a, T> ZeroColumnMut<'a, T> { +impl<'a, T> ParMut<'a, T> { pub(super) fn new(matrix: &'a mut M) -> Self where M: MatrixMut + ?Sized, { - debug_assert_eq!(matrix.ncols(), 0); + let layout = matrix.layout(); + let ptr = matrix.as_nonnull_mut(); Self { - ptr: matrix.as_nonnull_mut(), + ptr, + layout, _lifetime: PhantomData, } } - pub(super) fn row(&self) -> &'a mut [T] { - // SAFETY: The matrix pointer is valid for its zero-element layout. The returned - // slice preserves that pointer and cannot access or overlap any element. - unsafe { std::slice::from_raw_parts_mut(self.ptr.as_ptr(), 0) } + pub(super) fn nrows(&self) -> usize { + self.layout.nrows() } - pub(super) fn window(&self, nrows: usize) -> Mut<'a, T> { + /// # Safety + /// + /// * `row < self.nrows()`. + /// * No other live reference derived from this `ParMut` may overlap any element of + /// row `row`. Zero-column rows never overlap, even when their addresses match. + pub(super) unsafe fn row_disjoint_unchecked(&self, row: usize) -> &'a mut [T] { + debug_assert!(row < self.layout.nrows()); + let ncols = self.layout.ncols(); + + // SAFETY: The caller guarantees that `row` is in-bounds and does not overlap any + // other live view. The validated parent layout makes the offset representable and + // places the row within the initialized matrix span. + unsafe { std::slice::from_raw_parts_mut(self.ptr.as_ptr().add(row * ncols), ncols) } + } + + /// # Safety + /// + /// * `rows.start <= rows.end <= self.nrows()`. + /// * No other live reference derived from this `ParMut` may overlap any element in + /// `rows`. Zero-column windows never overlap. + pub(super) unsafe fn window_disjoint_unchecked( + &self, + rows: std::ops::Range, + ) -> Mut<'a, T> { + debug_assert!(rows.start <= rows.end); + debug_assert!(rows.end <= self.layout.nrows()); + + let ncols = self.layout.ncols(); + let nrows = rows.end - rows.start; + + // SAFETY: The caller guarantees an ordered, in-bounds range. The validated parent + // layout makes the offset representable and places it within or one past the matrix. + let ptr = unsafe { self.ptr.add(rows.start * ncols) }; + Mut { - ptr: self.ptr, - // SAFETY: Every zero-column layout has zero elements. - layout: unsafe { Layout::new_unchecked(nrows, 0) }, + ptr, + // SAFETY: This window has no more rows than the validated parent and keeps its + // column count, so its element count and byte span cannot exceed the parent. + layout: unsafe { Layout::new_unchecked(nrows, ncols) }, _lifetime: PhantomData, } } } + +#[cfg(all(test, feature = "rayon"))] +mod tests { + use super::ParMut; + use crate::views::rowmajor::{Matrix, MatrixMut, Owned}; + + #[test] + fn par_mut_zero_column_views_can_coexist() { + let mut matrix = Owned::from_element(usize::MAX, 0, 0); + let ptr = matrix.as_ptr(); + let matrix = ParMut::new(&mut matrix); + + // SAFETY: Empty views do not overlap any elements, even when their addresses match. + let rows = unsafe { + [ + matrix.row_disjoint_unchecked(0), + matrix.row_disjoint_unchecked(1), + matrix.row_disjoint_unchecked(usize::MAX - 1), + ] + }; + // SAFETY: Empty windows do not overlap any elements, including the live row views + // whose logical rows fall within these windows. + let windows = unsafe { + [ + matrix.window_disjoint_unchecked(0..2), + matrix.window_disjoint_unchecked(2..usize::MAX), + ] + }; + + assert!(rows.iter().all(|row| row.is_empty() && row.as_ptr() == ptr)); + assert!(windows.iter().all(|window| { + window.ncols() == 0 && window.as_slice().is_empty() && window.as_ptr() == ptr + })); + } + + #[test] + fn par_mut_disjoint_nonempty_views_can_coexist() { + let mut matrix = Owned::from_fn(6, 2, |rc| rc.row * 100 + rc.col); + { + let matrix = ParMut::new(&mut matrix); + + // SAFETY: These rows and windows are in-bounds and pairwise disjoint. + let (row0, row1, mut window2, mut window4) = unsafe { + ( + matrix.row_disjoint_unchecked(0), + matrix.row_disjoint_unchecked(1), + matrix.window_disjoint_unchecked(2..4), + matrix.window_disjoint_unchecked(4..6), + ) + }; + row0[0] = 10; + row1[1] = 11; + *window2.element_mut(0, 0) = 20; + *window2.element_mut(1, 1) = 31; + *window4.element_mut(0, 0) = 40; + *window4.element_mut(1, 1) = 51; + } + + assert_eq!( + matrix.as_slice(), + &[10, 1, 100, 11, 20, 201, 300, 31, 40, 401, 500, 51] + ); + } +}