diff --git a/benchs/bench_flat_l2_panorama.py b/benchs/bench_flat_l2_panorama.py new file mode 100644 index 000000000..660109ba6 --- /dev/null +++ b/benchs/bench_flat_l2_panorama.py @@ -0,0 +1,99 @@ +# 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 DatasetGIST1M +except ImportError: + from faiss.contrib.datasets import DatasetGIST1M + +ds = DatasetGIST1M() + +nq = 10 +xq = ds.get_queries()[:nq] +xb = ds.get_database() +gt = ds.get_groundtruth()[:nq] + +xt = ds.get_train() + +nb, d = xb.shape +nt, d = xt.shape + +k = 10 +gt = gt[:, :k] + + +def eval_qps(index): + faiss.cvar.indexPanorama_stats.reset() + t0 = time.time() + _, I = index.search(xq, k=k) + t = time.time() - t0 + speed = t * 1000 / nq # ms/query + qps = 1000 / speed + + corrects = (gt == I).sum() + recall = corrects / (nq * k) + ratio_dims_scanned = faiss.cvar.indexPanorama_stats.ratio_dims_scanned + print( + f"\tRecall@{k}: {recall:.6f}, speed: {speed:.6f} ms/query, " + f"dims scanned: {ratio_dims_scanned * 100:.2f}%" + ) + return recall, qps + + +def build_index(name): + 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) + return index + + +nlevels = 8 +batch_size = 512 + +plt.figure(figsize=(8, 6), dpi=80) + +names = [ + "Flat", + f"PCA{d},FlatL2Panorama{nlevels}_{batch_size}", +] + +labels = [] +qps_values = [] + +for name in names: + print(f"======{name}") + index = build_index(name) + recall, qps = eval_qps(index) + labels.append(f"{name}\n(r@{recall:.3f})") + qps_values.append(qps) + +x = np.arange(len(names)) +plt.bar(x, qps_values, color=['#1f77b4', '#ff7f0e']) +speedup = qps_values[1] / qps_values[0] +ax = plt.gca() +ax.text( + x[1], + qps_values[1] * 1.01, + f"{speedup:.2f}x", + ha="center", + va="bottom", +) +plt.xticks(x, labels, rotation=0) +plt.ylabel("QPS") +plt.title("Flat Indexes on GIST1M") + +plt.tight_layout() +plt.savefig("bench_flat_l2_panorama.png", bbox_inches="tight") diff --git a/benchs/bench_ivf_flat_panorama.py b/benchs/bench_ivf_flat_panorama.py index 17d8fee2c..a7fab3c03 100644 --- a/benchs/bench_ivf_flat_panorama.py +++ b/benchs/bench_ivf_flat_panorama.py @@ -92,7 +92,7 @@ eval_and_plot(f"IVF{nlist},Flat") # IVFFlatPanorama (with PCA transform to concentrate energy in early dimensions) eval_and_plot(f"PCA{d},IVF{nlist},FlatPanorama{nlevels}") -plt.title("Indexes on GIST1M") +plt.title("IVF Flat Indexes on GIST1M") plt.xlabel(f"Recall@{k}") plt.ylabel("QPS") plt.yscale("log") diff --git a/faiss/CMakeLists.txt b/faiss/CMakeLists.txt index ce22c223e..97d063e29 100644 --- a/faiss/CMakeLists.txt +++ b/faiss/CMakeLists.txt @@ -83,6 +83,7 @@ set(FAISS_SRC impl/residual_quantizer_encode_steps.cpp impl/zerocopy_io.cpp impl/NNDescent.cpp + impl/Panorama.cpp impl/PanoramaStats.cpp invlists/BlockInvertedLists.cpp invlists/DirectMap.cpp @@ -175,6 +176,7 @@ set(FAISS_HEADERS impl/maybe_owned_vector.h impl/NNDescent.h impl/NSG.h + impl/Panorama.h impl/PanoramaStats.h impl/PolysemousTraining.h impl/ProductQuantizer-inl.h diff --git a/faiss/IndexFlat.cpp b/faiss/IndexFlat.cpp index fcb96a93d..8693ff806 100644 --- a/faiss/IndexFlat.cpp +++ b/faiss/IndexFlat.cpp @@ -11,12 +11,15 @@ #include #include +#include #include #include #include #include #include +#include #include +#include namespace faiss { @@ -519,4 +522,166 @@ void IndexFlat1D::search( done:; } } + +/************************************************************** + * shared flat Panorama search code + **************************************************************/ + +namespace { + +template +inline void flat_pano_search_core( + const IndexFlatPanorama& index, + BlockHandler& handler, + idx_t n, + const float* x, + float radius, + const SearchParameters* params) { + using SingleResultHandler = typename BlockHandler::SingleResultHandler; + + IDSelector* sel = params ? params->sel : nullptr; + bool use_sel = sel != nullptr; + + [[maybe_unused]] int nt = std::min(int(n), omp_get_max_threads()); + size_t n_batches = (index.ntotal + index.batch_size - 1) / index.batch_size; + +#pragma omp parallel num_threads(nt) + { + SingleResultHandler res(handler); + + std::vector query_cum_norms(index.n_levels + 1); + std::vector exact_distances(index.batch_size); + std::vector active_indices(index.batch_size); + +#pragma omp for + for (int64_t i = 0; i < n; i++) { + const float* xi = x + i * index.d; + index.pano.compute_query_cum_sums(xi, query_cum_norms.data()); + + PanoramaStats local_stats; + local_stats.reset(); + + res.begin(i); + + for (size_t batch_no = 0; batch_no < n_batches; batch_no++) { + size_t batch_start = batch_no * index.batch_size; + + float threshold; + if constexpr (use_radius) { + threshold = radius; + } else { + threshold = res.heap_dis[0]; + } + + size_t num_active = + index.pano + .progressive_filter_batch>( + 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); + + for (size_t j = 0; j < num_active; j++) { + res.add_result( + exact_distances[active_indices[j]], + batch_start + active_indices[j]); + } + } + + res.end(); + indexPanorama_stats.add(local_stats); + } + } +} + +} // anonymous namespace + +/*************************************************** + * IndexFlatPanorama + ***************************************************/ + +void IndexFlatPanorama::add(idx_t n, const float* x) { + size_t offset = ntotal; + ntotal += n; + 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)); + + const uint8_t* code = reinterpret_cast(x); + pano.copy_codes_to_level_layout(codes.data(), offset, n, code); + pano.compute_cumulative_sums(cum_sums.data(), offset, n, x); +} + +void IndexFlatPanorama::search( + idx_t n, + const float* x, + idx_t k, + float* distances, + idx_t* labels, + const SearchParameters* params) const { + FAISS_THROW_IF_NOT(k > 0); + FAISS_THROW_IF_NOT(batch_size >= k); + + HeapBlockResultHandler, false> handler( + size_t(n), distances, labels, size_t(k), nullptr); + + flat_pano_search_core(*this, handler, n, x, 0.0f, params); +} + +void IndexFlatPanorama::range_search( + idx_t n, + const float* x, + float radius, + RangeSearchResult* result, + const SearchParameters* params) const { + RangeSearchBlockResultHandler, false> handler( + result, radius, nullptr); + + flat_pano_search_core(*this, handler, n, x, radius, params); +} + +void IndexFlatPanorama::reset() { + IndexFlat::reset(); + cum_sums.clear(); +} + +void IndexFlatPanorama::reconstruct(idx_t key, float* recons) const { + pano.reconstruct(key, recons, codes.data()); +} + +void IndexFlatPanorama::reconstruct_n(idx_t i, idx_t n, float* recons) const { + Index::reconstruct_n(i, n, recons); +} + +size_t IndexFlatPanorama::remove_ids(const IDSelector& /* sel */) { + FAISS_THROW_MSG("remove_ids not implemented for IndexFlatPanorama"); + return 0; +} + +void IndexFlatPanorama::merge_from( + Index& /* otherIndex */, + idx_t /* add_id */) { + FAISS_THROW_MSG("merge_from not implemented for IndexFlatPanorama"); +} + +void IndexFlatPanorama::add_sa_codes( + idx_t /* n */, + const uint8_t* /* codes_in */, + const idx_t* /* xids */) { + FAISS_THROW_MSG("add_sa_codes not implemented for IndexFlatPanorama"); +} + +void IndexFlatPanorama::permute_entries(const idx_t* /* perm */) { + FAISS_THROW_MSG("permute_entries not implemented for IndexFlatPanorama"); +} } // namespace faiss diff --git a/faiss/IndexFlat.h b/faiss/IndexFlat.h index 75c9c95c4..bed7ea281 100644 --- a/faiss/IndexFlat.h +++ b/faiss/IndexFlat.h @@ -13,6 +13,7 @@ #include #include +#include namespace faiss { @@ -99,6 +100,76 @@ struct IndexFlatL2 : IndexFlat { void clear_l2norms(); }; +struct IndexFlatPanorama : IndexFlat { + const size_t batch_size; + const size_t n_levels; + std::vector cum_sums; + Panorama pano; + + /** + * @param d dimensionality of the input vectors + * @param metric metric type + * @param n_levels number of Panorama levels + * @param batch_size batch size for Panorama storage + */ + explicit IndexFlatPanorama( + idx_t d, + MetricType metric, + size_t n_levels, + size_t batch_size) + : IndexFlat(d, metric), + batch_size(batch_size), + n_levels(n_levels), + pano(code_size, n_levels, batch_size) { + FAISS_THROW_IF_NOT(metric == METRIC_L2); + } + + void add(idx_t n, const float* x) override; + + 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; + + void reset() override; + + void reconstruct(idx_t key, float* recons) const override; + + void reconstruct_n(idx_t i, idx_t n, float* recons) const override; + + size_t remove_ids(const IDSelector& sel) override; + + void merge_from(Index& otherIndex, idx_t add_id) override; + + void add_sa_codes(idx_t n, const uint8_t* codes_in, const idx_t* xids) + override; + + void permute_entries(const idx_t* perm); +}; + +struct IndexFlatL2Panorama : IndexFlatPanorama { + /** + * @param d dimensionality of the input vectors + * @param n_levels number of Panorama levels + * @param batch_size batch size for Panorama storage + */ + explicit IndexFlatL2Panorama( + idx_t d, + size_t n_levels, + size_t batch_size = 512) + : IndexFlatPanorama(d, METRIC_L2, n_levels, batch_size) {} +}; + /// optimized version for 1D "vectors". struct IndexFlat1D : IndexFlatL2 { bool continuous_update = true; ///< is the permutation updated continuously? diff --git a/faiss/IndexIVFFlatPanorama.cpp b/faiss/IndexIVFFlatPanorama.cpp index 386e24f2c..547ed3eb2 100644 --- a/faiss/IndexIVFFlatPanorama.cpp +++ b/faiss/IndexIVFFlatPanorama.cpp @@ -70,29 +70,8 @@ struct IVFFlatScannerPanorama : InvertedListScanner { float q_norm = 0.0f; void set_query(const float* query) override { this->xi = query; - - const size_t d = vd.d; - const size_t level_width_floats = storage->level_width / sizeof(float); - - std::vector suffix_sums(d + 1); - suffix_sums[d] = 0.0f; - - for (int j = d - 1; j >= 0; j--) { - float squared_val = query[j] * query[j]; - suffix_sums[j] = suffix_sums[j + 1] + squared_val; - } - - for (size_t level = 0; level < storage->n_levels; level++) { - size_t start_idx = level * level_width_floats; - if (start_idx < d) { - cum_sums[level] = sqrt(suffix_sums[start_idx]); - } else { - cum_sums[level] = 0.0f; - } - } - - cum_sums[storage->n_levels] = 0.0f; - q_norm = suffix_sums[0]; + this->storage->pano.compute_query_cum_sums(query, cum_sums.data()); + q_norm = cum_sums[0] * cum_sums[0]; } void set_list(idx_t list_no, float /* coarse_dis */) override { @@ -107,95 +86,6 @@ struct IVFFlatScannerPanorama : InvertedListScanner { "IndexIVFFlatPanorama does not support distance_to_code"); } - /// Helper function for progressive filtering that both scan_codes and - /// scan_codes_range use. Processes a batch of vectors through all levels, - /// computing exact distances and pruning based on a threshold. - /// Returns the number of active survivors after all levels. - size_t progressive_filter_batch( - size_t batch_no, - size_t list_size, - const uint8_t* codes_base, - const float* cum_sums_data, - float threshold, - std::vector& exact_distances, - std::vector& active_indices, - const idx_t* ids, - PanoramaStats& local_stats) const { - const size_t d = vd.d; - const size_t level_width_floats = storage->level_width / sizeof(float); - - size_t batch_start = batch_no * storage->kBatchSize; - size_t curr_batch_size = - std::min(list_size - batch_start, storage->kBatchSize); - - size_t cumsum_batch_offset = - batch_no * storage->kBatchSize * (storage->n_levels + 1); - const float* batch_cum_sums = cum_sums_data + cumsum_batch_offset; - - size_t batch_offset = batch_no * storage->kBatchSize * code_size; - const uint8_t* storage_base = codes_base + batch_offset; - - // Initialize active set with ID-filtered vectors. - size_t num_active = 0; - for (size_t i = 0; i < curr_batch_size; i++) { - size_t global_idx = batch_start + i; - bool include = !use_sel || sel->is_member(ids[global_idx]); - - active_indices[num_active] = i; - float cum_sum = batch_cum_sums[i]; - exact_distances[i] = cum_sum * cum_sum + q_norm; - - num_active += include; - } - - if (num_active == 0) { - return 0; - } - - size_t total_active = num_active; - - const float* level_cum_sums = batch_cum_sums + storage->kBatchSize; - - // Progressive filtering through levels. - for (size_t level = 0; level < storage->n_levels; level++) { - local_stats.total_dims_scanned += num_active; - local_stats.total_dims += total_active; - - float query_cum_norm = cum_sums[level + 1]; - - size_t level_offset = - level * storage->level_width * storage->kBatchSize; - const float* level_storage = - (const float*)(storage_base + level_offset); - - size_t next_active = 0; - for (size_t i = 0; i < num_active; i++) { - uint32_t idx = active_indices[i]; - const float* yj = level_storage + idx * level_width_floats; - const float* query_level = xi + level * level_width_floats; - - size_t actual_level_width = std::min( - level_width_floats, d - level * level_width_floats); - float dot_product = - fvec_inner_product(query_level, yj, actual_level_width); - - exact_distances[idx] -= 2.0f * dot_product; - - float cum_sum = level_cum_sums[idx]; - float cauchy_schwarz_bound = 2.0f * cum_sum * query_cum_norm; - float lower_bound = exact_distances[idx] - cauchy_schwarz_bound; - - active_indices[next_active] = idx; - next_active += C::cmp(threshold, lower_bound) ? 1 : 0; - } - - num_active = next_active; - level_cum_sums += storage->kBatchSize; - } - - return num_active; - } - size_t scan_codes( size_t list_size, const uint8_t* codes, @@ -208,7 +98,6 @@ struct IVFFlatScannerPanorama : InvertedListScanner { const size_t n_batches = (list_size + storage->kBatchSize - 1) / storage->kBatchSize; - const uint8_t* codes_base = codes; const float* cum_sums_data = storage->get_cum_sums(list_no); std::vector exact_distances(storage->kBatchSize); @@ -217,34 +106,25 @@ struct IVFFlatScannerPanorama : InvertedListScanner { PanoramaStats local_stats; local_stats.reset(); - // Panorama's IVFFlat core progressive filtering algorithm: - // Process vectors in batches for cache efficiency. For each batch: - // 1. Apply ID selection filter and initialize distances - // (||y||^2 + ||x||^2). - // 2. Maintain an "active set" of candidate indices that haven't been - // pruned yet. - // 3. For each level, refine distances incrementally and compact the - // active set: - // - Compute dot product for current level: exact_dist -= 2*. - // - Use Cauchy-Schwarz bound on remaining levels to get lower bound - // - Prune candidates whose lower bound exceeds k-th best distance. - // - Compact active_indices to remove pruned candidates (branchless) - // 4. After all levels, survivors are exact distances; update heap. - // This achieves early termination while maintaining SIMD-friendly - // sequential access patterns in the level-oriented storage layout. for (size_t batch_no = 0; batch_no < n_batches; batch_no++) { size_t batch_start = batch_no * storage->kBatchSize; - size_t num_active = progressive_filter_batch( - batch_no, - list_size, - codes_base, - cum_sums_data, - simi[0], - exact_distances, - active_indices, - ids, - local_stats); + size_t num_active = + storage->pano + .progressive_filter_batch>( + codes, + cum_sums_data, + xi, + cum_sums.data(), + batch_no, + list_size, + sel, + ids, + use_sel, + active_indices, + exact_distances, + simi[0], + local_stats); // Add batch survivors to heap. for (size_t i = 0; i < num_active; i++) { @@ -274,7 +154,6 @@ struct IVFFlatScannerPanorama : InvertedListScanner { const size_t n_batches = (list_size + storage->kBatchSize - 1) / storage->kBatchSize; - const uint8_t* codes_base = codes; const float* cum_sums_data = storage->get_cum_sums(list_no); std::vector exact_distances(storage->kBatchSize); @@ -288,16 +167,22 @@ struct IVFFlatScannerPanorama : InvertedListScanner { for (size_t batch_no = 0; batch_no < n_batches; batch_no++) { size_t batch_start = batch_no * storage->kBatchSize; - size_t num_active = progressive_filter_batch( - batch_no, - list_size, - codes_base, - cum_sums_data, - radius, - exact_distances, - active_indices, - ids, - local_stats); + size_t num_active = + storage->pano + .progressive_filter_batch>( + codes, + cum_sums_data, + xi, + cum_sums.data(), + batch_no, + list_size, + sel, + ids, + use_sel, + active_indices, + exact_distances, + radius, + local_stats); // Add batch survivors to range result. for (size_t i = 0; i < num_active; i++) { diff --git a/faiss/clone_index.cpp b/faiss/clone_index.cpp index a921e3783..062165381 100644 --- a/faiss/clone_index.cpp +++ b/faiss/clone_index.cpp @@ -276,6 +276,7 @@ Index* Cloner::clone_Index(const Index* index) { // IndexFlat TRYCLONE(IndexFlat1D, index) TRYCLONE(IndexFlatL2, index) + TRYCLONE(IndexFlatL2Panorama, index) TRYCLONE(IndexFlatIP, index) TRYCLONE(IndexFlat, index) diff --git a/faiss/impl/Panorama.cpp b/faiss/impl/Panorama.cpp new file mode 100644 index 000000000..c37332d0e --- /dev/null +++ b/faiss/impl/Panorama.cpp @@ -0,0 +1,150 @@ +/* + * 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 + +#include +#include +#include +#include + +namespace faiss { + +/************************************************************** + * Panorama structure implementation + **************************************************************/ + +Panorama::Panorama(size_t code_size, size_t n_levels, size_t batch_size) + : code_size(code_size), n_levels(n_levels), batch_size(batch_size) { + set_derived_values(); +} + +void Panorama::set_derived_values() { + this->d = code_size / sizeof(float); + this->level_width_floats = ((d + n_levels - 1) / n_levels); + this->level_width = this->level_width_floats * sizeof(float); +} + +/** + * @brief Copy codes to level-oriented layout + * @param codes The base pointer to codes + * @param offset Where to start writing new data (in number of vectors) + * @param n_entry The number of new vectors to write + * @param code The new vector data + */ +void Panorama::copy_codes_to_level_layout( + uint8_t* codes, + size_t offset, + size_t n_entry, + const uint8_t* code) { + for (size_t entry_idx = 0; entry_idx < n_entry; entry_idx++) { + size_t current_pos = offset + entry_idx; + + // Determine which batch we're in and position within that batch. + size_t batch_no = current_pos / batch_size; + size_t pos_in_batch = current_pos % batch_size; + + // Copy entry into level-oriented layout for this batch. + size_t batch_offset = batch_no * batch_size * code_size; + for (size_t level = 0; level < n_levels; level++) { + size_t level_offset = level * level_width * batch_size; + size_t start_byte = level * level_width; + size_t actual_level_width = + std::min(level_width, code_size - level * level_width); + + const uint8_t* src = code + entry_idx * code_size + start_byte; + uint8_t* dest = codes + batch_offset + level_offset + + pos_in_batch * actual_level_width; + + memcpy(dest, src, actual_level_width); + } + } +} + +void Panorama::compute_cumulative_sums( + float* cumsum_base, + size_t offset, + size_t n_entry, + const float* vectors) { + std::vector suffix_sums(d + 1); + + for (size_t entry_idx = 0; entry_idx < n_entry; entry_idx++) { + size_t current_pos = offset + entry_idx; + size_t batch_no = current_pos / batch_size; + size_t pos_in_batch = current_pos % batch_size; + + const float* vector = vectors + entry_idx * d; + + // Compute suffix sums of squared values. + suffix_sums[d] = 0.0f; + for (int j = d - 1; j >= 0; j--) { + float squared_val = vector[j] * vector[j]; + suffix_sums[j] = suffix_sums[j + 1] + squared_val; + } + + // Store cumulative sums in batch-oriented layout. + size_t cumsum_batch_offset = batch_no * batch_size * (n_levels + 1); + + for (size_t level = 0; level < n_levels; level++) { + size_t start_idx = level * level_width_floats; + size_t cumsum_offset = + cumsum_batch_offset + level * batch_size + pos_in_batch; + if (start_idx < d) { + cumsum_base[cumsum_offset] = std::sqrt(suffix_sums[start_idx]); + } else { + cumsum_base[cumsum_offset] = 0.0f; + } + } + + // Last level sum is always 0. + size_t cumsum_offset = + cumsum_batch_offset + n_levels * batch_size + pos_in_batch; + cumsum_base[cumsum_offset] = 0.0f; + } +} + +void Panorama::compute_query_cum_sums(const float* query, float* query_cum_sums) + const { + std::vector suffix_sums(d + 1); + suffix_sums[d] = 0.0f; + + for (int j = d - 1; j >= 0; j--) { + float squared_val = query[j] * query[j]; + suffix_sums[j] = suffix_sums[j + 1] + squared_val; + } + + for (size_t level = 0; level < n_levels; level++) { + size_t start_idx = level * level_width_floats; + if (start_idx < d) { + query_cum_sums[level] = std::sqrt(suffix_sums[start_idx]); + } else { + query_cum_sums[level] = 0.0f; + } + } + + query_cum_sums[n_levels] = 0.0f; +} + +void Panorama::reconstruct(idx_t key, float* recons, const uint8_t* codes_base) + const { + uint8_t* recons_buffer = reinterpret_cast(recons); + + size_t batch_no = key / batch_size; + size_t pos_in_batch = key % batch_size; + size_t batch_offset = batch_no * batch_size * code_size; + + for (size_t level = 0; level < n_levels; level++) { + size_t level_offset = level * level_width * batch_size; + const uint8_t* src = codes_base + batch_offset + level_offset + + pos_in_batch * level_width; + uint8_t* dest = recons_buffer + level * level_width; + size_t copy_size = + std::min(level_width, code_size - level * level_width); + memcpy(dest, src, copy_size); + } +} +} // namespace faiss diff --git a/faiss/impl/Panorama.h b/faiss/impl/Panorama.h new file mode 100644 index 000000000..a3a687eeb --- /dev/null +++ b/faiss/impl/Panorama.h @@ -0,0 +1,195 @@ +/* + * 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. + */ + +// -*- c++ -*- + +#ifndef FAISS_PANORAMA_H +#define FAISS_PANORAMA_H + +#include +#include +#include + +#include +#include +#include +#include + +namespace faiss { + +/** + * Implements the core logic of Panorama-based refinement. + * arXiv: https://arxiv.org/abs/2510.00566 + * + * Panorama partitions the dimensions of all vectors into L contiguous levels. + * During the refinement stage of ANNS, it computes distances between the query + * and its candidates level-by-level. After processing each level, it prunes the + * candidates whose lower bound exceeds the k-th best distance. + * + * In order to enable speedups, the dimensions (or codes) of each vector are + * stored in a batched, level-major manner. Within each batch of b vectors, the + * dimensions corresponding to level 1 will be stored first (for all elements in + * that batch), followed by level 2, and so on. This allows for efficient memory + * access patterns. + * + * Coupled with the appropriate orthogonal PreTransform (e.g. PCA, Cayley, + * etc.), Panorama can prune the vast majority of dimensions, greatly + * accelerating the refinement stage. + */ +struct Panorama { + size_t d = 0; + size_t code_size = 0; + size_t n_levels = 0; + size_t level_width = 0; + size_t level_width_floats = 0; + size_t batch_size = 0; + + explicit Panorama(size_t code_size, size_t n_levels, size_t batch_size); + + void set_derived_values(); + + /// Helper method to copy codes into level-oriented batch layout at a given + /// offset in the list. + void copy_codes_to_level_layout( + uint8_t* codes, + size_t offset, + size_t n_entry, + const uint8_t* code); + + /// Helper method to compute the cumulative sums of the codes. + /// The cumsums also follow the level-oriented batch layout to minimize the + /// number of random memory accesses. + void compute_cumulative_sums( + float* cumsum_base, + size_t offset, + size_t n_entry, + const float* vectors); + + /// Compute the cumulative sums of the query vector. + void compute_query_cum_sums(const float* query, float* query_cum_sums) + const; + + /// Panorama's core progressive filtering algorithm: + /// Process vectors in batches for cache efficiency. For each batch: + /// 1. Apply ID selection filter and initialize distances + /// (||y||^2 + ||x||^2). + /// 2. Maintain an "active set" of candidate indices that haven't been + /// pruned yet. + /// 3. For each level, refine distances incrementally and compact the active + /// set: + /// - Compute dot product for current level: exact_dist -= 2*. + /// - Use Cauchy-Schwarz bound on remaining levels to get lower bound + /// - Prune candidates whose lower bound exceeds k-th best distance. + /// - Compact active_indices to remove pruned candidates (branchless) + /// 4. After all levels, survivors are exact distances; update heap. + /// This achieves early termination while maintaining SIMD-friendly + /// sequential access patterns in the level-oriented storage layout. + template + size_t progressive_filter_batch( + const uint8_t* codes_base, + const float* cum_sums, + const float* query, + const float* query_cum_sums, + size_t batch_no, + size_t list_size, + const IDSelector* sel, + const idx_t* ids, + bool use_sel, + std::vector& active_indices, + std::vector& exact_distances, + float threshold, + PanoramaStats& local_stats) const; + + void reconstruct(idx_t key, float* recons, const uint8_t* codes_base) const; +}; + +template +size_t Panorama::progressive_filter_batch( + const uint8_t* codes_base, + const float* cum_sums, + const float* query, + const float* query_cum_sums, + size_t batch_no, + size_t list_size, + const IDSelector* sel, + const idx_t* ids, + bool use_sel, + std::vector& active_indices, + std::vector& exact_distances, + float threshold, + PanoramaStats& local_stats) const { + size_t batch_start = batch_no * batch_size; + size_t curr_batch_size = std::min(list_size - batch_start, batch_size); + + size_t cumsum_batch_offset = batch_no * batch_size * (n_levels + 1); + const float* batch_cum_sums = cum_sums + cumsum_batch_offset; + const float* level_cum_sums = batch_cum_sums + batch_size; + float q_norm = query_cum_sums[0] * query_cum_sums[0]; + + size_t batch_offset = batch_no * batch_size * code_size; + const uint8_t* storage_base = codes_base + batch_offset; + + // Initialize active set with ID-filtered vectors. + size_t num_active = 0; + for (size_t i = 0; i < curr_batch_size; i++) { + size_t global_idx = batch_start + i; + idx_t id = (ids == nullptr) ? global_idx : ids[global_idx]; + bool include = !use_sel || sel->is_member(id); + + active_indices[num_active] = i; + float cum_sum = batch_cum_sums[i]; + exact_distances[i] = cum_sum * cum_sum + q_norm; + + num_active += include; + } + + if (num_active == 0) { + return 0; + } + + size_t total_active = num_active; + for (size_t level = 0; level < n_levels; level++) { + local_stats.total_dims_scanned += num_active; + local_stats.total_dims += total_active; + + float query_cum_norm = query_cum_sums[level + 1]; + + size_t level_offset = level * level_width * batch_size; + const float* level_storage = + (const float*)(storage_base + level_offset); + + size_t next_active = 0; + for (size_t i = 0; i < num_active; i++) { + uint32_t idx = active_indices[i]; + size_t actual_level_width = std::min( + level_width_floats, d - level * level_width_floats); + + const float* yj = level_storage + idx * actual_level_width; + const float* query_level = query + level * level_width_floats; + + float dot_product = + fvec_inner_product(query_level, yj, actual_level_width); + + exact_distances[idx] -= 2.0f * dot_product; + + float cum_sum = level_cum_sums[idx]; + float cauchy_schwarz_bound = 2.0f * cum_sum * query_cum_norm; + float lower_bound = exact_distances[idx] - cauchy_schwarz_bound; + + active_indices[next_active] = idx; + next_active += C::cmp(threshold, lower_bound) ? 1 : 0; + } + + num_active = next_active; + level_cum_sums += batch_size; + } + + return num_active; +} +} // namespace faiss + +#endif diff --git a/faiss/impl/index_read.cpp b/faiss/impl/index_read.cpp index fb4fd6663..0dcc8571a 100644 --- a/faiss/impl/index_read.cpp +++ b/faiss/impl/index_read.cpp @@ -713,6 +713,20 @@ Index* read_index(IOReader* f, int io_flags) { if (h == fourcc("null")) { // denotes a missing index, useful for some cases return nullptr; + } else if (h == fourcc("IxFP")) { + int d; + size_t n_levels, batch_size; + READ1(d); + READ1(n_levels); + READ1(batch_size); + IndexFlatL2Panorama* idxp = + new IndexFlatL2Panorama(d, n_levels, batch_size); + READ1(idxp->ntotal); + READ1(idxp->is_trained); + READVECTOR(idxp->codes); + READVECTOR(idxp->cum_sums); + idxp->verbose = false; + idx = idxp; } else if ( h == fourcc("IxFI") || h == fourcc("IxF2") || h == fourcc("IxFl")) { IndexFlat* idxf; diff --git a/faiss/impl/index_write.cpp b/faiss/impl/index_write.cpp index 0aeae1a0f..e5959d081 100644 --- a/faiss/impl/index_write.cpp +++ b/faiss/impl/index_write.cpp @@ -432,6 +432,18 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) { // eg. for a storage component of HNSW that is set to nullptr uint32_t h = fourcc("null"); WRITE1(h); + } else if ( + const IndexFlatL2Panorama* idxpan = + dynamic_cast(idx)) { + uint32_t h = fourcc("IxFP"); + WRITE1(h); + WRITE1(idxpan->d); + WRITE1(idxpan->n_levels); + WRITE1(idxpan->batch_size); + WRITE1(idxpan->ntotal); + WRITE1(idxpan->is_trained); + WRITEVECTOR(idxpan->codes); + WRITEVECTOR(idxpan->cum_sums); } else if (const IndexFlat* idxf = dynamic_cast(idx)) { uint32_t h = fourcc(idxf->metric_type == METRIC_INNER_PRODUCT ? "IxFI" diff --git a/faiss/index_factory.cpp b/faiss/index_factory.cpp index afbc9141d..8eb1bd9d8 100644 --- a/faiss/index_factory.cpp +++ b/faiss/index_factory.cpp @@ -569,6 +569,18 @@ Index* parse_other_indexes( return new IndexFlat(d, metric); } + // IndexFlatL2Panorama + if (match("FlatL2Panorama([0-9]+)(_[0-9]+)?")) { + FAISS_THROW_IF_NOT(metric == METRIC_L2); + int nlevels = std::stoi(sm[1].str()); + if (sm[2].length() > 0) { + int batch_size = std::stoi(sm[2].str().substr(1)); + return new IndexFlatL2Panorama(d, nlevels, (size_t)batch_size); + } else { + return new IndexFlatL2Panorama(d, nlevels); + } + } + // IndexLSH if (match("LSH([0-9]*)(r?)(t?)")) { int nbits = sm[1].length() > 0 ? std::stoi(sm[1].str()) : d; diff --git a/faiss/invlists/InvertedLists.cpp b/faiss/invlists/InvertedLists.cpp index 99ac67b90..e079b6c42 100644 --- a/faiss/invlists/InvertedLists.cpp +++ b/faiss/invlists/InvertedLists.cpp @@ -358,7 +358,8 @@ ArrayInvertedListsPanorama::ArrayInvertedListsPanorama( n_levels(n_levels), level_width( (((code_size / sizeof(float)) + n_levels - 1) / n_levels) * - sizeof(float)) { + sizeof(float)), + pano(code_size, n_levels, kBatchSize) { FAISS_THROW_IF_NOT(n_levels > 0); FAISS_THROW_IF_NOT(code_size % sizeof(float) == 0); FAISS_THROW_IF_NOT_MSG( @@ -390,8 +391,11 @@ size_t ArrayInvertedListsPanorama::add_entries( codes[list_no].resize(num_batches * kBatchSize * code_size); cum_sums[list_no].resize(num_batches * kBatchSize * (n_levels + 1)); - copy_codes_to_level_layout(list_no, o, n_entry, code); - compute_cumulative_sums(list_no, o, n_entry, code); + // Cast to float* is safe here as we guarantee codes are always float + // vectors for `IndexIVFFlatPanorama` (verified by the constructor). + const float* vectors = reinterpret_cast(code); + pano.copy_codes_to_level_layout(codes[list_no].data(), o, n_entry, code); + pano.compute_cumulative_sums(cum_sums[list_no].data(), o, n_entry, vectors); return o; } @@ -406,8 +410,14 @@ void ArrayInvertedListsPanorama::update_entries( assert(n_entry + offset <= ids[list_no].size()); memcpy(&ids[list_no][offset], ids_in, sizeof(ids_in[0]) * n_entry); - copy_codes_to_level_layout(list_no, offset, n_entry, code); - compute_cumulative_sums(list_no, offset, n_entry, code); + + // Cast to float* is safe here as we guarantee codes are always float + // vectors for `IndexIVFFlatPanorama` (verified by the constructor). + const float* vectors = reinterpret_cast(code); + pano.copy_codes_to_level_layout( + codes[list_no].data(), offset, n_entry, code); + pano.compute_cumulative_sums( + cum_sums[list_no].data(), offset, n_entry, vectors); } void ArrayInvertedListsPanorama::resize(size_t list_no, size_t new_size) { @@ -426,21 +436,8 @@ const uint8_t* ArrayInvertedListsPanorama::get_single_code( uint8_t* recons_buffer = new uint8_t[code_size]; - const uint8_t* codes_base = codes[list_no].data(); - - size_t batch_no = offset / kBatchSize; - size_t pos_in_batch = offset % kBatchSize; - size_t batch_offset = batch_no * kBatchSize * code_size; - - for (size_t level = 0; level < n_levels; level++) { - size_t level_offset = level * level_width * kBatchSize; - const uint8_t* src = codes_base + batch_offset + level_offset + - pos_in_batch * level_width; - uint8_t* dest = recons_buffer + level * level_width; - size_t copy_size = - std::min(level_width, code_size - level * level_width); - memcpy(dest, src, copy_size); - } + float* recons = reinterpret_cast(recons_buffer); + pano.reconstruct(offset, recons, codes[list_no].data()); return recons_buffer; } @@ -463,94 +460,6 @@ InvertedListsIterator* ArrayInvertedListsPanorama::get_iterator( return nullptr; } -void ArrayInvertedListsPanorama::compute_cumulative_sums( - size_t list_no, - size_t offset, - size_t n_entry, - const uint8_t* code) { - // Cast to float* is safe here as we guarantee codes are always float - // vectors for `IndexIVFFlatPanorama` (verified by the constructor). - const float* vectors = reinterpret_cast(code); - const size_t d = code_size / sizeof(float); - - std::vector suffix_sums(d + 1); - - for (size_t entry_idx = 0; entry_idx < n_entry; entry_idx++) { - size_t current_pos = offset + entry_idx; - size_t batch_no = current_pos / kBatchSize; - size_t pos_in_batch = current_pos % kBatchSize; - - const float* vector = vectors + entry_idx * d; - - // Compute suffix sums of squared values. - suffix_sums[d] = 0.0f; - for (int j = d - 1; j >= 0; j--) { - float squared_val = vector[j] * vector[j]; - suffix_sums[j] = suffix_sums[j + 1] + squared_val; - } - - // Store cumulative sums in batch-oriented layout. - size_t cumsum_batch_offset = batch_no * kBatchSize * (n_levels + 1); - float* cumsum_base = cum_sums[list_no].data(); - - const size_t level_width_floats = level_width / sizeof(float); - for (size_t level = 0; level < n_levels; level++) { - size_t start_idx = level * level_width_floats; - size_t cumsum_offset = - cumsum_batch_offset + level * kBatchSize + pos_in_batch; - if (start_idx < d) { - cumsum_base[cumsum_offset] = sqrt(suffix_sums[start_idx]); - } else { - cumsum_base[cumsum_offset] = 0.0f; - } - } - - // Last level sum is always 0. - size_t cumsum_offset = - cumsum_batch_offset + n_levels * kBatchSize + pos_in_batch; - cumsum_base[cumsum_offset] = 0.0f; - } -} - -// Helper method to copy codes into level-oriented batch layout at a given -// offset in the list. -void ArrayInvertedListsPanorama::copy_codes_to_level_layout( - size_t list_no, - size_t offset, - size_t n_entry, - const uint8_t* code) { - uint8_t* codes_base = codes[list_no].data(); - size_t current_pos = offset; - for (size_t entry_idx = 0; entry_idx < n_entry;) { - // Determine which batch we're in and position within that batch. - size_t batch_no = current_pos / kBatchSize; - size_t pos_in_batch = current_pos % kBatchSize; - size_t entries_in_this_batch = - std::min(n_entry - entry_idx, kBatchSize - pos_in_batch); - - // Copy entries into level-oriented layout for this batch. - size_t batch_offset = batch_no * kBatchSize * code_size; - for (size_t level = 0; level < n_levels; level++) { - size_t level_offset = level * level_width * kBatchSize; - size_t start_byte = level * level_width; - size_t copy_size = - std::min(level_width, code_size - level * level_width); - - for (size_t i = 0; i < entries_in_this_batch; i++) { - const uint8_t* src = - code + (entry_idx + i) * code_size + start_byte; - uint8_t* dest = codes_base + batch_offset + level_offset + - (pos_in_batch + i) * level_width; - - memcpy(dest, src, copy_size); - } - } - - entry_idx += entries_in_this_batch; - current_pos += entries_in_this_batch; - } -} - /***************************************************************** * Meta-inverted list implementations *****************************************************************/ diff --git a/faiss/invlists/InvertedLists.h b/faiss/invlists/InvertedLists.h index 019213c88..43c1ecc0c 100644 --- a/faiss/invlists/InvertedLists.h +++ b/faiss/invlists/InvertedLists.h @@ -18,6 +18,7 @@ #include #include +#include #include namespace faiss { @@ -283,6 +284,7 @@ struct ArrayInvertedListsPanorama : ArrayInvertedLists { std::vector> cum_sums; const size_t n_levels; const size_t level_width; // in code units + Panorama pano; ArrayInvertedListsPanorama(size_t nlist, size_t code_size, size_t n_levels); @@ -318,24 +320,6 @@ struct ArrayInvertedListsPanorama : ArrayInvertedLists { /// Frees codes returned by `get_single_code`. void release_codes(size_t list_no, const uint8_t* codes) const override; - - private: - /// Helper method to copy codes into level-oriented batch layout at a given - /// offset in the list. - void copy_codes_to_level_layout( - size_t list_no, - size_t offset, - size_t n_entry, - const uint8_t* code); - - /// Helper method to compute the cumulative sums of the codes. - /// The cumsums also follow the level-oriented batch layout to minimize the - /// number of random memory accesses. - void compute_cumulative_sums( - size_t list_no, - size_t offset, - size_t n_entry, - const uint8_t* code); }; /***************************************************************** diff --git a/faiss/python/swigfaiss.swig b/faiss/python/swigfaiss.swig index d780e4247..2ec52266b 100644 --- a/faiss/python/swigfaiss.swig +++ b/faiss/python/swigfaiss.swig @@ -171,6 +171,7 @@ typedef uint64_t size_t; #include #include #include +#include #include #include @@ -490,6 +491,9 @@ void gpu_sync_all_devices() %newobject *::get_FlatCodesDistanceComputer() const; %include + +%ignore faiss::IndexFlatPanorama::pano; + %include %include @@ -505,6 +509,7 @@ void gpu_sync_all_devices() %include %include %include +%include %include %include @@ -740,6 +745,7 @@ struct faiss::simd16uint16 {}; DOWNCAST ( IndexIVF ) DOWNCAST ( IndexFlatIP ) DOWNCAST ( IndexFlatL2 ) + DOWNCAST ( IndexFlatL2Panorama ) DOWNCAST ( IndexFlat ) DOWNCAST ( IndexRefineFlat ) DOWNCAST ( IndexRefine ) diff --git a/tests/test_flat_l2_panorama.py b/tests/test_flat_l2_panorama.py new file mode 100644 index 000000000..b0ea82683 --- /dev/null +++ b/tests/test_flat_l2_panorama.py @@ -0,0 +1,519 @@ +""" +Comprehensive test suite for IndexFlatL2Panorama. + +Panorama is an adaptation of IndexFlatL2 that uses level-oriented storage +and progressive filtering with Cauchy-Schwarz bounds to achieve significant +speedups (up to 45x when combined with PCA or Cayley transforms) with zero +loss in accuracy. + +Paper: https://www.arxiv.org/pdf/2510.00566 +""" + +import unittest +import tempfile +import os + +import faiss +import numpy as np +from faiss.contrib.datasets import SyntheticDataset + + +class TestIndexFlatL2Panorama(unittest.TestCase): + """Test Suite for IndexFlatL2Panorama.""" + + # Helper methods for index creation and data generation + + def generate_data(self, d, nt, nb, nq, seed=42): + ds = SyntheticDataset(d, nt, nb, nq, seed=seed) + return ds.get_train(), ds.get_database(), ds.get_queries() + + def create_flat(self, d, xb=None): + index = faiss.IndexFlatL2(d) + if xb is not None: + index.add(xb) + return index + + def create_panorama(self, d, nlevels, xb=None, batch_size=128): + """Create and initialize IndexFlatL2Panorama.""" + index = faiss.IndexFlatL2Panorama(d, nlevels, batch_size) + if xb is not None: + index.add(xb) + return index + + def assert_search_results_equal( + self, + D_regular, + I_regular, + D_panorama, + I_panorama, + rtol=1e-5, + atol=1e-7, + otol=1e-3, + ): + # Allow small tolerance in overlap rate to account for floating-point errors + # in distance computations that can affect ordering when distances are nearly equal. + # Faiss: (a - b) * (a - b) vs. Panorama: a * a + b * b - 2(a * b) + overlap_rate = np.mean(I_regular == I_panorama) + + self.assertGreater( + overlap_rate, + 1 - otol, + f"Overlap rate {overlap_rate:.6f} is not > {1-otol:.3f}. ", + ) + np.testing.assert_allclose( + D_regular, D_panorama, rtol=rtol, atol=atol, err_msg="Distances mismatch" + ) + + def assert_range_results_equal( + self, + lims_regular, + D_regular, + I_regular, + lims_panorama, + D_panorama, + I_panorama, + nq, + ): + np.testing.assert_array_equal( + lims_regular, lims_panorama, err_msg="Different result counts" + ) + + for i in range(nq): + n_results = lims_regular[i + 1] - lims_regular[i] + if n_results > 0: + ids_reg = I_regular[lims_regular[i] : lims_regular[i + 1]] + dist_reg = D_regular[lims_regular[i] : lims_regular[i + 1]] + ids_pan = I_panorama[lims_panorama[i] : lims_panorama[i + 1]] + dist_pan = D_panorama[lims_panorama[i] : lims_panorama[i + 1]] + + sort_reg, sort_pan = np.argsort(ids_reg), np.argsort(ids_pan) + np.testing.assert_array_equal(ids_reg[sort_reg], ids_pan[sort_pan]) + np.testing.assert_allclose( + dist_reg[sort_reg], dist_pan[sort_pan], rtol=1e-5 + ) + + # Core functionality tests + + def test_exact_match_with_flat(self): + """Core test: Panorama must return identical results to IndexFlatL2""" + d, nb, nt, nq, nlevels, k = 128, 50000, 60000, 10, 8, 20 + _, xb, xq = self.generate_data(d, nt, nb, nq, seed=42) + + index_regular = self.create_flat(d, xb) + index_panorama = self.create_panorama(d, nlevels, xb) + + D_regular, I_regular = index_regular.search(xq, k) + D_panorama, I_panorama = index_panorama.search(xq, k) + + self.assert_search_results_equal( + D_regular, I_regular, D_panorama, I_panorama + ) + + def test_exact_match_with_flat_medium(self): + """Core test: Medium scale version""" + d, nb, nt, nq, nlevels, k = 64, 10000, 15000, 10, 4, 10 + _, xb, xq = self.generate_data(d, nt, nb, nq, seed=42) + + index_regular = self.create_flat(d, xb) + index_panorama = self.create_panorama(d, nlevels, xb) + + D_regular, I_regular = index_regular.search(xq, k) + D_panorama, I_panorama = index_panorama.search(xq, k) + + self.assert_search_results_equal( + D_regular, I_regular, D_panorama, I_panorama + ) + + def test_range_search(self): + """Test range search returns correct results within radius""" + d, nb, nt, nq, nlevels = 128, 50000, 75000, 10, 8 + _, xb, xq = self.generate_data(d, nt, nb, nq, seed=123) + + index_regular = self.create_flat(d, xb) + index_panorama = self.create_panorama(d, nlevels, xb) + + for radius in [0.5, 1.0, 2.0, 5.0]: + with self.subTest(radius=radius): + lims_reg, D_reg, I_reg = index_regular.range_search(xq, radius) + lims_pan, D_pan, I_pan = index_panorama.range_search(xq, radius) + + self.assertTrue( + np.all(D_pan <= radius), f"Some distances exceed radius={radius}" + ) + self.assert_range_results_equal( + lims_reg, D_reg, I_reg, lims_pan, D_pan, I_pan, nq + ) + + # Parameter variation tests + + def test_different_n_levels(self): + """Test correctness with various n_levels parameter values""" + d, nb, nt, nq, k = 128, 25000, 40000, 10, 15 + _, xb, xq = self.generate_data(d, nt, nb, nq, seed=456) + + index_base = self.create_flat(d, xb) + D_base, I_base = index_base.search(xq, k) + + nt = faiss.omp_get_max_threads() + faiss.omp_set_num_threads(1) + + prev_ratio_dims_scanned = float("inf") + for nlevels in [1, 2, 4, 8, 16, 32]: + with self.subTest(nlevels=nlevels): + faiss.cvar.indexPanorama_stats.reset() + index = self.create_panorama(d, nlevels, xb) + D, I = index.search(xq, k) + self.assert_search_results_equal(D_base, I_base, D, I) + + ratio_dims_scanned = faiss.cvar.indexPanorama_stats.ratio_dims_scanned + self.assertLess(ratio_dims_scanned, prev_ratio_dims_scanned) + prev_ratio_dims_scanned = ratio_dims_scanned + + faiss.omp_set_num_threads(nt) + + def test_uneven_dimension_division(self): + """Test when n_levels doesn't evenly divide dimension""" + test_cases = [(65, 4), (63, 8), (100, 7)] + + for d, nlevels in test_cases: + with self.subTest(d=d, nlevels=nlevels): + nb, nt, nq, k = 5000, 7000, 10, 5 + _, xb, xq = self.generate_data(d, nt, nb, nq, seed=789) + + index_regular = self.create_flat(d, xb) + index_panorama = self.create_panorama(d, nlevels, xb) + + D_regular, I_regular = index_regular.search(xq, k) + D_panorama, I_panorama = index_panorama.search(xq, k) + + self.assert_search_results_equal( + D_regular, I_regular, D_panorama, I_panorama + ) + + def test_single_level(self): + """Test edge case with n_levels=1""" + d, nb, nt, nq, nlevels, k = 32, 500, 700, 10, 1, 5 + _, xb, xq = self.generate_data(d, nt, nb, nq, seed=333) + + index_regular = self.create_flat(d, xb) + index_panorama = self.create_panorama(d, nlevels, xb) + + D_regular, I_regular = index_regular.search(xq, k) + D_panorama, I_panorama = index_panorama.search(xq, k) + + self.assert_search_results_equal(D_regular, I_regular, D_panorama, I_panorama) + + def test_multiple_levels_small_dimension(self): + """Test edge case: more levels than dimension naturally supports""" + d, nb, nt, nq, nlevels, k = 16, 200, 300, 10, 8, 5 + _, xb, xq = self.generate_data(d, nt, nb, nq, seed=1212) + + index_regular = self.create_flat(d, xb) + index_panorama = self.create_panorama(d, nlevels, xb) + + D_regular, I_regular = index_regular.search(xq, k) + D_panorama, I_panorama = index_panorama.search(xq, k) + + self.assert_search_results_equal(D_regular, I_regular, D_panorama, I_panorama) + + # ID selector tests (disabled) + + def test_id_selector_range(self): + """Test ID filtering with range selector""" + d, nb, nt, nq, nlevels, k = 128, 80000, 120000, 500, 8, 20 + _, xb, xq = self.generate_data(d, nt, nb, nq, seed=321) + + index_regular = self.create_flat(d, xb) + index_panorama = self.create_panorama(d, nlevels, xb) + + params = faiss.SearchParameters() + params.sel = faiss.IDSelectorRange(10000, 50000) + + D_regular, I_regular = index_regular.search(xq, k, params=params) + D_panorama, I_panorama = index_panorama.search(xq, k, params=params) + + valid_ids = I_panorama >= 0 + self.assertTrue(np.all(I_panorama[valid_ids] >= 10000)) + self.assertTrue(np.all(I_panorama[valid_ids] < 50000)) + + np.testing.assert_array_equal(I_regular, I_panorama) + np.testing.assert_allclose( + D_regular[valid_ids], D_panorama[valid_ids], rtol=1e-5 + ) + + def test_id_selector_batch(self): + """Test ID filtering with batch selector""" + d, nb, nt, nq, nlevels, k = 128, 60000, 90000, 400, 8, 20 + _, xb, xq = self.generate_data(d, nt, nb, nq, seed=522) + + index_regular = self.create_flat(d, xb) + index_panorama = self.create_panorama(d, nlevels, xb) + + allowed_ids = np.array([i * 100 for i in range(500)], dtype=np.int64) + params = faiss.SearchParameters() + params.sel = faiss.IDSelectorBatch(allowed_ids) + + D_regular, I_regular = index_regular.search(xq, k, params=params) + D_panorama, I_panorama = index_panorama.search(xq, k, params=params) + + allowed_set = set(allowed_ids) + valid_panorama = I_panorama >= 0 + for id_val in I_panorama[valid_panorama]: + self.assertIn(int(id_val), allowed_set) + + np.testing.assert_array_equal(I_regular, I_panorama) + np.testing.assert_allclose( + D_regular[valid_panorama], D_panorama[valid_panorama], rtol=1e-5 + ) + + def test_selector_with_small_dataset(self): + """Test ID selectors with dataset smaller than batch size""" + d, nb, nt, nq, nlevels, k = 32, 100, 200, 10, 4, 5 + _, xb, xq = self.generate_data(d, nt, nb, nq, seed=888) + + index_regular = self.create_flat(d, xb) + index_panorama = self.create_panorama(d, nlevels, xb) + + params = faiss.SearchParameters() + params.sel = faiss.IDSelectorRange(20, 60) + + D_regular, I_regular = index_regular.search(xq, k, params=params) + D_panorama, I_panorama = index_panorama.search(xq, k, params=params) + + valid_ids = I_panorama >= 0 + if np.any(valid_ids): + self.assertTrue(np.all(I_panorama[valid_ids] >= 20)) + self.assertTrue(np.all(I_panorama[valid_ids] < 60)) + + np.testing.assert_array_equal(I_regular, I_panorama) + np.testing.assert_allclose( + D_regular[valid_ids], D_panorama[valid_ids], rtol=1e-5 + ) + + def test_selector_excludes_all(self): + """Test selector that excludes all results""" + d, nb, nt, nq, nlevels, k = 32, 300, 400, 5, 4, 10 + _, xb, xq = self.generate_data(d, nt, nb, nq, seed=999) + + index_panorama = self.create_panorama(d, nlevels, xb) + + params = faiss.SearchParameters() + params.sel = faiss.IDSelectorRange(nb + 100, nb + 200) + + D, I = index_panorama.search(xq, k, params=params) + self.assertTrue(np.all(I == -1)) + + # Batch size and edge case tests + + def test_batch_boundaries(self): + """Test correctness at various batch size boundaries""" + d, nt, nq, k = 128, 10000, 10, 15 + # random train not needed for Flat indices + xq = np.random.rand(nq, d).astype("float32") + + # Use index's batch_size + probe_index = self.create_panorama(d, nlevels=8) + batch_size = getattr(probe_index, "batch_size") + + for nb in [batch_size - 1, batch_size, batch_size * 2, batch_size * 3, batch_size * 5 - 1]: + with self.subTest(nb=nb): + np.random.seed(987) + xb = np.random.rand(nb, d).astype("float32") + + index_regular = self.create_flat(d, xb) + index_panorama = self.create_panorama(d, nlevels=8, xb=xb) + + D_regular, I_regular = index_regular.search(xq, k) + D_panorama, I_panorama = index_panorama.search(xq, k) + + self.assert_search_results_equal( + D_regular, I_regular, D_panorama, I_panorama + ) + + def test_batch_size(self): + """Test correctness across different internal batch sizes""" + d, nb, nt, nq, nlevels = 128, 10000, 15000, 10, 8 + _, xb, xq = self.generate_data(d, nt, nb, nq, seed=4242) + + index_regular = self.create_flat(d, xb) + + for k, bs in [(1, 1), (10, 128), (100, 4096)]: + with self.subTest(batch_size=bs): + index_panorama = self.create_panorama(d, nlevels, xb=xb, batch_size=bs) + D_regular, I_regular = index_regular.search(xq, k) + D_panorama, I_panorama = index_panorama.search(xq, k) + self.assert_search_results_equal( + D_regular, I_regular, D_panorama, I_panorama + ) + + def test_empty_result_handling(self): + """Test handling of empty search results (shapes only)""" + d, nb, nt, nq, nlevels, k = 32, 100, 200, 10, 4, 10 + _, xb, _ = self.generate_data(d, nt, nb, nq, seed=111) + xq = np.random.rand(nq, d).astype("float32") + 10.0 # Queries far from database + + index = self.create_panorama(d, nlevels, xb) + D, I = index.search(xq, k) + + self.assertEqual(D.shape, (nq, k)) + self.assertEqual(I.shape, (nq, k)) + + def test_very_small_dataset(self): + """Test with dataset smaller than batch size (< 128 vectors)""" + test_cases = [10, 50, 100] + + for nb in test_cases: + with self.subTest(nb=nb): + d, nt, nlevels, nq = 32, max(nb, 100), 4, 10 + k = min(3, nb) + _, xb, xq = self.generate_data(d, nt, nb, nq, seed=666 + nb) + + index_regular = self.create_flat(d, xb) + index_panorama = self.create_panorama(d, nlevels, xb) + + D_regular, I_regular = index_regular.search(xq, k) + D_panorama, I_panorama = index_panorama.search(xq, k) + + self.assert_search_results_equal( + D_regular, I_regular, D_panorama, I_panorama + ) + + def test_range_search_edge_cases(self): + """Test range search with extreme radius values""" + d, nb, nt, nq, nlevels = 32, 500, 700, 10, 4 + xt, xb, xq = self.generate_data(d, nt, nb, nq, seed=777) + + index_regular = self.create_flat(d) + index_panorama = self.create_panorama(d, nlevels) + + for radius in [0.01, 0.1, 100.0, 1000.0]: + with self.subTest(radius=radius): + lims_reg, D_reg, I_reg = index_regular.range_search(xq, radius) + lims_pan, D_pan, I_pan = index_panorama.range_search(xq, radius) + + np.testing.assert_array_equal(lims_reg, lims_pan) + + if len(I_reg) > 0: + self.assert_range_results_equal( + lims_reg, D_reg, I_reg, lims_pan, D_pan, I_pan, nq + ) + + # Dynamic operations tests + + def test_incremental_add(self): + """Test adding vectors incrementally in multiple batches""" + d, nlevels, k = 128, 8, 15 + + index_regular = self.create_flat(d) + index_panorama = self.create_panorama(d, nlevels) + + # Keep total nb under 100k + batch_sizes = [5000, 10000, 15000, 20000] + for batch_size in batch_sizes: + xb_batch = np.random.rand(batch_size, d).astype("float32") + index_regular.add(xb_batch) + index_panorama.add(xb_batch) + + nq = 10 + xq = np.random.rand(nq, d).astype("float32") + + D_regular, I_regular = index_regular.search(xq, k) + D_panorama, I_panorama = index_panorama.search(xq, k) + + self.assert_search_results_equal(D_regular, I_regular, D_panorama, I_panorama) + + def test_add_search_add_search(self): + """Test interleaved add and search operations""" + d, nlevels, k = 32, 4, 5 + np.random.seed(555) + + index_regular = self.create_flat(d) + index_panorama = self.create_panorama(d, nlevels) + + # First add and search + xb1 = np.random.rand(200, d).astype("float32") + index_regular.add(xb1) + index_panorama.add(xb1) + + xq1 = np.random.rand(10, d).astype("float32") + + D_reg_1, I_reg_1 = index_regular.search(xq1, k) + D_pan_1, I_pan_1 = index_panorama.search(xq1, k) + self.assert_search_results_equal(D_reg_1, I_reg_1, D_pan_1, I_pan_1) + + # Second add and search + xb2 = np.random.rand(300, d).astype("float32") + index_regular.add(xb2) + index_panorama.add(xb2) + + xq2 = np.random.rand(10, d).astype("float32") + D_reg_2, I_reg_2 = index_regular.search(xq2, k) + D_pan_2, I_pan_2 = index_panorama.search(xq2, k) + self.assert_search_results_equal(D_reg_2, I_reg_2, D_pan_2, I_pan_2) + + def test_reconstruct(self): + """Test reconstruct and reconstruct_n return original vectors""" + d, nb, nt, nq, nlevels = 128, 10000, 15000, 10, 8 + _, xb, _ = self.generate_data(d, nt, nb, nq, seed=2025) + + index_panorama = self.create_panorama(d, nlevels, xb) + + # Test reconstruct for single vector + idx = 123 + v_panorama = index_panorama.reconstruct(idx) + np.testing.assert_array_equal(xb[idx], v_panorama) + + # Test reconstruct_n for range of vectors + start_idx, n_vectors = 120, 10 + vn_panorama = index_panorama.reconstruct_n(start_idx, n_vectors) + np.testing.assert_array_equal(xb[start_idx:start_idx + n_vectors], vn_panorama) + + def test_serialization(self): + """Test that writing and reading Panorama indexes preserves search results""" + d, nb, nt, nq, nlevels, k = 128, 10000, 15000, 100, 8, 20 + _, xb, xq = self.generate_data(d, nt, nb, nq, seed=2024) + + index = self.create_panorama(d, nlevels, xb) + + D_before, I_before = index.search(xq, k) + faiss.write_index(index, "index.bin") + index_after = faiss.read_index("index.bin") + D_after, I_after = index_after.search(xq, k) + os.unlink("index.bin") + + np.testing.assert_array_equal(I_before, I_after) + np.testing.assert_array_equal(D_before, D_after) + + def test_ratio_dims_scanned(self): + """Test the correctness of the ratio of dimensions scanned""" + d, nb, nq, k = 128, 500000, 10, 1 + + # Setup: All vectors in the dataset are [1, 1, ..., 1], except for + # one, which is [0, 0, ..., 0]. The queries are [0, 0, ..., 0]. This + # ensures that pruning happens as early as it can. As a result, the + # ratio of dimensions scanned should be approximately (1 / nlevels). + xb = np.ones((nb, d)).astype("float32") + xb[-1] = 0 + xq = np.zeros((nq, d)).astype("float32") + + index_base = self.create_flat(d, xb) + D_base, I_base = index_base.search(xq, k) + + nt = faiss.omp_get_max_threads() + faiss.omp_set_num_threads(1) + + ratios = [] + nlevels_list = [1, 2, 4, 8, 16, 32] + for nlevels in nlevels_list: + with self.subTest(nlevels=nlevels): + faiss.cvar.indexPanorama_stats.reset() + index = self.create_panorama(d, nlevels, xb) + D, I = index.search(xq, k) + self.assert_search_results_equal(D_base, I_base, D, I) + + ratios.append(faiss.cvar.indexPanorama_stats.ratio_dims_scanned) + + expected_ratios = [1 / nlevels for nlevels in nlevels_list] + np.testing.assert_allclose(ratios, expected_ratios, atol=1e-3) + + faiss.omp_set_num_threads(nt) diff --git a/tests/test_index.py b/tests/test_index.py index b3954e9d0..a9d992ed9 100644 --- a/tests/test_index.py +++ b/tests/test_index.py @@ -145,6 +145,52 @@ class TestIndexFlatL2(unittest.TestCase): np.testing.assert_equal(D3, D1) +class TestIndexFlatL2Panorama(unittest.TestCase): + def test_indexflat_l2_panorama(self): + d = 32 + nlevels = 8 + batch_size = 128 + nq = 200 + nb = 1000 + nt = 0 + k = 10 + + (xt, xb, xq) = get_dataset_2(d, nt, nb, nq) + index = faiss.IndexFlatL2Panorama(d, nlevels, batch_size) + + ### k-NN search + + index.add(xb) + D1, I1 = index.search(xq, k) + + all_dis = ((xq.reshape(nq, 1, d) - xb.reshape(1, nb, d)) ** 2).sum(2) + Iref = all_dis.argsort(axis=1)[:, :k] + + Dref = all_dis[np.arange(nq)[:, None], Iref] + + # not too many elements are off. + self.assertLessEqual((Iref != I1).sum(), Iref.size * 0.0002) + np.testing.assert_almost_equal(Dref, D1, decimal=5) + + ### Range search + + radius = float(np.median(Dref[:, -1])) + + lims, D2, I2 = index.range_search(xq, radius) + + for i in range(nq): + l0, l1 = lims[i:i + 2] + _, Il = D2[l0:l1], I2[l0:l1] + Ilref, = np.where(all_dis[i] < radius) + Il.sort() + Ilref.sort() + np.testing.assert_equal(Il, Ilref) + np.testing.assert_almost_equal( + all_dis[i, Ilref], D2[l0:l1], + decimal=5 + ) + + class EvalIVFPQAccuracy(unittest.TestCase): def test_IndexIVFPQ(self): diff --git a/tests/test_panorama.py b/tests/test_ivf_flat_panorama.py similarity index 96% rename from tests/test_panorama.py rename to tests/test_ivf_flat_panorama.py index f645c9dde..ade917223 100644 --- a/tests/test_panorama.py +++ b/tests/test_ivf_flat_panorama.py @@ -179,6 +179,9 @@ class TestIndexIVFFlatPanorama(unittest.TestCase): index_base = self.create_ivf_flat(d, nlist, xt, xb, nprobe=32) D_base, I_base = index_base.search(xq, k) + nt = faiss.omp_get_max_threads() + faiss.omp_set_num_threads(1) + prev_ratio_dims_scanned = float('inf') for nlevels in [1, 2, 4, 8, 16, 32]: with self.subTest(nlevels=nlevels): @@ -194,13 +197,17 @@ class TestIndexIVFFlatPanorama(unittest.TestCase): self.assertLess(ratio_dims_scanned, prev_ratio_dims_scanned) prev_ratio_dims_scanned = ratio_dims_scanned + faiss.omp_set_num_threads(nt) + def test_uneven_dimension_division(self): """Test when n_levels doesn't evenly divide dimension""" test_cases = [(65, 4), (63, 8), (100, 7)] + # TODO(aknayar): Test functions like get_single_code(). + for d, nlevels in test_cases: with self.subTest(d=d, nlevels=nlevels): - nb, nt, nq, nlist, k = 500, 700, 20, 16, 5 + nb, nt, nq, nlist, k = 5000, 700, 20, 16, 5 xt, xb, xq = self.generate_data(d, nt, nb, nq, seed=789) index_regular = self.create_ivf_flat(d, nlist, xt, xb, nprobe=4) @@ -525,24 +532,16 @@ class TestIndexIVFFlatPanorama(unittest.TestCase): d, nb, nt, nq, nlist, nlevels, k = 128, 10000, 15000, 100, 128, 8, 20 xt, xb, xq = self.generate_data(d, nt, nb, nq, seed=2024) - indexes = [self.create_panorama(d, nlist, nlevels, xt, xb, nprobe=32)] + index = self.create_panorama(d, nlist, nlevels, xt, xb, nprobe=32) - for index in indexes: - D_before, I_before = index.search(xq, k) + D_before, I_before = index.search(xq, k) + faiss.write_index(index, "index.bin") + index_after = faiss.read_index("index.bin") + D_after, I_after = index_after.search(xq, k) + os.unlink("index.bin") - fd, fname = tempfile.mkstemp() - os.close(fd) - try: - faiss.write_index(index, fname) - index_after = faiss.read_index(fname) - finally: - if os.path.exists(fname): - os.unlink(fname) - - D_after, I_after = index_after.search(xq, k) - - np.testing.assert_array_equal(I_before, I_after) - np.testing.assert_array_equal(D_before, D_after) + np.testing.assert_array_equal(I_before, I_after) + np.testing.assert_array_equal(D_before, D_after) def test_ratio_dims_scanned(self): """Test the correctness of the ratio of dimensions scanned""" @@ -560,6 +559,9 @@ class TestIndexIVFFlatPanorama(unittest.TestCase): index_base = self.create_ivf_flat(d, nlist, xt, xb, nprobe=1) D_base, I_base = index_base.search(xq, k) + nt = faiss.omp_get_max_threads() + faiss.omp_set_num_threads(1) + ratios = [] nlevels_list = [1, 2, 4, 8, 16, 32] for nlevels in nlevels_list: @@ -577,3 +579,5 @@ class TestIndexIVFFlatPanorama(unittest.TestCase): expected_ratios = [1 / nlevels for nlevels in nlevels_list] np.testing.assert_allclose(ratios, expected_ratios, atol=1e-3) + + faiss.omp_set_num_threads(nt)