diff --git a/faiss/CMakeLists.txt b/faiss/CMakeLists.txt index fcf26ab10..b007e3296 100644 --- a/faiss/CMakeLists.txt +++ b/faiss/CMakeLists.txt @@ -139,6 +139,8 @@ set(FAISS_SRC IndexRefine.cpp IndexReplicas.cpp IndexRowwiseMinMax.cpp + IndexSQFastScan.cpp + IndexIVFSQFastScan.cpp IndexScalarQuantizer.cpp IndexShards.cpp IndexShardsIVF.cpp @@ -272,6 +274,8 @@ set(FAISS_HEADERS IndexRaBitQ.h IndexRaBitQFastScan.h IndexRowwiseMinMax.h + IndexSQFastScan.h + IndexIVFSQFastScan.h IndexScalarQuantizer.h IndexShards.h IndexShardsIVF.h diff --git a/faiss/IndexIVFSQFastScan.cpp b/faiss/IndexIVFSQFastScan.cpp new file mode 100644 index 000000000..855fed48d --- /dev/null +++ b/faiss/IndexIVFSQFastScan.cpp @@ -0,0 +1,900 @@ +/* + * 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 +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace faiss { + +namespace { + +size_t roundup(size_t a, size_t b) { + return (a + b - 1) / b * b; +} + +bool is_native_4bit(ScalarQuantizer::QuantizerType qtype) { + return qtype == ScalarQuantizer::QT_4bit || + qtype == ScalarQuantizer::QT_4bit_uniform; +} + +bool needs_rerank(ScalarQuantizer::QuantizerType qtype) { + return qtype == ScalarQuantizer::QT_6bit || + qtype == ScalarQuantizer::QT_8bit || + qtype == ScalarQuantizer::QT_8bit_uniform || + qtype == ScalarQuantizer::QT_8bit_direct || + qtype == ScalarQuantizer::QT_8bit_direct_signed; +} + +bool is_fallback(ScalarQuantizer::QuantizerType qtype) { + return !is_native_4bit(qtype) && !needs_rerank(qtype); +} + +bool is_uniform_range(ScalarQuantizer::QuantizerType qtype) { + return qtype == ScalarQuantizer::QT_8bit_uniform || + qtype == ScalarQuantizer::QT_8bit_direct || + qtype == ScalarQuantizer::QT_8bit_direct_signed; +} + +void get_uniform_range(const ScalarQuantizer& sq, float& vmin, float& vdiff) { + if (sq.qtype == ScalarQuantizer::QT_8bit_direct) { + vmin = 0; + vdiff = 255; + } else if (sq.qtype == ScalarQuantizer::QT_8bit_direct_signed) { + vmin = -128; + vdiff = 255; + } else { + vmin = sq.trained[0]; + vdiff = sq.trained[1]; + } +} + +void float_to_4bit_nibbles( + const float* x, + uint8_t* nibbles, + idx_t n, + int d, + int M2, + const ScalarQuantizer& sq) { + const bool is_uniform = is_uniform_range(sq.qtype) || + sq.qtype == ScalarQuantizer::QT_4bit_uniform; + const float* vmin_arr = nullptr; + const float* vdiff_arr = nullptr; + float vmin_s = 0, inv_vdiff_s = 0; + + if (is_uniform) { + float vdiff; + get_uniform_range(sq, vmin_s, vdiff); + inv_vdiff_s = (vdiff > 0) ? (1.0f / vdiff) : 0; + } else { + vmin_arr = sq.trained.data(); + vdiff_arr = sq.trained.data() + d; + } + + const int half = M2 / 2; + for (idx_t i = 0; i < n; i++) { + const float* xi = x + i * d; + uint8_t* dst = nibbles + i * half; + memset(dst, 0, half); + for (int m = 0; m + 1 < d; m += 2) { + float f0, f1; + if (is_uniform) { + f0 = (xi[m] - vmin_s) * inv_vdiff_s; + f1 = (xi[m + 1] - vmin_s) * inv_vdiff_s; + } else { + f0 = (vdiff_arr[m] > 0) ? (xi[m] - vmin_arr[m]) / vdiff_arr[m] + : 0; + f1 = (vdiff_arr[m + 1] > 0) + ? (xi[m + 1] - vmin_arr[m + 1]) / vdiff_arr[m + 1] + : 0; + } + uint8_t lo = (uint8_t)std::min(15, std::max(0, (int)(f0 * 15.0f))); + uint8_t hi = (uint8_t)std::min(15, std::max(0, (int)(f1 * 15.0f))); + dst[m / 2] = lo | (hi << 4); + } + if (d & 1) { + float f; + if (is_uniform) { + f = (xi[d - 1] - vmin_s) * inv_vdiff_s; + } else { + f = (vdiff_arr[d - 1] > 0) + ? (xi[d - 1] - vmin_arr[d - 1]) / vdiff_arr[d - 1] + : 0; + } + dst[(d - 1) / 2] = + (uint8_t)std::min(15, std::max(0, (int)(f * 15.0f))); + } + } +} + +} // anonymous namespace + +// ----------------------------------------------------------------------- +// Constructors +// ----------------------------------------------------------------------- + +IndexIVFSQFastScan::IndexIVFSQFastScan( + Index* quantizer_in, + size_t d_in, + size_t nlist_in, + ScalarQuantizer::QuantizerType qtype, + MetricType metric, + int bbs_in, + bool by_residual_in) + : IndexIVFFastScan(quantizer_in, d_in, nlist_in, 0, metric, false), + sq(d_in, qtype) { + if (is_fallback(qtype)) { + by_residual = by_residual_in; + if (ScalarQuantizer::TurboQuantRefine::is_turboq_full(qtype)) { + by_residual = false; + } + } else { + // FastScan path: by_residual=false because our 2D LUT cannot + // handle per-probe residuals (same rationale as IndexIVFPQFastScan). + by_residual = false; + } + + if (is_fallback(qtype)) { + // Fallback: use standard ArrayInvertedLists with full SQ codes. + // Search uses IndexIVF's InvertedListScanner path. + code_size = sq.code_size; + replace_invlists(new ArrayInvertedLists(nlist_in, code_size), true); + } else { + size_t M_in = d_in; + M = M_in; + nbits = 4; + bbs = bbs_in; + ksub = 16; + M2 = roundup(M_in, 2); + code_size = M2 / 2; + + replace_invlists( + new BlockInvertedLists(nlist_in, get_CodePacker()), true); + + if (needs_rerank(qtype)) { + orig_codes_invlists = + new ArrayInvertedLists(nlist_in, sq.code_size); + direct_map.set_type(DirectMap::Hashtable, invlists, 0); + } + } +} + +IndexIVFSQFastScan::IndexIVFSQFastScan() = default; + +IndexIVFSQFastScan::IndexIVFSQFastScan( + const IndexIVFScalarQuantizer& orig, + int bbs_in) + : IndexIVFFastScan( + orig.quantizer, + orig.d, + orig.nlist, + 0, + orig.metric_type, + false), + sq(orig.sq) { + ntotal = orig.ntotal; + is_trained = orig.is_trained; + nprobe = orig.nprobe; + + if (is_fallback(sq.qtype)) { + by_residual = orig.by_residual; + code_size = sq.code_size; + replace_invlists(new ArrayInvertedLists(nlist, code_size), true); + + for (size_t list_no = 0; list_no < nlist; list_no++) { + size_t list_size = orig.invlists->list_size(list_no); + if (list_size == 0) { + continue; + } + InvertedLists::ScopedCodes codes(orig.invlists, list_no); + InvertedLists::ScopedIds ids(orig.invlists, list_no); + invlists->add_entries(list_no, list_size, ids.get(), codes.get()); + } + return; + } + + FAISS_THROW_IF_NOT_MSG( + !orig.by_residual, + "IndexIVFSQFastScan: conversion from IndexIVFScalarQuantizer " + "with by_residual=true is not supported for native/rerank " + "types. Set orig.by_residual=false and retrain, or construct " + "IndexIVFSQFastScan directly."); + + by_residual = false; + + size_t M_in = d; + M = M_in; + nbits = 4; + bbs = bbs_in; + ksub = 16; + M2 = roundup(M_in, 2); + code_size = M2 / 2; + + replace_invlists(new BlockInvertedLists(nlist, get_CodePacker()), true); + + if (needs_rerank(sq.qtype)) { + orig_codes_invlists = new ArrayInvertedLists(nlist, sq.code_size); + direct_map.set_type(DirectMap::Hashtable, invlists, 0); + } + + for (size_t list_no = 0; list_no < nlist; list_no++) { + size_t list_size = orig.invlists->list_size(list_no); + if (list_size == 0) { + continue; + } + + InvertedLists::ScopedCodes orig_codes(orig.invlists, list_no); + InvertedLists::ScopedIds orig_ids(orig.invlists, list_no); + + if (needs_rerank(sq.qtype)) { + orig_codes_invlists->add_entries( + list_no, list_size, orig_ids.get(), orig_codes.get()); + for (size_t j = 0; j < list_size; j++) { + direct_map.add_single_id(orig_ids.get()[j], list_no, j); + } + } + + // Decode to float and re-quantize to 4-bit. + std::vector recon(list_size * d); + sq.decode(orig_codes.get(), recon.data(), list_size); + + std::vector flat4(list_size * M2 / 2, 0); + if (is_native_4bit(sq.qtype)) { + sq.compute_codes(recon.data(), flat4.data(), list_size); + } else { + float_to_4bit_nibbles( + recon.data(), flat4.data(), list_size, d, M2, sq); + } + + BlockInvertedLists* bil = dynamic_cast(invlists); + FAISS_THROW_IF_NOT(bil); + + size_t nb2 = roundup(list_size, bbs); + AlignedTable packed(nb2 * M2 / 2); + pq4_pack_codes(flat4.data(), list_size, M, nb2, bbs, M2, packed.get()); + bil->add_entries(list_no, list_size, orig_ids.get(), packed.get()); + } +} + +IndexIVFSQFastScan::~IndexIVFSQFastScan() { + delete orig_codes_invlists; +} + +// ----------------------------------------------------------------------- +// Virtual method implementations +// ----------------------------------------------------------------------- + +size_t IndexIVFSQFastScan::fast_scan_code_size() const { + return M2 / 2; +} + +void IndexIVFSQFastScan::train_encoder( + idx_t n, + const float* x, + const idx_t* /*assign*/) { + sq.train(n, x); +} + +idx_t IndexIVFSQFastScan::train_encoder_num_vectors() const { + return 100000; +} + +void IndexIVFSQFastScan::encode_vectors( + idx_t n, + const float* x, + const idx_t* list_nos, + uint8_t* codes_out, + bool include_listnos) const { + std::vector residuals; + const float* to_encode = x; + + if (by_residual) { + residuals.resize(n * d); + for (idx_t i = 0; i < n; i++) { + if (list_nos[i] >= 0) { + quantizer->compute_residual( + x + i * d, residuals.data() + i * d, list_nos[i]); + } else { + memset(residuals.data() + i * d, 0, sizeof(float) * d); + } + } + to_encode = residuals.data(); + } + + if (is_fallback(sq.qtype)) { + sq.compute_codes(to_encode, codes_out, n); + } else if (is_native_4bit(sq.qtype)) { + sq.compute_codes(to_encode, codes_out, n); + } else { + float_to_4bit_nibbles(to_encode, codes_out, n, d, M2, sq); + } + + if (include_listnos) { + size_t coarse_size = coarse_code_size(); + for (idx_t i = n - 1; i >= 0; i--) { + uint8_t* code = codes_out + i * (coarse_size + code_size); + memmove(code + coarse_size, codes_out + i * code_size, code_size); + encode_listno(list_nos[i], code); + } + } +} + +bool IndexIVFSQFastScan::lookup_table_is_3d() const { + // For SQ, the LUT doesn't depend on which list we're scanning + // (unlike PQ with residuals and precomputed tables). + return false; +} + +void IndexIVFSQFastScan::compute_LUT( + size_t n, + const float* x, + const CoarseQuantized& /*cq*/, + AlignedTable& dis_tables, + AlignedTable& /*biases*/, + const FastScanDistancePostProcessing&) const { + size_t dim = d; + size_t ksub_val = 16; + size_t dim12 = dim * ksub_val; + + // Build reconstruction table: for each dimension, the 16 centroid values. + bool is_uniform = is_uniform_range(sq.qtype) || + sq.qtype == ScalarQuantizer::QT_4bit_uniform; + + std::vector recon_table(dim * ksub_val); + + if (is_uniform) { + float vmin, vdiff; + get_uniform_range(sq, vmin, vdiff); + for (size_t c = 0; c < ksub_val; c++) { + float recon = vmin + ((c + 0.5f) / 15.0f) * vdiff; + for (size_t m = 0; m < dim; m++) { + recon_table[m * ksub_val + c] = recon; + } + } + } else { + const float* vmin = sq.trained.data(); + const float* vdiff = sq.trained.data() + dim; + for (size_t m = 0; m < dim; m++) { + for (size_t c = 0; c < ksub_val; c++) { + recon_table[m * ksub_val + c] = + vmin[m] + ((c + 0.5f) / 15.0f) * vdiff[m]; + } + } + } + + // by_residual is always false for the fast-scan path (non-fallback + // types), so codes store raw quantized values and the LUT computes + // distances directly from the raw query. No biases needed. + + dis_tables.resize(n * dim12); + + if (metric_type == METRIC_L2) { + for (size_t i = 0; i < n; i++) { + const float* xi = x + i * dim; + float* lut_i = dis_tables.get() + i * dim12; + for (size_t m = 0; m < dim; m++) { + float qi = xi[m]; + const float* recon_m = recon_table.data() + m * ksub_val; + float* lut_m = lut_i + m * ksub_val; + for (size_t c = 0; c < ksub_val; c++) { + float diff = qi - recon_m[c]; + lut_m[c] = diff * diff; + } + } + } + } else { + for (size_t i = 0; i < n; i++) { + const float* xi = x + i * dim; + float* lut_i = dis_tables.get() + i * dim12; + for (size_t m = 0; m < dim; m++) { + float qi = xi[m]; + const float* recon_m = recon_table.data() + m * ksub_val; + float* lut_m = lut_i + m * ksub_val; + for (size_t c = 0; c < ksub_val; c++) { + lut_m[c] = qi * recon_m[c]; + } + } + } + } +} + +void IndexIVFSQFastScan::sa_decode(idx_t n, const uint8_t* bytes, float* x) + const { + size_t coarse_size = coarse_code_size(); + +#pragma omp parallel if (n > 1) + { + std::vector residual(d); + +#pragma omp for + for (idx_t i = 0; i < n; i++) { + const uint8_t* code = bytes + i * (code_size + coarse_size); + int64_t list_no = decode_listno(code); + float* xi = x + i * d; + sq.decode(code + coarse_size, xi, 1); + if (by_residual) { + quantizer->reconstruct(list_no, residual.data()); + for (int j = 0; j < d; j++) { + xi[j] += residual[j]; + } + } + } + } +} + +void IndexIVFSQFastScan::reconstruct_from_offset( + int64_t list_no, + int64_t offset, + float* recons) const { + if (is_fallback(sq.qtype)) { + // Fallback: codes stored directly in invlists. + InvertedLists::ScopedCodes codes(invlists, list_no); + sq.decode(codes.get() + offset * sq.code_size, recons, 1); + } else if (orig_codes_invlists) { + // Rerank: use original full-precision codes for reconstruction. + InvertedLists::ScopedCodes codes(orig_codes_invlists, list_no); + sq.decode(codes.get() + offset * sq.code_size, recons, 1); + } else { + // Native 4-bit: unpack from block inverted lists. + std::vector code(M2 / 2, 0); + InvertedLists::ScopedCodes list_codes(invlists, list_no); + BitstringWriter bsw(code.data(), M2 / 2); + for (size_t m = 0; m < M; m++) { + uint8_t c = pq4_get_packed_element( + list_codes.get(), bbs, M2, offset, m); + bsw.write(c, 4); + } + sq.decode(code.data(), recons, 1); + } + + if (by_residual) { + std::vector centroid(d); + quantizer->reconstruct(list_no, centroid.data()); + for (int j = 0; j < d; j++) { + recons[j] += centroid[j]; + } + } +} + +InvertedListScanner* IndexIVFSQFastScan::get_InvertedListScanner( + bool store_pairs, + const IDSelector* sel, + const IVFSearchParameters*) const { + return sq.select_InvertedListScanner( + metric_type, quantizer, store_pairs, sel, by_residual); +} + +// ----------------------------------------------------------------------- +// search — fallback types use IndexIVF's scanner-based path +// ----------------------------------------------------------------------- + +void IndexIVFSQFastScan::search( + idx_t n, + const float* x, + idx_t k, + float* distances, + idx_t* labels, + const SearchParameters* params) const { + if (is_fallback(sq.qtype)) { + IndexIVF::search(n, x, k, distances, labels, params); + } else if (needs_rerank(sq.qtype)) { + idx_t k_coarse = idx_t(k * rerank_factor); + if (k_coarse < k) { + k_coarse = k; + } + + std::vector coarse_dis(n * k_coarse); + std::vector coarse_ids(n * k_coarse); + IndexIVFFastScan::search( + n, x, k_coarse, coarse_dis.data(), coarse_ids.data(), params); + + // Rerank using original full-precision codes +#pragma omp parallel + { + std::unique_ptr dc( + sq.get_distance_computer(metric_type)); + +#pragma omp for + for (idx_t i = 0; i < n; i++) { + const float* xi = x + i * d; + dc->set_query(xi); + + float* heap_dis = distances + i * k; + idx_t* heap_ids = labels + i * k; + + if (metric_type == METRIC_L2) { + maxheap_heapify(k, heap_dis, heap_ids); + } else { + minheap_heapify(k, heap_dis, heap_ids); + } + + for (idx_t j = 0; j < k_coarse; j++) { + idx_t id = coarse_ids[i * k_coarse + j]; + if (id < 0) { + continue; + } + + idx_t lo = direct_map.get(id); + idx_t list_no = lo_listno(lo); + idx_t offset = lo_offset(lo); + + InvertedLists::ScopedCodes codes( + orig_codes_invlists, list_no); + float dis = dc->distance_to_code( + codes.get() + offset * sq.code_size); + + if (metric_type == METRIC_L2) { + if (dis < heap_dis[0]) { + maxheap_replace_top(k, heap_dis, heap_ids, dis, id); + } + } else { + if (dis > heap_dis[0]) { + minheap_replace_top(k, heap_dis, heap_ids, dis, id); + } + } + } + + if (metric_type == METRIC_L2) { + maxheap_reorder(k, heap_dis, heap_ids); + } else { + minheap_reorder(k, heap_dis, heap_ids); + } + } + } + } else { + IndexIVFFastScan::search(n, x, k, distances, labels, params); + } +} + +void IndexIVFSQFastScan::search_preassigned( + idx_t n, + const float* x, + idx_t k, + const idx_t* assign, + const float* centroid_dis, + float* distances, + idx_t* labels, + bool store_pairs, + const IVFSearchParameters* params, + IndexIVFStats* stats) const { + if (is_fallback(sq.qtype)) { + IndexIVF::search_preassigned( + n, + x, + k, + assign, + centroid_dis, + distances, + labels, + store_pairs, + params, + stats); + } else if (needs_rerank(sq.qtype)) { + idx_t k_coarse = idx_t(k * rerank_factor); + if (k_coarse < k) { + k_coarse = k; + } + + std::vector coarse_dis(n * k_coarse); + std::vector coarse_ids(n * k_coarse); + IndexIVFFastScan::search_preassigned( + n, + x, + k_coarse, + assign, + centroid_dis, + coarse_dis.data(), + coarse_ids.data(), + false, + params, + stats); + +#pragma omp parallel + { + std::unique_ptr dc( + sq.get_distance_computer(metric_type)); + +#pragma omp for + for (idx_t i = 0; i < n; i++) { + const float* xi = x + i * d; + dc->set_query(xi); + + float* heap_dis = distances + i * k; + idx_t* heap_ids = labels + i * k; + + if (metric_type == METRIC_L2) { + maxheap_heapify(k, heap_dis, heap_ids); + } else { + minheap_heapify(k, heap_dis, heap_ids); + } + + for (idx_t j = 0; j < k_coarse; j++) { + idx_t id = coarse_ids[i * k_coarse + j]; + if (id < 0) { + continue; + } + + idx_t lo = direct_map.get(id); + idx_t list_no = lo_listno(lo); + idx_t offset = lo_offset(lo); + + InvertedLists::ScopedCodes codes( + orig_codes_invlists, list_no); + float dis = dc->distance_to_code( + codes.get() + offset * sq.code_size); + + if (metric_type == METRIC_L2) { + if (dis < heap_dis[0]) { + maxheap_replace_top(k, heap_dis, heap_ids, dis, id); + } + } else { + if (dis > heap_dis[0]) { + minheap_replace_top(k, heap_dis, heap_ids, dis, id); + } + } + } + + if (metric_type == METRIC_L2) { + maxheap_reorder(k, heap_dis, heap_ids); + } else { + minheap_reorder(k, heap_dis, heap_ids); + } + } + } + } else { + IndexIVFFastScan::search_preassigned( + n, + x, + k, + assign, + centroid_dis, + distances, + labels, + store_pairs, + params, + stats); + } +} + +void IndexIVFSQFastScan::range_search( + idx_t n, + const float* x, + float radius, + RangeSearchResult* result, + const SearchParameters* params) const { + if (is_fallback(sq.qtype)) { + IndexIVF::range_search(n, x, radius, result, params); + } else if (needs_rerank(sq.qtype)) { + const IVFSearchParameters* ivf_params = nullptr; + if (params) { + ivf_params = dynamic_cast(params); + } + size_t cur_nprobe = ivf_params ? ivf_params->nprobe : nprobe; + const IDSelector* sel = params ? params->sel : nullptr; + + std::vector assign(n * cur_nprobe); + std::vector centroid_dis(n * cur_nprobe); + quantizer->search(n, x, cur_nprobe, centroid_dis.data(), assign.data()); + + std::vector partial_results( + omp_get_max_threads()); + +#pragma omp parallel + { + int rank = omp_get_thread_num(); + RangeSearchPartialResult* pres = + new RangeSearchPartialResult(result); + partial_results[rank] = pres; + + std::unique_ptr dc( + sq.get_distance_computer(metric_type)); + +#pragma omp for + for (idx_t i = 0; i < n; i++) { + dc->set_query(x + i * d); + RangeQueryResult& qres = pres->new_result(i); + + for (size_t p = 0; p < cur_nprobe; p++) { + idx_t list_no = assign[i * cur_nprobe + p]; + if (list_no < 0) { + continue; + } + size_t list_size = orig_codes_invlists->list_size(list_no); + if (list_size == 0) { + continue; + } + InvertedLists::ScopedCodes codes( + orig_codes_invlists, list_no); + InvertedLists::ScopedIds ids(orig_codes_invlists, list_no); + + for (size_t j = 0; j < list_size; j++) { + if (sel && !sel->is_member(ids[j])) { + continue; + } + float dis = dc->distance_to_code( + codes.get() + j * sq.code_size); + if (metric_type == METRIC_L2) { + if (dis < radius) { + qres.add(dis, ids[j]); + } + } else { + if (dis > radius) { + qres.add(dis, ids[j]); + } + } + } + } + } + } + RangeSearchPartialResult::merge(partial_results); + } else { + IndexIVFFastScan::range_search(n, x, radius, result, params); + } +} + +// ----------------------------------------------------------------------- +// add_with_ids — custom implementation to store both packed and original codes +// ----------------------------------------------------------------------- + +void IndexIVFSQFastScan::add_with_ids( + idx_t n, + const float* x, + const idx_t* xids) { + FAISS_THROW_IF_NOT(is_trained); + + constexpr idx_t bs = 65536; + if (n > bs) { + for (idx_t i0 = 0; i0 < n; i0 += bs) { + idx_t i1 = std::min(n, i0 + bs); + add_with_ids(i1 - i0, x + i0 * d, xids ? xids + i0 : nullptr); + } + return; + } + + // Assign to coarse clusters + std::unique_ptr idx(new idx_t[n]); + quantizer->assign(n, x, idx.get()); + + if (is_fallback(sq.qtype)) { + // Fallback: store SQ codes directly in invlists (same as + // IndexIVFScalarQuantizer). + std::vector flat_codes(n * sq.code_size); + encode_vectors(n, x, idx.get(), flat_codes.data()); + + for (idx_t i = 0; i < n; i++) { + if (idx[i] < 0) { + continue; + } + idx_t id = xids ? xids[i] : ntotal + i; + invlists->add_entry( + idx[i], id, flat_codes.data() + i * sq.code_size); + } + ntotal += n; + return; + } + + // Produce 4-bit packed codes for the SIMD scan, + // and (for rerank types) also store original codes. + AlignedTable flat_codes(n * code_size); + encode_vectors(n, x, idx.get(), flat_codes.get()); + + // Store original SQ codes for rerank types + std::vector orig_flat_codes; + if (needs_rerank(sq.qtype)) { + orig_flat_codes.resize(n * sq.code_size); + for (idx_t i = 0; i < n; i++) { + if (idx[i] < 0) { + continue; + } + if (by_residual) { + std::vector residual(d); + quantizer->compute_residual(x + i * d, residual.data(), idx[i]); + sq.compute_codes( + residual.data(), + orig_flat_codes.data() + i * sq.code_size, + 1); + } else { + sq.compute_codes( + x + i * d, + orig_flat_codes.data() + i * sq.code_size, + 1); + } + } + } + + // Sort by list assignment + std::vector order(n); + for (idx_t i = 0; i < n; i++) { + order[i] = i; + } + std::stable_sort(order.begin(), order.end(), [&idx](idx_t a, idx_t b) { + return idx[a] < idx[b]; + }); + + BlockInvertedLists* bil = dynamic_cast(invlists); + FAISS_THROW_IF_NOT_MSG(bil, "only block inverted lists supported"); + + idx_t i0 = 0; + while (i0 < n) { + idx_t list_no = idx[order[i0]]; + idx_t i1 = i0 + 1; + while (i1 < n && idx[order[i1]] == list_no) { + i1++; + } + + if (list_no == -1) { + i0 = i1; + continue; + } + + size_t list_size = bil->list_size(list_no); + bil->resize(list_no, list_size + i1 - i0); + + AlignedTable list_codes((i1 - i0) * code_size); + for (idx_t i = i0; i < i1; i++) { + size_t ofs = list_size + i - i0; + idx_t id = xids ? xids[order[i]] : ntotal + order[i]; + bil->ids[list_no][ofs] = id; + memcpy(list_codes.data() + (i - i0) * code_size, + flat_codes.data() + order[i] * code_size, + code_size); + + if (needs_rerank(sq.qtype)) { + orig_codes_invlists->add_entry( + list_no, + id, + orig_flat_codes.data() + order[i] * sq.code_size); + direct_map.add_single_id(id, list_no, ofs); + } + } + + pq4_pack_codes_range( + list_codes.data(), + M, + list_size, + list_size + i1 - i0, + bbs, + M2, + bil->codes[list_no].data(), + 0, + get_block_stride()); + + i0 = i1; + } + + ntotal += n; +} + +void IndexIVFSQFastScan::reset() { + IndexIVFFastScan::reset(); + if (orig_codes_invlists) { + for (size_t i = 0; i < nlist; i++) { + orig_codes_invlists->resize(i, 0); + } + direct_map.clear(); + } +} + +} // namespace faiss diff --git a/faiss/IndexIVFSQFastScan.h b/faiss/IndexIVFSQFastScan.h new file mode 100644 index 000000000..a1a0b1397 --- /dev/null +++ b/faiss/IndexIVFSQFastScan.h @@ -0,0 +1,131 @@ +/* + * 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. + */ + +#pragma once + +#include +#include +#include + +namespace faiss { + +/** Fast scan version of IndexIVFScalarQuantizer. + * + * Uses the PQ4 FastScan SIMD infrastructure (vpshufb) for accelerated + * scanning within inverted lists. Supported quantizer types: + * + * - QT_4bit, QT_4bit_uniform: native 4-bit codes are packed directly + * into the SIMD block layout. No precision loss. + * + * - QT_6bit, QT_8bit, QT_8bit_uniform, QT_8bit_direct, + * QT_8bit_direct_signed: re-quantised to 4-bit for the SIMD scan, + * then reranked with exact original-precision distances. + * + * - All other types (QT_fp16, QT_bf16, TurboQuant): fall back to the + * ScalarQuantizer's own SIMD-optimised InvertedListScanner (same + * behavior as IndexIVFScalarQuantizer, no fast-scan acceleration). + * + * For reranked types, both the original full-precision codes and the + * packed 4-bit codes are stored in the inverted lists. The block + * inverted lists hold packed 4-bit codes; original codes are stored + * in a parallel ArrayInvertedLists (orig_codes_invlists). + */ +struct IndexIVFSQFastScan : IndexIVFFastScan { + ScalarQuantizer sq; + + /// Overselection ratio for reranking (default 2). + float rerank_factor = 2; + + /// Parallel inverted lists storing original SQ codes (for rerank + /// types). Owned by this index. nullptr for native 4-bit and + /// fallback types. + InvertedLists* orig_codes_invlists = nullptr; + + IndexIVFSQFastScan( + Index* quantizer, + size_t d, + size_t nlist, + ScalarQuantizer::QuantizerType qtype, + MetricType metric = METRIC_L2, + int bbs = 32, + bool by_residual = true); + + IndexIVFSQFastScan(); + + /// Build from an existing IndexIVFScalarQuantizer. + explicit IndexIVFSQFastScan( + const IndexIVFScalarQuantizer& orig, + int bbs = 32); + + ~IndexIVFSQFastScan() override; + + size_t fast_scan_code_size() const override; + + void train_encoder(idx_t n, const float* x, const idx_t* assign) override; + + idx_t train_encoder_num_vectors() const override; + + void encode_vectors( + idx_t n, + const float* x, + const idx_t* list_nos, + uint8_t* codes, + bool include_listno = false) const override; + + bool lookup_table_is_3d() const override; + + void compute_LUT( + size_t n, + const float* x, + const CoarseQuantized& cq, + AlignedTable& dis_tables, + AlignedTable& biases, + const FastScanDistancePostProcessing& context) const override; + + void sa_decode(idx_t n, const uint8_t* bytes, float* x) const override; + + void reconstruct_from_offset(int64_t list_no, int64_t offset, float* recons) + const override; + + InvertedListScanner* get_InvertedListScanner( + bool store_pairs, + const IDSelector* sel, + const IVFSearchParameters*) const override; + + void search( + idx_t n, + const float* x, + idx_t k, + float* distances, + idx_t* labels, + const SearchParameters* params = nullptr) const override; + + void search_preassigned( + idx_t n, + const float* x, + idx_t k, + const idx_t* assign, + const float* centroid_dis, + float* distances, + idx_t* labels, + bool store_pairs, + const IVFSearchParameters* params = nullptr, + IndexIVFStats* stats = nullptr) const override; + + void range_search( + idx_t n, + const float* x, + float radius, + RangeSearchResult* result, + const SearchParameters* params = nullptr) const override; + + void add_with_ids(idx_t n, const float* x, const idx_t* xids) override; + + void reset() override; +}; + +} // namespace faiss diff --git a/faiss/IndexSQFastScan.cpp b/faiss/IndexSQFastScan.cpp new file mode 100644 index 000000000..0fbf66e3c --- /dev/null +++ b/faiss/IndexSQFastScan.cpp @@ -0,0 +1,504 @@ +/* + * 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 +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace faiss { + +namespace { + +size_t roundup(size_t a, size_t b) { + return (a + b - 1) / b * b; +} + +bool is_native_4bit(ScalarQuantizer::QuantizerType qtype) { + return qtype == ScalarQuantizer::QT_4bit || + qtype == ScalarQuantizer::QT_4bit_uniform; +} + +} // anonymous namespace + +// ----------------------------------------------------------------------- +// Constructors +// ----------------------------------------------------------------------- + +IndexSQFastScan::IndexSQFastScan( + int d_in, + ScalarQuantizer::QuantizerType qtype, + MetricType metric, + int bbs_in) + : sq(d_in, qtype) { + FAISS_THROW_IF_NOT_MSG( + is_native_4bit(qtype), + "IndexSQFastScan only supports QT_4bit and QT_4bit_uniform. " + "For higher-precision types, use " + "IndexRefine(IndexSQFastScan(...), IndexScalarQuantizer(...))."); + init_fastscan(d_in, d_in, 4, metric, bbs_in); +} + +IndexSQFastScan::IndexSQFastScan() = default; + +IndexSQFastScan::IndexSQFastScan(const IndexScalarQuantizer& orig, int bbs_in) + : sq(orig.sq) { + FAISS_THROW_IF_NOT_MSG( + is_native_4bit(sq.qtype), + "IndexSQFastScan conversion constructor only supports " + "QT_4bit and QT_4bit_uniform."); + init_fastscan(orig.d, orig.d, 4, orig.metric_type, bbs_in); + ntotal = orig.ntotal; + is_trained = orig.is_trained; + + ntotal2 = roundup(ntotal, bbs_in); + codes.resize(ntotal2 * M2 / 2); + pq4_pack_codes( + orig.codes.data(), ntotal, M, ntotal2, bbs_in, M2, codes.get()); +} + +// ----------------------------------------------------------------------- +// train / add / reset +// ----------------------------------------------------------------------- + +void IndexSQFastScan::train(idx_t n, const float* x) { + if (is_trained) { + return; + } + sq.train(n, x); + is_trained = true; +} + +void IndexSQFastScan::add(idx_t n, const float* x) { + FAISS_THROW_IF_NOT(is_trained); + IndexFastScan::add(n, x); +} + +void IndexSQFastScan::reset() { + IndexFastScan::reset(); +} + +// ----------------------------------------------------------------------- +// search +// ----------------------------------------------------------------------- + +void IndexSQFastScan::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(is_trained); + + const IDSelector* sel = params ? params->sel : nullptr; + + if (!sel) { + IndexFastScan::search(n, x, k, distances, labels, nullptr); + return; + } + + // IDSelector path: unpack codes and use SQ scanner + std::vector flat_codes(ntotal * sq.code_size); + { + std::unique_ptr packer(get_CodePacker()); + for (idx_t i = 0; i < ntotal; i++) { + packer->unpack_1( + codes.data(), i, flat_codes.data() + i * sq.code_size); + } + } + +#pragma omp parallel + { + std::unique_ptr scanner( + sq.select_InvertedListScanner(metric_type, nullptr, true, sel)); + scanner->list_no = 0; + +#pragma omp for + for (idx_t i = 0; i < n; i++) { + float* D = distances + k * i; + idx_t* I = labels + k * i; + if (metric_type == METRIC_L2) { + maxheap_heapify(k, D, I); + } else { + minheap_heapify(k, D, I); + } + scanner->set_query(x + i * d); + scanner->scan_codes(ntotal, flat_codes.data(), nullptr, D, I, k); + if (metric_type == METRIC_L2) { + maxheap_reorder(k, D, I); + } else { + minheap_reorder(k, D, I); + } + } + } +} + +// ----------------------------------------------------------------------- +// compute_codes -- called by IndexFastScan::add +// ----------------------------------------------------------------------- + +void IndexSQFastScan::compute_codes(uint8_t* out_codes, idx_t n, const float* x) + const { + sq.compute_codes(x, out_codes, n); +} + +// ----------------------------------------------------------------------- +// compute_float_LUT -- builds the d x 16 distance LUT for vpshufb +// ----------------------------------------------------------------------- + +void IndexSQFastScan::compute_float_LUT( + float* lut, + idx_t n, + const float* x, + const FastScanDistancePostProcessing&) const { + const size_t dim = d; + const size_t ksub_val = 16; + + bool is_uniform = (sq.qtype == ScalarQuantizer::QT_4bit_uniform); + + std::vector recon_table(dim * ksub_val); + + if (is_uniform) { + float vmin = sq.trained[0]; + float vdiff = sq.trained[1]; + for (size_t c = 0; c < ksub_val; c++) { + float recon = vmin + ((c + 0.5f) / 15.0f) * vdiff; + for (size_t m = 0; m < dim; m++) { + recon_table[m * ksub_val + c] = recon; + } + } + } else { + // QT_4bit: per-dimension ranges + const float* vmin = sq.trained.data(); + const float* vdiff = sq.trained.data() + dim; + for (size_t m = 0; m < dim; m++) { + for (size_t c = 0; c < ksub_val; c++) { + recon_table[m * ksub_val + c] = + vmin[m] + ((c + 0.5f) / 15.0f) * vdiff[m]; + } + } + } + + if (metric_type == METRIC_L2) { + for (idx_t i = 0; i < n; i++) { + const float* xi = x + i * dim; + float* lut_i = lut + i * dim * ksub_val; + for (size_t m = 0; m < dim; m++) { + float qi = xi[m]; + const float* recon_m = recon_table.data() + m * ksub_val; + float* lut_m = lut_i + m * ksub_val; + for (size_t c = 0; c < ksub_val; c++) { + float diff = qi - recon_m[c]; + lut_m[c] = diff * diff; + } + } + } + } else { + for (idx_t i = 0; i < n; i++) { + const float* xi = x + i * dim; + float* lut_i = lut + i * dim * ksub_val; + for (size_t m = 0; m < dim; m++) { + float qi = xi[m]; + const float* recon_m = recon_table.data() + m * ksub_val; + float* lut_m = lut_i + m * ksub_val; + for (size_t c = 0; c < ksub_val; c++) { + lut_m[c] = qi * recon_m[c]; + } + } + } + } +} + +void IndexSQFastScan::sa_decode(idx_t n, const uint8_t* bytes, float* x) const { + sq.decode(bytes, x, n); +} + +void IndexSQFastScan::reconstruct(idx_t key, float* recons) const { + FAISS_THROW_IF_NOT(key >= 0 && key < ntotal); + std::unique_ptr packer(get_CodePacker()); + std::vector flat_code(sq.code_size); + packer->unpack_1(codes.data(), key, flat_code.data()); + sq.decode(flat_code.data(), recons, 1); +} + +size_t IndexSQFastScan::sa_code_size() const { + return sq.code_size; +} + +size_t IndexSQFastScan::fast_scan_code_size() const { + return M2 / 2; +} + +// ----------------------------------------------------------------------- +// sa_encode +// ----------------------------------------------------------------------- + +void IndexSQFastScan::sa_encode(idx_t n, const float* x, uint8_t* bytes) const { + FAISS_THROW_IF_NOT(is_trained); + sq.compute_codes(x, bytes, n); +} + +// ----------------------------------------------------------------------- +// reconstruct_n +// ----------------------------------------------------------------------- + +void IndexSQFastScan::reconstruct_n(idx_t i0, idx_t ni, float* recons) const { + FAISS_THROW_IF_NOT(i0 >= 0 && i0 + ni <= ntotal); + std::unique_ptr packer(get_CodePacker()); + std::vector flat_code(sq.code_size); + for (idx_t i = 0; i < ni; i++) { + packer->unpack_1(codes.data(), i0 + i, flat_code.data()); + sq.decode(flat_code.data(), recons + i * d, 1); + } +} + +// ----------------------------------------------------------------------- +// add_sa_codes +// ----------------------------------------------------------------------- + +void IndexSQFastScan::add_sa_codes( + idx_t n, + const uint8_t* code, + const idx_t* /*xids*/) { + FAISS_THROW_IF_NOT(is_trained); + // Decode to float then re-add via IndexFastScan for SIMD packing + std::vector recon(n * d); + sq.decode(code, recon.data(), n); + IndexFastScan::add(n, recon.data()); +} + +// ----------------------------------------------------------------------- +// permute_entries +// ----------------------------------------------------------------------- + +void IndexSQFastScan::permute_entries(const idx_t* perm) { + std::vector flat_old; + { + std::unique_ptr packer(get_CodePacker()); + size_t code_sz = code_size; + flat_old.resize(ntotal * code_sz, 0); + for (idx_t i = 0; i < ntotal; i++) { + packer->unpack_1(codes.data(), i, flat_old.data() + i * code_sz); + } + } + + size_t code_sz = code_size; + std::vector flat_new(ntotal2 * code_sz, 0); + for (idx_t i = 0; i < ntotal; i++) { + memcpy(flat_new.data() + i * code_sz, + flat_old.data() + perm[i] * code_sz, + code_sz); + } + pq4_pack_codes(flat_new.data(), ntotal, M, ntotal2, bbs, M2, codes.get()); +} + +// ----------------------------------------------------------------------- +// get_distance_computer / get_FlatCodesDistanceComputer +// ----------------------------------------------------------------------- + +namespace { + +/// Wrapper that owns an unpacked code buffer and delegates distance +/// computation to an underlying SQDistanceComputer. +struct OwningFlatCodesDistanceComputer : FlatCodesDistanceComputer { + std::vector owned_codes; + std::unique_ptr dc; + + OwningFlatCodesDistanceComputer( + std::vector&& codes_buf, + ScalarQuantizer::SQDistanceComputer* dc_in) + : owned_codes(std::move(codes_buf)), dc(dc_in) { + this->codes = owned_codes.data(); + this->code_size = dc->code_size; + dc->codes = owned_codes.data(); + } + + void set_query(const float* x) override { + dc->set_query(x); + } + + float distance_to_code(const uint8_t* code) override { + return dc->distance_to_code(code); + } + + void distance_to_code_batch_4( + const uint8_t* c0, + const uint8_t* c1, + const uint8_t* c2, + const uint8_t* c3, + float& d0, + float& d1, + float& d2, + float& d3) override { + dc->distance_to_code_batch_4(c0, c1, c2, c3, d0, d1, d2, d3); + } + + float symmetric_dis(idx_t i, idx_t j) override { + return dc->symmetric_dis(i, j); + } +}; + +} // anonymous namespace + +FlatCodesDistanceComputer* IndexSQFastScan::get_FlatCodesDistanceComputer() + const { + std::unique_ptr packer(get_CodePacker()); + size_t cs = sq.code_size; + std::vector flat(ntotal * cs); + for (idx_t i = 0; i < ntotal; i++) { + packer->unpack_1(codes.data(), i, flat.data() + i * cs); + } + + ScalarQuantizer::SQDistanceComputer* dc = + sq.get_distance_computer(metric_type); + dc->code_size = cs; + + return new OwningFlatCodesDistanceComputer(std::move(flat), dc); +} + +DistanceComputer* IndexSQFastScan::get_distance_computer() const { + return get_FlatCodesDistanceComputer(); +} + +// ----------------------------------------------------------------------- +// range_search +// ----------------------------------------------------------------------- + +void IndexSQFastScan::range_search( + idx_t n, + const float* x, + float radius, + RangeSearchResult* result, + const SearchParameters* params) const { + FAISS_THROW_IF_NOT(is_trained); + const IDSelector* sel = params ? params->sel : nullptr; + + // Unpack codes for range search + std::vector flat_codes(ntotal * sq.code_size); + { + std::unique_ptr packer(get_CodePacker()); + for (idx_t i = 0; i < ntotal; i++) { + packer->unpack_1( + codes.data(), i, flat_codes.data() + i * sq.code_size); + } + } + + std::vector partial_results( + omp_get_max_threads()); + + size_t cs = sq.code_size; + +#pragma omp parallel + { + int rank = omp_get_thread_num(); + RangeSearchPartialResult* pres = new RangeSearchPartialResult(result); + partial_results[rank] = pres; + + std::vector decoded(d); + +#pragma omp for + for (idx_t i = 0; i < n; i++) { + const float* qi = x + i * d; + RangeQueryResult& qres = pres->new_result(i); + + for (idx_t j = 0; j < ntotal; j++) { + if (sel && !sel->is_member(j)) { + continue; + } + sq.decode(flat_codes.data() + j * cs, decoded.data(), 1); + float dis; + if (metric_type == METRIC_INNER_PRODUCT) { + dis = fvec_inner_product(qi, decoded.data(), d); + if (dis > radius) { + qres.add(dis, j); + } + } else { + dis = fvec_L2sqr(qi, decoded.data(), d); + if (dis < radius) { + qres.add(dis, j); + } + } + } + } + } + + RangeSearchPartialResult::merge(partial_results); +} + +// ----------------------------------------------------------------------- +// remove_ids +// ----------------------------------------------------------------------- + +size_t IndexSQFastScan::remove_ids(const IDSelector& sel) { + return IndexFastScan::remove_ids(sel); +} + +// ----------------------------------------------------------------------- +// check_compatible_for_merge +// ----------------------------------------------------------------------- + +void IndexSQFastScan::check_compatible_for_merge( + const Index& otherIndex) const { + IndexFastScan::check_compatible_for_merge(otherIndex); + const IndexSQFastScan* other = + dynamic_cast(&otherIndex); + FAISS_THROW_IF_NOT_MSG(other, "merge requires IndexSQFastScan"); + FAISS_THROW_IF_NOT_MSG( + other->sq.qtype == sq.qtype, + "merge requires matching ScalarQuantizer types"); +} + +// ----------------------------------------------------------------------- +// merge_from +// ----------------------------------------------------------------------- + +void IndexSQFastScan::merge_from(Index& otherIndex, idx_t add_id) { + check_compatible_for_merge(otherIndex); + IndexFastScan::merge_from(otherIndex, add_id); +} + +// ----------------------------------------------------------------------- +// search1 +// ----------------------------------------------------------------------- + +void IndexSQFastScan::search1( + const float* x, + ResultHandler& handler, + SearchParameters* params) const { + const IDSelector* sel = params ? params->sel : nullptr; + std::unique_ptr dc(get_distance_computer()); + dc->set_query(x); + for (idx_t i = 0; i < ntotal; i++) { + if (sel && !sel->is_member(i)) { + continue; + } + float dis = (*dc)(i); + handler.add_result(dis, i); + } +} + +} // namespace faiss diff --git a/faiss/IndexSQFastScan.h b/faiss/IndexSQFastScan.h new file mode 100644 index 000000000..dca370291 --- /dev/null +++ b/faiss/IndexSQFastScan.h @@ -0,0 +1,117 @@ +/* + * 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. + */ + +#pragma once + +#include + +#include +#include +#include +#include + +namespace faiss { + +/** Fast scan version of IndexScalarQuantizer for native 4-bit types. + * + * Supported quantizer types: + * - QT_4bit, QT_4bit_uniform: native 4-bit codes mapped onto the + * PQ4 FastScan SIMD infrastructure (vpshufb / equivalent). + * No precision loss. + * + * For higher-precision types (QT_8bit, QT_6bit, etc.), use + * IndexRefine(IndexSQFastScan(QT_4bit), IndexScalarQuantizer(QT_8bit)) + * to get fast-scan with reranking. For fallback types (QT_fp16, + * QT_bf16), use IndexScalarQuantizer directly. + * + * M = d subquantizers with 16 levels and uint16 SIMD accumulators + * (safe for d <= 257). + */ +struct IndexSQFastScan : IndexFastScan { + ScalarQuantizer sq; + + /** Constructor. + * + * @param d dimensionality of input vectors + * @param qtype QT_4bit or QT_4bit_uniform only + * @param metric distance metric (METRIC_L2 or METRIC_INNER_PRODUCT) + * @param bbs block size for SIMD processing (multiple of 32) + */ + IndexSQFastScan( + int d, + ScalarQuantizer::QuantizerType qtype, + MetricType metric = METRIC_L2, + int bbs = 32); + + IndexSQFastScan(); + + /// Build from an existing IndexScalarQuantizer (must be QT_4bit*) + explicit IndexSQFastScan(const IndexScalarQuantizer& orig, int bbs = 32); + + void train(idx_t n, const float* x) override; + + void add(idx_t n, const float* x) override; + + void reset() override; + + void search( + idx_t n, + const float* x, + idx_t k, + float* distances, + idx_t* labels, + const SearchParameters* params = nullptr) const override; + + void compute_codes(uint8_t* codes, idx_t n, const float* x) const override; + + void check_compatible_for_merge(const Index& otherIndex) const override; + + void compute_float_LUT( + float* lut, + idx_t n, + const float* x, + const FastScanDistancePostProcessing& context) const override; + + void sa_decode(idx_t n, const uint8_t* bytes, float* x) const override; + + void sa_encode(idx_t n, const float* x, uint8_t* bytes) const override; + + void reconstruct(idx_t key, float* recons) const override; + + void reconstruct_n(idx_t i0, idx_t ni, float* recons) const override; + + void add_sa_codes(idx_t n, const uint8_t* codes, const idx_t* xids) + override; + + void permute_entries(const idx_t* perm); + + FlatCodesDistanceComputer* get_FlatCodesDistanceComputer() const; + + DistanceComputer* get_distance_computer() const override; + + void range_search( + idx_t n, + const float* x, + float radius, + RangeSearchResult* result, + const SearchParameters* params = nullptr) const override; + + size_t remove_ids(const IDSelector& sel) override; + + void merge_from(Index& otherIndex, idx_t add_id = 0) override; + + void search1( + const float* x, + ResultHandler& handler, + SearchParameters* params = nullptr) const override; + + size_t sa_code_size() const override; + + size_t fast_scan_code_size() const override; +}; + +} // namespace faiss diff --git a/faiss/impl/fast_scan/simd_result_handlers.h b/faiss/impl/fast_scan/simd_result_handlers.h index ecb748044..341004d69 100644 --- a/faiss/impl/fast_scan/simd_result_handlers.h +++ b/faiss/impl/fast_scan/simd_result_handlers.h @@ -410,7 +410,8 @@ struct SingleResultHandler : ResultHandlerCompare { if (!normalizers) { dis[q] = idis[q]; } else { - float one_a = 1 / normalizers[2 * q]; + float a = normalizers[2 * q]; + float one_a = a != 0 ? 1 / a : 0; float b = normalizers[2 * q + 1]; dis[q] = b + idis[q] * one_a; } @@ -549,7 +550,8 @@ struct HeapHandler : ResultHandlerCompare { float one_a = 1.0, b = 0.0; if (normalizers) { - one_a = 1 / normalizers[2 * q]; + float a = normalizers[2 * q]; + one_a = a != 0 ? 1 / a : 0; b = normalizers[2 * q + 1]; } for (int j = 0; j < k; j++) { @@ -679,7 +681,8 @@ struct ReservoirHandler : ResultHandlerCompare { float one_a = 1.0, b = 0.0; if (normalizers) { - one_a = 1 / normalizers[2 * q]; + float a = normalizers[2 * q]; + one_a = a != 0 ? 1 / a : 0; b = normalizers[2 * q + 1]; } for (size_t i = 0; i < res.i; i++) { @@ -807,7 +810,8 @@ struct RangeHandler : ResultHandlerCompare { rres.lims[0] = 0; for (int q = 0; q < nq; q++) { - float one_a = 1 / normalizers[2 * q]; + float a = normalizers[2 * q]; + float one_a = a != 0 ? 1 / a : 0; float b = normalizers[2 * q + 1]; for (size_t i = rres.lims[q]; i < rres.lims[q + 1]; i++) { rres.distances[i] = rres.distances[i] * one_a + b; @@ -873,7 +877,8 @@ struct PartialRangeHandler : RangeHandler { size_t* lims = n_per_query.data(); for (int q = 0; q < nq; q++) { - float one_a = 1 / normalizers[2 * q]; + float a = normalizers[2 * q]; + float one_a = a != 0 ? 1 / a : 0; float b = normalizers[2 * q + 1]; RangeQueryResult& qres = pres.new_result(q + q0); for (size_t i = lims[q]; i < lims[q + 1]; i++) { @@ -962,7 +967,8 @@ struct SingleQueryResultCollectHandler void end() override { if (normalizers) { - float one_a = 1 / normalizers[0]; + float a = normalizers[0]; + float one_a = a != 0 ? 1 / a : 0; float b = normalizers[1]; for (size_t i = 0; i < collect.size(); i++) { collect[i].second = collect[i].second * one_a + b; diff --git a/faiss/impl/index_read.cpp b/faiss/impl/index_read.cpp index e52a6f6e4..fd3a23e2c 100644 --- a/faiss/impl/index_read.cpp +++ b/faiss/impl/index_read.cpp @@ -45,6 +45,7 @@ #include #include #include +#include #include #include #include @@ -57,6 +58,7 @@ #include #include #include +#include #ifdef FAISS_ENABLE_SVS #include #include @@ -2760,6 +2762,33 @@ std::unique_ptr read_index_up(IOReader* f, int io_flags) { idxnnd->d); } idx = std::move(idxnnd); + } else if (h == fourcc("ISfs")) { + auto idxsqfs = std::make_unique(); + read_index_header(*idxsqfs, f); + read_ScalarQuantizer(&idxsqfs->sq, f, *idxsqfs); + READ1(idxsqfs->implem); + READ1(idxsqfs->bbs); + READ1(idxsqfs->qbs); + FAISS_THROW_IF_NOT_MSG(idxsqfs->qbs >= 0, "qbs must be non-negative"); + READ1(idxsqfs->ntotal2); + READ1(idxsqfs->M2); + READVECTOR(idxsqfs->codes); + + // Restore FastScan base-class fields from the SQ + idxsqfs->M = idxsqfs->sq.d; + idxsqfs->nbits = 4; + idxsqfs->ksub = 16; + idxsqfs->code_size = idxsqfs->M2 / 2; + + validate_fastscan_fields( + idxsqfs->M, + idxsqfs->M2, + idxsqfs->ksub, + idxsqfs->bbs, + "IndexSQFastScan"); + + idx = std::move(idxsqfs); + } else if (h == fourcc("IPfs")) { auto idxpqfs = std::make_unique(); read_index_header(*idxpqfs, f); @@ -2787,6 +2816,46 @@ std::unique_ptr read_index_up(IOReader* f, int io_flags) { idx = std::move(idxpqfs); + } else if (h == fourcc("IwSf")) { + auto ivfsqfs = std::make_unique(); + read_ivf_header(ivfsqfs.get(), f); + READ1_BOOL(ivfsqfs->by_residual); + READ1(ivfsqfs->code_size); + READ1(ivfsqfs->bbs); + READ1(ivfsqfs->M2); + READ1(ivfsqfs->implem); + READ1(ivfsqfs->rerank_factor); + read_ScalarQuantizer(&ivfsqfs->sq, f, *ivfsqfs); + read_InvertedLists(*ivfsqfs, f, io_flags); + + ivfsqfs->M = ivfsqfs->d; + ivfsqfs->nbits = 4; + ivfsqfs->ksub = 16; + ivfsqfs->init_code_packer(); + + bool has_orig; + READ1(has_orig); + if (has_orig) { + ivfsqfs->orig_codes_invlists = read_InvertedLists(f, io_flags); + // Rebuild direct_map from orig_codes_invlists for reranking + ivfsqfs->direct_map.set_type( + DirectMap::Hashtable, ivfsqfs->invlists, 0); + for (size_t list_no = 0; list_no < ivfsqfs->nlist; list_no++) { + size_t list_size = + ivfsqfs->orig_codes_invlists->list_size(list_no); + if (list_size == 0) { + continue; + } + InvertedLists::ScopedIds ids( + ivfsqfs->orig_codes_invlists, list_no); + for (size_t j = 0; j < list_size; j++) { + ivfsqfs->direct_map.add_single_id(ids[j], list_no, j); + } + } + } + + idx = std::move(ivfsqfs); + } else if (h == fourcc("IwPf")) { auto ivpq = std::make_unique(); read_ivf_header(ivpq.get(), f); diff --git a/faiss/impl/index_write.cpp b/faiss/impl/index_write.cpp index a0b7abc6b..0ffe46185 100644 --- a/faiss/impl/index_write.cpp +++ b/faiss/impl/index_write.cpp @@ -40,6 +40,7 @@ #include #include #include +#include #include #include #include @@ -52,6 +53,7 @@ #include #include #include +#include #ifdef FAISS_ENABLE_SVS #include #include @@ -951,6 +953,19 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) { write_index_header(idxnnd, f); write_NNDescent(&idxnnd->nndescent, f); write_index(idxnnd->storage, f); + } else if ( + const IndexSQFastScan* idxsqfs = + dynamic_cast(idx)) { + uint32_t h = fourcc("ISfs"); + WRITE1(h); + write_index_header(idxsqfs, f); + write_ScalarQuantizer(&idxsqfs->sq, f); + WRITE1(idxsqfs->implem); + WRITE1(idxsqfs->bbs); + WRITE1(idxsqfs->qbs); + WRITE1(idxsqfs->ntotal2); + WRITE1(idxsqfs->M2); + WRITEVECTOR(idxsqfs->codes); } else if ( const IndexPQFastScan* idxpqfs = dynamic_cast(idx)) { @@ -964,6 +979,26 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) { WRITE1(idxpqfs->ntotal2); WRITE1(idxpqfs->M2); WRITEVECTOR(idxpqfs->codes); + } else if ( + const IndexIVFSQFastScan* ivfsqfs = + dynamic_cast(idx)) { + uint32_t h = fourcc("IwSf"); + WRITE1(h); + write_ivf_header(ivfsqfs, f); + WRITE1(ivfsqfs->by_residual); + WRITE1(ivfsqfs->code_size); + WRITE1(ivfsqfs->bbs); + WRITE1(ivfsqfs->M2); + WRITE1(ivfsqfs->implem); + WRITE1(ivfsqfs->rerank_factor); + write_ScalarQuantizer(&ivfsqfs->sq, f); + write_InvertedLists(ivfsqfs->invlists, f); + // Write orig_codes_invlists if present + bool has_orig = (ivfsqfs->orig_codes_invlists != nullptr); + WRITE1(has_orig); + if (has_orig) { + write_InvertedLists(ivfsqfs->orig_codes_invlists, f); + } } else if ( const IndexIVFPQFastScan* ivpq_2 = dynamic_cast(idx)) { diff --git a/faiss/index_factory.cpp b/faiss/index_factory.cpp index cdbd0bea0..94fa1a45c 100644 --- a/faiss/index_factory.cpp +++ b/faiss/index_factory.cpp @@ -35,6 +35,7 @@ #include #include #include +#include #include #include #include @@ -46,6 +47,7 @@ #include #include #include +#include #include #include @@ -177,6 +179,11 @@ std::map sq_types = { }; const std::string sq_pattern = "(SQ0|SQ4|SQ8|SQ6|SQfp16|SQbf16|SQ8_direct_signed|SQ8_direct|SQtqmse1|SQtqmse2|SQtqmse3|SQtqmse4|SQtqmse8|SQtq2|SQtq3|SQtq4|SQtq5)"; +// Native 4-bit types with "fs" suffix for IndexSQFastScan, optional _bbs +const std::string sq_fs_pattern = "(SQ4)fs(_[0-9]+)?"; +// All SQ types with "fs" suffix for IndexIVFSQFastScan, optional _bbs +const std::string ivf_sq_fs_pattern = + "(SQ0|SQ4|SQ8|SQ6|SQfp16|SQbf16|SQ8_direct_signed|SQ8_direct|SQtqmse1|SQtqmse2|SQtqmse3|SQtqmse4|SQtqmse8|SQtq2|SQtq3|SQtq4|SQtq5)fs(_[0-9]+)?"; std::map aq_search_type = { {"_Nfloat", AdditiveQuantizer::ST_norm_float}, @@ -389,6 +396,17 @@ IndexIVF* parse_IndexIVF( return new IndexIVFFlatPanorama( get_q(), d, nlist, nlevels, mt, own_il, bs); } + if (match(ivf_sq_fs_pattern)) { + int bbs = mres_to_int(sm[2], 32, 1); + return new IndexIVFSQFastScan( + get_q(), + d, + nlist, + sq_types[sm[1].str()], + mt, + bbs, + /*by_residual=*/true); + } if (match(sq_pattern)) { return new IndexIVFScalarQuantizer( get_q(), @@ -831,6 +849,12 @@ Index* parse_other_indexes( return new IndexLattice(d, M, nbit, r2); } + // IndexSQFastScan (must be checked before IndexScalarQuantizer) + if (match(sq_fs_pattern)) { + int bbs = mres_to_int(sm[2], 32, 1); + return new IndexSQFastScan(d, sq_types[sm[1].str()], metric, bbs); + } + // IndexScalarQuantizer if (match(sq_pattern)) { return new IndexScalarQuantizer(d, sq_types[description], metric); diff --git a/faiss/python/__init__.py b/faiss/python/__init__.py index 1f2409d17..b6f54da9e 100644 --- a/faiss/python/__init__.py +++ b/faiss/python/__init__.py @@ -340,6 +340,7 @@ add_ref_in_method_explicit_own(IndexIVFSpectralHash, "replace_vt") add_ref_in_constructor(Index2Layer, 0) add_ref_in_constructor(Level1Quantizer, 0) add_ref_in_constructor(IndexIVFScalarQuantizer, 0) +add_ref_in_constructor(IndexIVFSQFastScan, 0) add_ref_in_constructor(IndexRowwiseMinMax, 0) add_ref_in_constructor(IndexRowwiseMinMaxFP16, 0) add_ref_in_constructor(IndexIDMap, 0) diff --git a/faiss/python/swigfaiss.swig b/faiss/python/swigfaiss.swig index d3d18e2f1..f43a29122 100644 --- a/faiss/python/swigfaiss.swig +++ b/faiss/python/swigfaiss.swig @@ -105,6 +105,8 @@ typedef uint64_t size_t; #include #include #include +#include +#include #include #include #include @@ -681,12 +683,14 @@ void gpu_sync_all_devices() %include %include %include +%include // NOTE(matthijs) simd_result_handlers.h uses simd16uint16 which is now a // class template. SWIG cannot handle this, and the Python API doesn't need // the internal SIMD handler types. // %include %include +%include %include %include @@ -868,6 +872,7 @@ typedef int cudaDataType_t; DOWNCAST ( IndexIVFPQR ) DOWNCAST ( IndexIVFPQ ) DOWNCAST ( IndexIVFPQFastScan ) + DOWNCAST ( IndexIVFSQFastScan ) DOWNCAST ( IndexIVFSpectralHash ) DOWNCAST ( IndexIVFScalarQuantizer ) DOWNCAST ( IndexIVFResidualQuantizer ) @@ -891,6 +896,8 @@ typedef int cudaDataType_t; DOWNCAST ( IndexRefineFlat ) DOWNCAST ( IndexRefine ) DOWNCAST ( IndexPQFastScan ) + DOWNCAST ( IndexSQFastScan ) + DOWNCAST ( IndexIVFSQFastScan ) DOWNCAST ( IndexPQ ) DOWNCAST ( IndexResidualQuantizer ) DOWNCAST ( IndexLocalSearchQuantizer ) diff --git a/tests/test_ivf_sq_fastscan.py b/tests/test_ivf_sq_fastscan.py new file mode 100644 index 000000000..0a066b0ce --- /dev/null +++ b/tests/test_ivf_sq_fastscan.py @@ -0,0 +1,372 @@ +#!/usr/bin/env python3 +# 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. + +"""Tests for IndexIVFSQFastScan.""" + +import numpy as np +import unittest + +import faiss +from faiss.contrib.datasets import SyntheticDataset + + +class TestIndexIVFSQFastScanBasic(unittest.TestCase): + """Basic construction, train, add, search tests.""" + + def test_construct_native_4bit(self): + d, nlist = 32, 8 + quantizer = faiss.IndexFlatL2(d) + idx = faiss.IndexIVFSQFastScan( + quantizer, d, nlist, faiss.ScalarQuantizer.QT_4bit + ) + self.assertEqual(idx.d, d) + self.assertEqual(idx.nlist, nlist) + + def test_construct_rerank_8bit(self): + d, nlist = 32, 8 + quantizer = faiss.IndexFlatL2(d) + idx = faiss.IndexIVFSQFastScan( + quantizer, d, nlist, faiss.ScalarQuantizer.QT_8bit + ) + self.assertEqual(idx.d, d) + + def test_train_add_search_4bit(self): + ds = SyntheticDataset(d=32, nt=1000, nb=1000, nq=10, seed=42) + + quantizer = faiss.IndexFlatL2(ds.d) + idx = faiss.IndexIVFSQFastScan( + quantizer, ds.d, 8, faiss.ScalarQuantizer.QT_4bit + ) + idx.train(ds.get_train()) + self.assertTrue(idx.is_trained) + idx.add(ds.get_database()) + self.assertEqual(idx.ntotal, ds.nb) + + idx.nprobe = 8 + D, I = idx.search(ds.get_queries(), 10) + self.assertEqual(D.shape, (10, 10)) + self.assertEqual(I.shape, (10, 10)) + self.assertTrue(np.all(I >= 0)) + + def test_train_add_search_8bit(self): + ds = SyntheticDataset(d=32, nt=1000, nb=1000, nq=10, seed=42) + + quantizer = faiss.IndexFlatL2(ds.d) + idx = faiss.IndexIVFSQFastScan( + quantizer, ds.d, 8, faiss.ScalarQuantizer.QT_8bit + ) + idx.train(ds.get_train()) + idx.add(ds.get_database()) + self.assertEqual(idx.ntotal, ds.nb) + + idx.nprobe = 8 + D, I = idx.search(ds.get_queries(), 10) + self.assertTrue(np.all(I >= 0)) + + def test_train_add_search_fp16_fallback(self): + ds = SyntheticDataset(d=32, nt=1000, nb=1000, nq=10, seed=42) + + quantizer = faiss.IndexFlatL2(ds.d) + idx = faiss.IndexIVFSQFastScan( + quantizer, ds.d, 8, faiss.ScalarQuantizer.QT_fp16 + ) + idx.train(ds.get_train()) + idx.add(ds.get_database()) + self.assertEqual(idx.ntotal, ds.nb) + + idx.nprobe = 8 + D, I = idx.search(ds.get_queries(), 10) + self.assertTrue(np.all(I >= 0)) + + +class TestIndexIVFSQFastScanRecall(unittest.TestCase): + """Recall parity tests against IndexIVFScalarQuantizer.""" + + def _compare_recall(self, qtype, nlist=16, nprobe=16): + ds = SyntheticDataset(d=32, nt=2000, nb=2000, nq=50, seed=42) + k = 10 + + # Ground truth + gt_idx = faiss.IndexFlatL2(ds.d) + gt_idx.add(ds.get_database()) + D_gt, I_gt = gt_idx.search(ds.get_queries(), k) + + # Shared coarse quantizer ensures same IVF assignment + quantizer = faiss.IndexFlatL2(ds.d) + + # Reference: IndexIVFScalarQuantizer + ref = faiss.IndexIVFScalarQuantizer( + quantizer, ds.d, nlist, qtype + ) + ref.train(ds.get_train()) + ref.add(ds.get_database()) + ref.nprobe = nprobe + D_ref, I_ref = ref.search(ds.get_queries(), k) + + # Test: IndexIVFSQFastScan with same quantizer and SQ + test = faiss.IndexIVFSQFastScan( + quantizer, ds.d, nlist, qtype + ) + test.sq = ref.sq + test.is_trained = True + test.add(ds.get_database()) + test.nprobe = nprobe + D_test, I_test = test.search(ds.get_queries(), k) + + # Compute recalls + recall_ref = np.mean([ + len(set(I_ref[i]) & set(I_gt[i])) / k + for i in range(ds.nq) + ]) + recall_test = np.mean([ + len(set(I_test[i]) & set(I_gt[i])) / k + for i in range(ds.nq) + ]) + + return recall_ref, recall_test + + def test_recall_4bit(self): + recall_ref, recall_test = self._compare_recall( + faiss.ScalarQuantizer.QT_4bit + ) + self.assertGreater(recall_test, 0.5) + + def test_recall_8bit(self): + recall_ref, recall_test = self._compare_recall( + faiss.ScalarQuantizer.QT_8bit + ) + self.assertGreater(recall_test, 0.7) + + def test_recall_6bit(self): + recall_ref, recall_test = self._compare_recall( + faiss.ScalarQuantizer.QT_6bit + ) + self.assertGreater(recall_test, 0.6) + + def test_recall_fp16(self): + recall_ref, recall_test = self._compare_recall( + faiss.ScalarQuantizer.QT_fp16 + ) + self.assertGreater(recall_test, 0.9) + + +class TestIndexIVFSQFastScanFactory(unittest.TestCase): + """Factory string tests.""" + + def test_factory_sq8fs(self): + d = 32 + idx = faiss.index_factory(d, "IVF16,SQ8fs") + self.assertIsInstance(idx, faiss.IndexIVFSQFastScan) + + def test_factory_sq4fs(self): + d = 32 + idx = faiss.index_factory(d, "IVF16,SQ4fs") + self.assertIsInstance(idx, faiss.IndexIVFSQFastScan) + + def test_factory_sqfp16fs(self): + d = 32 + idx = faiss.index_factory(d, "IVF16,SQfp16fs") + self.assertIsInstance(idx, faiss.IndexIVFSQFastScan) + + def test_factory_sq6fs(self): + d = 32 + idx = faiss.index_factory(d, "IVF16,SQ6fs") + self.assertIsInstance(idx, faiss.IndexIVFSQFastScan) + + def test_factory_with_bbs(self): + d = 32 + idx = faiss.index_factory(d, "IVF16,SQ8fs_64") + self.assertIsInstance(idx, faiss.IndexIVFSQFastScan) + self.assertEqual(idx.bbs, 64) + + def test_factory_train_search(self): + ds = SyntheticDataset(d=32, nt=1000, nb=1000, nq=10, seed=42) + + idx = faiss.index_factory(ds.d, "IVF16,SQ8fs") + idx.train(ds.get_train()) + idx.add(ds.get_database()) + idx.nprobe = 16 + D, I = idx.search(ds.get_queries(), 5) + self.assertTrue(np.all(I >= 0)) + + +class TestIndexIVFSQFastScanIP(unittest.TestCase): + """Inner product metric tests.""" + + def test_inner_product_4bit(self): + ds = SyntheticDataset(d=32, nt=1000, nb=1000, nq=10, seed=42) + + quantizer = faiss.IndexFlatIP(ds.d) + idx = faiss.IndexIVFSQFastScan( + quantizer, ds.d, 8, faiss.ScalarQuantizer.QT_4bit, + faiss.METRIC_INNER_PRODUCT + ) + idx.train(ds.get_train()) + idx.add(ds.get_database()) + idx.nprobe = 8 + D, I = idx.search(ds.get_queries(), 10) + self.assertTrue(np.all(I >= 0)) + + def test_inner_product_8bit(self): + ds = SyntheticDataset(d=32, nt=1000, nb=1000, nq=10, seed=42) + + quantizer = faiss.IndexFlatIP(ds.d) + idx = faiss.IndexIVFSQFastScan( + quantizer, ds.d, 8, faiss.ScalarQuantizer.QT_8bit, + faiss.METRIC_INNER_PRODUCT + ) + idx.train(ds.get_train()) + idx.add(ds.get_database()) + idx.nprobe = 8 + D, I = idx.search(ds.get_queries(), 10) + self.assertTrue(np.all(I >= 0)) + + +class TestIndexIVFSQFastScanIO(unittest.TestCase): + """Serialization round-trip tests.""" + + def test_io_roundtrip_4bit(self): + ds = SyntheticDataset(d=32, nt=500, nb=500, nq=5, seed=42) + + quantizer = faiss.IndexFlatL2(ds.d) + idx = faiss.IndexIVFSQFastScan( + quantizer, ds.d, 8, faiss.ScalarQuantizer.QT_4bit + ) + idx.train(ds.get_train()) + idx.add(ds.get_database()) + idx.nprobe = 8 + + D1, I1 = idx.search(ds.get_queries(), 5) + + idx2 = faiss.deserialize_index(faiss.serialize_index(idx)) + idx2.nprobe = 8 + D2, I2 = idx2.search(ds.get_queries(), 5) + np.testing.assert_array_equal(I1, I2) + np.testing.assert_allclose(D1, D2, rtol=1e-5) + + def test_io_roundtrip_8bit(self): + ds = SyntheticDataset(d=32, nt=500, nb=500, nq=5, seed=42) + + quantizer = faiss.IndexFlatL2(ds.d) + idx = faiss.IndexIVFSQFastScan( + quantizer, ds.d, 8, faiss.ScalarQuantizer.QT_8bit + ) + idx.train(ds.get_train()) + idx.add(ds.get_database()) + idx.nprobe = 8 + + D1, I1 = idx.search(ds.get_queries(), 5) + + idx2 = faiss.deserialize_index(faiss.serialize_index(idx)) + idx2.nprobe = 8 + D2, I2 = idx2.search(ds.get_queries(), 5) + np.testing.assert_array_equal(I1, I2) + np.testing.assert_allclose(D1, D2, rtol=1e-5) + + +class TestIndexIVFSQFastScanReset(unittest.TestCase): + """Reset and re-add tests.""" + + def test_reset(self): + ds = SyntheticDataset(d=32, nt=500, nb=500, nq=0, seed=42) + + quantizer = faiss.IndexFlatL2(ds.d) + idx = faiss.IndexIVFSQFastScan( + quantizer, ds.d, 8, faiss.ScalarQuantizer.QT_8bit + ) + idx.train(ds.get_train()) + idx.add(ds.get_database()) + self.assertEqual(idx.ntotal, ds.nb) + + idx.reset() + self.assertEqual(idx.ntotal, 0) + + idx.add(ds.get_database()) + self.assertEqual(idx.ntotal, ds.nb) + + +class TestIndexIVFSQFastScanReconstruct(unittest.TestCase): + """Reconstruction tests.""" + + def test_reconstruct_4bit(self): + ds = SyntheticDataset(d=32, nt=200, nb=200, nq=0, seed=42) + + quantizer = faiss.IndexFlatL2(ds.d) + idx = faiss.IndexIVFSQFastScan( + quantizer, ds.d, 8, faiss.ScalarQuantizer.QT_4bit + ) + idx.train(ds.get_train()) + idx.add(ds.get_database()) + idx.make_direct_map() + + recon = idx.reconstruct(0) + self.assertEqual(recon.shape, (ds.d,)) + err = np.linalg.norm(recon - ds.get_database()[0]) + self.assertLess(err, 5.0) + + def test_reconstruct_8bit(self): + ds = SyntheticDataset(d=32, nt=200, nb=200, nq=0, seed=42) + + quantizer = faiss.IndexFlatL2(ds.d) + idx = faiss.IndexIVFSQFastScan( + quantizer, ds.d, 8, faiss.ScalarQuantizer.QT_8bit + ) + idx.train(ds.get_train()) + idx.add(ds.get_database()) + idx.make_direct_map() + + recon = idx.reconstruct(0) + self.assertEqual(recon.shape, (ds.d,)) + err = np.linalg.norm(recon - ds.get_database()[0]) + self.assertLess(err, 1.0) + + +class TestIndexIVFSQFastScanEdgeCases(unittest.TestCase): + """Edge cases and special scenarios.""" + + def test_single_vector(self): + ds = SyntheticDataset(d=32, nt=100, nb=1, nq=1, seed=42) + + quantizer = faiss.IndexFlatL2(ds.d) + idx = faiss.IndexIVFSQFastScan( + quantizer, ds.d, 4, faiss.ScalarQuantizer.QT_4bit + ) + idx.train(ds.get_train()) + idx.add(ds.get_database()) + idx.nprobe = 4 + D, I = idx.search(ds.get_queries(), 1) + self.assertEqual(I[0, 0], 0) + + def test_k_equals_1(self): + ds = SyntheticDataset(d=32, nt=500, nb=500, nq=10, seed=42) + + quantizer = faiss.IndexFlatL2(ds.d) + idx = faiss.IndexIVFSQFastScan( + quantizer, ds.d, 8, faiss.ScalarQuantizer.QT_8bit + ) + idx.train(ds.get_train()) + idx.add(ds.get_database()) + idx.nprobe = 8 + D, I = idx.search(ds.get_queries(), 1) + self.assertEqual(D.shape, (10, 1)) + self.assertEqual(I.shape, (10, 1)) + + def test_odd_dimension(self): + ds = SyntheticDataset(d=33, nt=500, nb=500, nq=10, seed=42) + + quantizer = faiss.IndexFlatL2(ds.d) + idx = faiss.IndexIVFSQFastScan( + quantizer, ds.d, 8, faiss.ScalarQuantizer.QT_4bit + ) + idx.train(ds.get_train()) + idx.add(ds.get_database()) + idx.nprobe = 8 + D, I = idx.search(ds.get_queries(), 5) + self.assertTrue(np.all(I >= 0)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_sq_fastscan.py b/tests/test_sq_fastscan.py new file mode 100644 index 000000000..9248f70e8 --- /dev/null +++ b/tests/test_sq_fastscan.py @@ -0,0 +1,489 @@ +# 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. + +""" +Unit tests for IndexSQFastScan. + +Tests cover: + - Construction for native 4-bit types only + - Rejection of non-4-bit types + - Train / add / search round-trip + - Recall parity with IndexScalarQuantizer (within expected bounds) + - Conversion constructor from IndexScalarQuantizer + - Reset clears state + - Edge cases: k=1, single vector, zero vectors, odd dimensions + - Inner product metric + - sa_encode / sa_decode consistency + - I/O round-trip + - Factory strings + - get_distance_computer + - range_search +""" + +from __future__ import absolute_import, division, print_function + +import os +import tempfile +import unittest + +import numpy as np + +import faiss +from faiss.contrib.datasets import SyntheticDataset + + +def recall_at_k(I_gt, I_test, k): + """Fraction of queries whose true NN (gt[:,0]) appears in test[:,:k].""" + return np.mean([I_gt[i, 0] in I_test[i, :k] for i in range(len(I_gt))]) + + +# Native 4-bit types (only types supported by IndexSQFastScan) +SQ = faiss.ScalarQuantizer + +NATIVE_4BIT = [ + ("QT_4bit", SQ.QT_4bit), + ("QT_4bit_uniform", SQ.QT_4bit_uniform), +] + + +class TestSQFastScanConstruction(unittest.TestCase): + """IndexSQFastScan can be constructed for native 4-bit types only.""" + + def test_native_4bit_construct(self): + d = 64 + for name, qtype in NATIVE_4BIT: + with self.subTest(qtype=name): + index = faiss.IndexSQFastScan(d, qtype) + self.assertEqual(index.d, d) + self.assertFalse(index.is_trained) + self.assertEqual(index.ntotal, 0) + + def test_default_construct(self): + index = faiss.IndexSQFastScan() + self.assertEqual(index.d, 0) + self.assertEqual(index.ntotal, 0) + + def test_rejects_non_4bit(self): + d = 64 + with self.assertRaises(RuntimeError): + faiss.IndexSQFastScan(d, SQ.QT_8bit) + with self.assertRaises(RuntimeError): + faiss.IndexSQFastScan(d, SQ.QT_6bit) + with self.assertRaises(RuntimeError): + faiss.IndexSQFastScan(d, SQ.QT_fp16) + with self.assertRaises(RuntimeError): + faiss.IndexSQFastScan(d, SQ.QT_bf16) + + +class TestSQFastScanTrainAddSearch(unittest.TestCase): + """Basic train/add/search works for native 4-bit types.""" + + def setUp(self): + self.d = 64 + self.nb = 5000 + self.nq = 50 + self.k = 10 + self.ds = SyntheticDataset(d=self.d, nt=2000, nb=self.nb, + nq=self.nq, seed=42) + + def test_native_4bit(self): + for name, qtype in NATIVE_4BIT: + with self.subTest(qtype=name): + index = faiss.IndexSQFastScan(self.d, qtype) + index.train(self.ds.get_train()) + self.assertTrue(index.is_trained) + + index.add(self.ds.get_database()) + self.assertEqual(index.ntotal, self.nb) + + D, I = index.search(self.ds.get_queries(), self.k) + self.assertEqual(D.shape, (self.nq, self.k)) + self.assertEqual(I.shape, (self.nq, self.k)) + self.assertTrue(np.all(I[:, 0] >= 0)) + self.assertTrue(np.all(np.isfinite(D))) + for q in range(self.nq): + for j in range(1, self.k): + self.assertLessEqual(D[q, j - 1], D[q, j] + 1e-5) + + +class TestSQFastScanRecallParity(unittest.TestCase): + """IndexSQFastScan recall should match IndexScalarQuantizer for 4-bit.""" + + def setUp(self): + self.d = 64 + self.ds = SyntheticDataset(d=self.d, nt=2000, nb=10000, + nq=100, seed=42) + gt_index = faiss.IndexFlatL2(self.d) + gt_index.add(self.ds.get_database()) + self.D_gt, self.I_gt = gt_index.search(self.ds.get_queries(), 10) + + def test_native_4bit_parity(self): + """Native 4-bit: FastScan should match SQ recall exactly.""" + for name, qtype in NATIVE_4BIT: + with self.subTest(qtype=name): + sq = faiss.IndexScalarQuantizer(self.d, qtype) + sq.train(self.ds.get_train()) + sq.add(self.ds.get_database()) + + fs = faiss.IndexSQFastScan(self.d, qtype) + fs.train(self.ds.get_train()) + fs.add(self.ds.get_database()) + + _, I_sq = sq.search(self.ds.get_queries(), 10) + _, I_fs = fs.search(self.ds.get_queries(), 10) + r_sq = recall_at_k(self.I_gt, I_sq, 1) + r_fs = recall_at_k(self.I_gt, I_fs, 1) + self.assertAlmostEqual(r_sq, r_fs, delta=0.05, + msg=f"{name}: SQ={r_sq:.3f} FS={r_fs:.3f}") + + +class TestSQFastScanConversionConstructor(unittest.TestCase): + """Conversion from IndexScalarQuantizer preserves search results.""" + + def setUp(self): + self.d = 64 + self.ds = SyntheticDataset(d=self.d, nt=2000, nb=5000, + nq=50, seed=42) + + def test_conversion_native_4bit(self): + for name, qtype in NATIVE_4BIT: + with self.subTest(qtype=name): + sq = faiss.IndexScalarQuantizer(self.d, qtype) + sq.train(self.ds.get_train()) + sq.add(self.ds.get_database()) + + fs = faiss.IndexSQFastScan(sq) + self.assertEqual(fs.ntotal, sq.ntotal) + self.assertTrue(fs.is_trained) + + D_sq, I_sq = sq.search(self.ds.get_queries(), 10) + D_fs, I_fs = fs.search(self.ds.get_queries(), 10) + r1_sq = (I_sq[:, 0] >= 0).mean() + r1_fs = (I_fs[:, 0] >= 0).mean() + self.assertEqual(r1_sq, r1_fs) + + def test_conversion_rejects_non_4bit(self): + sq = faiss.IndexScalarQuantizer(self.d, SQ.QT_8bit) + sq.train(self.ds.get_train()) + with self.assertRaises(RuntimeError): + faiss.IndexSQFastScan(sq) + + +class TestSQFastScanReset(unittest.TestCase): + """Reset should clear all stored data.""" + + def test_reset(self): + d = 64 + ds = SyntheticDataset(d=d, nt=1000, nb=500, nq=0, seed=42) + for name, qtype in NATIVE_4BIT: + with self.subTest(qtype=name): + index = faiss.IndexSQFastScan(d, qtype) + index.train(ds.get_train()) + index.add(ds.get_database()) + self.assertEqual(index.ntotal, 500) + + index.reset() + self.assertEqual(index.ntotal, 0) + + index.add(ds.get_database()) + self.assertEqual(index.ntotal, 500) + + +class TestSQFastScanEdgeCases(unittest.TestCase): + """Edge cases: k=1, single vector, zero vectors, odd dimensions.""" + + def test_k_equals_1(self): + d = 32 + ds = SyntheticDataset(d=d, nt=500, nb=1000, nq=10, seed=42) + for name, qtype in NATIVE_4BIT: + with self.subTest(qtype=name): + index = faiss.IndexSQFastScan(d, qtype) + index.train(ds.get_train()) + index.add(ds.get_database()) + D, I = index.search(ds.get_queries(), 1) + self.assertEqual(D.shape, (10, 1)) + self.assertEqual(I.shape, (10, 1)) + self.assertTrue(np.all(I >= 0)) + + def test_single_vector(self): + d = 32 + xb = np.random.RandomState(42).randn(1, d).astype('float32') + xq = xb.copy() + for name, qtype in NATIVE_4BIT: + with self.subTest(qtype=name): + index = faiss.IndexSQFastScan(d, qtype) + index.train(xb) + index.add(xb) + D, I = index.search(xq, 1) + self.assertEqual(I[0, 0], 0) + self.assertLess(D[0, 0], 1.0) + + def test_zero_vectors(self): + d = 32 + xb = np.zeros((100, d), dtype='float32') + for name, qtype in NATIVE_4BIT: + with self.subTest(qtype=name): + index = faiss.IndexSQFastScan(d, qtype) + index.train(xb) + index.add(xb) + D, I = index.search(np.zeros((1, d), dtype='float32'), 10) + self.assertTrue(np.all(np.isfinite(D))) + + def test_odd_dimensions(self): + """Dimensions not aligned to SIMD width.""" + for d in [7, 9, 15, 17, 31, 33, 63, 65]: + with self.subTest(d=d): + ds = SyntheticDataset(d=d, nt=500, nb=500, nq=5, seed=42) + index = faiss.IndexSQFastScan(d, SQ.QT_4bit) + index.train(ds.get_train()) + index.add(ds.get_database()) + D, I = index.search(ds.get_queries(), 10) + self.assertEqual(D.shape, (5, 10)) + self.assertTrue(np.all(I[:, 0] >= 0)) + + +class TestSQFastScanInnerProduct(unittest.TestCase): + """Inner product metric works correctly.""" + + def test_ip_basic(self): + d = 64 + ds = SyntheticDataset(d=d, nt=2000, nb=5000, nq=50, seed=42) + xb = ds.get_database().copy() + xq = ds.get_queries().copy() + faiss.normalize_L2(xb) + faiss.normalize_L2(xq) + + for name, qtype in NATIVE_4BIT: + with self.subTest(qtype=name): + index = faiss.IndexSQFastScan( + d, qtype, faiss.METRIC_INNER_PRODUCT) + index.train(xb) + index.add(xb) + D, I = index.search(xq, 10) + self.assertEqual(D.shape, (50, 10)) + self.assertTrue(np.all(I[:, 0] >= 0)) + self.assertTrue(np.all(np.isfinite(D))) + for q in range(50): + for j in range(1, 10): + self.assertGreaterEqual( + D[q, j - 1] + 1e-5, D[q, j]) + + +class TestSQFastScanSaDecode(unittest.TestCase): + """sa_decode should reconstruct vectors close to the original.""" + + def test_sa_decode(self): + d = 64 + ds = SyntheticDataset(d=d, nt=1000, nb=100, nq=0, seed=42) + for name, qtype in NATIVE_4BIT: + with self.subTest(qtype=name): + index = faiss.IndexSQFastScan(d, qtype) + index.train(ds.get_train()) + index.add(ds.get_database()) + + xr = index.reconstruct_n(0, index.ntotal) + self.assertEqual(xr.shape, (100, d)) + self.assertTrue(np.all(np.isfinite(xr))) + + +class TestSQFastScanIncrementalAdd(unittest.TestCase): + """Adding vectors in batches should give same results as all-at-once.""" + + def test_incremental_add(self): + d = 64 + ds = SyntheticDataset(d=d, nt=1000, nb=1000, nq=20, seed=42) + xb = ds.get_database() + xq = ds.get_queries() + + for name, qtype in NATIVE_4BIT: + with self.subTest(qtype=name): + idx1 = faiss.IndexSQFastScan(d, qtype) + idx1.train(ds.get_train()) + idx1.add(xb) + + idx2 = faiss.IndexSQFastScan(d, qtype) + idx2.train(ds.get_train()) + idx2.add(xb[:500]) + idx2.add(xb[500:]) + + self.assertEqual(idx1.ntotal, idx2.ntotal) + + D1, I1 = idx1.search(xq, 10) + D2, I2 = idx2.search(xq, 10) + + np.testing.assert_array_equal(I1, I2) + np.testing.assert_allclose(D1, D2, atol=1e-5) + + +class TestSQFastScanSaEncode(unittest.TestCase): + """sa_encode / sa_decode round-trip matches IndexScalarQuantizer.""" + + def test_sa_encode_matches_sq(self): + d = 64 + ds = SyntheticDataset(d=d, nt=1000, nb=100, nq=0, seed=42) + for name, qtype in NATIVE_4BIT: + with self.subTest(qtype=name): + sq = faiss.IndexScalarQuantizer(d, qtype) + sq.train(ds.get_train()) + + fs = faiss.IndexSQFastScan(d, qtype) + fs.train(ds.get_train()) + + xb = ds.get_database() + cs = fs.sa_code_size() + + sq_codes = sq.sa_encode(xb) + fs_codes = fs.sa_encode(xb) + np.testing.assert_array_equal(sq_codes, fs_codes, + err_msg=f"{name}: sa_encode differs") + + sq_dec = sq.sa_decode(sq_codes) + fs_dec = fs.sa_decode(fs_codes) + np.testing.assert_allclose(sq_dec, fs_dec, atol=1e-6, + err_msg=f"{name}: sa_decode differs") + + +class TestSQFastScanAddSaCodes(unittest.TestCase): + """add_sa_codes should produce working indexes.""" + + def test_add_sa_codes(self): + d = 64 + ds = SyntheticDataset(d=d, nt=1000, nb=500, nq=20, seed=42) + for name, qtype in NATIVE_4BIT: + with self.subTest(qtype=name): + ref = faiss.IndexSQFastScan(d, qtype) + ref.train(ds.get_train()) + ref.add(ds.get_database()) + + test = faiss.IndexSQFastScan(d, qtype) + test.train(ds.get_train()) + codes = test.sa_encode(ds.get_database()) + test.add_sa_codes(codes) + + self.assertEqual(ref.ntotal, test.ntotal) + + D_ref, I_ref = ref.search(ds.get_queries(), 10) + D_test, I_test = test.search(ds.get_queries(), 10) + self.assertTrue(np.all(I_ref[:, 0] >= 0)) + self.assertTrue(np.all(I_test[:, 0] >= 0)) + + +class TestSQFastScanPermuteEntries(unittest.TestCase): + """permute_entries reorders the packed codes.""" + + def test_permute_entries(self): + d, nb = 64, 100 + ds = SyntheticDataset(d=d, nt=1000, nb=nb, nq=0, seed=42) + for name, qtype in NATIVE_4BIT: + with self.subTest(qtype=name): + index = faiss.IndexSQFastScan(d, qtype) + index.train(ds.get_train()) + index.add(ds.get_database()) + + before = index.reconstruct_n(0, nb) + perm = np.arange(nb - 1, -1, -1).astype("int64") + index.permute_entries(perm) + + np.testing.assert_allclose( + index.reconstruct_n(0, nb), before[perm], atol=1e-5) + + +class TestSQFastScanDistanceComputer(unittest.TestCase): + """get_distance_computer works for native 4-bit types.""" + + def test_distance_computer(self): + d = 64 + ds = SyntheticDataset(d=d, nt=1000, nb=500, nq=0, seed=42) + for name, qtype in NATIVE_4BIT: + with self.subTest(qtype=name): + index = faiss.IndexSQFastScan(d, qtype) + index.train(ds.get_train()) + index.add(ds.get_database()) + dc = index.get_distance_computer() + self.assertIsNotNone(dc) + + +class TestSQFastScanRangeSearch(unittest.TestCase): + """range_search produces results consistent with kNN search.""" + + def test_range_search(self): + d = 64 + ds = SyntheticDataset(d=d, nt=1000, nb=1000, nq=10, seed=42) + for name, qtype in NATIVE_4BIT: + with self.subTest(qtype=name): + sq = faiss.IndexScalarQuantizer(d, qtype) + sq.train(ds.get_train()) + sq.add(ds.get_database()) + + fs = faiss.IndexSQFastScan(d, qtype) + fs.train(ds.get_train()) + fs.add(ds.get_database()) + + D, _ = sq.search(ds.get_queries()[:1], 1000) + radius = float(D[0, 500]) + + lims_sq, D_sq, I_sq = sq.range_search( + ds.get_queries(), radius) + lims_fs, D_fs, I_fs = fs.range_search( + ds.get_queries(), radius) + + np.testing.assert_array_equal(lims_sq, lims_fs, + err_msg=f"{name}: lims differ") + + +class TestSQFastScanIO(unittest.TestCase): + """I/O round-trip preserves index state and search results.""" + + def test_io_roundtrip(self): + d = 64 + ds = SyntheticDataset(d=d, nt=1000, nb=500, nq=20, seed=42) + for name, qtype in NATIVE_4BIT: + with self.subTest(qtype=name): + index = faiss.IndexSQFastScan(d, qtype) + index.train(ds.get_train()) + index.add(ds.get_database()) + + D1, I1 = index.search(ds.get_queries(), 10) + + with tempfile.NamedTemporaryFile(suffix='.idx', delete=False) as f: + fname = f.name + try: + faiss.write_index(index, fname) + loaded = faiss.read_index(fname) + finally: + os.unlink(fname) + + self.assertIsInstance(loaded, faiss.IndexSQFastScan) + self.assertEqual(loaded.ntotal, index.ntotal) + self.assertEqual(loaded.d, d) + + D2, I2 = loaded.search(ds.get_queries(), 10) + np.testing.assert_array_equal(I1, I2) + np.testing.assert_allclose(D1, D2, atol=1e-5) + + +class TestSQFastScanFactory(unittest.TestCase): + """Factory string produces IndexSQFastScan.""" + + def test_factory_strings(self): + d = 64 + index = faiss.index_factory(d, "SQ4fs") + self.assertIsInstance(index, faiss.IndexSQFastScan) + self.assertEqual(index.d, d) + + def test_factory_with_bbs(self): + d = 64 + index = faiss.index_factory(d, "SQ4fs_64") + self.assertIsInstance(index, faiss.IndexSQFastScan) + self.assertEqual(index.bbs, 64) + + def test_factory_metric(self): + d = 64 + index = faiss.index_factory(d, "SQ4fs", faiss.METRIC_INNER_PRODUCT) + self.assertIsInstance(index, faiss.IndexSQFastScan) + + +if __name__ == '__main__': + unittest.main()