add hnsw unit test for PR 3840

mengdilin · facebook-github-bot · commit b8c3a9743fc0 · 2024-09-11T11:58:43.000-07:00
Summary: #3845 Add unit tests for helper search utilities for HNSW. These utility functions live inside an anonymous namespace and each has a reference version gated behind a const bool, I refactored them so the reference version is a flag for the function which defaults to false. If we are concerned about the performance overhead of the extra if branching (whether to use reference version or not) inside these utility functions, I'm happy to lift out the reference versions to their own functions inside the unit test Differential Revision: D62510014
diff --git a/faiss/impl/HNSW.cpp b/faiss/impl/HNSW.cpp
@@ -470,105 +470,6 @@ void search_neighbors_to_add(
     vt.advance();
 }
 
-/**************************************************************
- * Searching subroutines
- **************************************************************/
-
-/// greedily update a nearest vector at a given level
-HNSWStats greedy_update_nearest(
-        const HNSW& hnsw,
-        DistanceComputer& qdis,
-        int level,
-        storage_idx_t& nearest,
-        float& d_nearest) {
-    // selects a version
-    const bool reference_version = false;
-
-    HNSWStats stats;
-
-    for (;;) {
-        storage_idx_t prev_nearest = nearest;
-
-        size_t begin, end;
-        hnsw.neighbor_range(nearest, level, &begin, &end);
-
-        size_t ndis = 0;
-
-        // select a version, based on a flag
-        if (reference_version) {
-            // a reference version
-            for (size_t i = begin; i < end; i++) {
-                storage_idx_t v = hnsw.neighbors[i];
-                if (v < 0)
-                    break;
-                ndis += 1;
-                float dis = qdis(v);
-                if (dis < d_nearest) {
-                    nearest = v;
-                    d_nearest = dis;
-                }
-            }
-        } else {
-            // a faster version
-
-            // the following version processes 4 neighbors at a time
-            auto update_with_candidate = [&](const storage_idx_t idx,
-                                             const float dis) {
-                if (dis < d_nearest) {
-                    nearest = idx;
-                    d_nearest = dis;
-                }
-            };
-
-            int n_buffered = 0;
-            storage_idx_t buffered_ids[4];
-
-            for (size_t j = begin; j < end; j++) {
-                storage_idx_t v = hnsw.neighbors[j];
-                if (v < 0)
-                    break;
-                ndis += 1;
-
-                buffered_ids[n_buffered] = v;
-                n_buffered += 1;
-
-                if (n_buffered == 4) {
-                    float dis[4];
-                    qdis.distances_batch_4(
-                            buffered_ids[0],
-                            buffered_ids[1],
-                            buffered_ids[2],
-                            buffered_ids[3],
-                            dis[0],
-                            dis[1],
-                            dis[2],
-                            dis[3]);
-
-                    for (size_t id4 = 0; id4 < 4; id4++) {
-                        update_with_candidate(buffered_ids[id4], dis[id4]);
-                    }
-
-                    n_buffered = 0;
-                }
-            }
-
-            // process leftovers
-            for (size_t icnt = 0; icnt < n_buffered; icnt++) {
-                float dis = qdis(buffered_ids[icnt]);
-                update_with_candidate(buffered_ids[icnt], dis);
-            }
-        }
-
-        // update stats
-        stats.ndis += ndis;
-        stats.nhops += 1;
-
-        if (nearest == prev_nearest) {
-            return stats;
-        }
-    }
-}
-
 } // namespace
 
 /// Finds neighbors and builds links with them, starting from an entry
