faiss 0.5.3 → 0.6.1

data/vendor/faiss/faiss/IndexFlat.cpp

@@ -12,6 +12,7 @@
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/ResultHandler.h>
+#include <faiss/impl/simd_dispatch.h>
 #include <faiss/utils/Heap.h>
 #include <faiss/utils/distances.h>
 #include <faiss/utils/extra_distances.h>
@@ -19,12 +20,11 @@
 #include <faiss/utils/sorting.h>
 #include <omp.h>
 #include <cstring>
-#include <numeric>
 
 namespace faiss {
 
-IndexFlat::IndexFlat(idx_t d, MetricType metric)
-        : IndexFlatCodes(sizeof(float) * d, d, metric) {}
+IndexFlat::IndexFlat(idx_t d_, MetricType metric)
+        : IndexFlatCodes(sizeof(float) * d_, d_, metric) {}
 
 void IndexFlat::search(
         idx_t n,
@@ -44,7 +44,6 @@ void IndexFlat::search(
         float_maxheap_array_t res = {size_t(n), size_t(k), labels, distances};
         knn_L2sqr(x, get_xb(), d, n, ntotal, &res, nullptr, sel);
     } else {
-        FAISS_THROW_IF_NOT(!sel); // TODO implement with selector
         knn_extra_metrics(
                 x,
                 get_xb(),
@@ -55,7 +54,8 @@ void IndexFlat::search(
                 metric_arg,
                 k,
                 distances,
-                labels);
+                labels,
+                sel);
     }
 }
 
@@ -65,6 +65,7 @@ void IndexFlat::range_search(
         float radius,
         RangeSearchResult* result,
         const SearchParameters* params) const {
+    FAISS_THROW_IF_NOT_MSG(result, "RangeSearchResult object must not be null");
     IDSelector* sel = params ? params->sel : nullptr;
 
     switch (metric_type) {
@@ -86,6 +87,7 @@ void IndexFlat::compute_distance_subset(
         idx_t k,
         float* distances,
         const idx_t* labels) const {
+    FAISS_THROW_IF_NOT(k > 0);
     switch (metric_type) {
         case METRIC_INNER_PRODUCT:
             fvec_inner_products_by_idx(distances, x, get_xb(), labels, d, n, k);
@@ -100,6 +102,7 @@ void IndexFlat::compute_distance_subset(
 
 namespace {
 
+template <SIMDLevel SL>
 struct FlatL2Dis : FlatCodesDistanceComputer {
     size_t d;
     idx_t nb;
@@ -109,7 +112,7 @@ struct FlatL2Dis : FlatCodesDistanceComputer {
 
     float distance_to_code(const uint8_t* code) final {
         ndis++;
-        return fvec_L2sqr(q, (float*)code, d);
+        return fvec_L2sqr<SL>(q, (float*)code, d);
     }
 
     float partial_dot_product(
@@ -117,19 +120,19 @@ struct FlatL2Dis : FlatCodesDistanceComputer {
             const uint32_t offset,
             const uint32_t num_components) final override {
         npartial_dot_products++;
-        return fvec_inner_product(
+        return fvec_inner_product<SL>(
                 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);
+        return fvec_L2sqr<SL>(b + j * d, b + i * d, d);
     }
 
-    explicit FlatL2Dis(const IndexFlat& storage, const float* q = nullptr)
+    explicit FlatL2Dis(const IndexFlat& storage, const float* q_ = nullptr)
             : FlatCodesDistanceComputer(
                       storage.codes.data(),
                       storage.code_size,
-                      q),
+                      q_),
               d(storage.d),
               nb(storage.ntotal),
               b(storage.get_xb()),
@@ -166,7 +169,7 @@ struct FlatL2Dis : FlatCodesDistanceComputer {
         float dp1 = 0;
         float dp2 = 0;
         float dp3 = 0;
-        fvec_L2sqr_batch_4(q, y0, y1, y2, y3, d, dp0, dp1, dp2, dp3);
+        fvec_L2sqr_batch_4<SL>(q, y0, y1, y2, y3, d, dp0, dp1, dp2, dp3);
         dis0 = dp0;
         dis1 = dp1;
         dis2 = dp2;
@@ -200,7 +203,7 @@ struct FlatL2Dis : FlatCodesDistanceComputer {
         float dp1_ = 0;
         float dp2_ = 0;
         float dp3_ = 0;
-        fvec_inner_product_batch_4(
+        fvec_inner_product_batch_4<SL>(
                 q + offset,
                 y0 + offset,
                 y1 + offset,
@@ -218,6 +221,7 @@ struct FlatL2Dis : FlatCodesDistanceComputer {
     }
 };
 
+template <SIMDLevel SL>
 struct FlatIPDis : FlatCodesDistanceComputer {
     size_t d;
     idx_t nb;
@@ -226,21 +230,21 @@ struct FlatIPDis : FlatCodesDistanceComputer {
     size_t ndis;
 
     float symmetric_dis(idx_t i, idx_t j) final override {
-        return fvec_inner_product(b + j * d, b + i * d, d);
+        return fvec_inner_product<SL>(b + j * d, b + i * d, d);
     }
 
     float distance_to_code(const uint8_t* code) final override {
         ndis++;
-        return fvec_inner_product(q, (const float*)code, d);
+        return fvec_inner_product<SL>(q, (const float*)code, d);
     }
 
-    explicit FlatIPDis(const IndexFlat& storage, const float* q = nullptr)
+    explicit FlatIPDis(const IndexFlat& storage, const float* q_in = nullptr)
             : FlatCodesDistanceComputer(
                       storage.codes.data(),
                       storage.code_size),
               d(storage.d),
               nb(storage.ntotal),
-              q(q),
+              q(q_in),
               b(storage.get_xb()),
               ndis(0) {}
 
@@ -274,7 +278,8 @@ struct FlatIPDis : FlatCodesDistanceComputer {
         float dp1 = 0;
         float dp2 = 0;
         float dp3 = 0;
-        fvec_inner_product_batch_4(q, y0, y1, y2, y3, d, dp0, dp1, dp2, dp3);
+        fvec_inner_product_batch_4<SL>(
+                q, y0, y1, y2, y3, d, dp0, dp1, dp2, dp3);
         dis0 = dp0;
         dis1 = dp1;
         dis2 = dp2;
@@ -285,14 +290,16 @@ struct FlatIPDis : FlatCodesDistanceComputer {
 } // namespace
 
 FlatCodesDistanceComputer* IndexFlat::get_FlatCodesDistanceComputer() const {
+    FlatCodesDistanceComputer* dc = nullptr;
     if (metric_type == METRIC_L2) {
-        return new FlatL2Dis(*this);
+        with_simd_level([&]<SIMDLevel SL>() { dc = new FlatL2Dis<SL>(*this); });
     } else if (metric_type == METRIC_INNER_PRODUCT) {
-        return new FlatIPDis(*this);
+        with_simd_level([&]<SIMDLevel SL>() { dc = new FlatIPDis<SL>(*this); });
     } else {
-        return get_extra_distance_computer(
+        dc = get_extra_distance_computer(
                 d, metric_type, metric_arg, ntotal, get_xb());
     }
+    return dc;
 }
 
 void IndexFlat::reconstruct(idx_t key, float* recons) const {
@@ -317,6 +324,7 @@ void IndexFlat::sa_decode(idx_t n, const uint8_t* bytes, float* x) const {
  ***************************************************/
 
 namespace {
+template <SIMDLevel SL>
 struct FlatL2WithNormsDis : FlatCodesDistanceComputer {
     size_t d;
     idx_t nb;
@@ -329,7 +337,7 @@ struct FlatL2WithNormsDis : FlatCodesDistanceComputer {
 
     float distance_to_code(const uint8_t* code) final override {
         ndis++;
-        return fvec_L2sqr(q, (float*)code, d);
+        return fvec_L2sqr<SL>(q, (float*)code, d);
     }
 
     float operator()(const idx_t i) final override {
@@ -337,7 +345,7 @@ struct FlatL2WithNormsDis : FlatCodesDistanceComputer {
                 reinterpret_cast<const float*>(codes + i * code_size);
 
         prefetch_L2(l2norms + i);
-        const float dp0 = fvec_inner_product(q, y, d);
+        const float dp0 = fvec_inner_product<SL>(q, y, d);
         return query_l2norm + l2norms[i] - 2 * dp0;
     }
 
@@ -349,19 +357,19 @@ struct FlatL2WithNormsDis : FlatCodesDistanceComputer {
 
         prefetch_L2(l2norms + i);
         prefetch_L2(l2norms + j);
-        const float dp0 = fvec_inner_product(yi, yj, d);
+        const float dp0 = fvec_inner_product<SL>(yi, yj, d);
         return l2norms[i] + l2norms[j] - 2 * dp0;
     }
 
     explicit FlatL2WithNormsDis(
             const IndexFlatL2& storage,
-            const float* q = nullptr)
+            const float* q_in = nullptr)
             : FlatCodesDistanceComputer(
                       storage.codes.data(),
                       storage.code_size),
               d(storage.d),
               nb(storage.ntotal),
-              q(q),
+              q(q_in),
               b(storage.get_xb()),
               ndis(0),
               l2norms(storage.cached_l2norms.data()),
@@ -369,7 +377,7 @@ struct FlatL2WithNormsDis : FlatCodesDistanceComputer {
 
     void set_query(const float* x) override {
         q = x;
-        query_l2norm = fvec_norm_L2sqr(q, d);
+        query_l2norm = fvec_norm_L2sqr<SL>(q, d);
     }
 
     // compute four distances
@@ -403,7 +411,8 @@ struct FlatL2WithNormsDis : FlatCodesDistanceComputer {
         float dp1 = 0;
         float dp2 = 0;
         float dp3 = 0;
-        fvec_inner_product_batch_4(q, y0, y1, y2, y3, d, dp0, dp1, dp2, dp3);
+        fvec_inner_product_batch_4<SL>(
+                q, y0, y1, y2, y3, d, dp0, dp1, dp2, dp3);
         dis0 = query_l2norm + l2norms[idx0] - 2 * dp0;
         dis1 = query_l2norm + l2norms[idx1] - 2 * dp1;
         dis2 = query_l2norm + l2norms[idx2] - 2 * dp2;
@@ -430,7 +439,11 @@ void IndexFlatL2::clear_l2norms() {
 FlatCodesDistanceComputer* IndexFlatL2::get_FlatCodesDistanceComputer() const {
     if (metric_type == METRIC_L2) {
         if (!cached_l2norms.empty()) {
-            return new FlatL2WithNormsDis(*this);
+            FlatCodesDistanceComputer* dc = nullptr;
+            with_simd_level([&]<SIMDLevel SL>() {
+                dc = new FlatL2WithNormsDis<SL>(*this);
+            });
+            return dc;
         }
     }
 
@@ -441,8 +454,8 @@ FlatCodesDistanceComputer* IndexFlatL2::get_FlatCodesDistanceComputer() const {
  * IndexFlat1D
  ***************************************************/
 
-IndexFlat1D::IndexFlat1D(bool continuous_update)
-        : IndexFlatL2(1), continuous_update(continuous_update) {}
+IndexFlat1D::IndexFlat1D(bool continuous_update_in)
+        : IndexFlatL2(1), continuous_update(continuous_update_in) {}
 
 /// if not continuous_update, call this between the last add and
 /// the first search
@@ -478,7 +491,8 @@ void IndexFlat1D::search(
             !params, "search params not supported for this index");
     FAISS_THROW_IF_NOT(k > 0);
     FAISS_THROW_IF_NOT_MSG(
-            perm.size() == ntotal, "Call update_permutation before search");
+            perm.size() == static_cast<size_t>(ntotal),
+            "Call update_permutation before search");
     const float* xb = get_xb();
 
 #pragma omp parallel for if (n > 10000)
@@ -583,7 +597,17 @@ void IndexFlat1D::search(
 
 namespace {
 
-template <bool use_radius, typename BlockHandler>
+template <typename Fn>
+inline auto dispatch_metric_compare(MetricType metric, Fn&& fn) {
+    if (is_similarity_metric(metric)) {
+        using C = CMin<float, int64_t>;
+        return fn.template operator()<C>();
+    }
+    using C = CMax<float, int64_t>;
+    return fn.template operator()<C>();
+}
+
+template <bool use_radius, typename C, typename BlockHandler>
 inline void flat_pano_search_core(
         const IndexFlatPanorama& index,
         BlockHandler& handler,
@@ -603,9 +627,11 @@ inline void flat_pano_search_core(
     {
         SingleResultHandler res(handler);
 
-        std::vector<float> query_cum_norms(index.n_levels + 1);
-        std::vector<float> exact_distances(index.batch_size);
+        std::vector<float> query_cum_norms(index.pano.n_levels + 1);
         std::vector<uint32_t> active_indices(index.batch_size);
+        std::vector<uint8_t> active_byteset(index.batch_size);
+        std::vector<float> exact_distances(index.batch_size);
+        std::vector<float> dot_buffer(index.batch_size);
 
 #pragma omp for
         for (int64_t i = 0; i < n; i++) {
@@ -627,22 +653,25 @@ inline void flat_pano_search_core(
                     threshold = res.heap_dis[0];
                 }
 
-                size_t num_active =
-                        index.pano
-                                .progressive_filter_batch<CMax<float, int64_t>>(
-                                        index.codes.data(),
-                                        index.cum_sums.data(),
-                                        xi,
-                                        query_cum_norms.data(),
-                                        batch_no,
-                                        index.ntotal,
-                                        sel,
-                                        nullptr,
-                                        use_sel,
-                                        active_indices,
-                                        exact_distances,
-                                        threshold,
-                                        local_stats);
+                size_t num_active = with_metric_type(
+                        index.metric_type, [&]<MetricType M>() {
+                            return index.pano.progressive_filter_batch<C, M>(
+                                    index.codes.data(),
+                                    index.cum_sums.data(),
+                                    xi,
+                                    query_cum_norms.data(),
+                                    batch_no,
+                                    index.ntotal,
+                                    sel,
+                                    nullptr,
+                                    use_sel,
+                                    active_indices,
+                                    active_byteset,
+                                    exact_distances,
+                                    dot_buffer,
+                                    threshold,
+                                    local_stats);
+                        });
 
                 for (size_t j = 0; j < num_active; j++) {
                     res.add_result(
@@ -669,7 +698,7 @@ void IndexFlatPanorama::add(idx_t n, const float* x) {
     size_t num_batches = (ntotal + batch_size - 1) / batch_size;
 
     codes.resize(num_batches * batch_size * code_size);
-    cum_sums.resize(num_batches * batch_size * (n_levels + 1));
+    cum_sums.resize(num_batches * batch_size * (pano.n_levels + 1));
 
     const uint8_t* code = reinterpret_cast<const uint8_t*>(x);
     pano.copy_codes_to_level_layout(codes.data(), offset, n, code);
@@ -684,12 +713,13 @@ void IndexFlatPanorama::search(
         idx_t* labels,
         const SearchParameters* params) const {
     FAISS_THROW_IF_NOT(k > 0);
-    FAISS_THROW_IF_NOT(batch_size >= k);
+    FAISS_THROW_IF_NOT(batch_size >= static_cast<size_t>(k));
 
-    HeapBlockResultHandler<CMax<float, int64_t>, false> handler(
-            size_t(n), distances, labels, size_t(k), nullptr);
-
-    flat_pano_search_core<false>(*this, handler, n, x, 0.0f, params);
+    dispatch_metric_compare(metric_type, [&]<typename C>() {
+        HeapBlockResultHandler<C, false> handler(
+                size_t(n), distances, labels, size_t(k), nullptr);
+        flat_pano_search_core<false, C>(*this, handler, n, x, 0.0f, params);
+    });
 }
 
 void IndexFlatPanorama::range_search(
@@ -698,10 +728,11 @@ void IndexFlatPanorama::range_search(
         float radius,
         RangeSearchResult* result,
         const SearchParameters* params) const {
-    RangeSearchBlockResultHandler<CMax<float, int64_t>, false> handler(
-            result, radius, nullptr);
-
-    flat_pano_search_core<true>(*this, handler, n, x, radius, params);
+    dispatch_metric_compare(metric_type, [&]<typename C>() {
+        RangeSearchBlockResultHandler<C, false> handler(
+                result, radius, nullptr);
+        flat_pano_search_core<true, C>(*this, handler, n, x, radius, params);
+    });
 }
 
 void IndexFlatPanorama::reset() {
@@ -740,7 +771,7 @@ size_t IndexFlatPanorama::remove_ids(const IDSelector& sel) {
         ntotal = j;
         size_t num_batches = (ntotal + batch_size - 1) / batch_size;
         codes.resize(num_batches * batch_size * code_size);
-        cum_sums.resize(num_batches * batch_size * (n_levels + 1));
+        cum_sums.resize(num_batches * batch_size * (pano.n_levels + 1));
     }
     return nremove;
 }
@@ -790,103 +821,136 @@ void IndexFlatPanorama::search_subset(
         idx_t k,
         float* distances,
         idx_t* labels) const {
-    using SingleResultHandler =
-            HeapBlockResultHandler<CMax<float, int64_t>, false>::
-                    SingleResultHandler;
-    HeapBlockResultHandler<CMax<float, int64_t>, false> handler(
-            size_t(n), distances, labels, size_t(k), nullptr);
-
-    FAISS_THROW_IF_NOT(k > 0);
-    FAISS_THROW_IF_NOT(batch_size == 1);
-
-    [[maybe_unused]] int nt = std::min(int(n), omp_get_max_threads());
+    with_simd_level([&]<SIMDLevel SL>() {
+        with_metric_type(metric_type, [&]<MetricType M>() {
+            constexpr bool is_sim = is_similarity_metric(M);
+            using C = std::conditional_t<
+                    is_sim,
+                    CMin<float, int64_t>,
+                    CMax<float, int64_t>>;
+            using SingleResultHandler =
+                    typename HeapBlockResultHandler<C, false>::
+                            SingleResultHandler;
+            HeapBlockResultHandler<C, false> handler(
+                    size_t(n), distances, labels, size_t(k), nullptr);
+
+            FAISS_THROW_IF_NOT(k > 0);
+            FAISS_THROW_IF_NOT(batch_size == 1);
+
+            [[maybe_unused]] int nt = std::min(int(n), omp_get_max_threads());
 
 #pragma omp parallel num_threads(nt)
-    {
-        SingleResultHandler res(handler);
-
-        std::vector<float> query_cum_norms(n_levels + 1);
-
-        // Panorama's optimized point-wise refinement (Algorithm 2):
-        // Batch-wise Panorama, as implemented in Panorama.h, incurs overhead
-        // from maintaining active_indices and exact_distances. This optimized
-        // implementation has minimal overhead and is thus preferred for
-        // IndexRefine's use case.
-        // 1. Initialize exact distance as ||y||^2 + ||x||^2.
-        // 2. For each level, refine distance incrementally:
-        //    - Compute dot product for current level: exact_dist -= 2*<x,y>.
-        //    - Use Cauchy-Schwarz bound on remaining levels to get lower bound.
-        //    - If there are less than k points in the heap, add the point to
-        //    the heap.
-        //    - Else, prune if lower bound exceeds k-th best distance.
-        // 3. After all levels, update heap if the point survived.
+            {
+                SingleResultHandler res(handler);
+
+                std::vector<float> query_cum_norms(pano.n_levels + 1);
+
+                // Panorama's optimized point-wise refinement (Algorithm 2):
+                // Batch-wise Panorama, as implemented in Panorama.h, incurs
+                // overhead from maintaining active_indices and exact_distances.
+                // This optimized implementation has minimal overhead and is
+                // thus preferred for IndexRefine's use case.
+                // 1. Initialize exact distance as ||y||^2 + ||x||^2.
+                // 2. For each level, refine distance incrementally:
+                //    - Compute dot product for current level: exact_dist -=
+                //    2*<x,y>.
+                //    - Use Cauchy-Schwarz bound on remaining levels to get
+                //    lower bound.
+                //    - If there are less than k points in the heap, add the
+                //    point to the heap.
+                //    - Else, prune if lower bound exceeds k-th best distance.
+                // 3. After all levels, update heap if the point survived.
 #pragma omp for
-        for (idx_t i = 0; i < n; i++) {
-            const idx_t* __restrict idsi = base_labels + i * k_base;
-            const float* xi = x + i * d;
-
-            PanoramaStats local_stats;
-            local_stats.reset();
-
-            pano.compute_query_cum_sums(xi, query_cum_norms.data());
-            float query_cum_norm = query_cum_norms[0] * query_cum_norms[0];
-
-            res.begin(i);
-
-            for (size_t j = 0; j < k_base; j++) {
-                idx_t idx = idsi[j];
-
-                if (idx < 0) {
-                    continue;
-                }
-
-                size_t cum_sum_offset = (n_levels + 1) * idx;
-                float cum_sum = cum_sums[cum_sum_offset];
-                float exact_distance = cum_sum * cum_sum + query_cum_norm;
-                cum_sum_offset++;
-
-                const float* x_ptr = xi;
-                const float* p_ptr =
-                        reinterpret_cast<const float*>(codes.data()) + d * idx;
-
-                local_stats.total_dims += d;
-
-                bool pruned = false;
-                for (size_t level = 0; level < n_levels; level++) {
-                    local_stats.total_dims_scanned += pano.level_width_floats;
-
-                    // Refine distance
-                    size_t actual_level_width = std::min(
-                            pano.level_width_floats,
-                            d - level * pano.level_width_floats);
-                    float dot_product = fvec_inner_product(
-                            x_ptr, p_ptr, actual_level_width);
-                    exact_distance -= 2 * dot_product;
-
-                    float cum_sum = cum_sums[cum_sum_offset];
-                    float cauchy_schwarz_bound =
-                            2.0f * cum_sum * query_cum_norms[level + 1];
-                    float lower_bound = exact_distance - cauchy_schwarz_bound;
-
-                    // Prune using Cauchy-Schwarz bound
-                    if (lower_bound > res.heap_dis[0]) {
-                        pruned = true;
-                        break;
+                for (idx_t i = 0; i < n; i++) {
+                    const idx_t* __restrict idsi = base_labels + i * k_base;
+                    const float* xi = x + i * d;
+
+                    PanoramaStats local_stats;
+                    local_stats.reset();
+
+                    pano.compute_query_cum_sums(xi, query_cum_norms.data());
+                    float query_cum_norm =
+                            query_cum_norms[0] * query_cum_norms[0];
+
+                    res.begin(i);
+
+                    for (idx_t j = 0; j < k_base; j++) {
+                        idx_t idx = idsi[j];
+
+                        if (idx < 0) {
+                            continue;
+                        }
+
+                        size_t cum_sum_offset = (pano.n_levels + 1) * idx;
+                        float cum_sum = cum_sums[cum_sum_offset];
+                        float exact_distance = 0.0f;
+                        if constexpr (!is_sim) {
+                            exact_distance = cum_sum * cum_sum + query_cum_norm;
+                        }
+                        cum_sum_offset++;
+
+                        const float* x_ptr = xi;
+                        const float* p_ptr =
+                                reinterpret_cast<const float*>(codes.data()) +
+                                d * idx;
+
+                        local_stats.total_dims += d;
+
+                        bool pruned = false;
+                        for (size_t level = 0; level < pano.n_levels; level++) {
+                            local_stats.total_dims_scanned +=
+                                    pano.level_width_floats;
+
+                            // Refine distance
+                            size_t actual_level_width = std::min(
+                                    pano.level_width_floats,
+                                    d - level * pano.level_width_floats);
+                            float dot_product = fvec_inner_product<SL>(
+                                    x_ptr, p_ptr, actual_level_width);
+                            if constexpr (is_sim) {
+                                exact_distance += dot_product;
+                            } else {
+                                exact_distance -= 2 * dot_product;
+                            }
+
+                            float level_cum_sum = cum_sums[cum_sum_offset];
+                            float cauchy_schwarz_bound;
+                            if constexpr (is_sim) {
+                                cauchy_schwarz_bound = -level_cum_sum *
+                                        query_cum_norms[level + 1];
+                            } else {
+                                cauchy_schwarz_bound = 2.0f * level_cum_sum *
+                                        query_cum_norms[level + 1];
+                            }
+                            float bound = exact_distance - cauchy_schwarz_bound;
+
+                            // Prune using Cauchy-Schwarz bound
+                            bool should_prune = false;
+                            if constexpr (is_sim) {
+                                should_prune = bound < res.heap_dis[0];
+                            } else {
+                                should_prune = bound > res.heap_dis[0];
+                            }
+                            if (should_prune) {
+                                pruned = true;
+                                break;
+                            }
+
+                            cum_sum_offset++;
+                            x_ptr += pano.level_width_floats;
+                            p_ptr += pano.level_width_floats;
+                        }
+
+                        if (!pruned) {
+                            res.add_result(exact_distance, idx);
+                        }
                     }
 
-                    cum_sum_offset++;
-                    x_ptr += pano.level_width_floats;
-                    p_ptr += pano.level_width_floats;
-                }
-
-                if (!pruned) {
-                    res.add_result(exact_distance, idx);
+                    res.end();
+                    indexPanorama_stats.add(local_stats);
                 }
             }
-
-            res.end();
-            indexPanorama_stats.add(local_stats);
-        }
-    }
+        });
+    });
 }
 } // namespace faiss