remove alignedslice and replace with direct poly calls

Cargo.toml

@@ -85,7 +85,6 @@ iai-callgrind = "0.14.0"
 itertools = "0.13.0"
 num-traits = "0.2.15"
 num_cpus = "1.16.0"
-once_cell = "1.19.0"
 opentelemetry = "0.30.0"
 opentelemetry_sdk = "0.30.0"
 paste = "1.0.15"

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_quantization::{alloc::aligned_slice, num::PowerOfTwo};
+use diskann_quantization::alloc::{AlignedAllocator, Poly};
 
 pub const TEST_INDEX_PATH: &str =
     "../test_data/disk_index_misc/disk_index_siftsmall_learn_256pts_R4_L50_A1.2_aligned_reader_test.index";
@@ -34,7 +34,7 @@ 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 =
-        aligned_slice::<u8>(read_length * num_read, PowerOfTwo::new(512).unwrap()).unwrap();
+        Poly::broadcast(0u8, read_length * num_read, AlignedAllocator::A512).unwrap();
 
     // create and add AlignedReads to the vector
     let mut mem_slices: Vec<&mut [u8]> = aligned_mem.chunks_mut(read_length).collect();

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_quantization::{alloc::aligned_slice, num::PowerOfTwo};
+use diskann_quantization::alloc::{AlignedAllocator, Poly};
 
 pub const TEST_INDEX_PATH: &str =
     "../test_data/disk_index_misc/disk_index_siftsmall_learn_256pts_R4_L50_A1.2_aligned_reader_test.index";
@@ -25,7 +25,7 @@ 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 =
-        aligned_slice::<u8>(read_length * num_read, PowerOfTwo::new(512).unwrap()).unwrap();
+        Poly::broadcast(0u8, read_length * num_read, AlignedAllocator::A512).unwrap();
 
     // create and add AlignedReads to the vector
     let mut mem_slices: Vec<&mut [u8]> = aligned_mem.chunks_mut(read_length).collect();

diskann-disk/src/search/pq/pq_scratch.rs

@@ -6,38 +6,29 @@
 
 use diskann::{error::IntoANNResult, utils::VectorRepr, ANNError, ANNResult};
 
-use diskann_quantization::{
-    alloc::{aligned_slice, AlignedSlice},
-    num::PowerOfTwo,
-};
+use diskann_quantization::alloc::{AlignedAllocator, Poly};
 
 #[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: AlignedSlice<f32>,
+    pub aligned_pqtable_dist_scratch: Poly<[f32], AlignedAllocator>,
 
     /// 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: AlignedSlice<f32>,
+    pub aligned_dist_scratch: Poly<[f32], AlignedAllocator>,
 
     /// 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: AlignedSlice<u8>,