@@ -644,7 +545,8 @@ void HNSW::add_with_locks(
     float d_nearest = ptdis(nearest);
 
     for (; level > pt_level; level--) {
-        greedy_update_nearest(*this, ptdis, level, nearest, d_nearest);
+        hnsw_utils::greedy_update_nearest(
+                *this, ptdis, level, nearest, d_nearest);
     }
 
     for (; level >= 0; level--) {
@@ -667,16 +569,16 @@ void HNSW::add_with_locks(
     }
 }
 
+namespace hnsw_utils {
+
 /**************************************************************
  * Searching
  **************************************************************/
 
-namespace {
 using MinimaxHeap = HNSW::MinimaxHeap;
 using Node = HNSW::Node;
 using C = HNSW::C;
 /** Do a BFS on the candidates list */
-
 int search_from_candidates(
         const HNSW& hnsw,
         DistanceComputer& qdis,
@@ -685,11 +587,9 @@ int search_from_candidates(
         VisitedTable& vt,
         HNSWStats& stats,
         int level,
-        int nres_in = 0,
-        const SearchParametersHNSW* params = nullptr) {
-    // selects a version
-    const bool reference_version = false;
-
+        int nres_in,
+        const SearchParametersHNSW* params,
+        const bool reference_version) {
     int nres = nres_in;
     int ndis = 0;
 
@@ -851,10 +751,8 @@ std::priority_queue<HNSW::Node> search_from_candidate_unbounded(
         DistanceComputer& qdis,
         int ef,
         VisitedTable* vt,
-        HNSWStats& stats) {
-    // selects a version
-    const bool reference_version = false;
-
+        HNSWStats& stats,
+        const bool reference_version) {
     int ndis = 0;
     std::priority_queue<Node> top_candidates;
     std::priority_queue<Node, std::vector<Node>, std::greater<Node>> candidates;
@@ -984,6 +882,104 @@ std::priority_queue<HNSW::Node> search_from_candidate_unbounded(
     return top_candidates;
 }
 
+/// greedily update a nearest vector at a given level
+HNSWStats greedy_update_nearest(
+        const HNSW& hnsw,
+        DistanceComputer& qdis,
+        int level,
+        storage_idx_t& nearest,
+        float& d_nearest,
+        const bool reference_version) {
+    HNSWStats stats;
+
+    for (;;) {
+        storage_idx_t prev_nearest = nearest;
+
+        size_t begin, end;
+        hnsw.neighbor_range(nearest, level, &begin, &end);
+
+        size_t ndis = 0;
+
+        // select a version, based on a flag
+        if (reference_version) {
+            // a reference version
+            for (size_t i = begin; i < end; i++) {
+                storage_idx_t v = hnsw.neighbors[i];
+                if (v < 0)
+                    break;
+                ndis += 1;
+                float dis = qdis(v);
+                if (dis < d_nearest) {
+                    nearest = v;
+                    d_nearest = dis;
+                }
+            }
+        } else {
+            // a faster version
+
+            // the following version processes 4 neighbors at a time
+            auto update_with_candidate = [&](const storage_idx_t idx,
+                                             const float dis) {
+                if (dis < d_nearest) {
+                    nearest = idx;
+                    d_nearest = dis;
+                }
+            };
+
+            int n_buffered = 0;
+            storage_idx_t buffered_ids[4];
+
+            for (size_t j = begin; j < end; j++) {
+                storage_idx_t v = hnsw.neighbors[j];
+                if (v < 0)
+                    break;
+                ndis += 1;
+
+                buffered_ids[n_buffered] = v;
+                n_buffered += 1;
+
+                if (n_buffered == 4) {
+                    float dis[4];
+                    qdis.distances_batch_4(
+                            buffered_ids[0],
+                            buffered_ids[1],
+                            buffered_ids[2],
+                            buffered_ids[3],
+                            dis[0],
+                            dis[1],
+                            dis[2],
+                            dis[3]);
+
+                    for (size_t id4 = 0; id4 < 4; id4++) {
+                        update_with_candidate(buffered_ids[id4], dis[id4]);
+                    }
+
+                    n_buffered = 0;
+                }
+            }
+
+            // process leftovers
+            for (size_t icnt = 0; icnt < n_buffered; icnt++) {
+                float dis = qdis(buffered_ids[icnt]);
+                update_with_candidate(buffered_ids[icnt], dis);
+            }
+        }
+
+        // update stats
+        stats.ndis += ndis;
+        stats.nhops += 1;
+
+        if (nearest == prev_nearest) {
+            return stats;
+        }
+    }
+}
+} // namespace hnsw_utils
+namespace {
+using MinimaxHeap = HNSW::MinimaxHeap;
+using Node = HNSW::Node;
+using C = HNSW::C;
+
 // just used as a lower bound for the minmaxheap, but it is set for heap search
 int extract_k_from_ResultHandler(ResultHandler<C>& res) {
     using RH = HeapBlockResultHandler<C>;
@@ -993,7 +989,7 @@ int extract_k_from_ResultHandler(ResultHandler<C>& res) {
     return 1;
 }
 
-} // anonymous namespace
+} // namespace
 
 HNSWStats HNSW::search(
         DistanceComputer& qdis,
@@ -1014,8 +1010,8 @@ HNSWStats HNSW::search(
     float d_nearest = qdis(nearest);
 
     for (int level = max_level; level >= 1; level--) {
-        HNSWStats local_stats =
-                greedy_update_nearest(*this, qdis, level, nearest, d_nearest);
+        HNSWStats local_stats = hnsw_utils::greedy_update_nearest(
+                *this, qdis, level, nearest, d_nearest);
         stats.combine(local_stats);
     }
 
@@ -1025,11 +1021,11 @@ HNSWStats HNSW::search(
 
         candidates.push(nearest, d_nearest);
 
-        search_from_candidates(
+        hnsw_utils::search_from_candidates(
                 *this, qdis, res, candidates, vt, stats, 0, 0, params);
     } else {
         std::priority_queue<Node> top_candidates =
-                search_from_candidate_unbounded(
+                hnsw_utils::search_from_candidate_unbounded(
                         *this, Node(d_nearest, nearest), qdis, ef, &vt, stats);
 
         while (top_candidates.size() > k) {
@@ -1081,7 +1077,7 @@ void HNSW::search_level_0(
 
             candidates.push(cj, nearest_d[j]);
 
-            nres = search_from_candidates(
+            nres = hnsw_utils::search_from_candidates(
                     hnsw,
                     qdis,
                     res,
@@ -1106,7 +1102,7 @@ void HNSW::search_level_0(
             candidates.push(cj, nearest_d[j]);
         }
 
-        search_from_candidates(
+        hnsw_utils::search_from_candidates(
                 hnsw, qdis, res, candidates, vt, search_stats, 0, 0, params);
     }
 }
diff --git a/faiss/impl/HNSW.h b/faiss/impl/HNSW.h
@@ -47,7 +47,6 @@ struct SearchParametersHNSW : SearchParameters {
     int efSearch = 16;
     bool check_relative_distance = true;
     bool bounded_queue = true;
-
     ~SearchParametersHNSW() {}
 };
 
@@ -255,4 +254,33 @@ struct HNSWStats {
 // global var that collects them all
 FAISS_API extern HNSWStats hnsw_stats;
 
+namespace hnsw_utils {
+int search_from_candidates(
+        const HNSW& hnsw,
+        DistanceComputer& qdis,
+        ResultHandler<HNSW::C>& res,
+        HNSW::MinimaxHeap& candidates,
+        VisitedTable& vt,
+        HNSWStats& stats,
+        int level,
+        int nres_in = 0,
+        const SearchParametersHNSW* params = nullptr,
+        const bool reference_version = false);
+
+HNSWStats greedy_update_nearest(
+        const HNSW& hnsw,
+        DistanceComputer& qdis,
+        int level,
+        HNSW::storage_idx_t& nearest,
+        float& d_nearest,
+        const bool reference_version = false);
+std::priority_queue<HNSW::Node> search_from_candidate_unbounded(
+        const HNSW& hnsw,
+        const HNSW::Node& node,
+        DistanceComputer& qdis,
+        int ef,
+        VisitedTable* vt,
+        HNSWStats& stats,
+        const bool reference_version = false);
+} // namespace hnsw_utils
 } // namespace faiss
diff --git a/tests/test_hnsw.cpp b/tests/test_hnsw.cpp
