Integrate Panorama into IndexHNSWFlatPanorama

benchs/bench_hnsw_flat_panorama.py

@@ -0,0 +1,154 @@
+# 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.
+
+import multiprocessing as mp
+import time
+
+import faiss
+import matplotlib.pyplot as plt
+import numpy as np
+
+try:
+    from faiss.contrib.datasets_fb import (
+        DatasetSIFT1M,
+        DatasetGIST1M,
+        SyntheticDataset,
+    )
+except ImportError:
+    from faiss.contrib.datasets import (
+        DatasetSIFT1M,
+        DatasetGIST1M,
+        SyntheticDataset,
+    )
+
+
+def eval_recall(index, efSearch_val, xq, gt, k):
+    """Evaluate recall and QPS for a given efSearch value."""
+    t0 = time.time()
+    _, I = index.search(xq, k=k)
+    t = time.time() - t0
+    speed = t * 1000 / len(xq)
+    qps = 1000 / speed
+
+    corrects = (gt == I).sum()
+    recall = corrects / (len(xq) * k)
+    print(
+        f"\tefSearch {efSearch_val:3d}, Recall@{k}: "
+        f"{recall:.6f}, speed: {speed:.6f} ms/query, QPS: {qps:.2f}"
+    )
+
+    return recall, qps
+
+
+def get_hnsw_index(index):
+    """Extract the underlying HNSW index from a PreTransform index."""
+    if isinstance(index, faiss.IndexPreTransform):
+        return faiss.downcast_index(index.index)
+    return index
+
+
+def eval_and_plot(name, ds, k=10, nlevels=8, plot_data=None):
+    """Evaluate an index configuration and collect data for plotting."""
+    xq = ds.get_queries()
+    xb = ds.get_database()
+    gt = ds.get_groundtruth()
+
+    if hasattr(ds, "get_train"):
+        xt = ds.get_train()
+    else:
+        # Use database as training data if no separate train set
+        xt = xb
+
+    nb, d = xb.shape
+    nq, d = xq.shape
+    gt = gt[:, :k]
+
+    print(f"\n======{name} on {ds.__class__.__name__}======")
+    print(f"Database: {nb} vectors, {d} dimensions")
+    print(f"Queries: {nq} vectors")
+
+    # Create index
+    index = faiss.index_factory(d, name)
+
+    faiss.omp_set_num_threads(mp.cpu_count())
+    index.train(xt)
+    index.add(xb)
+
+    faiss.omp_set_num_threads(1)
+
+    # Get the underlying HNSW index for setting efSearch
+    hnsw_index = get_hnsw_index(index)
+
+    data = []
+    for efSearch in [16, 32, 64, 128, 256, 512]:
+        hnsw_index.hnsw.efSearch = efSearch
+        recall, qps = eval_recall(index, efSearch, xq, gt, k)
+        data.append((recall, qps))
+
+    if plot_data is not None:
+        data = np.array(data)
+        plot_data.append((name, data))
+
+
+def benchmark_dataset(ds, dataset_name, k=10, nlevels=8, M=32):
+    """Benchmark both regular HNSW and HNSW Panorama on a dataset."""
+    d = ds.d
+
+    plot_data = []
+
+    # HNSW Flat (baseline)
+    eval_and_plot(f"HNSW{M},Flat", ds, k=k, nlevels=nlevels, plot_data=plot_data)
+
+    # HNSW Flat Panorama (with PCA to concentrate energy)
+    eval_and_plot(
+        f"PCA{d},HNSW{M},FlatPanorama{nlevels}",
+        ds,
+        k=k,
+        nlevels=nlevels,
+        plot_data=plot_data,
+    )
+
+    # Plot results
+    plt.figure(figsize=(8, 6), dpi=80)
+    for name, data in plot_data:
+        plt.plot(data[:, 0], data[:, 1], marker="o", label=name)
+
+    plt.title(f"HNSW Indexes on {dataset_name}")
+    plt.xlabel(f"Recall@{k}")
+    plt.ylabel("QPS")
+    plt.yscale("log")
+    plt.legend(bbox_to_anchor=(1.02, 0.1), loc="upper left", borderaxespad=0)
+    plt.grid(True, alpha=0.3)
+
+    output_file = f"bench_hnsw_flat_panorama_{dataset_name}.png"
+    plt.savefig(output_file, bbox_inches="tight")
+    print(f"Saved plot to {output_file}")
+    plt.close()
+
+
+if __name__ == "__main__":
+    k = 10
+    nlevels = 8
+    M = 32
+
+    # Test on 3 datasets with varying dimensionality:
+    # SIFT1M (128d), GIST1M (960d), and Synthetic high-dim (2048d)
+    datasets = [
+        (DatasetSIFT1M(), "SIFT1M"),
+        (DatasetGIST1M(), "GIST1M"),
+        # Synthetic high-dimensional dataset: 2048d, 100k train, 1M database, 10k queries
+        (SyntheticDataset(2048, 100000, 1000000, 10000), "Synthetic2048D"),
+    ]
+
+    for ds, name in datasets:
+        print(f"\n{'='*60}")
+        print(f"Benchmarking on {name}")
+        print(f"{'='*60}")
+        benchmark_dataset(ds, name, k=k, nlevels=nlevels, M=M)
+
+    print("\n" + "="*60)
+    print("All benchmarks completed!")
+    print("="*60)
+

faiss/IndexAdditiveQuantizer.cpp

@@ -58,6 +58,10 @@ struct AQDistanceComputerDecompress : FlatCodesDistanceComputer {
         q = x;
     }
 