+    pub aligned_pq_coord_scratch: Poly<[u8], AlignedAllocator>,
 
     /// 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.
     pub rotated_query: Vec<f32>,
 }
 
 impl PQScratch {
-    /// 128 bytes alignment to optimize for the L2 Adjacent Cache Line Prefetcher.
-    const ALIGNED_ALLOC_128: PowerOfTwo = match PowerOfTwo::new(128) {
-        Ok(v) => v,
-        Err(_) => unreachable!(),
-    };
-
     /// Create a new pq scratch
     pub fn new(
         graph_degree: usize,
@@ -46,12 +37,12 @@ impl PQScratch {
         num_centers: usize,
     ) -> ANNResult<Self> {
         let aligned_pq_coord_scratch =
-            aligned_slice(graph_degree * num_pq_chunks, PQScratch::ALIGNED_ALLOC_128)
+            Poly::broadcast(0u8, graph_degree * num_pq_chunks, AlignedAllocator::A128)
                 .map_err(ANNError::log_index_error)?;
         let aligned_pqtable_dist_scratch =
-            aligned_slice(num_centers * num_pq_chunks, PQScratch::ALIGNED_ALLOC_128)
+            Poly::broadcast(0f32, num_centers * num_pq_chunks, AlignedAllocator::A128)
                 .map_err(ANNError::log_index_error)?;
-        let aligned_dist_scratch = aligned_slice(graph_degree, PQScratch::ALIGNED_ALLOC_128)
+        let aligned_dist_scratch = Poly::broadcast(0f32, graph_degree, AlignedAllocator::A128)
             .map_err(ANNError::log_index_error)?;
         let rotated_query = vec![0.0f32; dim];
 
@@ -94,6 +85,7 @@ impl PQScratch {
 
 #[cfg(test)]
 mod tests {
+    use diskann_quantization::num::PowerOfTwo;
     use rstest::rstest;
 
     use super::PQScratch;
@@ -110,20 +102,16 @@ mod tests {
         let mut pq_scratch: PQScratch =
             PQScratch::new(graph_degree, dim, num_pq_chunks, num_centers).unwrap();
 
-        // Check alignment of the remaining AlignedSlice buffers.
         assert_eq!(
-            (pq_scratch.aligned_pqtable_dist_scratch.as_ptr() as usize)
-                % PQScratch::ALIGNED_ALLOC_128.raw(),
+            (pq_scratch.aligned_pqtable_dist_scratch.as_ptr() as usize) % PowerOfTwo::V128.raw(),
             0
         );
         assert_eq!(
-            (pq_scratch.aligned_dist_scratch.as_ptr() as usize)
-                % PQScratch::ALIGNED_ALLOC_128.raw(),
+            (pq_scratch.aligned_dist_scratch.as_ptr() as usize) % PowerOfTwo::V128.raw(),
             0
         );
         assert_eq!(
-            (pq_scratch.aligned_pq_coord_scratch.as_ptr() as usize)
-                % PQScratch::ALIGNED_ALLOC_128.raw(),
+            (pq_scratch.aligned_pq_coord_scratch.as_ptr() as usize) % PowerOfTwo::V128.raw(),
             0
         );
 

diskann-disk/src/search/provider/disk_sector_graph.rs

@@ -9,7 +9,7 @@ use std::ops::Deref;
 
 use diskann::{ANNError, ANNResult};
 use diskann_quantization::{
-    alloc::{aligned_slice, AlignedSlice},
+    alloc::{AlignedAllocator, Poly},
     num::PowerOfTwo,
 };
 
@@ -37,7 +37,7 @@ pub struct DiskSectorGraph<AlignedReaderType: AlignedFileReader> {
     /// 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: AlignedSlice<u8>,
+    sectors_data: Poly<[u8], AlignedAllocator>,
     /// Current sector index into which the next read reads data
     cur_sector_idx: u64,
 
@@ -76,9 +76,12 @@ impl<AlignedReaderType: AlignedFileReader> DiskSectorGraph<AlignedReaderType> {
 
         Ok(Self {
             sector_reader,
-            sectors_data: aligned_slice(
+            sectors_data: Poly::broadcast(
+                0u8,
                 max_n_batch_sector_read * num_sectors_per_node * block_size,
-                PowerOfTwo::new(block_size).map_err(ANNError::log_index_error)?,
+                AlignedAllocator::new(
+                    PowerOfTwo::new(block_size).map_err(ANNError::log_index_error)?,
+                ),
             )
             .map_err(ANNError::log_index_error)?,
             cur_sector_idx: 0,
@@ -94,9 +97,12 @@ impl<AlignedReaderType: AlignedFileReader> DiskSectorGraph<AlignedReaderType> {
     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 = aligned_slice(
+            self.sectors_data = Poly::broadcast(
+                0u8,
                 max_n_batch_sector_read * self.num_sectors_per_node * self.block_size,
-                PowerOfTwo::new(self.block_size).map_err(ANNError::log_index_error)?,
+                AlignedAllocator::new(
+                    PowerOfTwo::new(self.block_size).map_err(ANNError::log_index_error)?,
+                ),
             )
             .map_err(ANNError::log_index_error)?;
         }
@@ -220,7 +226,7 @@ mod disk_sector_graph_test {
     ) -> DiskSectorGraph<<AlignedFileReaderFactory as AlignedReaderFactory>::AlignedReaderType>
     {
         DiskSectorGraph {
-            sectors_data: aligned_slice(512, PowerOfTwo::new(512).unwrap()).unwrap(),
+            sectors_data: Poly::broadcast(0u8, 512, AlignedAllocator::A512).unwrap(),
             sector_reader,
             cur_sector_idx: 0,
             num_nodes_per_sector,

diskann-disk/src/search/provider/disk_vertex_provider_factory.rs

@@ -6,7 +6,10 @@ 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_quantization::{alloc::aligned_slice, num::PowerOfTwo};
+use diskann_quantization::{
+    alloc::{AlignedAllocator, Poly},
+    num::PowerOfTwo,
+};
 use hashbrown::HashSet;
 use tracing::info;
 
@@ -52,9 +55,10 @@ 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 = aligned_slice::<u8>(
+        let mut read_buf = Poly::broadcast(
+            0u8,
             buffer_len,
-            PowerOfTwo::new(buffer_len).map_err(ANNError::log_index_error)?,
+            AlignedAllocator::new(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)?;

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_quantization::{alloc::aligned_slice, num::PowerOfTwo};
+    use diskann_quantization::alloc::{AlignedAllocator, Poly};
     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 =
@@ -171,7 +171,7 @@ mod tests {
         let read_length = 512; // adjust according to your logic
         let num_read = 10;
         let mut aligned_mem =
-            aligned_slice::<u8>(read_length * num_read, PowerOfTwo::new(512).unwrap()).unwrap();
+            Poly::broadcast(0u8, read_length * num_read, AlignedAllocator::A512).unwrap();
 
         // create and add AlignedReads to the vector
         let mut mem_slices: Vec<&mut [u8]> = aligned_mem.chunks_mut(read_length).collect();
@@ -215,7 +215,7 @@ 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 =
-            aligned_slice::<u8>(read_length * num_read, PowerOfTwo::new(512).unwrap()).unwrap();
+            Poly::broadcast(0u8, read_length * num_read, AlignedAllocator::A512).unwrap();
 
         // create and add AlignedReads to the vector
         let mut mem_slices: Vec<&mut [u8]> = aligned_mem.chunks_mut(read_length).collect();
@@ -255,7 +255,7 @@ mod tests {
         let read_length = 512; // adjust according to your logic
         let num_sector = 10;
         let mut aligned_mem =
-            aligned_slice::<u8>(read_length * num_sector, PowerOfTwo::new(512).unwrap()).unwrap();
+            Poly::broadcast(0u8, read_length * num_sector, AlignedAllocator::A512).unwrap();
 
         // Each slice will be used as the buffer for a read request of a sector.
         let mut mem_slices: Vec<&mut [u8]> = aligned_mem.chunks_mut(read_length).collect();

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_quantization::{alloc::aligned_slice, num::PowerOfTwo};
+    use diskann_quantization::alloc::{AlignedAllocator, Poly};
 
     fn test_index_path() -> String {
         "/disk_index_misc/disk_index_siftsmall_learn_256pts_R4_L50_A1.2_aligned_reader_test.index"
@@ -79,7 +79,7 @@ mod tests {
         let read_length = 512;
         let num_read = 10;
         let mut aligned_mem =
-            aligned_slice::<u8>(read_length * num_read, PowerOfTwo::new(512).unwrap()).unwrap();
+            Poly::broadcast(0u8, read_length * num_read, AlignedAllocator::A512).unwrap();
 
         // create and add AlignedReads to the vector
         let mut mem_slices: Vec<&mut [u8]> = aligned_mem.chunks_mut(read_length).collect();

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_quantization::{alloc::aligned_slice, num::PowerOfTwo};
+    use diskann_quantization::alloc::{AlignedAllocator, Poly};
 
     fn test_index_path() -> String {
         test_data_root()
@@ -170,7 +170,7 @@ mod tests {
         let read_length = 512; // adjust according to your logic
         let num_read = 10;
         let mut aligned_mem =
-            aligned_slice::<u8>(read_length * num_read, PowerOfTwo::new(512).unwrap()).unwrap();
+            Poly::broadcast(0u8, read_length * num_read, AlignedAllocator::A512).unwrap();
 
         // create and add AlignedReads to the vector
         let mut mem_slices: Vec<&mut [u8]> = aligned_mem.chunks_mut(read_length).collect();
@@ -209,7 +209,7 @@ mod tests {
         let read_length = DEFAULT_DISK_SECTOR_LEN; // adjust according to your logic
         let num_sector = 10;
         let mut aligned_mem =
-            aligned_slice::<u8>(read_length * num_sector, PowerOfTwo::new(512).unwrap()).unwrap();
+            Poly::broadcast(0u8, read_length * num_sector, AlignedAllocator::A512).unwrap();
 
         // Each slice will be used as the buffer for a read request of a sector.
         let mut mem_slices: Vec<&mut [u8]> = aligned_mem.chunks_mut(read_length).collect();

diskann-quantization/src/alloc/aligned.rs

@@ -20,6 +20,19 @@ pub struct AlignedAllocator {
 }
 
 impl AlignedAllocator {
+    /// Aligned allocators for commonly specified boundaries in the codebase (4..4096)
+    pub const A4: Self = Self::new(PowerOfTwo::V4);
+    pub const A8: Self = Self::new(PowerOfTwo::V8);
+    pub const A16: Self = Self::new(PowerOfTwo::V16);
+    pub const A32: Self = Self::new(PowerOfTwo::V32);
+    pub const A64: Self = Self::new(PowerOfTwo::V64);
+    pub const A128: Self = Self::new(PowerOfTwo::V128);
+    pub const A256: Self = Self::new(PowerOfTwo::V256);
+    pub const A512: Self = Self::new(PowerOfTwo::V512);
+    pub const A1024: Self = Self::new(PowerOfTwo::V1024);
+    pub const A2048: Self = Self::new(PowerOfTwo::V2048);
+    pub const A4096: Self = Self::new(PowerOfTwo::V4096);
+
     /// Construct a new allocator that uses the given alignment.
     #[inline]
     pub const fn new(alignment: PowerOfTwo) -> Self {