Add prefetch for vector code

faiss/impl/DistanceComputer.h

@@ -8,6 +8,7 @@
 #pragma once
 
 #include <faiss/Index.h>
+#include <faiss/utils/prefetch.h>
 
 namespace faiss {
 
@@ -30,6 +31,9 @@ struct DistanceComputer {
     /// compute distance of vector i to current query
     virtual float operator()(idx_t i) = 0;
 
+    /// prefetch data for vector i, if the implementation supports it
+    virtual void prefetch(idx_t i) {}
+
     /// compute distances of current query to 4 stored vectors.
     /// certain DistanceComputer implementations may benefit
     /// heavily from this.
@@ -73,6 +77,10 @@ struct NegativeDistanceComputer : DistanceComputer {
         basedis->set_query(x);
     }
 
+    void prefetch(idx_t i) override {
+        basedis->prefetch(i);
+    }
+
     /// compute distance of vector i to current query
     float operator()(idx_t i) override {
         return -(*basedis)(i);
@@ -150,6 +158,27 @@ struct FlatCodesDistanceComputer : DistanceComputer {
                 dis3);
     }
 
+    void prefetch(idx_t i) override {
+        if (codes == nullptr || code_size == 0) {
+            return;
+        }
+
+        const uint8_t* code = codes + i * code_size;
+
+        // Prefetch a few cache lines. Do not prefetch the whole vector/code;
+        // the distance loop itself will stream through the rest.
+        prefetch_L2(code);
+        if (code_size > 64) {
+            prefetch_L2(code + 64);
+        }
+        if (code_size > 128) {
+            prefetch_L2(code + 128);
+        }
+        if (code_size > 192) {
+            prefetch_L2(code + 192);
+        }
+    }
+
     /// Computes a partial dot product over a slice of the query vector.
     /// The slice is defined by the following parameters:
     ///   — `offset`: the starting index of the first component to include

faiss/impl/HNSW.cpp

@@ -869,6 +869,8 @@ void HNSW::add_with_locks(
 
 namespace {
 
+static constexpr size_t kHNSWCodePrefetchAhead = 1;
+
 /** Helper to extract search parameters from HNSW and SearchParameters */
 inline void extract_search_params(
         const HNSW& hnsw,
@@ -980,6 +982,13 @@ int search_from_candidates_fixVT(
         };
 
         for (size_t j = begin; j < jmax; j++) {
+            if (j + kHNSWCodePrefetchAhead < jmax) {
+                int vp = hnsw.neighbors[j + kHNSWCodePrefetchAhead];
+                if (vp >= 0) {
+                    qdis.prefetch(vp);
+                }
+            }
+
             int v1 = hnsw.neighbors[j];
 
             saved_j[counter] = v1;
@@ -1420,6 +1429,13 @@ TopCandidatesQueue<C> search_from_candidate_unbounded_fixVT(
         };
 
         for (size_t j = begin; j < jmax; j++) {
+            if (j + kHNSWCodePrefetchAhead < jmax) {
+                int vp = hnsw.neighbors[j + kHNSWCodePrefetchAhead];
+                if (vp >= 0) {
+                    qdis.prefetch(vp);
+                }
+            }
+
             int v1 = hnsw.neighbors[j];
 
             saved_j[counter] = v1;