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
Files
7f6bf10b1c83b1634c3cfd3367ba1e76907c79b9
faiss/faiss/IndexIVFSQFastScan.cpp
T
Michael Norris 7f6bf10b1c Halve IndexSQFastScan rerank memory with a split-nibble layout (#5662) 
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
2026-09-25 13:47:19 -07:00

856 lines
28 KiB
C++
Raw Blame History

/*
* 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 <faiss/IndexIVFSQFastScan.h>
#include <faiss/impl/sq_fastscan_utils.h>
#include <omp.h>
#include <algorithm>
#include <cstring>
#include <memory>
#include <vector>
#include <faiss/impl/AuxIndexStructures.h>
#include <faiss/impl/FaissAssert.h>
#include <faiss/impl/IDSelector.h>
#include <faiss/impl/ScalarQuantizer.h>
#include <faiss/impl/fast_scan/fast_scan.h>
#include <faiss/invlists/BlockInvertedLists.h>
#include <faiss/invlists/DirectMap.h>
#include <faiss/invlists/InvertedLists.h>
#include <faiss/utils/Heap.h>
#include <faiss/utils/distances.h>
#include <faiss/utils/hamming.h>
#include <faiss/utils/utils.h>
namespace faiss {
using namespace faiss::sq_fastscan;
// -----------------------------------------------------------------------
// 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)) {
lo_codes_invlists =
new ArrayInvertedLists(nlist_in, sq_rerank_size(d, qtype));
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)) {
lo_codes_invlists = new ArrayInvertedLists(nlist, lo_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);
// 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);
std::vector<uint8_t> flat4(list_size * M2 / 2, 0);
if (is_native_4bit(sq.qtype)) {
sq.compute_codes(recon.data(), flat4.data(), list_size);
} else {
// 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);
FAISS_THROW_IF_NOT(bil);
size_t nb2 = roundup(list_size, bbs);
AlignedTable<uint8_t> 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 lo_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<float> 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 {
// 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) {
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<float>& dis_tables,
AlignedTable<float>& /*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<float> 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<float> 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::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,
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 (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);
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<float> 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(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<float> coarse_dis(n * k_coarse);
std::vector<idx_t> 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<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++) {
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);
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]) {
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(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<float> coarse_dis(n * k_coarse);
std::vector<idx_t> 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<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++) {
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);
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]) {
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<const IVFSearchParameters*>(params);
}
size_t cur_nprobe = ivf_params ? ivf_params->nprobe : nprobe;
const IDSelector* sel = params ? params->sel : nullptr;
std::vector<idx_t> assign(n * cur_nprobe);
std::vector<float> centroid_dis(n * cur_nprobe);
quantizer->search(n, x, cur_nprobe, centroid_dis.data(), assign.data());
std::vector<RangeSearchPartialResult*> 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<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++) {
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 = lo_codes_invlists->list_size(list_no);
if (list_size == 0) {
continue;
}
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;
}
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]);
}
} 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_t[]> 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<uint8_t> 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<uint8_t> flat_codes(n * code_size);
encode_vectors(n, x, idx.get(), flat_codes.get());
// Same quantisation as encode_vectors, or the halves disagree.
std::vector<uint8_t> lo_flat_codes;
if (needs_rerank(sq.qtype)) {
std::vector<float> to_encode(size_t(n) * d);
for (idx_t i = 0; i < n; i++) {
if (by_residual && idx[i] >= 0) {
quantizer->compute_residual(
x + i * d, to_encode.data() + i * d, idx[i]);
} else {
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
std::vector<idx_t> 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<BlockInvertedLists*>(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<uint8_t> 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)) {
lo_codes_invlists->add_entry(
list_no,
id,
lo_flat_codes.data() + order[i] * lo_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 (lo_codes_invlists) {
for (size_t i = 0; i < nlist; i++) {
lo_codes_invlists->resize(i, 0);
}
direct_map.clear();
}
}
} // namespace faiss
Reference in New Issue View Git Blame Copy Permalink