+    const float* get_query() const override {
+        return q;
+    }
+
     float symmetric_dis(idx_t i, idx_t j) final {
         aq.decode(codes + i * d, tmp.data(), 1);
         aq.decode(codes + j * d, tmp.data() + d, 1);
@@ -77,15 +81,18 @@ struct AQDistanceComputerLUT : FlatCodesDistanceComputer {
     std::vector<float> LUT;
     const AdditiveQuantizer& aq;
     size_t d;
+    const float* q;
 
     explicit AQDistanceComputerLUT(const IndexAdditiveQuantizer& iaq)
             : FlatCodesDistanceComputer(iaq.codes.data(), iaq.code_size),
               LUT(iaq.aq->total_codebook_size + iaq.d * 2),
               aq(*iaq.aq),
-              d(iaq.d) {}
+              d(iaq.d),
+              q(nullptr) {}
 
     float bias;
     void set_query(const float* x) final {
+        q = x;
         // this is quite sub-optimal for multiple queries
         aq.compute_LUT(1, x, LUT.data());
         if (is_IP) {
@@ -95,6 +102,10 @@ struct AQDistanceComputerLUT : FlatCodesDistanceComputer {
         }
     }
 
+    const float* get_query() const override {
+        return q;
+    }
+
     float symmetric_dis(idx_t i, idx_t j) final {
         float* tmp = LUT.data();
         aq.decode(codes + i * d, tmp, 1);

faiss/IndexBinaryHNSW.cpp

@@ -227,7 +227,11 @@ void IndexBinaryHNSW::search(
         for (idx_t i = 0; i < n; i++) {
             res.begin(i);
             dis->set_query((float*)(x + i * code_size));
-            hnsw.search(*dis, res, vt);
+            // Given that IndexBinaryHNSW is not an IndexHNSW, we pass nullptr
+            // as the index parameter. This state does not get used in the
+            // search function, as it is merely there to to enable Panorama
+            // execution for IndexHNSWFlatPanorama.
+            hnsw.search(*dis, nullptr, res, vt);
             res.end();
         }
     }

faiss/IndexFlat.cpp

@@ -103,12 +103,22 @@ struct FlatL2Dis : FlatCodesDistanceComputer {
     const float* q;
     const float* b;
     size_t ndis;
+    size_t npartial_dot_products;
 
     float distance_to_code(const uint8_t* code) final {
         ndis++;
         return fvec_L2sqr(q, (float*)code, d);
     }
 
+    float partial_dot_product(
+            const idx_t i,
+            const uint32_t offset,
+            const uint32_t num_components) final override {
+        npartial_dot_products++;
+        return fvec_inner_product(
+                q + offset, b + i * d + offset, num_components);
+    }
+
     float symmetric_dis(idx_t i, idx_t j) override {
         return fvec_L2sqr(b + j * d, b + i * d, d);
     }
@@ -121,12 +131,17 @@ struct FlatL2Dis : FlatCodesDistanceComputer {
               nb(storage.ntotal),
               q(q),
               b(storage.get_xb()),
-              ndis(0) {}
+              ndis(0),
+              npartial_dot_products(0) {}
 
     void set_query(const float* x) override {
         q = x;
     }
 
+    const float* get_query() const override {
+        return q;
+    }
+
     // compute four distances
     void distances_batch_4(
             const idx_t idx0,
@@ -159,6 +174,50 @@ struct FlatL2Dis : FlatCodesDistanceComputer {
         dis2 = dp2;
         dis3 = dp3;
     }
+
+    void partial_dot_product_batch_4(
+            const idx_t idx0,
+            const idx_t idx1,
+            const idx_t idx2,
+            const idx_t idx3,
+            float& dp0,
+            float& dp1,
+            float& dp2,
+            float& dp3,
+            const uint32_t offset,
+            const uint32_t num_components) final override {
+        npartial_dot_products += 4;
+
+        // compute first, assign next
+        const float* __restrict y0 =
+                reinterpret_cast<const float*>(codes + idx0 * code_size);
+        const float* __restrict y1 =
+                reinterpret_cast<const float*>(codes + idx1 * code_size);
+        const float* __restrict y2 =
+                reinterpret_cast<const float*>(codes + idx2 * code_size);
+        const float* __restrict y3 =
+                reinterpret_cast<const float*>(codes + idx3 * code_size);
+
+        float dp0_ = 0;
+        float dp1_ = 0;
+        float dp2_ = 0;
+        float dp3_ = 0;
+        fvec_inner_product_batch_4(
+                q + offset,
+                y0 + offset,
+                y1 + offset,
+                y2 + offset,
+                y3 + offset,
+                num_components,
+                dp0_,
+                dp1_,
+                dp2_,
+                dp3_);
+        dp0 = dp0_;
+        dp1 = dp1_;
+        dp2 = dp2_;
+        dp3 = dp3_;
+    }
 };
 
 struct FlatIPDis : FlatCodesDistanceComputer {
@@ -191,6 +250,10 @@ struct FlatIPDis : FlatCodesDistanceComputer {
         q = x;
     }
 
+    const float* get_query() const override {
+        return q;
+    }
+
     // compute four distances
     void distances_batch_4(
             const idx_t idx0,
@@ -315,6 +378,10 @@ struct FlatL2WithNormsDis : FlatCodesDistanceComputer {
         query_l2norm = fvec_norm_L2sqr(q, d);
     }
 
+    const float* get_query() const override {
+        return q;
+    }
+
     // compute four distances
     void distances_batch_4(
             const idx_t idx0,

faiss/IndexFlatCodes.cpp

@@ -142,6 +142,10 @@ struct GenericFlatCodesDistanceComputer : FlatCodesDistanceComputer {
         query = x;
     }
 
+    const float* get_query() const override {
+        return query;
+    }
+
     float operator()(idx_t i) override {
         codec.sa_decode(1, codes + i * code_size, vec_buffer.data());
         return vd(query, vec_buffer.data());