Skip to content
Open
Show file tree
Hide file tree
Changes from 5 commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

12 changes: 6 additions & 6 deletions diskann-benchmark/src/multi_vector/driver.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ use diskann_benchmark_runner::{
},
Checker, Input,
};
use diskann_quantization::multi_vector::{Mat, MatRef, MaxSimKernel, Overflow, Standard};
use diskann_quantization::multi_vector::{Mat, MatRef, MaxSimKernel, Overflow, RowMajor};
use rand::{
distr::{Distribution, StandardUniform},
rngs::StdRng,
Expand Down Expand Up @@ -68,8 +68,8 @@ impl Input for MultiVectorTolerance {

/// Random query / doc fixture for a single benchmark run.
pub(super) struct Data<T: Copy> {
pub(super) queries: Mat<Standard<T>>,
pub(super) docs: Mat<Standard<T>>,
pub(super) queries: Mat<RowMajor<T>>,
pub(super) docs: Mat<RowMajor<T>>,
}

impl<T: Copy> Data<T>
Expand All @@ -79,11 +79,11 @@ where
pub(super) fn new(run: &Run) -> Result<Self, Overflow> {
let mut rng = StdRng::seed_from_u64(0x12345);
let queries = Mat::from_fn(
Standard::new(run.num_query_vectors.get(), run.dim.get())?,
RowMajor::new(run.num_query_vectors.get(), run.dim.get())?,
|| StandardUniform.sample(&mut rng),
);
Comment on lines +80 to 86
let docs = Mat::from_fn(
Standard::new(run.num_doc_vectors.get(), run.dim.get())?,
RowMajor::new(run.num_doc_vectors.get(), run.dim.get())?,
|| StandardUniform.sample(&mut rng),
);
Ok(Self { queries, docs })
Expand All @@ -96,7 +96,7 @@ where

pub(super) fn run_with_kernel<T: Copy>(
run: &Run,
doc: MatRef<'_, Standard<T>>,
doc: MatRef<'_, RowMajor<T>>,
kernel: &dyn MaxSimKernel<T>,
) -> RunResult {
let mut scores = vec![0.0f32; run.num_query_vectors.get()];
Expand Down
2 changes: 1 addition & 1 deletion diskann-disk/src/search/pq/quantizer_preprocess.rs
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ pub fn quantizer_preprocess(
) -> 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::MatrixViewMut::try_from(
&mut (*pq_scratch.aligned_pqtable_dist_scratch)[..expected_len],
table.nchunks(),
table.ncenters(),
Expand Down
4 changes: 2 additions & 2 deletions diskann-disk/src/storage/quant/compressor.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
*/

use diskann::{utils::VectorRepr, ANNResult};
use diskann_utils::views::{MatrixView, MutMatrixView};
use diskann_utils::views::{MatrixView, MatrixViewMut};

/// [`QuantCompressor`] defines the interface for quantizer with [`QuantDataGenerator`]
///
Expand All @@ -30,6 +30,6 @@ where
type CompressorContext;

fn new(context: &Self::CompressorContext) -> ANNResult<Self>;
fn compress(&self, vector: MatrixView<f32>, output: MutMatrixView<u8>) -> ANNResult<()>;
fn compress(&self, vector: MatrixView<f32>, output: MatrixViewMut<u8>) -> ANNResult<()>;
fn compressed_bytes(&self) -> usize;
}
4 changes: 2 additions & 2 deletions diskann-disk/src/storage/quant/generator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -147,7 +147,7 @@ where

// Wrap the data in `MatrixViews` 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::MatrixViewMut::try_from(
block_compressed_base,
cur_block_size,
compressed_size,
Expand Down Expand Up @@ -221,7 +221,7 @@ mod generator_tests {
fn compress(
&self,
_vector: views::MatrixView<f32>,
mut output: views::MutMatrixView<u8>,
mut output: views::MatrixViewMut<u8>,
) -> ANNResult<()> {
output
.row_iter_mut()
Expand Down
12 changes: 6 additions & 6 deletions diskann-disk/src/storage/quant/pq/pq_generation.rs
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ use diskann_providers::{
utils::{BridgeErr, RayonThreadPoolRef},
};
use diskann_quantization::{product::TransposedTable, CompressInto};
use diskann_utils::views::MatrixBase;
use diskann_utils::views::{MatrixView, MatrixViewMut};
use diskann_vector::distance::Metric;
use tracing::info;

Expand Down Expand Up @@ -119,7 +119,7 @@ where
context.storage_provider,
)?;

let mut full_pivot_data_mat = diskann_utils::views::MutMatrixView::try_from(
let mut full_pivot_data_mat = diskann_utils::views::MatrixViewMut::try_from(
full_pivot_data.as_mut_slice(),
context.num_centers,
full_dim,
Expand All @@ -146,8 +146,8 @@ where

fn compress(
&self,
vector: MatrixBase<&[f32]>,
output: MatrixBase<&mut [u8]>,
vector: MatrixView<'_, f32>,
output: MatrixViewMut<'_, u8>,
) -> Result<(), diskann::ANNError> {
self.table
.compress_into(vector, output)
Expand Down Expand Up @@ -175,7 +175,7 @@ mod pq_generation_tests {
use diskann_utils::{
io::{read_bin, write_bin},
test_data_root,
views::{MatrixView, MutMatrixView},
views::{MatrixView, MatrixViewMut},
};
use diskann_vector::distance::Metric;
use rstest::rstest;
Expand Down Expand Up @@ -329,7 +329,7 @@ mod pq_generation_tests {
let mut compressed_mat = vec![0_u8; num_chunks * npts];
let result = compressor.unwrap().compress(
data_matrix.as_view(),
MutMatrixView::try_from(&mut compressed_mat, npts, num_chunks).unwrap(),
MatrixViewMut::try_from(&mut compressed_mat, npts, num_chunks).unwrap(),
);
assert!(result.is_ok());

Expand Down
2 changes: 1 addition & 1 deletion diskann-garnet/src/provider.rs
Original file line number Diff line number Diff line change
Expand Up @@ -499,7 +499,7 @@ impl<T: VectorRepr> GarnetProvider<T> {
Ok(v) => v,
Err(_) => return false,
};
let view = match MatrixView::try_from(&*converted, view.nrows(), view.ncols()) {
let view = match MatrixView::try_from(&converted, view.nrows(), view.ncols()) {
Ok(v) => v,
Err(_) => return false,
};
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,7 @@

use std::fmt;

use diskann_utils::views::MutMatrixView;
use diskann_utils::views::MatrixViewMut;
use diskann_vector::{MathematicalValue, PureDistanceFunction, distance::InnerProduct};

/// Parameters for Determinant-Diversity post-processor with validation.
Expand Down Expand Up @@ -199,7 +199,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: MatrixViewMut<'_, f32>,
distances: &[f32],
query: &[f32],
k: usize,
Expand Down Expand Up @@ -338,7 +338,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: MatrixViewMut<'_, f32>,
distances: &[f32],
k: usize,
power: f32,
Expand Down
4 changes: 2 additions & 2 deletions diskann-providers/src/model/pq/fixed_chunk_pq_table.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ use diskann_quantization::{
product::{self, BasicTable},
views::ChunkOffsetsBase,
};
use diskann_utils::views::{self, MatrixBase, MatrixView};
use diskann_utils::views::{self, Matrix, MatrixView};
use diskann_vector::{PureDistanceFunction, distance};
use diskann_wide::ARCH;

Expand Down Expand Up @@ -133,7 +133,7 @@ impl FixedChunkPQTable {
pub fn new(dim: usize, pq_table: Box<[f32]>, chunk_offsets: Box<[usize]>) -> ANNResult<Self> {
let len = pq_table.len();
let table = BasicTable::new(
MatrixBase::try_from(pq_table, len / dim, dim).bridge_err()?,
Matrix::try_from(pq_table, len / dim, dim).bridge_err()?,
ChunkOffsetsBase::new(chunk_offsets).bridge_err()?,
)
.map_err(|err| ANNError::log_pq_error(diskann_quantization::error::format(&err)))?;
Expand Down
8 changes: 4 additions & 4 deletions diskann-providers/src/model/pq/pq_construction.rs
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ use diskann_quantization::{
};
use diskann_utils::{
io::Metadata,
views::{MatrixView, MutMatrixView},
views::{MatrixView, MatrixViewMut},
};
use rand::{Rng, distr::Distribution};
use rayon::prelude::*;
Expand Down Expand Up @@ -293,7 +293,7 @@ pub fn move_train_data_by_centroid(
/// # Panics
///
/// Panics if `y.len() != x.ncols()`.
pub fn accum_row_inplace<T>(mut x: MutMatrixView<T>, y: &[T])
pub fn accum_row_inplace<T>(mut x: MatrixViewMut<T>, y: &[T])
where
T: Copy + std::ops::AddAssign,
{
Expand Down Expand Up @@ -362,7 +362,7 @@ where
// Instead of subtracting the center from each data set component, we instead
// add it to each center.
let mut full_pivot_data_mat =
MutMatrixView::try_from(full_pivot_data.as_mut_slice(), num_centers, full_dim)
MatrixViewMut::try_from(full_pivot_data.as_mut_slice(), num_centers, full_dim)
.bridge_err()?;
accum_row_inplace(full_pivot_data_mat.as_mut_view(), centroid.as_slice());

Expand Down Expand Up @@ -430,7 +430,7 @@ where
// Wrap the data in `MatrixViews` 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)
MatrixViewMut::try_from(&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()?;

Expand Down
2 changes: 1 addition & 1 deletion diskann-providers/src/model/pq/views.rs
Original file line number Diff line number Diff line change
Expand Up @@ -106,7 +106,7 @@ mod tests {
let data = vec![0; ncols * nrows];

test_error(|| {
views::MatrixView::try_from(&*data, nrows, ncols + 1)
views::MatrixView::try_from(&data, nrows, ncols + 1)
.bridge_err()
.unwrap_err()
});
Expand Down
4 changes: 2 additions & 2 deletions diskann-providers/src/storage/pq_storage.rs
Original file line number Diff line number Diff line change
Expand Up @@ -523,7 +523,7 @@ mod pq_storage_tests {
#[test]
fn write_read_roundtrip_with_legacy_centroid() {
use crate::model::pq::accum_row_inplace;
use diskann_utils::views::MutMatrixView;
use diskann_utils::views::MatrixViewMut;

let storage_provider = VirtualStorageProvider::new_memory();
let pivot_path = "/roundtrip_legacy_centroid_pivots.bin";
Expand Down Expand Up @@ -563,7 +563,7 @@ mod pq_storage_tests {

// Fold the centroid into the pivots — this is what production callers do.
let mut pivot_mat =
MutMatrixView::try_from(loaded_pivots.as_mut_slice(), num_centers, dim).unwrap();
MatrixViewMut::try_from(loaded_pivots.as_mut_slice(), num_centers, dim).unwrap();
accum_row_inplace(pivot_mat.as_mut_view(), &loaded_centroid);

// Each pivot row should have the centroid added element-wise.
Expand Down
8 changes: 4 additions & 4 deletions diskann-quantization/src/algorithms/kmeans/lloyds.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ use super::common::square_norm;
use crate::multi_vector::{BlockTransposed, BlockTransposedRef};
use diskann_utils::{
strided::StridedView,
views::{Matrix, MatrixView, MutMatrixView},
views::{Matrix, MatrixView, MatrixViewMut},
};

////////////////////////////////
Expand Down Expand Up @@ -342,7 +342,7 @@ fn update((d0, i0): (f32s, u32s), (d1, i1): (f32s, u32s)) -> (f32s, u32s) {
// Update Step //
/////////////////

fn update_centroids(mut centers: MutMatrixView<'_, f32>, data: StridedView<'_, f32>, map: &[u32]) {
fn update_centroids(mut centers: MatrixViewMut<'_, f32>, data: StridedView<'_, f32>, map: &[u32]) {
let mut sums = Matrix::<f64>::new(0.0, centers.nrows(), centers.ncols());
let mut counts: Vec<u32> = vec![0; centers.nrows()];
data.row_iter().zip(map.iter()).for_each(|(row, &center)| {
Expand Down Expand Up @@ -373,7 +373,7 @@ pub(crate) fn lloyds_inner(
data: StridedView<'_, f32>,
square_norms: &[f32],
transpose: BlockTransposedRef<'_, f32, 16>,
mut centers: MutMatrixView<'_, f32>,
mut centers: MatrixViewMut<'_, f32>,
max_reps: usize,
) -> (Vec<u32>, f32) {
// Check our requirements.
Expand Down Expand Up @@ -440,7 +440,7 @@ pub(crate) fn lloyds_inner(
/// dimension.
pub fn lloyds(
data: MatrixView<'_, f32>,
centers: MutMatrixView<'_, f32>,
centers: MatrixViewMut<'_, f32>,
max_reps: usize,
) -> (Vec<u32>, f32) {
assert_eq!(
Expand Down
8 changes: 4 additions & 4 deletions diskann-quantization/src/algorithms/kmeans/plusplus.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ use std::{collections::HashSet, fmt};

use diskann_utils::{
strided::StridedView,
views::{MatrixView, MutMatrixView},
views::{MatrixView, MatrixViewMut},
};
use diskann_wide::{SIMDMulAdd, SIMDPartialOrd, SIMDSelect, SIMDVector};
use rand::{
Expand Down Expand Up @@ -379,7 +379,7 @@ impl KMeansPlusPlusError {
}

pub(crate) fn kmeans_plusplus_into_inner<const N: usize>(
mut points: MutMatrixView<'_, f32>,
mut points: MatrixViewMut<'_, f32>,
data: StridedView<'_, f32>,
transpose: BlockTransposedRef<'_, f32, N>,
norms: &[f32],
Expand Down Expand Up @@ -498,7 +498,7 @@ where
}

pub fn kmeans_plusplus_into(
centers: MutMatrixView<'_, f32>,
centers: MatrixViewMut<'_, f32>,
data: MatrixView<'_, f32>,
rng: &mut dyn RngCore,
) -> Result<(), KMeansPlusPlusError> {
Expand Down Expand Up @@ -581,7 +581,7 @@ mod tests {
/// ...
/// K-1, K, K+1, K+3 ... N+K-2
/// ```
fn set_default_values(mut x: MutMatrixView<'_, f32>) {
fn set_default_values(mut x: MatrixViewMut<'_, f32>) {
for (i, row) in x.row_iter_mut().enumerate() {
for (j, r) in row.iter_mut().enumerate() {
*r = (i + j) as f32;
Expand Down
1 change: 0 additions & 1 deletion diskann-quantization/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -216,7 +216,6 @@ mod tests {
t.pass("tests/compile-fail/bootstrap/bootstrap.rs");
t.compile_fail("tests/compile-fail/*.rs");
t.compile_fail("tests/compile-fail/error/*.rs");
t.compile_fail("tests/compile-fail/multi/*.rs");
}
}

Expand Down
10 changes: 5 additions & 5 deletions diskann-quantization/src/minmax/multi/max_sim.rs
Original file line number Diff line number Diff line change
Expand Up @@ -144,7 +144,7 @@ mod tests {
use crate::algorithms::transforms::NullTransform;
use crate::bits::{Representation, Unsigned};
use crate::minmax::{Data, MinMaxQuantizer};
use crate::multi_vector::{Defaulted, Mat, Standard};
use crate::multi_vector::{Mat, RowMajor};
use crate::num::Positive;
use diskann_utils::ReborrowMut;
use std::num::NonZeroUsize;
Expand Down Expand Up @@ -202,9 +202,9 @@ mod tests {
where
Unsigned: Representation<NBITS>,
{
let input_mat = MatRef::new(Standard::<f32>::new(n, dim).unwrap(), input).unwrap();
let input_mat = MatRef::from_repr(RowMajor::<f32>::new(n, dim).unwrap(), input).unwrap();
let mut output: Mat<MinMaxMeta<NBITS>> =
Mat::new(MinMaxMeta::new(n, dim), Defaulted).unwrap();
Mat::from_default(MinMaxMeta::new(n, dim)).unwrap();
quantizer
.compress_into(input_mat, output.reborrow_mut())
.unwrap();
Expand Down Expand Up @@ -301,10 +301,10 @@ mod tests {
let query_data = vec![0u8; 2 * row_bytes];
let doc_data = vec![0u8; 3 * row_bytes];

let query: QueryMatRef<_> = MatRef::new(MinMaxMeta::<8>::new(2, dim), &query_data)
let query: QueryMatRef<_> = MatRef::from_repr(MinMaxMeta::<8>::new(2, dim), &query_data)
.unwrap()
.into();
let doc = MatRef::new(MinMaxMeta::<8>::new(3, dim), &doc_data).unwrap();
let doc = MatRef::from_repr(MinMaxMeta::<8>::new(3, dim), &doc_data).unwrap();

let mut scores = vec![0.0f32; 5]; // Wrong size
MaxSim::new(&mut scores).evaluate(query, doc);
Expand Down
Loading
Loading