Integrate Panorama into IndexHNSWFlatPanorama

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();
         }
     }
@@ -290,7 +294,9 @@ struct FlatHammingDis : DistanceComputer {
 
     ~FlatHammingDis() override {
 #pragma omp critical
-        { hnsw_stats.ndis += ndis; }
+        {
+            hnsw_stats.ndis += ndis;
+        }
     }
 };
 

faiss/IndexFlat.cpp

@@ -103,12 +103,21 @@ 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, num_components);
+    }
+
     float symmetric_dis(idx_t i, idx_t j) override {
         return fvec_L2sqr(b + j * d, b + i * d, d);
     }
@@ -121,7 +130,8 @@ 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;
@@ -159,6 +169,41 @@ 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, y1, y2, y3, d, dp0_, dp1_, dp2_, dp3_);
+        dp0 = dp0_;
+        dp1 = dp1_;
+        dp2 = dp2_;
+        dp3 = dp3_;
+    }
 };
 
 struct FlatIPDis : FlatCodesDistanceComputer {

faiss/IndexHNSW.cpp

@@ -276,7 +276,7 @@ void hnsw_search(
                 res.begin(i);
                 dis->set_query(x + i * index->d);
 
-                HNSWStats stats = hnsw.search(*dis, res, vt, params);
+                HNSWStats stats = hnsw.search(*dis, index, res, vt, params);
                 n1 += stats.n1;
                 n2 += stats.n2;
                 ndis += stats.ndis;
@@ -647,6 +647,155 @@ IndexHNSWFlat::IndexHNSWFlat(int d, int M, MetricType metric)
     is_trained = true;
 }
 
+IndexHNSWFlat::IndexHNSWFlat(int d, int M, IndexFlat* storage)
+        : IndexHNSW(storage, M) {
+    own_fields = true;
+    is_trained = true;
+}
+
+/**************************************************************
+ * IndexHNSWFlatPanorama implementation
+ **************************************************************/
+
+namespace {
+void compute_cum_sums(
+        const float* x,
+        float* dst_cum_sums,
+        int d,
+        int n_levels,
+        int level_width) {
+    // Iterate backwards through levels, accumulating sum as we go.
+    // This avoids computing the suffix sum for each vector, which takes
+    // extra memory.
+
+    float sum = 0.0f;
+    dst_cum_sums[n_levels] = 0.0f;
+    for (int level = n_levels - 1; level >= 0; level--) {
+        int start_idx = level * level_width;
+        int end_idx = std::min(start_idx + level_width, d);
+        for (int j = start_idx; j < end_idx; j++) {
+            sum += x[j] * x[j];
+        }
+        dst_cum_sums[level] = sqrt(sum);
+    }
+}
+} // namespace
+
+IndexHNSWFlatPanorama::IndexHNSWFlatPanorama()
+        : IndexHNSWFlat(),
+          query_cum_sums(),
+          query_norm_sq(0.0f),
+          cum_sums(),
+          level_width(0),
+          n_levels(0) {}
+
+IndexHNSWFlatPanorama::IndexHNSWFlatPanorama(
+        int d,
+        int M,
+        int n_levels,
+        MetricType metric)
+        : IndexHNSWFlat(d, M, new IndexFlatL2(d)),
+          query_cum_sums(n_levels + 1),
+          query_norm_sq(0.0f),
+          cum_sums(),
+          level_width(d + (n_levels - 1) / n_levels),
+          n_levels(n_levels) {
+    // For now, we only support L2 distance.
+    // Supporting dot product and cosine distance is a trivial addition
+    // left for future work.
+    FAISS_THROW_IF_NOT(metric == METRIC_L2);
+
+    // Enable Panorama search mode.
+    // This is not ideal, but is still more simple than making a subclass of
+    // HNSW and overriding the search logic.
+    hnsw.is_panorama = true;
+}
+
+void IndexHNSWFlatPanorama::add(idx_t n, const float* x) {
+    idx_t n0 = ntotal;
+    cum_sums.resize((ntotal + n) * (n_levels + 1));
+
+    for (size_t idx = 0; idx < n; idx++) {
+        const float* vector = x + idx * d;
+        compute_cum_sums(
+                vector,
+                &cum_sums[(n0 + idx) * (n_levels + 1)],
+                d,
+                n_levels,
+                level_width);
+    }
+
+    IndexHNSWFlat::add(n, x);
+}
+
+void IndexHNSWFlatPanorama::reset() {
+    cum_sums.clear();
+    IndexHNSWFlat::reset();
+}
+
+size_t IndexHNSWFlatPanorama::remove_ids(const IDSelector& sel) {
+    idx_t j = 0;
+    for (idx_t i = 0; i < ntotal; i++) {
+        if (!sel.is_member(i)) {
+            if (i > j) {
+                memmove(&cum_sums[j * (n_levels + 1)],
+                        &cum_sums[i * (n_levels + 1)],
+                        (n_levels + 1) * sizeof(float));
+            }
+            j++;
+        }
+    }
+
+    idx_t n_removed = IndexHNSWFlat::remove_ids(sel);
+    if (n_removed > 0) {
+        // `ntotal` is updated by the base class.
+        cum_sums.resize(ntotal * (n_levels + 1));
+    }
+
+    return n_removed;
+}
+
+void IndexHNSWFlatPanorama::permute_entries(const idx_t* perm) {
+    std::vector<float> new_cum_sums(ntotal * (n_levels + 1));
+
+    for (idx_t i = 0; i < ntotal; i++) {
+        idx_t src = perm[i];
+        memcpy(&new_cum_sums[i * (n_levels + 1)],
+               &cum_sums[src * (n_levels + 1)],
+               (n_levels + 1) * sizeof(float));
+    }
+
+    std::swap(cum_sums, new_cum_sums);
+    IndexHNSWFlat::permute_entries(perm);
+}
+
+const float* IndexHNSWFlatPanorama::get_cum_sum(idx_t i) const {
+    return cum_sums.data() + i * (n_levels + 1);
+}
+
+void IndexHNSWFlatPanorama::search(
+        idx_t n,
+        const float* x,
+        idx_t k,
+        float* distances,
+        idx_t* labels,
+        const SearchParameters* params) const {
+    compute_cum_sums(x, query_cum_sums.data(), d, n_levels, level_width);
+    query_norm_sq = query_cum_sums[0] * query_cum_sums[0];
+    IndexHNSWFlat::search(n, x, k, distances, labels, params);
+}
+
+void IndexHNSWFlatPanorama::range_search(
+        idx_t n,
+        const float* x,
+        float radius,
+        RangeSearchResult* result,
+        const SearchParameters* params) const {
+    compute_cum_sums(x, query_cum_sums.data(), d, n_levels, level_width);
+    query_norm_sq = query_cum_sums[0] * query_cum_sums[0];
+    IndexHNSWFlat::range_search(n, x, radius, result, params);
+}
+
 /**************************************************************
  * IndexHNSWPQ implementation
  **************************************************************/

faiss/IndexHNSW.h

@@ -111,7 +111,7 @@ struct IndexHNSW : Index {
 
     void link_singletons();
 
-    void permute_entries(const idx_t* perm);
+    virtual void permute_entries(const idx_t* perm);
 
     DistanceComputer* get_distance_computer() const override;
 };
@@ -123,6 +123,69 @@ struct IndexHNSW : Index {
 struct IndexHNSWFlat : IndexHNSW {
     IndexHNSWFlat();
     IndexHNSWFlat(int d, int M, MetricType metric = METRIC_L2);
+    IndexHNSWFlat(int d, int M, IndexFlat* storage);
+};
+
+/** Panorama implementation of IndexHNSWFlat following
+ * https://www.arxiv.org/pdf/2510.00566.
+ *
+ * Unlike cluster-based Panorama, the vectors have to
+ * be higher dimensional (i.e. typically d > 512) and/or
+ * be able to compress a lot of their energy in the early
+ * dimensions to be effective. This is because HNSW accesses
+ * vectors in a random order, which makes cache misses dominate
+ * the distance computation time.
+ *
+ * The num_levels parameter controls the granularity of progressive distance
+ * refinement, allowing candidates to be eliminated early using partial
+ * distance computations rather than computing full distances.
+ *
+ * NOTE: This version of HNSW handles search slightly differently than the
+ * vanilla HNSW, as it uses partial distance computations with progressive
+ * refinement bounds. Instead of computing full distances immediately for all
+ * candidates, Panorama maintains lower and upper bounds that are incrementally
+ * tightened across refinement levels. Candidates are inserted into the search
+ * beam using approximate distance estimates (LB+UB)/2 and are only fully
+ * evaluated when they survive pruning and enter the result heap. This allows
+ * the algorithm to prune unpromising candidates early using Cauchy-Schwarz
+ * bounds on partial inner products. Hence, recall is not guaranteed to be the
+ * same as vanilla HNSW due to the heterogeneous precision within the search
+ * beam (exact vs. partial distance estimates affecting traversal order).
+ */
+struct IndexHNSWFlatPanorama : IndexHNSWFlat {
+    IndexHNSWFlatPanorama();
+    IndexHNSWFlatPanorama(
+            int d,
+            int M,
+            int num_levels,
+            MetricType metric = METRIC_L2);
+
+    void add(idx_t n, const float* x) override;
+    void reset() override;
+    size_t remove_ids(const IDSelector& sel) override;
+    void permute_entries(const idx_t* perm) override;
+
+    const float* get_cum_sum(idx_t i) const;
+
+    void search(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
+    void range_search(
+            idx_t n,
+            const float* x,
+            float radius,
+            RangeSearchResult* result,
+            const SearchParameters* params = nullptr) const override;
+
+    mutable std::vector<float> query_cum_sums;
+    mutable float query_norm_sq;
+    std::vector<float> cum_sums;
+    const size_t level_width;
+    const size_t n_levels;
 };
 
 /** PQ index topped with with a HNSW structure to access elements