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

Add InvertedListScanner for IVFPQFastScan (#4537)

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

Description
-
Adds InvertedListScanner for IVFPQFastScan. Right now IVFPQFastScan does not support a Scanner because the codes are PQFS packed. This is required from Unicorn to simplify the integration layer.

We can implement a scanner based on the search_implem_10. We can't do distance_to_code but we can do the whole scan. (that is what is done in this diff in the new `IVFPQFastScanScanner.scan_codes`).

Summary of changes (v10)
-
- High level: there's 2 heaps at play during a `scan_codes` call. One is the uint16_t heap inside the HeapHandler, which tracks the best results for a given iteration of nprobe. The other tracks the results across iterations and is in float: heap_distances. At the end of each call, we find the best results from the current iteration (`curr_dists`) then `heap_addn` those to the output distances `distances`).
- Moved `make_knn_handler` to the IndexIVFFastScan.h declaration so it can be used outside just IndexIVFFastScan.cpp. It is moved out of the anonymous namespace in the .cpp file too.
- Add the IVFPQFastScanScanner to IndexIVFPQFastScan. This impl for scan_codes is directly translated from search_implem_10 in in IndexIVFPQFastScan, because both search_implem_10 and scan_codes process 1 query and 1 invlist at a time (technically, search_implem_10 processes `n` queries at a time, just not batched. search_implem_10 has a simple loop over `n`, so search_implem_10 can be factored out to 1 query for this Scanner easily.). The only changes are removal of extraneous stuff from search_implem_10.
- Changes to test_lowlevel_ivf.cpp to test this. Refactored the test because IVFPQFastScanScanner does not behave like other scanners.

v15 vs v16
-
v15 matches search_implem_10 exactly. The unit tests pass with bit-exact match. They were not passing before due to the biases and LUT not matching exactly to search_implem_10.

Problems:
1. It requires quantizer search. Otherwise the biases and LUT will not match search_implem_10. The LUT is built from several inverted lists.
2. It requires mutable state across scan_codes calls. I called them lut_counter and bias_counter.
(I think even though this matches results, these problems make this option a no-go. It is mostly here to show the investigation.)

v16 does not use quantizer search and has no mutable state.
- cq.dis == distance from query (xi) to centroid represented by list_no.
- cq.ids == list_no

Both are of length 1 since we just have 1 list which is set explicitly.
Then we can call compute_LUT as normal.

The LUTs and biases are completely different from an existing approach, so results will not match an existing approach.

Match percentage between search_implem_10 and various by_residual approaches (I called it `recall` in the test):
by_residual_IP: 0.9745
by_residual_L2: 0.936
Thus I set a baseline of 0.9 in the test.

Reference impl:
https://github.com/facebookresearch/faiss/blob/514b44fca8542bafe8640adcbf1cccce1900f74c/faiss/IndexIVFFastScan.cpp#L918-L992

ntotal vs number of blocks?
-
- index.bbs is often 32 https://github.com/facebookresearch/faiss/wiki/Fast-accumulation-of-PQ-and-AQ-codes-(FastScan)#implementation. This is the number of codes per block.

~~Why do we need this io_simd_dis and io_simd_ids?~~
-
Old, but keeping it struck out. This describes the uint16_t heap kept over multiple scans.

~~Other scanners can scan over multiple lists with `set_query` then `set_list`, then calling `scan_codes`. Why? Because they take the `float* distances`. The heap can operate on these floats to re-use values from last iteration.~~

~~IVFPQFastScan is different. Codes are pq4 packed. The HeapHandler has to get initialized with values from scanning the prior list.~~

~~If the user is only setting 1 list and 1 query for the lifetime of this scanner (which I believe is Unicorn's use case), it does not matter because there's no need to use prior results and the user can follow the test named IVFPQFS_null_params. But if the user wants to use like the other Scanners, they can follow the other 2 tests in test_lowlevel_ivf.cpp that pass the heap size.~~

Reviewed By: ddrcoder, mdouze

Differential Revision: D80114737

fbshipit-source-id: a28392d444bc4dbe5d12e0fd84884e691cf00cf0
This commit is contained in:
Michael Norris
2025-10-15 20:39:45 -07:00
committed by meta-codesync[bot]
parent dbc7550615
commit f58fd4c8ca
7 changed files with 324 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files
faiss/IndexIVF.h
+3 -3
View File
@@ -497,12 +497,12 @@ struct InvertedListScanner {
/// compute a single query-to-code distance
virtual float distance_to_code(const uint8_t* code) const = 0;
/** scan a set of codes, compute distances to current query and
/** scan a set of codes, compute distances to current query, and
* update heap of results if necessary. Default implementation
* calls distance_to_code.
*
* @param n number of codes to scan
* @param codes codes to scan (n * code_size)
* @param n number of codes to scan
* @param codes codes to scan (n * code_size)
* @param ids corresponding ids (ignored if store_pairs)
* @param distances heap distances (size k)
* @param labels heap labels (size k)
faiss/IndexIVFFastScan.cpp
+22 -20
View File
@@ -379,7 +379,8 @@ ResultHandlerCompare<C, true>* make_knn_handler_fixC(
idx_t k,
float* distances,
idx_t* labels,
const IDSelector* sel) {
const IDSelector* sel,
const float* normalizers) {
using HeapHC = HeapHandler<C, true>;
using ReservoirHC = ReservoirHandler<C, true>;
using SingleResultHC = SingleResultHandler<C, true>;
@@ -387,29 +388,12 @@ ResultHandlerCompare<C, true>* make_knn_handler_fixC(
if (k == 1) {
return new SingleResultHC(n, 0, distances, labels, sel);
} else if (impl % 2 == 0) {
return new HeapHC(n, 0, k, distances, labels, sel);
return new HeapHC(n, 0, k, distances, labels, sel, normalizers);
} else /* if (impl % 2 == 1) */ {
return new ReservoirHC(n, 0, k, 2 * k, distances, labels, sel);
}
}
SIMDResultHandlerToFloat* make_knn_handler(
bool is_max,
int impl,
idx_t n,
idx_t k,
float* distances,
idx_t* labels,
const IDSelector* sel) {
if (is_max) {
return make_knn_handler_fixC<CMax<uint16_t, int64_t>>(
impl, n, k, distances, labels, sel);
} else {
return make_knn_handler_fixC<CMin<uint16_t, int64_t>>(
impl, n, k, distances, labels, sel);
}
}
using CoarseQuantized = IndexIVFFastScan::CoarseQuantized;
struct CoarseQuantizedWithBuffer : CoarseQuantized {
@@ -443,7 +427,7 @@ struct CoarseQuantizedWithBuffer : CoarseQuantized {
};
struct CoarseQuantizedSlice : CoarseQuantizedWithBuffer {
size_t i0, i1;
const size_t i0, i1;
CoarseQuantizedSlice(const CoarseQuantized& cq, size_t i0, size_t i1)
: CoarseQuantizedWithBuffer(cq), i0(i0), i1(i1) {
if (done()) {
@@ -486,6 +470,24 @@ int compute_search_nslice(
} // namespace
SIMDResultHandlerToFloat* IndexIVFFastScan::make_knn_handler(
bool is_max,
int impl,
idx_t n,
idx_t k,
float* distances,
idx_t* labels,
const IDSelector* sel,
const float* normalizers) const {
if (is_max) {
return make_knn_handler_fixC<CMax<uint16_t, int64_t>>(
impl, n, k, distances, labels, sel, normalizers);
} else {
return make_knn_handler_fixC<CMin<uint16_t, int64_t>>(
impl, n, k, distances, labels, sel, normalizers);
}
}
void IndexIVFFastScan::search_dispatch_implem(
idx_t n,
const float* x,
faiss/IndexIVFFastScan.h
+10 -1
View File
@@ -99,7 +99,7 @@ struct IndexIVFFastScan : IndexIVF {
// compact way of conveying coarse quantization results
struct CoarseQuantized {
size_t nprobe;
size_t nprobe = 0;
const float* dis = nullptr;
const idx_t* ids = nullptr;
};
@@ -147,6 +147,15 @@ struct IndexIVFFastScan : IndexIVF {
const SearchParameters* params = nullptr) const override;
// internal search funcs
SIMDResultHandlerToFloat* make_knn_handler(
bool is_max,
int impl,
idx_t n,
idx_t k,
float* distances,
idx_t* labels,
const IDSelector* sel,
const float* normalizers = nullptr) const;
// dispatch to implementations and parallelize
void search_dispatch_implem(
faiss/IndexIVFPQFastScan.cpp
+129 -1
View File
@@ -7,6 +7,7 @@
#include <faiss/IndexIVFPQFastScan.h>
#include <array>
#include <cassert>
#include <cstdio>
@@ -14,6 +15,7 @@
#include <faiss/impl/AuxIndexStructures.h>
#include <faiss/impl/FaissAssert.h>
#include <faiss/utils/Heap.h>
#include <faiss/utils/distances.h>
#include <faiss/utils/simdlib.h>
@@ -213,7 +215,7 @@ void IndexIVFPQFastScan::compute_LUT(
AlignedTable<float>& biases) const {
size_t dim12 = pq.ksub * pq.M;
size_t d = pq.d;
size_t nprobe = this->nprobe;
size_t nprobe = cq.nprobe;
if (by_residual) {
if (metric_type == METRIC_L2) {
@@ -292,4 +294,130 @@ void IndexIVFPQFastScan::compute_LUT(
}
}
/*********************************************************
* InvertedListScanner for IVFPQFS
*********************************************************/
namespace {
struct IVFPQFastScanScanner : InvertedListScanner {
static constexpr int impl = 10; // based on search_implem_10
static constexpr size_t nq = 1; // 1 query at a time.
const IndexIVFPQFastScan& index;
AlignedTable<uint8_t> dis_tables;
AlignedTable<uint16_t> biases;
std::array<float, 2> normalizers{};
const float* xi = nullptr;
IVFPQFastScanScanner(
const IndexIVFPQFastScan& index,
bool store_pairs,
const IDSelector* sel)
: InvertedListScanner(store_pairs, sel), index(index) {
this->keep_max = is_similarity_metric(index.metric_type);
}
void set_query(const float* query) override {
this->xi = query;
}
void set_list(idx_t list_no, float coarse_dis) override {
this->list_no = list_no;
IndexIVFFastScan::CoarseQuantized cq{
.nprobe = 1, // 1 due to explicitly passing in list_no
.dis = &coarse_dis, // dis from query to list_no centroid.
.ids = &list_no, // id of the current list we are scanning
};
index.compute_LUT_uint8(1, xi, cq, dis_tables, biases, &normalizers[0]);
}
float distance_to_code(const uint8_t* /* code */) const override {
// It's not really possible to implement a distance_to_code since codes
// for 32 database vectors are intermixed.
FAISS_THROW_MSG("not implemented");
}
// Based on IVFFastScan search_implem_10, since it also deals with 1 query
// at a time.
size_t scan_codes(
size_t ntotal,
const uint8_t* codes,
const idx_t* ids,
float* distances,
idx_t* labels,
size_t k) const override {
// initialize the current iteration heap to the worst possible value of
// the prior loop
std::vector<float> curr_dists(k, distances[0]);
std::vector<idx_t> curr_labels(k, labels[0]);
std::unique_ptr<SIMDResultHandlerToFloat> handler(
index.make_knn_handler(
!keep_max,
impl,
nq,
k,
curr_dists.data(),
curr_labels.data(),
sel,
&normalizers[0]));
// This does not quite match search_implem_10, but it is fine because
// the scanner operates on a single query at a time, and this value is
// used as the query index. For a single query, the value is always 0.
int qmap1[1] = {0};
handler->q_map = qmap1;
handler->begin(&normalizers[0]);
const uint8_t* LUT = dis_tables.get();
handler->dbias = biases.get();
handler->ntotal = ntotal;
handler->id_map = ids;
pq4_accumulate_loop(
1,
roundup(ntotal, index.bbs),
index.bbs,
static_cast<int>(index.M2),
codes,
LUT,
*handler,
nullptr);
// The handler is for the results of this iteration.
// Then we need a second heap to combine across iterations.
handler->end();
if (keep_max) {
minheap_addn(
k,
distances,
labels,
curr_dists.data(),
curr_labels.data(),
k);
} else {
maxheap_addn(
k,
distances,
labels,
curr_dists.data(),
curr_labels.data(),
k);
}
return handler->num_updates();
}
};
} // anonymous namespace
InvertedListScanner* IndexIVFPQFastScan::get_InvertedListScanner(
bool store_pairs,
const IDSelector* sel,
const IVFSearchParameters*) const {
return new IVFPQFastScanScanner(*this, store_pairs, sel);
}
} // namespace faiss
faiss/IndexIVFPQFastScan.h
+5
View File
@@ -81,6 +81,11 @@ struct IndexIVFPQFastScan : IndexIVFFastScan {
const CoarseQuantized& cq,
AlignedTable<float>& dis_tables,
AlignedTable<float>& biases) const override;
InvertedListScanner* get_InvertedListScanner(
bool store_pairs,
const IDSelector* sel,
const IVFSearchParameters*) const override;
};
} // namespace faiss
faiss/impl/simd_result_handlers.h
+36 -8
View File
@@ -46,7 +46,7 @@ struct SIMDResultHandler {
virtual ~SIMDResultHandler() {}
};
/* Result handler that will return float resutls eventually */
/* Result handler that will return float results eventually */
struct SIMDResultHandlerToFloat : SIMDResultHandler {
size_t nq; // number of queries
size_t ntotal; // ignore excess elements after ntotal
@@ -70,6 +70,12 @@ struct SIMDResultHandlerToFloat : SIMDResultHandler {
virtual void end() {
normalizers = nullptr;
}
// Number of updates made to the underlying data structure.
// For example: number of heap updates.
virtual size_t num_updates() {
return 0;
}
};
FAISS_API extern bool simd_result_handlers_accept_virtual;
@@ -318,8 +324,8 @@ struct HeapHandler : ResultHandlerCompare<C, with_id_map> {
std::vector<TI> iids;
float* dis;
int64_t* ids;
int64_t k; // number of results to keep
size_t k; // number of results to keep
size_t nup = 0; // number of heap updates
HeapHandler(
size_t nq,
@@ -327,14 +333,30 @@ struct HeapHandler : ResultHandlerCompare<C, with_id_map> {
int64_t k,
float* dis,
int64_t* ids,
const IDSelector* sel_in)
const IDSelector* sel_in,
const float* normalizers = nullptr)
: RHC(nq, ntotal, sel_in),
idis(nq * k),
iids(nq * k),
idis(nq * k, threshold_idis(dis, normalizers)),
iids(nq * k, -1),
dis(dis),
ids(ids),
k(k) {
heap_heapify<C>(k * nq, idis.data(), iids.data());
k(k) {}
static uint16_t threshold_idis(float* dis_in, const float* normalizers) {
if (dis_in[0] == std::numeric_limits<float>::max()) {
return std::numeric_limits<uint16_t>::max();
}
if (dis_in[0] == std::numeric_limits<float>::lowest()) {
return 0;
}
if (normalizers) {
float one_a = 1 / normalizers[0], b = normalizers[1];
float f = (dis_in[0] - b) / one_a;
f = C::is_max ? std::ceil(f) : std::floor(f);
return std::clamp<float>(
f, 0, std::numeric_limits<uint16_t>::max());
}
return C::neutral();
}
void handle(size_t q, size_t b, simd16uint16 d0, simd16uint16 d1) final {
@@ -372,6 +394,7 @@ struct HeapHandler : ResultHandlerCompare<C, with_id_map> {
if (C::cmp(heap_dis[0], dis_2)) {
heap_replace_top<C>(
k, heap_dis, heap_ids, dis_2, real_idx);
nup++;
}
}
}
@@ -384,6 +407,7 @@ struct HeapHandler : ResultHandlerCompare<C, with_id_map> {
if (C::cmp(heap_dis[0], dis_2)) {
int64_t idx = this->adjust_id(b, j);
heap_replace_top<C>(k, heap_dis, heap_ids, dis_2, idx);
nup++;
}
}
}
@@ -408,6 +432,10 @@ struct HeapHandler : ResultHandlerCompare<C, with_id_map> {
}
}
}
size_t num_updates() override {
return nup;
}
};
/** Simple top-N implementation using a reservoir.
tests/test_lowlevel_ivf.cpp
+119 -14
View File
@@ -5,7 +5,6 @@
* LICENSE file in the root directory of this source tree.
*/
#include <cinttypes>
#include <cstdio>
#include <cstdlib>
@@ -20,6 +19,7 @@
#include <faiss/IVFlib.h>
#include <faiss/IndexBinaryIVF.h>
#include <faiss/IndexIVF.h>
#include <faiss/IndexIVFPQFastScan.h>
#include <faiss/IndexPreTransform.h>
#include <faiss/index_factory.h>
#include <faiss/utils/distances.h>
@@ -64,17 +64,31 @@ std::unique_ptr<Index> make_trained_index(
return index;
}
std::vector<idx_t> search_index(Index* index, const float* xq) {
std::pair<std::vector<idx_t>, std::vector<float>> search_index(
Index* index,
const float* xq) {
std::vector<idx_t> I(k * nq);
std::vector<float> D(k * nq);
index->search(nq, xq, k, D.data(), I.data());
return I;
return {I, D};
}
/*************************************************************
* Test functions for a given index type
*************************************************************/
void test_get_InvertedListScanner(
IndexIVF* index_ivf,
const IndexPreTransform* index_pt,
std::vector<uint8_t> codes,
std::unique_ptr<InvertedListScanner> scanner,
std::vector<float> xq,
std::vector<idx_t> ref_I,
std::vector<float> ref_D,
MetricType metric,
bool distance_to_code_supported = true,
float accuracy_requirement = 1.0);
void test_lowlevel_access(const char* index_key, MetricType metric) {
std::unique_ptr<Index> index = make_trained_index(index_key, metric);
@@ -152,8 +166,30 @@ void test_lowlevel_access(const char* index_key, MetricType metric) {
// sample some example queries and get reference search results.
auto xq = make_data(nq);
auto ref_I = search_index(index.get(), xq.data());
auto res = search_index(index.get(), xq.data());
test_get_InvertedListScanner(
index_ivf,
index_pt,
codes,
nullptr,
xq,
res.first,
res.second,
metric);
}
void test_get_InvertedListScanner(
IndexIVF* index_ivf,
const IndexPreTransform* index_pt,
std::vector<uint8_t> codes,
std::unique_ptr<InvertedListScanner> scanner,
std::vector<float> xq,
std::vector<idx_t> ref_I,
std::vector<float> ref_D,
MetricType metric,
bool distance_to_code_supported,
float accuracy_requirement) {
// handle preprocessing
const float* xqt = xq.data();
std::unique_ptr<float[]> del_xqt;
@@ -167,22 +203,27 @@ void test_lowlevel_access(const char* index_key, MetricType metric) {
// quantize the queries to get the inverted list ids to visit.
int nprobe = index_ivf->nprobe;
const InvertedLists* il = index_ivf->invlists;
std::vector<idx_t> q_lists(nq * nprobe);
std::vector<float> q_dis(nq * nprobe);
index_ivf->quantizer->search(nq, xqt, nprobe, q_dis.data(), q_lists.data());
// object that does the scanning and distance computations.
std::unique_ptr<InvertedListScanner> scanner(
index_ivf->get_InvertedListScanner());
if (scanner == nullptr) {
// standard flow
scanner = std::unique_ptr<InvertedListScanner>(
index_ivf->get_InvertedListScanner());
}
float recall = 0.0;
for (int i = 0; i < nq; i++) {
std::vector<idx_t> I(k, -1);
float default_dis = metric == METRIC_L2 ? HUGE_VAL : -HUGE_VAL;
float default_dis = metric == METRIC_L2
? std::numeric_limits<float>::max()
: std::numeric_limits<float>::lowest();
std::vector<float> D(k, default_dis);
scanner->set_query(xqt + i * dt);
scanner->set_query(xqt + i * index_ivf->d);
for (int j = 0; j < nprobe; j++) {
int list_no = q_lists[i * nprobe + j];
@@ -202,7 +243,7 @@ void test_lowlevel_access(const char* index_key, MetricType metric) {
I.data(),
k);
if (j == 0) {
if (distance_to_code_supported && j == 0) {
// all results so far come from list_no, so let's check if
// the distance function works
for (int jj = 0; jj < k; jj++) {
@@ -229,9 +270,57 @@ void test_lowlevel_access(const char* index_key, MetricType metric) {
// check that we have the same results as the reference search
for (int j = 0; j < k; j++) {
EXPECT_EQ(I[j], ref_I[i * k + j]);
if (I[j] != ref_I[i * k + j]) {
if (accuracy_requirement == 1.0) {
EXPECT_EQ(D[j], ref_D[i * k + j]);
}
if (D[j] == ref_D[i * k + j]) {
recall += 1.0;
}
} else {
recall += 1.0;
}
}
}
recall = recall / (k * nq);
EXPECT_GE(recall, accuracy_requirement);
}
void test_ivfpqfs_scanner(MetricType metric, bool by_residual = false) {
auto index =
std::unique_ptr<Index>(index_factory(d, "IVF32,PQ4x4fs", metric));
ParameterSpace().set_index_parameter(index.get(), "nprobe", 4);
IndexIVFPQFastScan* index_ivf = static_cast<IndexIVFPQFastScan*>(
ivflib::extract_index_ivf(index.get()));
index_ivf->by_residual = by_residual;
// implem_10 also processes one query at a time, so compare with that.
index_ivf->implem = 10;
auto xt = make_data(nt);
index_ivf->train(nt, xt.data());
auto xb = make_data(nb);
index_ivf->add(nb, xb.data());
// Initialize scanner with context in params so heap results can persist
// across multiple inverted list scans.
auto scanner = std::unique_ptr<InvertedListScanner>(
index_ivf->get_InvertedListScanner(true, nullptr, nullptr));
// ref data
auto xq = make_data(nq);
auto res = search_index(index_ivf, xq.data());
// distance_to_code_supported = false because codes are intermixed.
test_get_InvertedListScanner(
index_ivf,
nullptr,
{},
std::move(scanner),
xq,
res.first,
res.second,
metric,
false,
by_residual ? 0.9 : 1.0 // recall is not 100% for by_residual
);
}
} // anonymous namespace
@@ -276,6 +365,22 @@ TEST(TestLowLevelIVF, IVFRQ) {
test_lowlevel_access("IVF32,RQ16x8", METRIC_L2);
}
TEST(TestLowLevelIVF, IVFPQFS_L2) {
test_ivfpqfs_scanner(METRIC_L2);
}
TEST(TestLowLevelIVF, IVFPQFS_IP) {
test_ivfpqfs_scanner(METRIC_INNER_PRODUCT);
}
TEST(TestLowLevelIVF, IVFPQFSr_L2) {
test_ivfpqfs_scanner(METRIC_L2, true);
}
TEST(TestLowLevelIVF, IVFPQFSr_IP) {
test_ivfpqfs_scanner(METRIC_INNER_PRODUCT, true);
}
/*************************************************************
* Same for binary (a bit simpler)
*************************************************************/
@@ -459,7 +564,7 @@ void test_threaded_search(const char* index_key, MetricType metric) {
// sample some example queries and get reference search results.
auto xq = make_data(nq);
auto ref_I = search_index(index.get(), xq.data());
auto res = search_index(index.get(), xq.data());
// handle preprocessing
const float* xqt = xq.data();
@@ -570,7 +675,7 @@ void test_threaded_search(const char* index_key, MetricType metric) {
// check that we have the same results as the reference search
for (int j = 0; j < k; j++) {
EXPECT_EQ(I[j], ref_I[i * k + j]);
EXPECT_EQ(I[j], res.first[i * k + j]);
}
}
}