Extend DistanceComputer batching to PQ, SQ, and PreTransform (#5309)

Summary:
Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5309

Implement `DistanceComputer` batching interfaces in SQ and PQ:
* Add `FlatCodesDistanceComputer::distance_to_code_batch_4` virtual , override in ScalarQuantizer distance computer to call `query_to_codes_batch_4`.
* Implement `PQDistanceComputer::distances_batch_4` using `PQCodeDist::distance_four_codes` for 4-way PQ lookup table evaluation.
* Propagate batching through `IndexPreTransform` wrapper (`PreTransformDistanceComputer`).

This allows HNSW graph traversal to issue batched distance calls instead of scalar `operator()` per neighbor, improving cache locality and enabling SIMD-friendly PQ four-code path, leading to more efficient search.

This is tested with a randomly-connected HNSW graph (only testing throughput, not recall) using `HNSW::fill_with_random_links()`, which has been tweaked to remove debug logging and properly initialized `entry_point` and `max_level`.

Reviewed By: mnorris11

Differential Revision: D108383062

fbshipit-source-id: 8153c43ba27e88fdf4114324cd00355c600d492b

benchs/bench_distance_computer.cpp

@@ -0,0 +1,118 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <memory>
+#include <vector>
+
+#include <benchmark/benchmark.h>
+#include <faiss/IndexHNSW.h>
+#include <faiss/index_factory.h>
+#include <faiss/utils/random.h>
+
+namespace faiss {
+
+const int d = 256;
+const size_t nb = 1000000;
+const size_t nq = 100000;
+
+std::vector<float> random_data(size_t n, int seed) {
+    std::vector<float> data(n);
+    float_rand(data.data(), n, seed);
+    return data;
+}
+
+const auto xb = random_data(nb * d, 1234);
+const auto xq = random_data(nq * d, 5678);
+
+std::unique_ptr<Index> make_index(const std::string& factory) {
+    std::unique_ptr<Index> index{faiss::index_factory(d, factory.c_str())};
+    index->train(100000, xb.data()); // train on subset for speed.
+    index->add(nb, xb.data());
+    return index;
+}
+
+std::unique_ptr<IndexHNSW> make_hollow_hnsw() {
+    auto index = std::make_unique<IndexHNSW>();
+    index->d = d;
+    index->own_fields = false;
+    auto& hnsw = index->hnsw;
+    hnsw.efSearch = 32;
+    hnsw.fill_with_random_links(nb);
+    index->ntotal = nb;
+    return index;
+}
+
+using benchmark::Counter;
+
+void bench_index(benchmark::State& state, const Index& index) {
+    std::vector<float> distances(nq * 10);
+    std::vector<faiss::idx_t> labels(nq * 10);
+    for (auto _ : state) {
+        index.search(nq, xq.data(), 10, distances.data(), labels.data());
+        benchmark::DoNotOptimize(distances[1]);
+    }
+    state.counters["nq"] = Counter(nq, Counter::kAvgThreadsRate);
+}
+
+void bench_hnsw(benchmark::State& state, Index& storage) {
+    static auto index = make_hollow_hnsw();
+    index->storage = &storage;
+    hnsw_stats.reset();
+    bench_index(state, *index);
+    state.counters["ndis"] = Counter(hnsw_stats.ndis, Counter::kAvgThreadsRate);
+}
+
+void bench_ivf(benchmark::State& state, Index& index) {
+    indexIVF_stats.reset();
+    bench_index(state, index);
+    state.counters["ndis"] =
+            Counter(indexIVF_stats.ndis, Counter::kAvgThreadsRate);
+}
+
+// Define and register a benchmark for a quantizer. The function name and the
+// index_factory string are both derived from it. Each index is built lazily in
+// a function-local static, so filtering to a single benchmark only builds the
+// index it actually uses.
+#define BENCH_IVF(Q)                                  \
+    void BM_IVF_##Q(benchmark::State& state) {        \
+        static auto index = make_index("IVF100," #Q); \
+        bench_ivf(state, *index);                     \
+    }                                                 \
+    BENCHMARK(BM_IVF_##Q)                             \
+            ->MeasureProcessCPUTime()                 \
+            ->Unit(benchmark::kMillisecond)           \
+            ->Name("IVF1K," #Q)
+
+#define BENCH_HNSW(Q)                           \
+    void BM_HNSW_##Q(benchmark::State& state) { \
+        static auto storage = make_index(#Q);   \
+        bench_hnsw(state, *storage);            \
+    }                                           \
+    BENCHMARK(BM_HNSW_##Q)                      \
+            ->MeasureProcessCPUTime()           \
+            ->Unit(benchmark::kMillisecond)     \
+            ->Name("HNSW32," #Q)
+
+BENCH_HNSW(Flat);
+BENCH_HNSW(SQ8);
+BENCH_HNSW(SQ6);
+BENCH_HNSW(SQ4);
+BENCH_HNSW(PQ64x8np);
+BENCH_HNSW(PQ64x4np);
+BENCH_HNSW(RaBitQ1);
+BENCH_HNSW(RaBitQ2);
+
+BENCH_IVF(SQ4);
+BENCH_IVF(SQ8);
+BENCH_IVF(PQ64x8np);
+
+#undef BENCH_IVF
+#undef BENCH_HNSW
+
+} // namespace faiss
+
+BENCHMARK_MAIN();

faiss/IndexPreTransform.cpp

@@ -345,6 +345,19 @@ struct PreTransformDistanceComputer : DistanceComputer {
     float operator()(idx_t i) override {
         return (*sub_dc)(i);
     }
+
+    void distances_batch_4(
+            const idx_t idx0,
+            const idx_t idx1,
+            const idx_t idx2,
+            const idx_t idx3,
+            float& dis0,
+            float& dis1,
+            float& dis2,
+            float& dis3) override {
+        sub_dc->distances_batch_4(
+                idx0, idx1, idx2, idx3, dis0, dis1, dis2, dis3);
+    }
 };
 
 } // anonymous namespace

faiss/impl/DistanceComputer.h

@@ -130,6 +130,26 @@ struct FlatCodesDistanceComputer : DistanceComputer {
         return distance_to_code(codes + i * code_size);
     }
 
+    void distances_batch_4(
+            const idx_t idx0,
+            const idx_t idx1,
+            const idx_t idx2,
+            const idx_t idx3,
+            float& dis0,
+            float& dis1,
+            float& dis2,
+            float& dis3) override {
+        distance_to_code_batch_4(
+                codes + idx0 * code_size,
+                codes + idx1 * code_size,
+                codes + idx2 * code_size,
+                codes + idx3 * code_size,
+                dis0,
+                dis1,
+                dis2,
+                dis3);
+    }
+
     /// 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
@@ -167,6 +187,20 @@ struct FlatCodesDistanceComputer : DistanceComputer {
 
     /// compute distance of current query to an encoded vector
     virtual float distance_to_code(const uint8_t* code) = 0;
+    virtual void distance_to_code_batch_4(
+            const uint8_t* c1,
+            const uint8_t* c2,
+            const uint8_t* c3,
+            const uint8_t* c4,
+            float& d1,
+            float& d2,
+            float& d3,
+            float& d4) {
+        d1 = distance_to_code(c1);
+        d2 = distance_to_code(c2);
+        d3 = distance_to_code(c3);
+        d4 = distance_to_code(c4);
+    }
 
     /// Compute partial dot products of current query to 4 stored vectors.
     /// See `partial_dot_product` for more details.

faiss/impl/HNSW.cpp

@@ -172,17 +172,21 @@ void HNSW::print_neighbor_stats(int level) const {
 }
 
 void HNSW::fill_with_random_links(size_t n) {
-    int max_level_2 = prepare_level_tab(n);
+    if (n == 0) {
+        return;
+    }
+    max_level = prepare_level_tab(n);
+    entry_point = 0;
+
     RandomGenerator rng2(456);
 
-    for (int level = max_level_2 - 1; level >= 0; --level) {
+    for (int level = max_level - 1; level >= 0; --level) {
         std::vector<int> elts;
         for (size_t i = 0; i < n; i++) {
             if (levels[i] > level) {
                 elts.push_back(i);
             }
         }
-        printf("linking %zd elements in level %d\n", elts.size(), level);
 
         if (elts.size() == 1) {
             continue;

faiss/impl/ScalarQuantizer.h

@@ -135,6 +135,18 @@ struct ScalarQuantizer : Quantizer {
         float distance_to_code(const uint8_t* code) final {
             return query_to_code(code);
         }
+
+        void distance_to_code_batch_4(
+                const uint8_t* c1,
+                const uint8_t* c2,
+                const uint8_t* c3,
+                const uint8_t* c4,
+                float& d1,
+                float& d2,
+                float& d3,
+                float& d4) override {
+            query_to_codes_batch_4(c1, c2, c3, c4, d1, d2, d3, d4);
+        }
     };
 
     /// TurboQuant full (QT_*_tq) refinement state, isolated from the

faiss/impl/pq_code_distance/PQDistanceComputer_impl.h

@@ -20,23 +20,39 @@ namespace faiss {
 namespace pq_code_distance {
 
 template <class PQCodeDist>
-struct PQDistanceComputer : FlatCodesDistanceComputer {
+struct PQDistanceComputer final : FlatCodesDistanceComputer {
     using PQDecoder = typename PQCodeDist::PQDecoder;
-    size_t d;
     MetricType metric;
-    idx_t nb;
     const ProductQuantizer& pq;
     const float* sdc;
     std::vector<float> precomputed_table;
-    size_t ndis;
-    const float* q;
 
     float distance_to_code(const uint8_t* code) final {
-        ndis++;
-
-        float dis = PQCodeDist::distance_single_code(
+        return PQCodeDist::distance_single_code(
                 pq.M, pq.nbits, precomputed_table.data(), code);
-        return dis;
+    }
+
+    void distance_to_code_batch_4(
+            const uint8_t* c1,
+            const uint8_t* c2,
+            const uint8_t* c3,
+            const uint8_t* c4,
+            float& d1,
+            float& d2,
+            float& d3,
+            float& d4) override {
+        PQCodeDist::distance_four_codes(
+                pq.M,
+                pq.nbits,
+                precomputed_table.data(),
+                c1,
+                c2,
+                c3,
+                c4,
+                d1,
+                d2,
+                d3,
+                d4);
     }
 
     float symmetric_dis(idx_t i, idx_t j) override {
@@ -50,7 +66,6 @@ struct PQDistanceComputer : FlatCodesDistanceComputer {
             accu += sdci[codei.decode() + (codej.decode() << codei.nbits)];
             sdci += uint64_t(1) << (2 * codei.nbits);
         }
-        ndis++;
         return accu;
     }
 
@@ -58,18 +73,14 @@ struct PQDistanceComputer : FlatCodesDistanceComputer {
             : FlatCodesDistanceComputer(
                       storage.codes.data(),
                       storage.code_size),
-              pq(storage.pq),
-              q(nullptr) {
+              pq(storage.pq) {
         precomputed_table.resize(pq.M * pq.ksub);
-        nb = storage.ntotal;
-        d = storage.d;
         metric = storage.metric_type;
         if (pq.sdc_table.size() == pq.ksub * pq.ksub * pq.M) {
             sdc = pq.sdc_table.data();
         } else {
             sdc = nullptr;
         }
-        ndis = 0;
     }
 
     void set_query(const float* x) override {