facebookresearch/faiss
1
0
Fork 0
mirror of https://github.com/facebookresearch/faiss.git synced 2026-10-12 02:13:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

Halve IndexSQFastScan rerank memory with a split-nibble layout (#5662)

Browse Source 
Summary:
Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5662

`IndexIVFSQFastScan` kept whole original codes in a parallel inverted list, so
an 8-bit quantizer cost 12 bits per dimension against 8 for plain SQ8. It now
keeps only the bits the scan does not read, which is the same split the flat
index uses in the diff below.

## The split is bit-width generic

An 8-bit code keeps 4 leftover bits. A 6-bit code keeps 2. `sq_fastscan_utils.h` packs the remainder with `BitstringWriter`, which lays bits out exactly as `Codec6bit` and `Codec8bit` do.

The LUT reconstruction collapses to one formula: `step = 2^(bits-4)`, `offset = step / 2`, `scale = 2^bits - 1`. At 4 bits this reduces to the native `(v + 0.5) / 15`, so the native and split paths no longer need separate constants. This also corrects a small bias, because the 8-bit case previously used an offset of 7.5 where the true band midpoint is 8.0.

## `IndexIVFSQFastScan` splits as well

The IVF form kept whole codes in a parallel inverted list. It now keeps only the leftover bits.

## Measured results, flat index, 200k vectors

| index | before | after |
| --- | --- | --- |
| `SQ8fs` | 12 bits/dim | **8.00 bits/dim** |
| `SQ6fs` | 10 bits/dim | **6.00 bits/dim** |

The IVF form reaches 1.37x `SQ8` at d=64 and 1.10x at d=512, against 1.79x before. The remainder is the ids duplicated into the parallel list and the block padding. Both are per-vector, so the ratio falls as `d` grows.

Recall does not drop. The split reaches 0.914 to 0.932 for `SQ6fs` and 0.969 to 0.980 for `SQ8fs` at `rerank_factor=4`, measured against exact neighbours restricted to the admitted set.

Also fixes a double-append in `add()`. `IndexFastScan::add` blocks adds over 65536 and calls `add()` virtually. The override therefore re-enters once per chunk, and every side-table entry was written twice for a large index.

## Serialization

The `ISfs` payload now carries the leftover bits instead of whole original
codes. The fourcc does not change, because `ISfs` has not been released with
rerank state and no stored index is affected.

Reviewed By: junjieqi

Differential Revision: D121092769

fbshipit-source-id: d6e8a4f541ce73057eba86215a83a0a80f7e106d
This commit is contained in:
Michael Norris
2026-09-25 13:47:19 -07:00
committed by meta-codesync[bot]
parent 94a4331fee
commit 7f6bf10b1c
4 changed files with 138 additions and 70 deletions
Download Patch File Download Diff File Expand all files Collapse all files
faiss/IndexIVFSQFastScan.cpp
+109 -55
View File
@@ -75,8 +75,8 @@ IndexIVFSQFastScan::IndexIVFSQFastScan(
new BlockInvertedLists(nlist_in, get_CodePacker()), true);
if (needs_rerank(qtype)) {
orig_codes_invlists =
new ArrayInvertedLists(nlist_in, sq.code_size);
lo_codes_invlists =
new ArrayInvertedLists(nlist_in, sq_rerank_size(d, qtype));
direct_map.set_type(DirectMap::Hashtable, invlists, 0);
}
}
@@ -136,7 +136,7 @@ IndexIVFSQFastScan::IndexIVFSQFastScan(
replace_invlists(new BlockInvertedLists(nlist, get_CodePacker()), true);
if (needs_rerank(sq.qtype)) {
orig_codes_invlists = new ArrayInvertedLists(nlist, sq.code_size);
lo_codes_invlists = new ArrayInvertedLists(nlist, lo_code_size());
direct_map.set_type(DirectMap::Hashtable, invlists, 0);
}
@@ -149,15 +149,7 @@ IndexIVFSQFastScan::IndexIVFSQFastScan(
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.
// Decode to float and re-encode for the scan.
std::vector<float> recon(list_size * d);
sq.decode(orig_codes.get(), recon.data(), list_size);
@@ -165,8 +157,21 @@ IndexIVFSQFastScan::IndexIVFSQFastScan(
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);
// Top 4 bits feed the scan, the rest are kept for rerank.
std::vector<uint8_t> lo(list_size * lo_code_size(), 0);
float_to_split_codes(
recon.data(),
flat4.data(),
lo.data(),
list_size,
d,
M2,
sq);
lo_codes_invlists->add_entries(
list_no, list_size, orig_ids.get(), lo.data());
for (size_t j = 0; j < list_size; j++) {
direct_map.add_single_id(orig_ids.get()[j], list_no, j);
}
}
BlockInvertedLists* bil = dynamic_cast<BlockInvertedLists*>(invlists);
@@ -180,7 +185,7 @@ IndexIVFSQFastScan::IndexIVFSQFastScan(
}
IndexIVFSQFastScan::~IndexIVFSQFastScan() {
delete orig_codes_invlists;
delete lo_codes_invlists;
}
// -----------------------------------------------------------------------
@@ -229,7 +234,10 @@ void IndexIVFSQFastScan::encode_vectors(
} 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);
// add_core keeps the leftover bits in lo_codes_invlists.
std::vector<uint8_t> lo_scratch(size_t(n) * lo_code_size());
float_to_split_codes(
to_encode, codes_out, lo_scratch.data(), n, d, M2, sq);
}
if (include_listnos) {
@@ -345,6 +353,34 @@ void IndexIVFSQFastScan::sa_decode(idx_t n, const uint8_t* bytes, float* x)
}
}
void IndexIVFSQFastScan::recombine_code(
int64_t list_no,
int64_t offset,
uint8_t* scratch,
int* values,
uint8_t* code_out) const {
// Interleaved for SIMD, so read one sub-quantizer at a time.
memset(scratch, 0, size_t(M2) / 2);
InvertedLists::ScopedCodes list_codes(invlists, list_no);
BitstringWriter bsw(scratch, M2 / 2);
for (size_t m = 0; m < M; m++) {
bsw.write(
pq4_get_packed_element(list_codes.get(), bbs, M2, offset, m),
4);
}
InvertedLists::ScopedCodes lo(lo_codes_invlists, list_no);
split_to_codes(
scratch, lo.get() + offset * lo_code_size(), d, sq.qtype, values);
// BitstringWriter matches Codec6bit/Codec8bit, so distance_to_code reads
// it.
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 IndexIVFSQFastScan::reconstruct_from_offset(
int64_t list_no,
int64_t offset,
@@ -353,10 +389,13 @@ void IndexIVFSQFastScan::reconstruct_from_offset(
// 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 if (lo_codes_invlists) {
// Rerank: rebuild the original code from its two nibble halves.
std::vector<uint8_t> scratch(M2 / 2), code(sq.code_size);
std::vector<int> values(d);
recombine_code(
list_no, offset, scratch.data(), values.data(), code.data());
sq.decode(code.data(), recons, 1);
} else {
// Native 4-bit: unpack from block inverted lists.
std::vector<uint8_t> code(M2 / 2, 0);
@@ -416,6 +455,8 @@ void IndexIVFSQFastScan::search(
{
std::unique_ptr<ScalarQuantizer::SQDistanceComputer> dc(
sq.get_distance_computer(metric_type));
std::vector<uint8_t> scratch(M2 / 2), code(sq.code_size);
std::vector<int> values(d);
#pragma omp for
for (idx_t i = 0; i < n; i++) {
@@ -441,10 +482,13 @@ void IndexIVFSQFastScan::search(
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);
recombine_code(
list_no,
offset,
scratch.data(),
values.data(),
code.data());
float dis = dc->distance_to_code(code.data());
if (metric_type == METRIC_L2) {
if (dis < heap_dis[0]) {
@@ -516,6 +560,8 @@ void IndexIVFSQFastScan::search_preassigned(
{
std::unique_ptr<ScalarQuantizer::SQDistanceComputer> dc(
sq.get_distance_computer(metric_type));
std::vector<uint8_t> scratch(M2 / 2), code(sq.code_size);
std::vector<int> values(d);
#pragma omp for
for (idx_t i = 0; i < n; i++) {
@@ -541,10 +587,13 @@ void IndexIVFSQFastScan::search_preassigned(
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);
recombine_code(
list_no,
offset,
scratch.data(),
values.data(),
code.data());
float dis = dc->distance_to_code(code.data());
if (metric_type == METRIC_L2) {
if (dis < heap_dis[0]) {
@@ -611,6 +660,8 @@ void IndexIVFSQFastScan::range_search(
std::unique_ptr<ScalarQuantizer::SQDistanceComputer> dc(
sq.get_distance_computer(metric_type));
std::vector<uint8_t> scratch(M2 / 2), code(sq.code_size);
std::vector<int> values(d);
#pragma omp for
for (idx_t i = 0; i < n; i++) {
@@ -622,20 +673,23 @@ void IndexIVFSQFastScan::range_search(
if (list_no < 0) {
continue;
}
size_t list_size = orig_codes_invlists->list_size(list_no);
size_t list_size = lo_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);
InvertedLists::ScopedIds ids(lo_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);
recombine_code(
list_no,
j,
scratch.data(),
values.data(),
code.data());
float dis = dc->distance_to_code(code.data());
if (metric_type == METRIC_L2) {
if (dis < radius) {
qres.add(dis, ids[j]);
@@ -701,28 +755,28 @@ void IndexIVFSQFastScan::add_with_ids(
AlignedTable<uint8_t> flat_codes(n * code_size);
encode_vectors(n, x, idx.get(), flat_codes.get());
// Store original SQ codes for rerank types
std::vector<uint8_t> orig_flat_codes;
// Same quantisation as encode_vectors, or the halves disagree.
std::vector<uint8_t> lo_flat_codes;
if (needs_rerank(sq.qtype)) {
orig_flat_codes.resize(n * sq.code_size);
std::vector<float> to_encode(size_t(n) * d);
for (idx_t i = 0; i < n; i++) {
if (idx[i] < 0) {
continue;
}
if (by_residual) {
std::vector<float> 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);
if (by_residual && idx[i] >= 0) {
quantizer->compute_residual(
x + i * d, to_encode.data() + i * d, idx[i]);
} else {
sq.compute_codes(
x + i * d,
orig_flat_codes.data() + i * sq.code_size,
1);
memcpy(to_encode.data() + i * d, x + i * d, sizeof(float) * d);
}
}
lo_flat_codes.assign(size_t(n) * lo_code_size(), 0);
std::vector<uint8_t> hi_discard(size_t(n) * (M2 / 2));
float_to_split_codes(
to_encode.data(),
hi_discard.data(),
lo_flat_codes.data(),
n,
d,
M2,
sq);
}
// Sort by list assignment
@@ -763,10 +817,10 @@ void IndexIVFSQFastScan::add_with_ids(
code_size);
if (needs_rerank(sq.qtype)) {
orig_codes_invlists->add_entry(
lo_codes_invlists->add_entry(
list_no,
id,
orig_flat_codes.data() + order[i] * sq.code_size);
lo_flat_codes.data() + order[i] * lo_code_size());
direct_map.add_single_id(id, list_no, ofs);
}
}
@@ -790,9 +844,9 @@ void IndexIVFSQFastScan::add_with_ids(
void IndexIVFSQFastScan::reset() {
IndexIVFFastScan::reset();
if (orig_codes_invlists) {
if (lo_codes_invlists) {
for (size_t i = 0; i < nlist; i++) {
orig_codes_invlists->resize(i, 0);
lo_codes_invlists->resize(i, 0);
}
direct_map.clear();
}
faiss/IndexIVFSQFastScan.h
+22 -8
View File
@@ -10,6 +10,7 @@
#include <faiss/IndexIVFFastScan.h>
#include <faiss/IndexScalarQuantizer.h>
#include <faiss/impl/ScalarQuantizer.h>
#include <faiss/impl/sq_fastscan_utils.h>
namespace faiss {
@@ -29,10 +30,8 @@ namespace faiss {
* 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).
* For reranked types the block inverted lists hold the packed high nibbles
* and a parallel ArrayInvertedLists (lo_codes_invlists) holds the low ones.
*/
struct IndexIVFSQFastScan : IndexIVFFastScan {
ScalarQuantizer sq;
@@ -40,10 +39,25 @@ struct IndexIVFSQFastScan : IndexIVFFastScan {
/// 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;
/// Parallel inverted lists holding the leftover low bits of each code.
/// The top 4 bits of every code live in the packed scan blocks.
/// Owned by this index. nullptr for native 4-bit and fallback types.
InvertedLists* lo_codes_invlists = nullptr;
/// Bytes per vector in lo_codes_invlists: the bits of each code left
/// over once the scan has taken the top 4.
size_t lo_code_size() const {
return sq_fastscan::sq_rerank_size(d, sq.qtype);
}
/// Rebuild the code at (list_no, offset) from its two halves. `scratch`
/// needs M2 / 2 bytes, `code_out` needs sq.code_size.
void recombine_code(
int64_t list_no,
int64_t offset,
uint8_t* scratch,
int* values,
uint8_t* code_out) const;
IndexIVFSQFastScan(
Index* quantizer,
faiss/impl/index_read.cpp
+4 -4
View File
@@ -2838,18 +2838,18 @@ std::unique_ptr<Index> read_index_up(IOReader* f, int io_flags) {
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->lo_codes_invlists = read_InvertedLists(f, io_flags);
// Rebuild direct_map from lo_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);
ivfsqfs->lo_codes_invlists->list_size(list_no);
if (list_size == 0) {
continue;
}
InvertedLists::ScopedIds ids(
ivfsqfs->orig_codes_invlists, list_no);
ivfsqfs->lo_codes_invlists, list_no);
for (size_t j = 0; j < list_size; j++) {
ivfsqfs->direct_map.add_single_id(ids[j], list_no, j);
}
faiss/impl/index_write.cpp
+3 -3
View File
@@ -995,11 +995,11 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) {
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);
// Write lo_codes_invlists if present
bool has_orig = (ivfsqfs->lo_codes_invlists != nullptr);
WRITE1(has_orig);
if (has_orig) {
write_InvertedLists(ivfsqfs->orig_codes_invlists, f);
write_InvertedLists(ivfsqfs->lo_codes_invlists, f);
}
} else if (
const IndexIVFPQFastScan* ivpq_2 =