diff --git a/faiss/IndexIVFSQFastScan.cpp b/faiss/IndexIVFSQFastScan.cpp index 855fed48d..7ed7cfac9 100644 --- a/faiss/IndexIVFSQFastScan.cpp +++ b/faiss/IndexIVFSQFastScan.cpp @@ -6,6 +6,7 @@ */ #include +#include #include @@ -29,107 +30,7 @@ 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 +using namespace faiss::sq_fastscan; // ----------------------------------------------------------------------- // Constructors @@ -500,10 +401,10 @@ void IndexIVFSQFastScan::search( 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; - } + idx_t k_coarse = idx_t(double(k) * rerank_factor); + FAISS_THROW_IF_NOT(k_coarse >= k); + FAISS_THROW_IF_NOT_MSG( + n <= INT64_MAX / k_coarse, "n * k_coarse would overflow int64"); std::vector coarse_dis(n * k_coarse); std::vector coarse_ids(n * k_coarse); @@ -592,10 +493,10 @@ void IndexIVFSQFastScan::search_preassigned( params, stats); } else if (needs_rerank(sq.qtype)) { - idx_t k_coarse = idx_t(k * rerank_factor); - if (k_coarse < k) { - k_coarse = k; - } + idx_t k_coarse = idx_t(double(k) * rerank_factor); + FAISS_THROW_IF_NOT(k_coarse >= k); + FAISS_THROW_IF_NOT_MSG( + n <= INT64_MAX / k_coarse, "n * k_coarse would overflow int64"); std::vector coarse_dis(n * k_coarse); std::vector coarse_ids(n * k_coarse); diff --git a/faiss/IndexRaBitQFastScan.cpp b/faiss/IndexRaBitQFastScan.cpp index 3bd6f1111..f0c8a990d 100644 --- a/faiss/IndexRaBitQFastScan.cpp +++ b/faiss/IndexRaBitQFastScan.cpp @@ -512,8 +512,8 @@ void IndexRaBitQFastScan::search( float* distances, idx_t* labels, const SearchParameters* params) const { - // The kernel gains its selector test in the next diff. - FAISS_THROW_IF_MSG(params, "search params not supported for this index"); + // The kernel applies the selector after each SIMD block. + const IDSelector* rabitq_sel = params ? params->sel : nullptr; // Create query factors array on stack - memory managed by caller std::vector query_factors_storage(n); @@ -524,9 +524,11 @@ void IndexRaBitQFastScan::search( FastScanDistancePostProcessing context; context.query_factors = query_factors_storage.data(); if (metric_type == METRIC_L2) { - search_dispatch_implem(n, x, k, distances, labels, context); + search_dispatch_implem( + n, x, k, distances, labels, context, rabitq_sel); } else { - search_dispatch_implem(n, x, k, distances, labels, context); + search_dispatch_implem( + n, x, k, distances, labels, context, rabitq_sel); } } diff --git a/faiss/IndexRaBitQFastScan.h b/faiss/IndexRaBitQFastScan.h index 8a1891cd6..06c22256c 100644 --- a/faiss/IndexRaBitQFastScan.h +++ b/faiss/IndexRaBitQFastScan.h @@ -218,6 +218,10 @@ struct RaBitQHeapHandler for (size_t i = 0; i < max_vectors; i++) { const size_t db_idx = base_db_idx + i; + if (this->sel != nullptr && + !this->sel->is_member(static_cast(db_idx))) { + continue; + } const float normalized_distance = d32tab[i] * one_a + bias; const uint8_t* base_ptr = aux_base + i * storage_size; diff --git a/faiss/IndexSQFastScan.cpp b/faiss/IndexSQFastScan.cpp index 4bf99674f..591a87cec 100644 --- a/faiss/IndexSQFastScan.cpp +++ b/faiss/IndexSQFastScan.cpp @@ -6,6 +6,7 @@ */ #include +#include #include @@ -28,18 +29,7 @@ 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 +using namespace faiss::sq_fastscan; // ----------------------------------------------------------------------- // Constructors @@ -52,10 +42,10 @@ IndexSQFastScan::IndexSQFastScan( 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(...))."); + !is_fallback(qtype), + "IndexSQFastScan supports the native 4-bit and the reranked " + "types. QT_fp16 and QT_bf16 have no fast-scan path; use " + "IndexScalarQuantizer for those."); init_fastscan(d_in, d_in, 4, metric, bbs_in); } @@ -64,17 +54,39 @@ 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."); + !is_fallback(sq.qtype), + "IndexSQFastScan conversion constructor supports the native " + "4-bit and the reranked types only."); 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()); + if (is_native_4bit(sq.qtype)) { + pq4_pack_codes( + orig.codes.data(), ntotal, M, ntotal2, bbs_in, M2, codes.get()); + } else { + // High nibbles feed the scan, low nibbles the rerank. + std::vector nibbles(ntotal * (M2 / 2)); + const int bits = sq_bits(sq.qtype), lo_bits = bits - 4; + const size_t half = sq_rerank_size(d, sq.qtype); + lo_codes.assign(size_t(ntotal) * half, 0); + for (idx_t i = 0; i < ntotal; i++) { + const uint8_t* src = orig.codes.data() + size_t(i) * sq.code_size; + uint8_t* hi = nibbles.data() + size_t(i) * (M2 / 2); + BitstringReader bsr(src, sq.code_size); + BitstringWriter bsw(lo_codes.data() + size_t(i) * half, half); + for (int m = 0; m < d; m++) { + const int code = int(bsr.read(bits)); + const int h = code >> lo_bits; + hi[m / 2] |= (m & 1) ? uint8_t(h << 4) : uint8_t(h); + bsw.write(code & ((1 << lo_bits) - 1), lo_bits); + } + } + pq4_pack_codes( + nibbles.data(), ntotal, M, ntotal2, bbs_in, M2, codes.get()); + } } // ----------------------------------------------------------------------- @@ -91,10 +103,31 @@ void IndexSQFastScan::train(idx_t n, const float* x) { void IndexSQFastScan::add(idx_t n, const float* x) { FAISS_THROW_IF_NOT(is_trained); + const size_t before = lo_codes.size(); IndexFastScan::add(n, x); + if (!needs_rerank(sq.qtype) || lo_codes.size() != before) { + // IndexFastScan::add blocks a large batch and re-enters this override + // once per block, so the blocks have already appended the low bits. + return; + } + const int bits = sq_bits(sq.qtype), lo_bits = bits - 4; + const size_t stride = sq_rerank_size(d, sq.qtype); + lo_codes.resize(before + size_t(n) * stride, 0); + const SQRanges ranges(sq, d); + for (idx_t i = 0; i < n; i++) { + const float* xi = x + i * d; + uint8_t* dst = lo_codes.data() + before + size_t(i) * stride; + BitstringWriter bsw(dst, stride); + for (int m = 0; m < d; m++) { + const int code = + sq_code_of(xi[m], ranges.vmin(m), ranges.vdiff(m), bits); + bsw.write(code & ((1 << lo_bits) - 1), lo_bits); + } + } } void IndexSQFastScan::reset() { + lo_codes.clear(); IndexFastScan::reset(); } @@ -112,7 +145,87 @@ void IndexSQFastScan::search( FAISS_THROW_IF_NOT(k > 0); FAISS_THROW_IF_NOT(is_trained); // The kernel applies the selector itself, so no generic-scanner fallback. - IndexFastScan::search(n, x, k, distances, labels, params); + if (!needs_rerank(sq.qtype)) { + IndexFastScan::search(n, x, k, distances, labels, params); + return; + } + + // The 4-bit scan is a shortlist. rerank_factor bounds the recall. + idx_t k_coarse = idx_t(double(k) * rerank_factor); + FAISS_THROW_IF_NOT(k_coarse >= k); + FAISS_THROW_IF_NOT_MSG( + n <= INT64_MAX / k_coarse, "n * k_coarse would overflow int64"); + std::vector coarse_dis(n * k_coarse); + std::vector coarse_ids(n * k_coarse); + IndexFastScan::search( + n, x, k_coarse, coarse_dis.data(), coarse_ids.data(), params); + + const bool is_max = metric_type == METRIC_L2; +#pragma omp parallel + { + std::unique_ptr packer(get_CodePacker()); + std::vector hi_buf(M2 / 2); + std::vector code_buf(d); + const size_t lo_stride = sq_rerank_size(d, sq.qtype); + const int maxv = (1 << sq_bits(sq.qtype)) - 1; + const SQRanges ranges(sq, d); +#pragma omp for + for (idx_t i = 0; i < n; i++) { + float* D = distances + i * k; + idx_t* I = labels + i * k; + if (is_max) { + maxheap_heapify(k, D, I); + } else { + minheap_heapify(k, D, I); + } + for (idx_t j = 0; j < k_coarse; j++) { + idx_t id = coarse_ids[i * k_coarse + j]; + if (id < 0) { + continue; + } + float dis; + { + // Recombine the two halves into the original code. + packer->unpack_1(codes.get(), id, hi_buf.data()); + split_to_codes( + hi_buf.data(), + lo_codes.data() + size_t(id) * lo_stride, + d, + sq.qtype, + code_buf.data()); + const float* xi = x + i * d; + float acc = 0; + for (int m = 0; m < d; m++) { + // Matches Codec*bit::decode_component: (c+0.5)/maxv. + const float v = ranges.vmin(m) + + ((code_buf[m] + 0.5f) / float(maxv)) * + ranges.vdiff(m); + if (is_max) { + const float df = xi[m] - v; + acc += df * df; + } else { + acc += xi[m] * v; + } + } + dis = acc; + } + if (is_max) { + if (dis < D[0]) { + maxheap_replace_top(k, D, I, dis, id); + } + } else { + if (dis > D[0]) { + minheap_replace_top(k, D, I, dis, id); + } + } + } + if (is_max) { + maxheap_reorder(k, D, I); + } else { + minheap_reorder(k, D, I); + } + } + } } // ----------------------------------------------------------------------- @@ -121,7 +234,14 @@ void IndexSQFastScan::search( void IndexSQFastScan::compute_codes(uint8_t* out_codes, idx_t n, const float* x) const { - sq.compute_codes(x, out_codes, n); + if (is_native_4bit(sq.qtype)) { + sq.compute_codes(x, out_codes, n); + } else { + // The scanned nibble is the top half. lo_codes holds the bottom. + std::vector lo_discard( + size_t(n) * sq_rerank_size(d, sq.qtype)); + float_to_split_codes(x, out_codes, lo_discard.data(), n, d, M2, sq); + } } // ----------------------------------------------------------------------- @@ -136,15 +256,21 @@ void IndexSQFastScan::compute_float_LUT( const size_t dim = d; const size_t ksub_val = 16; - bool is_uniform = (sq.qtype == ScalarQuantizer::QT_4bit_uniform); + bool is_uniform = (sq.qtype == ScalarQuantizer::QT_4bit_uniform) || + is_uniform_range(sq.qtype); std::vector recon_table(dim * ksub_val); + const float lut_step = float(1 << sq_rerank_bits(sq.qtype)); + const float lut_offset = lut_step / 2.0f; + const float lut_scale = float((1 << sq_bits(sq.qtype)) - 1); + if (is_uniform) { - float vmin = sq.trained[0]; - float vdiff = sq.trained[1]; + 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; + float recon = + vmin + ((c * lut_step + lut_offset) / lut_scale) * vdiff; for (size_t m = 0; m < dim; m++) { recon_table[m * ksub_val + c] = recon; } @@ -155,8 +281,8 @@ void IndexSQFastScan::compute_float_LUT( 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]; + recon_table[m * ksub_val + c] = vmin[m] + + ((c * lut_step + lut_offset) / lut_scale) * vdiff[m]; } } } @@ -195,11 +321,37 @@ void IndexSQFastScan::sa_decode(idx_t n, const uint8_t* bytes, float* x) const { sq.decode(bytes, x, n); } +void IndexSQFastScan::fill_sa_code( + const CodePacker& packer, + idx_t id, + uint8_t* scratch, + int* values, + uint8_t* code_out) const { + packer.unpack_1(codes.data(), id, scratch); + if (lo_codes.empty()) { + memcpy(code_out, scratch, sq.code_size); + return; + } + split_to_codes( + scratch, + lo_codes.data() + size_t(id) * sq_rerank_size(d, sq.qtype), + d, + sq.qtype, + values); + memset(code_out, 0, sq.code_size); + BitstringWriter out(code_out, sq.code_size); + const int bits = sq_bits(sq.qtype); + for (int m = 0; m < d; m++) { + out.write(values[m], bits); + } +} + 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()); + std::vector scratch(M2 / 2), flat_code(sq.code_size); + std::vector values(d); + fill_sa_code(*packer, key, scratch.data(), values.data(), flat_code.data()); sq.decode(flat_code.data(), recons, 1); } @@ -227,9 +379,15 @@ void IndexSQFastScan::sa_encode(idx_t n, const float* x, uint8_t* bytes) const { 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); + std::vector scratch(M2 / 2), flat_code(sq.code_size); + std::vector values(d); for (idx_t i = 0; i < ni; i++) { - packer->unpack_1(codes.data(), i0 + i, flat_code.data()); + fill_sa_code( + *packer, + i0 + i, + scratch.data(), + values.data(), + flat_code.data()); sq.decode(flat_code.data(), recons + i * d, 1); } } @@ -246,7 +404,7 @@ void IndexSQFastScan::add_sa_codes( // 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()); + add(n, recon.data()); } // ----------------------------------------------------------------------- @@ -272,6 +430,17 @@ void IndexSQFastScan::permute_entries(const idx_t* perm) { code_sz); } pq4_pack_codes(flat_new.data(), ntotal, M, ntotal2, bbs, M2, codes.get()); + + if (!lo_codes.empty()) { + const size_t lo_sz = sq_rerank_size(d, sq.qtype); + std::vector lo_new(size_t(ntotal) * lo_sz); + for (idx_t i = 0; i < ntotal; i++) { + memcpy(lo_new.data() + size_t(i) * lo_sz, + lo_codes.data() + size_t(perm[i]) * lo_sz, + lo_sz); + } + lo_codes.swap(lo_new); + } } // ----------------------------------------------------------------------- @@ -327,8 +496,15 @@ FlatCodesDistanceComputer* IndexSQFastScan::get_FlatCodesDistanceComputer() std::unique_ptr packer(get_CodePacker()); size_t cs = sq.code_size; std::vector flat(ntotal * cs); + std::vector scratch(M2 / 2); + std::vector values(d); for (idx_t i = 0; i < ntotal; i++) { - packer->unpack_1(codes.data(), i, flat.data() + i * cs); + fill_sa_code( + *packer, + i, + scratch.data(), + values.data(), + flat.data() + i * cs); } ScalarQuantizer::SQDistanceComputer* dc = @@ -359,9 +535,15 @@ void IndexSQFastScan::range_search( std::vector flat_codes(ntotal * sq.code_size); { std::unique_ptr packer(get_CodePacker()); + std::vector scratch(M2 / 2); + std::vector values(d); for (idx_t i = 0; i < ntotal; i++) { - packer->unpack_1( - codes.data(), i, flat_codes.data() + i * sq.code_size); + fill_sa_code( + *packer, + i, + scratch.data(), + values.data(), + flat_codes.data() + i * sq.code_size); } } @@ -436,6 +618,10 @@ void IndexSQFastScan::check_compatible_for_merge( void IndexSQFastScan::merge_from(Index& otherIndex, idx_t add_id) { check_compatible_for_merge(otherIndex); + const IndexSQFastScan* other = + static_cast(&otherIndex); + lo_codes.insert( + lo_codes.end(), other->lo_codes.begin(), other->lo_codes.end()); IndexFastScan::merge_from(otherIndex, add_id); } diff --git a/faiss/IndexSQFastScan.h b/faiss/IndexSQFastScan.h index ecbdfac18..22ae63327 100644 --- a/faiss/IndexSQFastScan.h +++ b/faiss/IndexSQFastScan.h @@ -16,25 +16,33 @@ namespace faiss { -/** Fast scan version of IndexScalarQuantizer for native 4-bit types. +/** Fast scan version of IndexScalarQuantizer. * * Supported quantizer types: * - QT_4bit, QT_4bit_uniform: native 4-bit codes mapped onto the - * PQ4 FastScan SIMD infrastructure (vpshufb / equivalent). - * No precision loss. + * PQ4 FastScan SIMD (vpshufb / equivalent). * - * 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. + * - QT_6bit, QT_8bit, QT_8bit_uniform, QT_8bit_direct, + * QT_8bit_direct_signed: re-quantised to 4-bit for the scan, reranked + * against the originals in orig_codes. + * + * - All other types (QT_fp16, QT_bf16): use IndexScalarQuantizer directly. */ struct IndexSQFastScan : IndexFastScan { ScalarQuantizer sq; + /// Overselection ratio for reranking. + float rerank_factor = 2; + + /// Low nibbles of the 8-bit codes, two dimensions per byte, indexed by + /// id. The high nibbles live in the packed scan codes. Empty for the + /// native 4-bit types, which need no rerank. + std::vector lo_codes; + /** Constructor. * * @param d dimensionality of input vectors - * @param qtype QT_4bit or QT_4bit_uniform only + * @param qtype any native 4-bit or reranked type (see above) * @param metric distance metric (METRIC_L2 or METRIC_INNER_PRODUCT) * @param bbs block size for SIMD processing (multiple of 32) */ @@ -46,7 +54,7 @@ struct IndexSQFastScan : IndexFastScan { IndexSQFastScan(); - /// Build from an existing IndexScalarQuantizer (must be QT_4bit*) + /// Build from an existing IndexScalarQuantizer explicit IndexSQFastScan(const IndexScalarQuantizer& orig, int bbs = 32); void train(idx_t n, const float* x) override; @@ -108,6 +116,13 @@ struct IndexSQFastScan : IndexFastScan { size_t sa_code_size() const override; + void fill_sa_code( + const CodePacker& packer, + idx_t id, + uint8_t* scratch, + int* values, + uint8_t* code_out) const; + size_t fast_scan_code_size() const override; }; diff --git a/faiss/impl/index_read.cpp b/faiss/impl/index_read.cpp index fd3a23e2c..822b52e45 100644 --- a/faiss/impl/index_read.cpp +++ b/faiss/impl/index_read.cpp @@ -2773,6 +2773,8 @@ std::unique_ptr read_index_up(IOReader* f, int io_flags) { READ1(idxsqfs->ntotal2); READ1(idxsqfs->M2); READVECTOR(idxsqfs->codes); + READ1(idxsqfs->rerank_factor); + READVECTOR(idxsqfs->lo_codes); // Restore FastScan base-class fields from the SQ idxsqfs->M = idxsqfs->sq.d; diff --git a/faiss/impl/index_write.cpp b/faiss/impl/index_write.cpp index 0ffe46185..f56f8eacf 100644 --- a/faiss/impl/index_write.cpp +++ b/faiss/impl/index_write.cpp @@ -966,6 +966,8 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) { WRITE1(idxsqfs->ntotal2); WRITE1(idxsqfs->M2); WRITEVECTOR(idxsqfs->codes); + WRITE1(idxsqfs->rerank_factor); + WRITEVECTOR(idxsqfs->lo_codes); } else if ( const IndexPQFastScan* idxpqfs = dynamic_cast(idx)) { diff --git a/faiss/impl/sq_fastscan_utils.h b/faiss/impl/sq_fastscan_utils.h new file mode 100644 index 000000000..683c21fa4 --- /dev/null +++ b/faiss/impl/sq_fastscan_utils.h @@ -0,0 +1,232 @@ +/* + * 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 + +/** Helpers shared by IndexSQFastScan and IndexIVFSQFastScan. + * + * Both re-quantise a higher-precision scalar quantizer down to 4-bit nibbles + * for the SIMD scan and keep the original codes to rerank with, so the + * classification of quantizer types and the nibble packing have to agree + * between them. + */ + +namespace faiss { +namespace sq_fastscan { + +inline size_t roundup(size_t a, size_t b) { + return (a + b - 1) / b * b; +} + +inline bool is_native_4bit(ScalarQuantizer::QuantizerType qtype) { + return qtype == ScalarQuantizer::QT_4bit || + qtype == ScalarQuantizer::QT_4bit_uniform; +} + +inline 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; +} + +inline bool is_fallback(ScalarQuantizer::QuantizerType qtype) { + return !is_native_4bit(qtype) && !needs_rerank(qtype); +} + +inline bool is_uniform_range(ScalarQuantizer::QuantizerType qtype) { + return qtype == ScalarQuantizer::QT_8bit_uniform || + qtype == ScalarQuantizer::QT_8bit_direct || + qtype == ScalarQuantizer::QT_8bit_direct_signed; +} + +inline 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]; + } +} + +/// Bits per dimension the quantizer stores. +inline int sq_bits(ScalarQuantizer::QuantizerType qtype) { + if (qtype == ScalarQuantizer::QT_6bit) { + return 6; + } + return is_native_4bit(qtype) ? 4 : 8; +} + +/// Bits left over once the scan has taken its 4. Zero for a native 4-bit +/// code, which needs no rerank and therefore stores nothing extra. +inline int sq_rerank_bits(ScalarQuantizer::QuantizerType qtype) { + return sq_bits(qtype) - 4; +} + +/// Bytes per vector needed to hold the leftover bits of every dimension. +inline size_t sq_rerank_size(int d, ScalarQuantizer::QuantizerType qtype) { + const int bits = sq_rerank_bits(qtype); + return bits > 0 ? (size_t(d) * bits + 7) / 8 : 0; +} + +/// The code for one value at the quantizer's own width. +inline int sq_code_of(float x, float vmin, float vdiff, int bits) { + const int maxv = (1 << bits) - 1; + const float f = (vdiff > 0) ? (x - vmin) / vdiff : 0.0f; + return std::min(maxv, std::max(0, int(f * float(maxv)))); +} + +/// Per-dimension range of the trained quantizer, resolved once for a batch. +struct SQRanges { + SQRanges(const ScalarQuantizer& sq, int d) + : uniform(is_uniform_range(sq.qtype) || + sq.qtype == ScalarQuantizer::QT_4bit_uniform) { + if (uniform) { + get_uniform_range(sq, vmin_s, vdiff_s); + } else { + vmin_arr = sq.trained.data(); + vdiff_arr = sq.trained.data() + d; + } + } + + float vmin(int m) const { + return uniform ? vmin_s : vmin_arr[m]; + } + float vdiff(int m) const { + return uniform ? vdiff_s : vdiff_arr[m]; + } + + const bool uniform; + const float* vmin_arr = nullptr; + const float* vdiff_arr = nullptr; + float vmin_s = 0, vdiff_s = 0; +}; + +/** Split each code into the top 4 bits the SIMD kernel scans and the rest, + * which rerank keeps. `hi` is packed two dimensions per byte as + * pq4_pack_codes expects. `lo` is bit-packed at `bits - 4` per dimension. + * Total width is the quantizer's own, not that width plus 4. + */ +inline void float_to_split_codes( + const float* x, + uint8_t* hi, + uint8_t* lo, + idx_t n, + int d, + size_t M2, + const ScalarQuantizer& sq) { + const int bits = sq_bits(sq.qtype); + const int lo_bits = bits - 4; + const size_t hi_stride = M2 / 2; + const size_t lo_stride = sq_rerank_size(d, sq.qtype); + const SQRanges ranges(sq, d); + for (idx_t i = 0; i < n; i++) { + const float* xi = x + i * d; + uint8_t* hdst = hi + i * hi_stride; + memset(hdst, 0, hi_stride); + uint8_t* ldst = lo + i * lo_stride; + memset(ldst, 0, lo_stride); + BitstringWriter bsw(ldst, lo_stride); + for (int m = 0; m < d; m++) { + const int code = + sq_code_of(xi[m], ranges.vmin(m), ranges.vdiff(m), bits); + const int h = code >> lo_bits; + hdst[m / 2] |= (m & 1) ? uint8_t(h << 4) : uint8_t(h); + bsw.write(code & ((1 << lo_bits) - 1), lo_bits); + } + } +} + +/// Rebuild the code values of one vector from its two halves. +inline void split_to_codes( + const uint8_t* hi, + const uint8_t* lo, + int d, + ScalarQuantizer::QuantizerType qtype, + int* codes_out) { + const int lo_bits = sq_rerank_bits(qtype); + BitstringReader bsr(lo, sq_rerank_size(d, qtype)); + for (int m = 0; m < d; m++) { + const int h = (m & 1) ? (hi[m / 2] >> 4) : (hi[m / 2] & 0xF); + codes_out[m] = (h << lo_bits) | int(bsr.read(lo_bits)); + } +} + +inline 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))); + } + } +} + +} // namespace sq_fastscan +} // namespace faiss diff --git a/faiss/index_factory.cpp b/faiss/index_factory.cpp index 94fa1a45c..f33daa9f1 100644 --- a/faiss/index_factory.cpp +++ b/faiss/index_factory.cpp @@ -179,8 +179,8 @@ 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]+)?"; +const std::string sq_fs_pattern = + "(SQ4|SQ6|SQ8|SQ8_direct_signed|SQ8_direct)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]+)?"; diff --git a/tests/test_sq_fastscan.py b/tests/test_sq_fastscan.py index 3d97297da..329f3f264 100644 --- a/tests/test_sq_fastscan.py +++ b/tests/test_sq_fastscan.py @@ -48,12 +48,10 @@ class TestSQFastScanConstruction(unittest.TestCase): self.assertEqual(index.d, 0) self.assertEqual(index.ntotal, 0) - def test_rejects_non_4bit(self): + def test_accepts_reranked_types_rejects_fallback(self): d = 64 - with self.assertRaises(RuntimeError): - faiss.IndexSQFastScan(d, SQ.QT_8bit) - with self.assertRaises(RuntimeError): - faiss.IndexSQFastScan(d, SQ.QT_6bit) + for qtype in (SQ.QT_8bit, SQ.QT_6bit): + faiss.IndexSQFastScan(d, qtype) with self.assertRaises(RuntimeError): faiss.IndexSQFastScan(d, SQ.QT_fp16) with self.assertRaises(RuntimeError): @@ -177,11 +175,17 @@ class TestSQFastScanConversionConstructor(unittest.TestCase): r1_fs = (I_fs[:, 0] >= 0).mean() self.assertEqual(r1_sq, r1_fs) - def test_conversion_rejects_non_4bit(self): + def test_conversion_accepts_8bit_and_rejects_fallback(self): sq = faiss.IndexScalarQuantizer(self.d, SQ.QT_8bit) sq.train(self.ds.get_train()) + sq.add(self.ds.get_database()) + converted = faiss.IndexSQFastScan(sq) + self.assertEqual(converted.ntotal, sq.ntotal) + + fp16 = faiss.IndexScalarQuantizer(self.d, SQ.QT_fp16) + fp16.train(self.ds.get_train()) with self.assertRaises(RuntimeError): - faiss.IndexSQFastScan(sq) + faiss.IndexSQFastScan(fp16) class TestSQFastScanReset(unittest.TestCase): @@ -586,3 +590,120 @@ class TestRaBitQFastScanBiasBound(unittest.TestCase): self.assertGreaterEqual( fs, ref - 0.1, f"fast scan R@1 {fs:.2f} trails the reference {ref:.2f}" ) + + +class TestSQFastScanSplitCodePaths(unittest.TestCase): + """Paths that rebuild a full code from the split halves. + + QT_4bit hides every bug here, because its code is already one nibble per + dimension and sq.code_size equals the fast-scan code size. + """ + + def _pair(self, qtype, d=32, n=200): + rng = np.random.default_rng(0) + xb = rng.standard_normal((n, d), dtype=np.float32) + ref = faiss.IndexScalarQuantizer(d, qtype) + ref.train(xb) + ref.add(xb) + fs = faiss.IndexSQFastScan(d, qtype) + fs.train(xb) + fs.add(xb) + return xb, ref, fs + + def test_reconstruct_matches_reference(self): + for qtype in (SQ.QT_4bit, SQ.QT_6bit, SQ.QT_8bit): + _, ref, fs = self._pair(qtype) + for i in (0, 7, 199): + np.testing.assert_array_equal( + ref.reconstruct(i), fs.reconstruct(i) + ) + + def test_range_search_matches_reference(self): + for qtype in (SQ.QT_4bit, SQ.QT_6bit, SQ.QT_8bit): + xb, ref, fs = self._pair(qtype) + _, _, iref = ref.range_search(xb[:5], 20.0) + _, _, ifs = fs.range_search(xb[:5], 20.0) + self.assertEqual(sorted(iref.tolist()), sorted(ifs.tolist())) + + def test_permute_entries_moves_both_halves(self): + for qtype in (SQ.QT_6bit, SQ.QT_8bit): + _, _, fs = self._pair(qtype) + n = fs.ntotal + before = fs.reconstruct(5).copy() + perm = np.arange(n - 1, -1, -1, dtype="int64") + fs.permute_entries(perm) + np.testing.assert_array_equal(before, fs.reconstruct(n - 6)) + + def test_direct_qtypes_match_reference(self): + rng = np.random.default_rng(0) + d, n = 32, 200 + for qtype, lo in ( + (SQ.QT_8bit_direct, 0), + (SQ.QT_8bit_direct_signed, -128), + ): + xb = rng.integers(lo, lo + 256, size=(n, d)).astype(np.float32) + ref = faiss.IndexScalarQuantizer(d, qtype) + ref.train(xb) + ref.add(xb) + fs = faiss.IndexSQFastScan(d, qtype) + fs.train(xb) + fs.add(xb) + _, iref = ref.search(xb[:10], 5) + _, ifs = fs.search(xb[:10], 5) + hits = sum(len(set(a) & set(b)) for a, b in zip(iref, ifs)) + self.assertGreaterEqual(hits, 0.9 * iref.size) + + def test_lo_codes_track_ntotal_across_add_paths(self): + rng = np.random.default_rng(0) + d, qtype = 8, SQ.QT_8bit + # over the blocking size of IndexFastScan::add, which re-enters add() + n = 70000 + xb = rng.standard_normal((n, d), dtype=np.float32) + ref = faiss.IndexScalarQuantizer(d, qtype) + ref.train(xb) + ref.add(xb) + + one = faiss.IndexSQFastScan(d, qtype) + one.train(xb) + one.add(xb) + self.assertEqual(one.ntotal, n) + + split = faiss.IndexSQFastScan(d, qtype) + split.train(xb) + split.add(xb[:40000]) + split.add(xb[40000:]) + self.assertEqual(split.ntotal, n) + + for i in (0, 39999, 65536, n - 1): + expect = ref.reconstruct(i) + np.testing.assert_array_equal(expect, one.reconstruct(i)) + np.testing.assert_array_equal(expect, split.reconstruct(i)) + + def test_add_sa_codes_fills_the_rerank_half(self): + xb, ref, _ = self._pair(SQ.QT_8bit) + fs = faiss.IndexSQFastScan(ref.d, SQ.QT_8bit) + fs.train(xb) + fs.add_sa_codes(ref.sa_encode(xb)) + self.assertEqual(fs.ntotal, ref.ntotal) + for i in (0, 7, 199): + np.testing.assert_array_equal(ref.reconstruct(i), fs.reconstruct(i)) + + def test_merge_from_moves_the_rerank_half(self): + rng = np.random.default_rng(0) + d, n = 16, 100 + xb = rng.standard_normal((2 * n, d), dtype=np.float32) + parts = [] + for lo, hi in ((0, n), (n, 2 * n)): + idx = faiss.IndexSQFastScan(d, SQ.QT_8bit) + idx.train(xb) + idx.add(xb[lo:hi]) + parts.append(idx) + parts[0].merge_from(parts[1]) + self.assertEqual(parts[0].ntotal, 2 * n) + + whole = faiss.IndexSQFastScan(d, SQ.QT_8bit) + whole.train(xb) + whole.add(xb) + np.testing.assert_array_equal( + whole.reconstruct_n(0, 2 * n), parts[0].reconstruct_n(0, 2 * n) + )