diff --git a/diskann-disk/benches/benchmarks/aligned_file_reader_bench.rs b/diskann-disk/benches/benchmarks/aligned_file_reader_bench.rs index 4988b0b19..d301af5a0 100644 --- a/diskann-disk/benches/benchmarks/aligned_file_reader_bench.rs +++ b/diskann-disk/benches/benchmarks/aligned_file_reader_bench.rs @@ -9,7 +9,7 @@ use diskann_disk::utils::aligned_file_reader::{ traits::{AlignedFileReader, AlignedReaderFactory}, AlignedFileReaderFactory, AlignedRead, }; -use diskann_providers::common::AlignedBoxWithSlice; +use diskann_quantization::{alloc::aligned_slice, num::PowerOfTwo}; pub const TEST_INDEX_PATH: &str = "../test_data/disk_index_misc/disk_index_siftsmall_learn_256pts_R4_L50_A1.2_aligned_reader_test.index"; @@ -33,12 +33,11 @@ pub fn benchmark_aligned_file_reader(c: &mut Criterion) { let read_length = 512; let num_read = MAX_IO_CONCURRENCY * 100; // The LinuxAlignedFileReader batches reads according to MAX_IO_CONCURRENCY. Make sure we have many batches to handle. - let mut aligned_mem = AlignedBoxWithSlice::::new(read_length * num_read, 512).unwrap(); + let mut aligned_mem = + aligned_slice::(read_length * num_read, PowerOfTwo::new(512).unwrap()).unwrap(); // create and add AlignedReads to the vector - let mut mem_slices = aligned_mem - .split_into_nonoverlapping_mut_slices(0..aligned_mem.len(), read_length) - .unwrap(); + let mut mem_slices: Vec<&mut [u8]> = aligned_mem.chunks_mut(read_length).collect(); // Read the same data from disk over and over again. We guarantee that it is not all zeros. let mut aligned_reads: Vec> = mem_slices diff --git a/diskann-disk/benches/benchmarks_iai/aligned_file_reader_bench_iai.rs b/diskann-disk/benches/benchmarks_iai/aligned_file_reader_bench_iai.rs index 2261e1729..8ff50f167 100644 --- a/diskann-disk/benches/benchmarks_iai/aligned_file_reader_bench_iai.rs +++ b/diskann-disk/benches/benchmarks_iai/aligned_file_reader_bench_iai.rs @@ -7,7 +7,7 @@ use diskann_disk::utils::aligned_file_reader::{ traits::{AlignedFileReader, AlignedReaderFactory}, AlignedFileReaderFactory, AlignedRead, }; -use diskann_providers::common::AlignedBoxWithSlice; +use diskann_quantization::{alloc::aligned_slice, num::PowerOfTwo}; pub const TEST_INDEX_PATH: &str = "../test_data/disk_index_misc/disk_index_siftsmall_learn_256pts_R4_L50_A1.2_aligned_reader_test.index"; @@ -24,12 +24,11 @@ pub fn benchmark_aligned_file_reader_iai() { let read_length = 512; let num_read = MAX_IO_CONCURRENCY * 100; // The LinuxAlignedFileReader batches reads according to MAX_IO_CONCURRENCY. Make sure we have many batches to handle. - let mut aligned_mem = AlignedBoxWithSlice::::new(read_length * num_read, 512).unwrap(); + let mut aligned_mem = + aligned_slice::(read_length * num_read, PowerOfTwo::new(512).unwrap()).unwrap(); // create and add AlignedReads to the vector - let mut mem_slices = aligned_mem - .split_into_nonoverlapping_mut_slices(0..aligned_mem.len(), read_length) - .unwrap(); + let mut mem_slices: Vec<&mut [u8]> = aligned_mem.chunks_mut(read_length).collect(); // Read the same data from disk over and over again. We guarantee that it is not all zeros. let mut aligned_reads: Vec> = mem_slices diff --git a/diskann-disk/src/search/pq/pq_scratch.rs b/diskann-disk/src/search/pq/pq_scratch.rs index 49a6fa2f9..bd7a8dfc3 100644 --- a/diskann-disk/src/search/pq/pq_scratch.rs +++ b/diskann-disk/src/search/pq/pq_scratch.rs @@ -6,22 +6,25 @@ use diskann::{error::IntoANNResult, utils::VectorRepr, ANNError, ANNResult}; -use diskann_providers::common::AlignedBoxWithSlice; +use diskann_quantization::{ + alloc::{aligned_slice, AlignedSlice}, + num::PowerOfTwo, +}; #[derive(Debug)] /// PQ scratch pub struct PQScratch { /// Aligned pq table distance scratch, the length must be at least [256 * NCHUNKS]. 256 is the number of PQ centroids. /// This is used to store the distance between each chunk in the query vector to each centroid, which is why the length is num of centroids * num of chunks - pub aligned_pqtable_dist_scratch: AlignedBoxWithSlice, + pub aligned_pqtable_dist_scratch: AlignedSlice, /// Aligned dist scratch, must be at least diskann MAX_DEGREE /// This is used to temporarily save the pq distance between query vector to the candidate vectors. - pub aligned_dist_scratch: AlignedBoxWithSlice, + pub aligned_dist_scratch: AlignedSlice, /// Aligned pq coord scratch, must be at least [N_CHUNKS * MAX_DEGREE] /// This is used to store the pq coordinates of the candidate vectors. - pub aligned_pq_coord_scratch: AlignedBoxWithSlice, + pub aligned_pq_coord_scratch: AlignedSlice, /// Query scratch buffer stored as `f32`. `set` initializes it by copying/converting the /// raw query values; `PQTable.PreprocessQuery` can then rotate or otherwise preprocess it. @@ -30,7 +33,10 @@ pub struct PQScratch { impl PQScratch { /// 128 bytes alignment to optimize for the L2 Adjacent Cache Line Prefetcher. - const ALIGNED_ALLOC_128: usize = 128; + const ALIGNED_ALLOC_128: PowerOfTwo = match PowerOfTwo::new(128) { + Ok(v) => v, + Err(_) => unreachable!(), + }; /// Create a new pq scratch pub fn new( @@ -40,11 +46,13 @@ impl PQScratch { num_centers: usize, ) -> ANNResult { let aligned_pq_coord_scratch = - AlignedBoxWithSlice::new(graph_degree * num_pq_chunks, PQScratch::ALIGNED_ALLOC_128)?; + aligned_slice(graph_degree * num_pq_chunks, PQScratch::ALIGNED_ALLOC_128) + .map_err(ANNError::log_index_error)?; let aligned_pqtable_dist_scratch = - AlignedBoxWithSlice::new(num_centers * num_pq_chunks, PQScratch::ALIGNED_ALLOC_128)?; - let aligned_dist_scratch = - AlignedBoxWithSlice::new(graph_degree, PQScratch::ALIGNED_ALLOC_128)?; + aligned_slice(num_centers * num_pq_chunks, PQScratch::ALIGNED_ALLOC_128) + .map_err(ANNError::log_index_error)?; + let aligned_dist_scratch = aligned_slice(graph_degree, PQScratch::ALIGNED_ALLOC_128) + .map_err(ANNError::log_index_error)?; let rotated_query = vec![0.0f32; dim]; Ok(Self { @@ -102,18 +110,20 @@ mod tests { let mut pq_scratch: PQScratch = PQScratch::new(graph_degree, dim, num_pq_chunks, num_centers).unwrap(); - // Check alignment of the remaining AlignedBoxWithSlice buffers. + // Check alignment of the remaining AlignedSlice buffers. assert_eq!( (pq_scratch.aligned_pqtable_dist_scratch.as_ptr() as usize) - % PQScratch::ALIGNED_ALLOC_128, + % PQScratch::ALIGNED_ALLOC_128.raw(), 0 ); assert_eq!( - (pq_scratch.aligned_dist_scratch.as_ptr() as usize) % PQScratch::ALIGNED_ALLOC_128, + (pq_scratch.aligned_dist_scratch.as_ptr() as usize) + % PQScratch::ALIGNED_ALLOC_128.raw(), 0 ); assert_eq!( - (pq_scratch.aligned_pq_coord_scratch.as_ptr() as usize) % PQScratch::ALIGNED_ALLOC_128, + (pq_scratch.aligned_pq_coord_scratch.as_ptr() as usize) + % PQScratch::ALIGNED_ALLOC_128.raw(), 0 ); diff --git a/diskann-disk/src/search/pq/quantizer_preprocess.rs b/diskann-disk/src/search/pq/quantizer_preprocess.rs index 4e75e778f..cc454ea7b 100644 --- a/diskann-disk/src/search/pq/quantizer_preprocess.rs +++ b/diskann-disk/src/search/pq/quantizer_preprocess.rs @@ -26,7 +26,7 @@ pub fn quantizer_preprocess( let dim = table.dim(); let expected_len = table.ncenters() * table.nchunks(); let dst = diskann_utils::views::MutMatrixView::try_from( - &mut pq_scratch.aligned_pqtable_dist_scratch.as_mut_slice()[..expected_len], + &mut (*pq_scratch.aligned_pqtable_dist_scratch)[..expected_len], table.nchunks(), table.ncenters(), ) @@ -69,13 +69,13 @@ pub fn quantizer_preprocess( // Compute the distance between each chunk of the query to each pq centroids. table.populate_chunk_distances( pq_scratch.rotated_query.as_slice(), - pq_scratch.aligned_pqtable_dist_scratch.as_mut_slice(), + &mut pq_scratch.aligned_pqtable_dist_scratch, )?; } Metric::InnerProduct => { table.populate_chunk_inner_products( pq_scratch.rotated_query.as_slice(), - pq_scratch.aligned_pqtable_dist_scratch.as_mut_slice(), + &mut pq_scratch.aligned_pqtable_dist_scratch, )?; } } diff --git a/diskann-disk/src/search/provider/disk_provider.rs b/diskann-disk/src/search/provider/disk_provider.rs index 076ffbcb6..0403a5019 100644 --- a/diskann-disk/src/search/provider/disk_provider.rs +++ b/diskann-disk/src/search/provider/disk_provider.rs @@ -426,7 +426,6 @@ where .scratch .pq_scratch .aligned_pqtable_dist_scratch - .as_slice() .to_vec(), }) } diff --git a/diskann-disk/src/search/provider/disk_sector_graph.rs b/diskann-disk/src/search/provider/disk_sector_graph.rs index 1f00ad6db..cd3412b57 100644 --- a/diskann-disk/src/search/provider/disk_sector_graph.rs +++ b/diskann-disk/src/search/provider/disk_sector_graph.rs @@ -8,7 +8,10 @@ use std::ops::Deref; use diskann::{ANNError, ANNResult}; -use diskann_providers::common::AlignedBoxWithSlice; +use diskann_quantization::{ + alloc::{aligned_slice, AlignedSlice}, + num::PowerOfTwo, +}; use crate::{ data_model::GraphHeader, @@ -34,7 +37,7 @@ pub struct DiskSectorGraph { /// index info for multi-sector nodes /// node `i` is in sector: [i * max_node_len.div_ceil(block_size)] /// offset in sector: [0] - sectors_data: AlignedBoxWithSlice, + sectors_data: AlignedSlice, /// Current sector index into which the next read reads data cur_sector_idx: u64, @@ -73,10 +76,11 @@ impl DiskSectorGraph { Ok(Self { sector_reader, - sectors_data: AlignedBoxWithSlice::new( + sectors_data: aligned_slice( max_n_batch_sector_read * num_sectors_per_node * block_size, - block_size, - )?, + PowerOfTwo::new(block_size).map_err(ANNError::log_index_error)?, + ) + .map_err(ANNError::log_index_error)?, cur_sector_idx: 0, num_nodes_per_sector, node_len, @@ -90,10 +94,11 @@ impl DiskSectorGraph { pub fn reconfigure(&mut self, max_n_batch_sector_read: usize) -> ANNResult<()> { if max_n_batch_sector_read > self.max_n_batch_sector_read { self.max_n_batch_sector_read = max_n_batch_sector_read; - self.sectors_data = AlignedBoxWithSlice::new( + self.sectors_data = aligned_slice( max_n_batch_sector_read * self.num_sectors_per_node * self.block_size, - self.block_size, - )?; + PowerOfTwo::new(self.block_size).map_err(ANNError::log_index_error)?, + ) + .map_err(ANNError::log_index_error)?; } Ok(()) } @@ -116,11 +121,22 @@ impl DiskSectorGraph { } let len_per_node = self.num_sectors_per_node * self.block_size; - let mut sector_slices = self.sectors_data.split_into_nonoverlapping_mut_slices( - cur_sector_idx_usize * len_per_node - ..(cur_sector_idx_usize + sectors_to_fetch.len()) * len_per_node, - len_per_node, - )?; + if len_per_node == 0 { + return Err(ANNError::log_index_error(format_args!( + "len_per_node is 0 (num_sectors_per_node={}, block_size={})", + self.num_sectors_per_node, self.block_size, + ))); + } + let range = cur_sector_idx_usize * len_per_node + ..(cur_sector_idx_usize + sectors_to_fetch.len()) * len_per_node; + debug_assert!( + range.len() % len_per_node == 0, + "range length {} is not divisible by {}", + range.len(), + len_per_node + ); + let mut sector_slices: Vec<&mut [u8]> = + self.sectors_data[range].chunks_mut(len_per_node).collect(); let mut read_requests = Vec::with_capacity(sector_slices.len()); for (local_sector_idx, slice) in sector_slices.iter_mut().enumerate() { let sector_id = sectors_to_fetch[local_sector_idx]; @@ -204,7 +220,7 @@ mod disk_sector_graph_test { ) -> DiskSectorGraph<::AlignedReaderType> { DiskSectorGraph { - sectors_data: AlignedBoxWithSlice::new(512, 512).unwrap(), + sectors_data: aligned_slice(512, PowerOfTwo::new(512).unwrap()).unwrap(), sector_reader, cur_sector_idx: 0, num_nodes_per_sector, diff --git a/diskann-disk/src/search/provider/disk_vertex_provider_factory.rs b/diskann-disk/src/search/provider/disk_vertex_provider_factory.rs index 116608610..7b5647b87 100644 --- a/diskann-disk/src/search/provider/disk_vertex_provider_factory.rs +++ b/diskann-disk/src/search/provider/disk_vertex_provider_factory.rs @@ -6,7 +6,7 @@ use std::{cmp::min, collections::VecDeque, sync::Arc, time::Instant}; use crate::data_model::GraphDataType; use diskann::{graph::AdjacencyList, utils::TryIntoVectorId, ANNError, ANNResult}; -use diskann_providers::common::AlignedBoxWithSlice; +use diskann_quantization::{alloc::aligned_slice, num::PowerOfTwo}; use hashbrown::HashSet; use tracing::info; @@ -52,14 +52,18 @@ where // since this is the implementation for the disk vertex provider, there're only two kinds of sector lengths: 4096 and 512. // it's okay to hardcoded at this place. let buffer_len = GraphHeader::get_size().next_multiple_of(DEFAULT_DISK_SECTOR_LEN); - let mut read_buf = AlignedBoxWithSlice::::new(buffer_len, buffer_len)?; - let aligned_read = AlignedRead::new(0_u64, read_buf.as_mut_slice())?; + let mut read_buf = aligned_slice::( + buffer_len, + PowerOfTwo::new(buffer_len).map_err(ANNError::log_index_error)?, + ) + .map_err(ANNError::log_index_error)?; + let aligned_read = AlignedRead::new(0_u64, &mut read_buf)?; self.aligned_reader_factory .build()? .read(&mut [aligned_read])?; // Create a GraphHeader from the buffer. - GraphHeader::try_from(&read_buf.as_slice()[8..]) + GraphHeader::try_from(&read_buf[8..]) } fn create_vertex_provider( diff --git a/diskann-disk/src/utils/aligned_file_reader/aligned_read.rs b/diskann-disk/src/utils/aligned_file_reader/aligned_read.rs index 1c55e15a8..782a8d696 100644 --- a/diskann-disk/src/utils/aligned_file_reader/aligned_read.rs +++ b/diskann-disk/src/utils/aligned_file_reader/aligned_read.rs @@ -7,7 +7,7 @@ use diskann::{ANNError, ANNResult}; pub const DISK_IO_ALIGNMENT: usize = 512; -/// Aligned read struct for disk IO, it takes the ownership of the AlignedBoxedSlice and returns the AlignedBoxWithSlice data immutably. +/// Aligned read struct for disk IO. pub struct AlignedRead<'a, T> { /// where to read from /// offset needs to be aligned with DISK_IO_ALIGNMENT diff --git a/diskann-disk/src/utils/aligned_file_reader/linux_aligned_file_reader.rs b/diskann-disk/src/utils/aligned_file_reader/linux_aligned_file_reader.rs index 457563b88..9ad5dba16 100644 --- a/diskann-disk/src/utils/aligned_file_reader/linux_aligned_file_reader.rs +++ b/diskann-disk/src/utils/aligned_file_reader/linux_aligned_file_reader.rs @@ -135,7 +135,7 @@ mod tests { use serde::{Deserialize, Serialize}; use super::*; - use diskann_providers::common::AlignedBoxWithSlice; + use diskann_quantization::{alloc::aligned_slice, num::PowerOfTwo}; pub const TEST_INDEX_PATH: &str = "../test_data/disk_index_misc/disk_index_siftsmall_learn_256pts_R4_L50_A1.2_aligned_reader_test.index"; pub const TRUTH_NODE_DATA_PATH: &str = @@ -170,12 +170,11 @@ mod tests { let read_length = 512; // adjust according to your logic let num_read = 10; - let mut aligned_mem = AlignedBoxWithSlice::::new(read_length * num_read, 512).unwrap(); + let mut aligned_mem = + aligned_slice::(read_length * num_read, PowerOfTwo::new(512).unwrap()).unwrap(); // create and add AlignedReads to the vector - let mut mem_slices = aligned_mem - .split_into_nonoverlapping_mut_slices(0..aligned_mem.len(), read_length) - .unwrap(); + let mut mem_slices: Vec<&mut [u8]> = aligned_mem.chunks_mut(read_length).collect(); let mut aligned_reads: Vec> = mem_slices .iter_mut() @@ -215,12 +214,11 @@ mod tests { let read_length = 512; let num_read = MAX_IO_CONCURRENCY * 100; // The LinuxAlignedFileReader batches reads according to MAX_IO_CONCURRENCY. Make sure we have many batches to handle. - let mut aligned_mem = AlignedBoxWithSlice::::new(read_length * num_read, 512).unwrap(); + let mut aligned_mem = + aligned_slice::(read_length * num_read, PowerOfTwo::new(512).unwrap()).unwrap(); // create and add AlignedReads to the vector - let mut mem_slices = aligned_mem - .split_into_nonoverlapping_mut_slices(0..aligned_mem.len(), read_length) - .unwrap(); + let mut mem_slices: Vec<&mut [u8]> = aligned_mem.chunks_mut(read_length).collect(); // Read the same data from disk over and over again. We guarantee that it is not all zeros. let mut aligned_reads: Vec> = mem_slices @@ -257,12 +255,10 @@ mod tests { let read_length = 512; // adjust according to your logic let num_sector = 10; let mut aligned_mem = - AlignedBoxWithSlice::::new(read_length * num_sector, 512).unwrap(); + aligned_slice::(read_length * num_sector, PowerOfTwo::new(512).unwrap()).unwrap(); // Each slice will be used as the buffer for a read request of a sector. - let mut mem_slices = aligned_mem - .split_into_nonoverlapping_mut_slices(0..aligned_mem.len(), read_length) - .unwrap(); + let mut mem_slices: Vec<&mut [u8]> = aligned_mem.chunks_mut(read_length).collect(); let mut aligned_reads: Vec> = mem_slices .iter_mut() diff --git a/diskann-disk/src/utils/aligned_file_reader/storage_provider_aligned_file_reader.rs b/diskann-disk/src/utils/aligned_file_reader/storage_provider_aligned_file_reader.rs index 2e0b88253..2e296eaf7 100644 --- a/diskann-disk/src/utils/aligned_file_reader/storage_provider_aligned_file_reader.rs +++ b/diskann-disk/src/utils/aligned_file_reader/storage_provider_aligned_file_reader.rs @@ -54,7 +54,7 @@ mod tests { use diskann_utils::test_data_root; use super::*; - use diskann_providers::common::AlignedBoxWithSlice; + use diskann_quantization::{alloc::aligned_slice, num::PowerOfTwo}; fn test_index_path() -> String { "/disk_index_misc/disk_index_siftsmall_learn_256pts_R4_L50_A1.2_aligned_reader_test.index" @@ -78,12 +78,11 @@ mod tests { let read_length = 512; let num_read = 10; - let mut aligned_mem = AlignedBoxWithSlice::::new(read_length * num_read, 512).unwrap(); + let mut aligned_mem = + aligned_slice::(read_length * num_read, PowerOfTwo::new(512).unwrap()).unwrap(); // create and add AlignedReads to the vector - let mut mem_slices = aligned_mem - .split_into_nonoverlapping_mut_slices(0..aligned_mem.len(), read_length) - .unwrap(); + let mut mem_slices: Vec<&mut [u8]> = aligned_mem.chunks_mut(read_length).collect(); let mut aligned_reads: Vec> = mem_slices .iter_mut() diff --git a/diskann-disk/src/utils/aligned_file_reader/windows_aligned_file_reader.rs b/diskann-disk/src/utils/aligned_file_reader/windows_aligned_file_reader.rs index 09b802271..0ac6887f5 100644 --- a/diskann-disk/src/utils/aligned_file_reader/windows_aligned_file_reader.rs +++ b/diskann-disk/src/utils/aligned_file_reader/windows_aligned_file_reader.rs @@ -123,7 +123,7 @@ mod tests { use super::*; use crate::utils::aligned_file_reader::AlignedRead; - use diskann_providers::common::AlignedBoxWithSlice; + use diskann_quantization::{alloc::aligned_slice, num::PowerOfTwo}; fn test_index_path() -> String { test_data_root() @@ -169,12 +169,11 @@ mod tests { let read_length = 512; // adjust according to your logic let num_read = 10; - let mut aligned_mem = AlignedBoxWithSlice::::new(read_length * num_read, 512).unwrap(); + let mut aligned_mem = + aligned_slice::(read_length * num_read, PowerOfTwo::new(512).unwrap()).unwrap(); // create and add AlignedReads to the vector - let mut mem_slices = aligned_mem - .split_into_nonoverlapping_mut_slices(0..aligned_mem.len(), read_length) - .unwrap(); + let mut mem_slices: Vec<&mut [u8]> = aligned_mem.chunks_mut(read_length).collect(); let mut aligned_reads: Vec> = mem_slices .iter_mut() @@ -210,12 +209,10 @@ mod tests { let read_length = DEFAULT_DISK_SECTOR_LEN; // adjust according to your logic let num_sector = 10; let mut aligned_mem = - AlignedBoxWithSlice::::new(read_length * num_sector, 512).unwrap(); + aligned_slice::(read_length * num_sector, PowerOfTwo::new(512).unwrap()).unwrap(); // Each slice will be used as the buffer for a read request of a sector. - let mut mem_slices = aligned_mem - .split_into_nonoverlapping_mut_slices(0..aligned_mem.len(), read_length) - .unwrap(); + let mut mem_slices: Vec<&mut [u8]> = aligned_mem.chunks_mut(read_length).collect(); let mut aligned_reads: Vec> = mem_slices .iter_mut() diff --git a/diskann-providers/src/common/aligned_allocator.rs b/diskann-providers/src/common/aligned_allocator.rs deleted file mode 100644 index c70d8ded6..000000000 --- a/diskann-providers/src/common/aligned_allocator.rs +++ /dev/null @@ -1,292 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ -#![warn(missing_debug_implementations, missing_docs)] - -//! Aligned allocator - -use std::{ - ops::{Deref, DerefMut, Range}, - ptr::copy_nonoverlapping, -}; - -use diskann::{ANNError, ANNResult}; -use diskann_quantization::{ - alloc::{AlignedAllocator, Poly}, - num::PowerOfTwo, -}; - -/// A box that holds a slice but is aligned to the specified layout for potential -/// cache efficiency improvements. -#[derive(Debug)] -pub struct AlignedBoxWithSlice { - val: Poly<[T], AlignedAllocator>, -} - -impl AlignedBoxWithSlice -where - T: Default, -{ - /// Creates a new `AlignedBoxWithSlice` with the given capacity and alignment. - /// The allocated memory are set to `T::default()`.. - /// - /// # Error - /// - /// Return IndexError if the alignment is not a power of two or if the layout is invalid. - pub fn new(capacity: usize, alignment: usize) -> ANNResult { - let allocator = - AlignedAllocator::new(PowerOfTwo::new(alignment).map_err(ANNError::log_index_error)?); - let val = Poly::from_iter((0..capacity).map(|_| T::default()), allocator) - .map_err(ANNError::log_index_error)?; - Ok(Self { val }) - } - - /// Returns a reference to the slice. - pub fn as_slice(&self) -> &[T] { - &self.val - } - - /// Returns a mutable reference to the slice. - pub fn as_mut_slice(&mut self) -> &mut [T] { - &mut self.val - } - - /// Copies data from the source slice to the destination box. - pub fn memcpy(&mut self, src: &[T]) -> ANNResult<()> { - if src.len() > self.val.len() { - return Err(ANNError::log_index_error(format_args!( - "source slice is too large (src:{}, dst:{})", - src.len(), - self.val.len() - ))); - } - - // Check that they don't overlap - let src_ptr = src.as_ptr(); - let src_end = unsafe { src_ptr.add(src.len()) }; - let dst_ptr = self.val.as_mut_ptr(); - let dst_end = unsafe { dst_ptr.add(self.val.len()) }; - - if src_ptr < dst_end && src_end > dst_ptr { - return Err(ANNError::log_index_error("Source and destination overlap")); - } - - // Call is safe because we checked that - // 1. the destination is large enough. - // 2. the source and destination don't overlap. - unsafe { - copy_nonoverlapping(src.as_ptr(), self.val.as_mut_ptr(), src.len()); - } - - Ok(()) - } - - /// Split the range of memory into nonoverlapping mutable slices. - /// The number of returned slices is (range length / slice_len) and each has a length of slice_len. - pub fn split_into_nonoverlapping_mut_slices( - &mut self, - range: Range, - slice_len: usize, - ) -> ANNResult> { - if !range.len().is_multiple_of(slice_len) || range.end > self.len() { - return Err(ANNError::log_index_error(format_args!( - "Cannot split range ({:?}) of AlignedBoxWithSlice (len: {}) into nonoverlapping mutable slices with length {}", - range, - self.len(), - slice_len, - ))); - } - - let mut slices = Vec::with_capacity(range.len() / slice_len); - let mut remaining_slice = &mut self.val[range]; - - while remaining_slice.len() >= slice_len { - let (left, right) = remaining_slice.split_at_mut(slice_len); - slices.push(left); - remaining_slice = right; - } - - Ok(slices) - } -} - -impl Deref for AlignedBoxWithSlice { - type Target = [T]; - - fn deref(&self) -> &Self::Target { - &self.val - } -} - -impl DerefMut for AlignedBoxWithSlice { - fn deref_mut(&mut self) -> &mut Self::Target { - &mut self.val - } -} - -#[cfg(test)] -mod tests { - #[cfg(not(miri))] - use std::{cell::RefCell, rc::Rc}; - - use diskann::ANNErrorKind; - use rand::Rng; - - use super::*; - - cfg_if::cfg_if! { - if #[cfg(miri)] { - // Use smaller allocations for Miri. - const TEST_SIZE: usize = 100; - } else { - const TEST_SIZE: usize = 1_000_000; - } - } - - #[test] - fn create_alignedvec_works_32() { - (0..100).for_each(|_| { - let size = TEST_SIZE; - let data = AlignedBoxWithSlice::::new(size, 32).unwrap(); - assert_eq!(data.len(), size, "Capacity should match"); - - let ptr = data.as_ptr() as usize; - assert_eq!(ptr % 32, 0, "Ptr should be aligned to 32"); - - // assert that the slice is initialized. - (0..size).for_each(|i| { - assert_eq!(data[i], f32::default()); - }); - - drop(data); - }); - } - - #[test] - fn create_alignedvec_works_256() { - let mut rng = crate::utils::create_rnd_in_tests(); - - (0..100).for_each(|_| { - let n = rng.random::(); - let size = usize::from(n) + 1; - let data = AlignedBoxWithSlice::::new(size, 256).unwrap(); - assert_eq!(data.len(), size, "Capacity should match"); - - let ptr = data.as_ptr() as usize; - assert_eq!(ptr % 256, 0, "Ptr should be aligned to 32"); - - // assert that the slice is initialized. - (0..size).for_each(|i| { - assert_eq!(data[i], u8::default()); - }); - - drop(data); - }); - } - - #[test] - fn create_zero_length_box() { - let x = AlignedBoxWithSlice::::new(0, 16).unwrap(); - assert_eq!(x.len(), 0); - } - - #[test] - fn as_slice_test() { - let size = TEST_SIZE; - let data = AlignedBoxWithSlice::::new(size, 32).unwrap(); - // assert that the slice is initialized. - (0..size).for_each(|i| { - assert_eq!(data[i], f32::default()); - }); - - let slice = data.as_slice(); - (0..size).for_each(|i| { - assert_eq!(slice[i], f32::default()); - }); - } - - #[test] - fn as_mut_slice_test() { - let size = TEST_SIZE; - let mut data = AlignedBoxWithSlice::::new(size, 32).unwrap(); - let mut_slice = data.as_mut_slice(); - (0..size).for_each(|i| { - assert_eq!(mut_slice[i], f32::default()); - }); - } - - #[test] - fn memcpy_test() { - let size = TEST_SIZE; - let mut data = AlignedBoxWithSlice::::new(size, 32).unwrap(); - let mut destination = AlignedBoxWithSlice::::new(size - 2, 32).unwrap(); - let mut_destination = destination.as_mut_slice(); - data.memcpy(mut_destination).unwrap(); - (0..size - 2).for_each(|i| { - assert_eq!(data[i], mut_destination[i]); - }); - } - - #[test] - #[should_panic(expected = "source slice is too large")] - fn memcpy_panic_test() { - let size = TEST_SIZE; - let mut data = AlignedBoxWithSlice::::new(size - 2, 32).unwrap(); - let mut destination = AlignedBoxWithSlice::::new(size, 32).unwrap(); - let mut_destination = destination.as_mut_slice(); - data.memcpy(mut_destination).unwrap(); - } - - // NOTE: This function is wildly unsafe and miri complains about it before we even - // get to the internal check. - // - // In safe Rust, it is impossible to obtain overlapping mutable and immutable slices, - // but the check in this function is a safe-guard against UB as a result of unsafe code, - // so we still need to exercise it. - #[cfg(not(miri))] - #[test] - fn test_memcpy_overlap() { - let aligned_box = Rc::new(RefCell::new(AlignedBoxWithSlice::new(4, 16).unwrap())); - aligned_box - .borrow_mut() - .as_mut_slice() - .copy_from_slice(&[1, 2, 3, 4]); - - let src_ptr = aligned_box.as_ptr(); - let src = &unsafe { src_ptr.as_ref().unwrap() }.as_slice()[0..3]; - - let result = aligned_box.borrow_mut().memcpy(src); - - assert!(result.is_err()); - assert_eq!(result.unwrap_err().kind(), ANNErrorKind::IndexError); - } - - #[test] - fn split_into_nonoverlapping_mut_slices_test() { - let size = 10; - let slice_len = 2; - let mut data = AlignedBoxWithSlice::::new(size, 32).unwrap(); - let slices = data - .split_into_nonoverlapping_mut_slices(2..8, slice_len) - .unwrap(); - assert_eq!(slices.len(), 3); - for (i, slice) in slices.into_iter().enumerate() { - assert_eq!(slice.len(), slice_len); - slice[0] = i as f32 + 1.0; - slice[1] = i as f32 + 1.0; - } - let expected_arr = [0.0f32, 0.0, 1.0, 1.0, 2.0, 2.0, 3.0, 3.0, 0.0, 0.0]; - assert_eq!(data.as_ref(), &expected_arr); - } - - #[test] - fn split_into_nonoverlapping_mut_slices_error_when_indivisible() { - let size = 10; - let slice_len = 2; - let range = 2..7; - let mut data = AlignedBoxWithSlice::::new(size, 32).unwrap(); - let result = data.split_into_nonoverlapping_mut_slices(range.clone(), slice_len); - assert!(result.is_err_and(|e| e.kind() == ANNErrorKind::IndexError)); - } -} diff --git a/diskann-providers/src/common/mod.rs b/diskann-providers/src/common/mod.rs index 7ba437722..b5c4a42a4 100644 --- a/diskann-providers/src/common/mod.rs +++ b/diskann-providers/src/common/mod.rs @@ -2,9 +2,6 @@ * Copyright (c) Microsoft Corporation. * Licensed under the MIT license. */ -mod aligned_allocator; -pub use aligned_allocator::AlignedBoxWithSlice; - mod minmax_repr; pub use minmax_repr::{MinMax4, MinMax8, MinMaxElement}; diff --git a/diskann-quantization/src/alloc/aligned_slice.rs b/diskann-quantization/src/alloc/aligned_slice.rs new file mode 100644 index 000000000..4e15f114f --- /dev/null +++ b/diskann-quantization/src/alloc/aligned_slice.rs @@ -0,0 +1,76 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use super::Poly; +use crate::alloc::{AlignedAllocator, AllocatorError}; + +use crate::num::PowerOfTwo; + +/// Type alias for an aligned, heap-allocated slice +/// +/// Shorthand for `Poly<[T], AlignedAllocator>` which is intended +/// for allocations requiring specific alignment (e.g. cache-line or disk-sector alignment) +pub type AlignedSlice = Poly<[T], AlignedAllocator>; + +/// Create a new [`AlignedSlice`] with the given capacity and alignment +/// initialized to `T::default()` +pub fn aligned_slice( + capacity: usize, + alignment: PowerOfTwo, +) -> Result, AllocatorError> { + let allocator = AlignedAllocator::new(alignment); + Poly::from_iter((0..capacity).map(|_| T::default()), allocator) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn aligned_slice_alignment_32() { + let data = aligned_slice::(1000, PowerOfTwo::new(32).unwrap()).unwrap(); + assert_eq!(data.len(), 1000); + assert_eq!(data.as_ptr() as usize % 32, 0); + } + + #[test] + fn aligned_slice_alignment_256() { + let data = aligned_slice::(500, PowerOfTwo::new(256).unwrap()).unwrap(); + assert_eq!(data.len(), 500); + assert_eq!(data.as_ptr() as usize % 256, 0); + } + + #[test] + fn aligned_slice_alignment_512() { + let data = aligned_slice::(4096, PowerOfTwo::new(512).unwrap()).unwrap(); + assert_eq!(data.len(), 4096); + assert_eq!(data.as_ptr() as usize % 512, 0); + } + + #[test] + fn aligned_slice_zero_length() { + // Zero-length aligned slices should succeed: `Poly::from_iter` + // special-cases empty iterators and returns an empty slice. + let data = aligned_slice::(0, PowerOfTwo::new(16).unwrap()).unwrap(); + assert!(data.is_empty()); + assert_eq!(data.len(), 0); + } + + #[test] + fn aligned_slice_default_initialized() { + let data = aligned_slice::(100, PowerOfTwo::new(64).unwrap()).unwrap(); + for &val in data.iter() { + assert_eq!(val, 0.0); + } + } + + #[test] + fn aligned_slice_deref_mut() { + let mut data = aligned_slice::(4, PowerOfTwo::new(32).unwrap()).unwrap(); + data[0] = 1.0; + data[1] = 2.0; + assert_eq!(&data[..2], &[1.0, 2.0]); + } +} diff --git a/diskann-quantization/src/alloc/mod.rs b/diskann-quantization/src/alloc/mod.rs index da58f8b77..e4aaae317 100644 --- a/diskann-quantization/src/alloc/mod.rs +++ b/diskann-quantization/src/alloc/mod.rs @@ -6,11 +6,13 @@ use std::{alloc::Layout, ptr::NonNull}; mod aligned; +mod aligned_slice; mod bump; mod poly; mod traits; pub use aligned::{AlignedAllocator, NotPowerOfTwo}; +pub use aligned_slice::{AlignedSlice, aligned_slice}; pub use bump::BumpAllocator; pub use poly::{CompoundError, Poly, TrustedIter, poly}; pub use traits::{Allocator, AllocatorCore, AllocatorError};