PiPNN 6/6: add HashPrune candidate merging

.github/workflows/nightly.yml

@@ -145,3 +145,12 @@ jobs:
           cargo +nightly miri test --locked -p diskann --features pipnn --lib -- \
             graph::pipnn::topk::tests::selection_returns_distinct_candidates_when_distances_tie \
             graph::pipnn::topk::tests::a_single_pair_updates_exactly_its_two_endpoints::case_1_finite
+
+      - name: PiPNN HashPrune pointer boundaries
+        env:
+          MIRIFLAGS: -Zmiri-disable-isolation -Zmiri-strict-provenance
+        run: |
+          cargo +nightly miri test --locked -p diskann --features pipnn --lib \
+            graph::pipnn::lsh::tests
+          cargo +nightly miri test --locked -p diskann --features pipnn --lib \
+            graph::pipnn::hash_prune::tests

diskann-benchmark/src/index/build.rs

@@ -143,12 +143,15 @@ where
 
     let started = std::time::Instant::now();
     let adjacency = {
-        let context = diskann::graph::pipnn::PiPNNBuildContext::new(
+        let mut context = diskann::graph::pipnn::PiPNNBuildContext::new(
             parameters.into(),
             &graph,
             metric,
             &pool,
         )?;
+        if let Some(hash_prune) = &parameters.hash_prune {
+            context = context.with_hash_prune(hash_prune.into())?;
+        }
         diskann::graph::pipnn::build_graph(data.as_view(), &context)?
     };
     let start_points = input

diskann-disk/src/build/builder/build.rs

@@ -79,8 +79,12 @@ where
         index_writer: DiskIndexWriter,
     ) -> ANNResult<Self> {
         #[cfg(feature = "pipnn")]
-        if let Some(config) = disk_build_param.pipnn_config() {
-            config.validate()?;
+        if let Some(parameters) = disk_build_param.pipnn_parameters() {
+            diskann::graph::pipnn::PiPNNConfig::from(parameters).validate()?;
+            if let Some(hash_prune) = &parameters.hash_prune {
+                diskann::graph::pipnn::HashPruneConfig::from(hash_prune)
+                    .validate(index_configuration.config.pruned_degree().get())?;
+            }
         }
 
         let pq_storage = PQStorage::new(
@@ -184,8 +188,8 @@ where
 
     async fn build_graph(&mut self, pool: RayonThreadPoolRef<'_>) -> ANNResult<()> {
         #[cfg(feature = "pipnn")]
-        if let Some(config) = self.disk_build_param.pipnn_config() {
-            return pipnn::build_graph(self, pool, config);
+        if let Some(parameters) = self.disk_build_param.pipnn_parameters().cloned() {
+            return pipnn::build_graph(self, pool, &parameters);
         }
 
         match determine_build_strategy::<Data>(

diskann-disk/src/build/builder/build/pipnn.rs

@@ -11,7 +11,7 @@
 //!
 //! PiPNN and Vamana use the same disk graph format.
 
-use diskann::graph::pipnn::{PiPNNBuildContext, PiPNNConfig};
+use diskann::graph::pipnn::PiPNNBuildContext;
 use diskann::{utils::VectorRepr, ANNError, ANNResult};
 use diskann_providers::{
     storage::{save_adjacency_graph, StorageReadProvider, StorageWriteProvider},
@@ -20,13 +20,13 @@ use diskann_providers::{
 use diskann_utils::io::{read_bin, Metadata};
 
 use super::{u32_try_from, DiskIndexBuilder};
-use crate::data_model::GraphDataType;
+use crate::{data_model::GraphDataType, PiPNNParameters};
 
 /// Build PiPNN adjacency and persist it through the canonical disk graph writer.
 pub(super) fn build_graph<Data, StorageProvider>(
     builder: &DiskIndexBuilder<'_, Data, StorageProvider>,
     pool: RayonThreadPoolRef<'_>,
-    config: PiPNNConfig,
+    parameters: &PiPNNParameters,
 ) -> ANNResult<()>
 where
     Data: GraphDataType<VectorIdType = u32>,
@@ -55,12 +55,15 @@ where
     // supplied Rayon pool.
     let data =
         read_bin::<Data::VectorDataType>(&mut builder.storage_provider.open_reader(&data_path)?)?;
-    let context = PiPNNBuildContext::new(
-        config,
+    let mut context = PiPNNBuildContext::new(
+        parameters.into(),
         &builder.index_configuration.config,
         builder.index_configuration.dist_metric,
         pool.as_rayon(),
     )?;
+    if let Some(hash_prune) = &parameters.hash_prune {
+        context = context.with_hash_prune(hash_prune.into())?;
+    }
     let adjacency = diskann::graph::pipnn::build_graph(data.as_view(), &context)?;
 
     // The disk header requires a start point. Use the same sampled medoid policy
@@ -116,6 +119,7 @@ mod tests {
             fanout: vec![nz(10), nz(3)],
             k: nz(2),
             replicas: nz(1),
+            hash_prune: Some(crate::HashPruneParameters::default()),
         }
     }
 
@@ -199,7 +203,7 @@ mod tests {
         let builder = builder(&storage, 3, 8, 1.0, 1.2, parameters.clone());
         let pool = create_thread_pool(1).unwrap();
 
-        let error = super::build_graph(&builder, pool.as_ref(), (&parameters).into()).unwrap_err();
+        let error = super::build_graph(&builder, pool.as_ref(), &parameters).unwrap_err();
         assert!(format!("{error:?}").contains("configured point count 3"));
         assert!(!storage.exists(&builder.index_writer.get_mem_index_file()));
     }
@@ -213,7 +217,7 @@ mod tests {
         let builder = builder(&storage, points, dimensions, 1.0, 1.2, parameters.clone());
         let pool = create_thread_pool(1).unwrap();
 
-        super::build_graph(&builder, pool.as_ref(), (&parameters).into()).unwrap();
+        super::build_graph(&builder, pool.as_ref(), &parameters).unwrap();
 
         let mut header = [0_u8; 24];
         std::io::Read::read_exact(
@@ -274,4 +278,34 @@ mod tests {
 
         assert!(format!("{error:?}").contains("p_samp (0) must be in (0, 1]"));
     }
+
+    #[test]
+    fn builder_rejects_hash_prune_capacity_before_quantizer_artifacts() {
+        let storage = VirtualStorageProvider::new_memory();
+        let parameters = PiPNNParameters {
+            hash_prune: Some(crate::HashPruneParameters {
+                num_hash_planes: 12,
+                l_max: 16,
+                final_prune: true,
+            }),
+            ..PiPNNParameters::default()
+        };
+        let params = DiskIndexBuildParameters::new_pipnn(
+            MemoryBudget::try_from_gb(1.0).unwrap(),
+            NumPQChunks::new_with(1, 1).unwrap(),
+            parameters,
+        );
+        let config = IndexConfiguration::new(Metric::L2, 1, 1, ONE, 1, graph_config(32, 1.2));
+        let writer =
+            DiskIndexWriter::new("/data.fbin".into(), "/index".into(), None, 4096).unwrap();
+
+        let error = match DiskIndexBuilder::<AdHoc<f32>, _>::new(&storage, params, config, writer) {
+            Ok(_) => panic!("HashPrune capacity below graph degree must be rejected"),
+            Err(error) => error,
+        };
+
+        assert!(format!("{error:?}").contains("must be at least the graph degree (32)"));
+        assert!(!storage.exists("/index_pq_pivots.bin"));
+        assert!(!storage.exists("/index_pq_compressed.bin"));
+    }
 }

diskann-disk/src/build/configuration/build_algorithm.rs

@@ -29,6 +29,43 @@ pub struct PiPNNParameters {
     pub k: NonZeroUsize,
     /// Number of independent partition passes.
     pub replicas: NonZeroUsize,
+    /// HashPrune policy. `None` keeps all unique direct candidates.
+    pub hash_prune: Option<HashPruneParameters>,
+}
+
+/// HashPrune parameters in the JSON build configuration.
+#[cfg(feature = "pipnn")]
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+#[serde(default, deny_unknown_fields)]
+pub struct HashPruneParameters {
+    /// Number of random-hyperplane sketch dimensions.
+    pub num_hash_planes: usize,
+    /// Maximum number of candidates retained per point.
+    pub l_max: usize,
+    /// Apply Vamana RobustPrune after reservoir extraction.
+    pub final_prune: bool,
+}
+
+#[cfg(feature = "pipnn")]
+impl Default for HashPruneParameters {
+    fn default() -> Self {
+        Self {
+            num_hash_planes: 12,
+            l_max: 64,
+            final_prune: true,
+        }
+    }
+}
+
+#[cfg(feature = "pipnn")]
+impl From<&HashPruneParameters> for diskann::graph::pipnn::HashPruneConfig {
+    fn from(config: &HashPruneParameters) -> Self {
+        Self {
+            num_hash_planes: config.num_hash_planes,
+            l_max: config.l_max,
+            final_prune: config.final_prune,
+        }
+    }
 }
 
 #[cfg(feature = "pipnn")]
@@ -44,6 +81,7 @@ impl Default for PiPNNParameters {
             fanout: FANOUT.to_vec(),
             k: K,
             replicas: NonZeroUsize::MIN,
+            hash_prune: None,
         }
     }
 }
@@ -119,6 +157,15 @@ mod tests {
         );
         assert_eq!(config.k.get(), 3);
         assert_eq!(config.replicas.get(), 1);
+        assert_eq!(config.hash_prune, None);
+
+        let explicit: BuildAlgorithm =
+            serde_json::from_str(r#"{"algorithm":"PiPNN","hash_prune":{}}"#).unwrap();
+        let BuildAlgorithm::PiPNN(explicit) = explicit else {
+            panic!("expected PiPNN");
+        };
+        assert_eq!(explicit.hash_prune, Some(HashPruneParameters::default()));
+
         assert!(
             serde_json::from_str::<BuildAlgorithm>(r#"{"algorithm":"PiPNN","l_max":72}"#).is_err()
         );

diskann-disk/src/build/configuration/disk_index_build_parameter.rs

@@ -198,9 +198,9 @@ impl DiskIndexBuildParameters {
     }
 
     #[cfg(feature = "pipnn")]
-    pub(crate) fn pipnn_config(&self) -> Option<diskann::graph::pipnn::PiPNNConfig> {
+    pub(crate) fn pipnn_parameters(&self) -> Option<&PiPNNParameters> {
         match &self.build_algorithm {
-            BuildAlgorithm::PiPNN(config) => Some(config.into()),
+            BuildAlgorithm::PiPNN(config) => Some(config),
             BuildAlgorithm::Vamana => None,
         }
     }

diskann-disk/src/build/configuration/mod.rs

@@ -5,7 +5,7 @@
 pub mod build_algorithm;
 pub use build_algorithm::BuildAlgorithm;
 #[cfg(feature = "pipnn")]
-pub use build_algorithm::PiPNNParameters;
+pub use build_algorithm::{HashPruneParameters, PiPNNParameters};
 
 pub mod disk_index_build_parameter;
 pub use disk_index_build_parameter::{DiskIndexBuildParameters, MemoryBudget, NumPQChunks};

diskann-disk/src/build/mod.rs

@@ -12,9 +12,9 @@ pub mod builder;
 pub mod configuration;
 
 // Re-export key types for convenience
-#[cfg(feature = "pipnn")]
-pub use configuration::PiPNNParameters;
 pub use configuration::{
     disk_index_build_parameter, filter_parameter, BuildAlgorithm, DiskIndexBuildParameters,
     QuantizationType, SphericalBits,
 };
+#[cfg(feature = "pipnn")]
+pub use configuration::{HashPruneParameters, PiPNNParameters};

diskann-disk/src/lib.rs

@@ -14,12 +14,12 @@ pub(crate) mod test_utils;
 pub mod error;
 
 pub mod build;
-#[cfg(feature = "pipnn")]
-pub use build::PiPNNParameters;
 pub use build::{
     disk_index_build_parameter, filter_parameter, BuildAlgorithm, DiskIndexBuildParameters,
     QuantizationType, SphericalBits,
 };
+#[cfg(feature = "pipnn")]
+pub use build::{HashPruneParameters, PiPNNParameters};
 
 pub mod data_model;
 pub mod search;

diskann-vector/src/lib.rs

@@ -38,14 +38,17 @@ pub mod distance;
 pub mod norm;
 
 cfg_if::cfg_if! {
-    if #[cfg(all(target_arch = "x86_64", target_feature = "avx2"))] {
+    // x86-64 guarantees SSE2; `_mm_prefetch` needs only SSE.
+    if #[cfg(target_arch = "x86_64")] {
         const CACHE_LINE_SIZE: usize = 64;
 
         #[inline(always)]
         unsafe fn prefetch_exactly<const N: usize>(ptr: *const i8) {
             use std::arch::x86_64::*;
             for i in 0..N {
-                _mm_prefetch(ptr.add(i * CACHE_LINE_SIZE), _MM_HINT_T0);
+                // SAFETY: the caller guarantees that all `N` computed addresses are
+                // inside the allocation.
+                unsafe { _mm_prefetch(ptr.add(i * CACHE_LINE_SIZE), _MM_HINT_T0) };
             }
         }
 
@@ -56,7 +59,8 @@ cfg_if::cfg_if! {
                 if CACHE_LINE_SIZE * i >= bytes {
                     break;
                 }
-                _mm_prefetch(ptr.add(i * CACHE_LINE_SIZE), _MM_HINT_T0);
+                // SAFETY: the loop uses only offsets below `bytes`.
+                unsafe { _mm_prefetch(ptr.add(i * CACHE_LINE_SIZE), _MM_HINT_T0) };
             }
         }
 
@@ -66,32 +70,48 @@ cfg_if::cfg_if! {
         pub fn prefetch_hint_max<const MAX_CACHE_LINES: usize, T>(vec: &[T]) {
             let vecsize = std::mem::size_of_val(vec);
             if vecsize >= MAX_CACHE_LINES * 64 {
-                // SAFETY: Pointer is in-bounds and use of the intrinsic is cfg gated.
+                // SAFETY: the slice contains every address passed to prefetch.
                 unsafe { prefetch_exactly::<MAX_CACHE_LINES>(vec.as_ptr().cast()) }
             } else {
-                // SAFETY: Pointer is in-bounds and use of the intrinsic is cfg gated.
+                // SAFETY: the slice covers `vecsize` bytes.
                 unsafe { prefetch_at_most::<MAX_CACHE_LINES>(vec.as_ptr().cast(), vecsize) }
             }
         }
 
+        /// Prefetch a raw byte range without creating a slice.
+        ///
+        /// # Safety
+        ///
+        /// `ptr` must identify an allocation of at least `bytes` bytes. The allocation
+        /// must remain live for this call. The function creates no Rust reference.
+        /// The caller controls concurrent mutation of the range.
+        #[inline]
+        pub unsafe fn prefetch_hint_all_raw(ptr: *const u8, bytes: usize) {
+            use std::arch::x86_64::*;
+
+            for offset in (0..bytes).step_by(CACHE_LINE_SIZE) {
+                // SAFETY: the caller guarantees the byte range, and `offset < bytes`.
+                unsafe { _mm_prefetch(ptr.add(offset).cast(), _MM_HINT_T0) };
+            }
+        }
+
         /// Prefetch the given vector in chunks of 64 bytes, which is a cache line size.
         /// The entire vector will be prefetched.
         #[inline]
         pub fn prefetch_hint_all<T>(vec: &[T]) {
-            use std::arch::x86_64::*;
-
-            let vecsize = std::mem::size_of_val(vec);
-            let num_prefetch_blocks = vecsize.div_ceil(64);
-            let vec_ptr = vec.as_ptr() as *const i8;
-            for d in 0..num_prefetch_blocks {
-                // SAFETY: Pointer is in-bounds and use of the intrinsic is gated by the
-                // `cfg`-guard on this function.
-                unsafe {
-                    std::arch::x86_64::_mm_prefetch(vec_ptr.add(d * CACHE_LINE_SIZE), _MM_HINT_T0);
-                }
-            }        }
+            // SAFETY: the slice remains live and covers exactly `size_of_val(vec)` bytes.
+            unsafe { prefetch_hint_all_raw(vec.as_ptr().cast(), std::mem::size_of_val(vec)) }
+        }
     } else {
         pub fn prefetch_hint_max<const MAX_CACHE_LINES: usize, T>(_vec: &[T]) {}
+
+        /// Accept a raw prefetch range and do nothing.
+        ///
+        /// # Safety
+        ///
+        /// The pointer contract is the same as the x86-64 implementation.
+        pub unsafe fn prefetch_hint_all_raw(_ptr: *const u8, _bytes: usize) {}
+
         pub fn prefetch_hint_all<T>(_vec: &[T]) {}
     }
 }