facebookresearch/faiss
1
0
Fork 0
mirror of https://github.com/facebookresearch/faiss.git synced 2026-10-11 22:50:00 +00:00
Code Issues Packages Projects Releases Wiki Activity

Inline PQ code distance kernels into scanner TUs (#5159)

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

After the SIMD dispatch refactoring, the PQ code distance implementations
(pq_code_distance_single_impl<SL>, pq_code_distance_four_impl<SL>) lived in
separate translation units from the scanner loops (IVFPQScannerT::scan_list_with_table,
PQDistanceComputer::distance_to_code). The compiler could not inline the SIMD
gather/accumulate code into the hot inner loops.

This diff converts the per-SIMD pq_code_distance .cpp files to .h headers and
includes them in the corresponding scanner TUs before the scanner _impl.h
includes. This puts the kernel definitions in the same TU, enabling the compiler
to inline the AVX2/AVX512 vgatherdps code directly into scan_list_with_table and
scan_list_polysemous_hc.

Changes:
- pq_code_distance-avx2.cpp → pq_code_distance-avx2.h (header, #pragma once)
- pq_code_distance-avx512.cpp → pq_code_distance-avx512.h (same)
- New pq_code_distance-generic.h with inline NONE/ARM_NEON specializations
- Scanner TUs (avx2.cpp, avx512.cpp, neon.cpp) include the PQ distance headers
- PQCodeDistance gains static constexpr simd_level member
- scan_list_polysemous uses PQCodeDist::simd_level for Hamming computer dispatch
  (was hardcoded to SIMDLevel::NONE)
- Scanner TUs include per-ISA hamming_computer headers for SIMD Hamming dispatch
- Build files (xplat.bzl, CMakeLists.txt) updated

Verified via objdump that scan_list_with_table and scan_list_polysemous_hc contain
zero calls to pq_code_distance_*_impl — all AVX2/AVX512 gather code is fully inlined.

Reviewed By: algoriddle

Differential Revision: D102942787

fbshipit-source-id: d2979f9cd629652ac1de4886afd3c335fbb4fac8
This commit is contained in:
Matthijs Douze
2026-05-04 07:35:56 -07:00
committed by meta-codesync[bot]
parent 417c53e0ce
commit 28b2b66cf2
19 changed files with 1953 additions and 1036 deletions
Download Patch File Download Diff File Expand all files Collapse all files
faiss/CMakeLists.txt
+10 -3
View File
@@ -11,7 +11,7 @@
set(FAISS_SIMD_AVX2_SRC
impl/fast_scan/impl-avx2.cpp
impl/hnsw/avx2.cpp
impl/pq_code_distance/pq_code_distance-avx2.cpp
impl/pq_code_distance/avx2.cpp
impl/scalar_quantizer/sq-avx2.cpp
impl/approx_topk/avx2.cpp
impl/binary_hamming/avx2.cpp
@@ -25,7 +25,7 @@ set(FAISS_SIMD_AVX2_SRC
set(FAISS_SIMD_AVX512_SRC
impl/fast_scan/impl-avx512.cpp
impl/hnsw/avx512.cpp
impl/pq_code_distance/pq_code_distance-avx512.cpp
impl/pq_code_distance/avx512.cpp
impl/scalar_quantizer/sq-avx512.cpp
impl/binary_hamming/avx512.cpp
utils/simd_impl/distances_avx512.cpp
@@ -39,6 +39,7 @@ set(FAISS_SIMD_NEON_SRC
impl/scalar_quantizer/sq-neon.cpp
impl/approx_topk/neon.cpp
impl/binary_hamming/neon.cpp
impl/pq_code_distance/neon.cpp
utils/simd_impl/distances_aarch64.cpp
utils/hamming_distance/hamming_neon.cpp
utils/simd_impl/partitioning_neon.cpp
@@ -133,6 +134,7 @@ set(FAISS_SRC
impl/PolysemousTraining.cpp
impl/ProductQuantizer.cpp
impl/pq_code_distance/pq_code_distance-generic.cpp
impl/pq_code_distance/IVFPQ_QueryTables.cpp
impl/AdditiveQuantizer.cpp
impl/RaBitQuantizer.cpp
impl/RaBitQuantizerMultiBit.cpp
@@ -311,8 +313,13 @@ set(FAISS_HEADERS
impl/simd_dispatch.h
impl/fast_scan/simd_result_handlers.h
impl/zerocopy_io.h
utils/pq_code_distance.h
impl/pq_code_distance/pq_code_distance-inl.h
impl/pq_code_distance/pq_code_distance-avx2.h
impl/pq_code_distance/pq_code_distance-avx512.h
impl/pq_code_distance/pq_code_distance-generic.h
impl/pq_code_distance/IVFPQ_QueryTables.h
impl/pq_code_distance/IVFPQScanner_impl.h
impl/pq_code_distance/PQDistanceComputer_impl.h
invlists/BlockInvertedLists.h
invlists/DirectMap.h
invlists/InvertedLists.h
faiss/IndexIVFPQ.cpp
+9 -809
View File
@@ -28,9 +28,15 @@
#include <faiss/impl/IDSelector.h>
#include <faiss/impl/ProductQuantizer.h>
#include <faiss/impl/ResultHandler.h>
#include <faiss/impl/pq_code_distance/pq_code_distance-inl.h>
#include <faiss/impl/pq_code_distance/pq_code_distance-generic.h>
#include <faiss/impl/simd_dispatch.h>
// Scalar (NONE) fallback for dynamic dispatch
#define THE_SIMD_LEVEL SIMDLevel::NONE
// NOLINTNEXTLINE(facebook-hte-InlineHeader)
#include <faiss/impl/pq_code_distance/IVFPQScanner_impl.h>
#undef THE_SIMD_LEVEL
namespace faiss {
/*****************************************
@@ -491,819 +497,13 @@ void IndexIVFPQ::precompute_table() {
verbose);
}
namespace {
#define TIC t0 = get_cycles()
#define TOC get_cycles() - t0
/** QueryTables manages the various ways of searching an
* IndexIVFPQ. The code contains a lot of branches, depending on:
* - metric_type: are we computing L2 or Inner product similarity?
* - by_residual: do we encode raw vectors or residuals?
* - use_precomputed_table: are x_R|x_C tables precomputed?
* - polysemous_ht: are we filtering with polysemous codes?
*/
struct QueryTables {
/*****************************************************
* General data from the IVFPQ
*****************************************************/
const IndexIVFPQ& ivfpq;
const IVFSearchParameters* params;
// copied from IndexIVFPQ for easier access
int d;
const ProductQuantizer& pq;
MetricType metric_type;
bool by_residual;
int use_precomputed_table;
int polysemous_ht;
// pre-allocated data buffers
float *sim_table, *sim_table_2;
float *residual_vec, *decoded_vec;
// single data buffer
std::vector<float> mem;
// for table pointers
std::vector<const float*> sim_table_ptrs;
explicit QueryTables(
const IndexIVFPQ& ivfpq_in,
const IVFSearchParameters* params_in)
: ivfpq(ivfpq_in),
params(params_in),
d(ivfpq_in.d),
pq(ivfpq_in.pq),
metric_type(ivfpq_in.metric_type),
by_residual(ivfpq_in.by_residual),
use_precomputed_table(ivfpq_in.use_precomputed_table) {
mem.resize(pq.ksub * pq.M * 2 + d * 2);
sim_table = mem.data();
sim_table_2 = sim_table + pq.ksub * pq.M;
residual_vec = sim_table_2 + pq.ksub * pq.M;
decoded_vec = residual_vec + d;
// for polysemous
polysemous_ht = ivfpq_in.polysemous_ht;
if (auto ivfpq_params =
dynamic_cast<const IVFPQSearchParameters*>(params_in)) {
polysemous_ht = ivfpq_params->polysemous_ht;
}
if (polysemous_ht != 0) {
q_code.resize(pq.code_size);
}
init_list_cycles = 0;
sim_table_ptrs.resize(pq.M);
}
/*****************************************************
* What we do when query is known
*****************************************************/
// field specific to query
const float* qi = nullptr;
// query-specific initialization
void init_query(const float* qi_in) {
this->qi = qi_in;
if (metric_type == METRIC_INNER_PRODUCT) {
init_query_IP();
} else {
init_query_L2();
}
if (!by_residual && polysemous_ht != 0) {
pq.compute_code(qi_in, q_code.data());
}
}
void init_query_IP() {
// precompute some tables specific to the query qi
pq.compute_inner_prod_table(qi, sim_table);
}
void init_query_L2() {
if (!by_residual) {
pq.compute_distance_table(qi, sim_table);
} else if (use_precomputed_table) {
pq.compute_inner_prod_table(qi, sim_table_2);
}
}
/*****************************************************
* When inverted list is known: prepare computations
*****************************************************/
// fields specific to list
idx_t key = 0;
float coarse_dis = 0.0f;
std::vector<uint8_t> q_code;
uint64_t init_list_cycles;
/// once we know the query and the centroid, we can prepare the
/// sim_table that will be used for accumulation
/// and dis0, the initial value
float precompute_list_tables() {
float dis0 = 0;
uint64_t t0;
TIC;
if (by_residual) {
if (metric_type == METRIC_INNER_PRODUCT) {
dis0 = precompute_list_tables_IP();
} else {
dis0 = precompute_list_tables_L2();
}
}
init_list_cycles += TOC;
return dis0;
}
float precompute_list_table_pointers() {
float dis0 = 0;
uint64_t t0;
TIC;
if (by_residual) {
if (metric_type == METRIC_INNER_PRODUCT) {
FAISS_THROW_MSG("not implemented");
} else {
dis0 = precompute_list_table_pointers_L2();
}
}
init_list_cycles += TOC;
return dis0;
}
/*****************************************************
* compute tables for inner prod
*****************************************************/
float precompute_list_tables_IP() {
// prepare the sim_table that will be used for accumulation
// and dis0, the initial value
ivfpq.quantizer->reconstruct(key, decoded_vec);
// decoded_vec = centroid
float dis0 = fvec_inner_product_dispatch(qi, decoded_vec, d);
if (polysemous_ht) {
for (int i = 0; i < d; i++) {
residual_vec[i] = qi[i] - decoded_vec[i];
}
pq.compute_code(residual_vec, q_code.data());
}
return dis0;
}
/*****************************************************
* compute tables for L2 distance
*****************************************************/
float precompute_list_tables_L2() {
float dis0 = 0;
if (use_precomputed_table == 0 || use_precomputed_table == -1) {
ivfpq.quantizer->compute_residual(qi, residual_vec, key);
pq.compute_distance_table(residual_vec, sim_table);
if (polysemous_ht != 0) {
pq.compute_code(residual_vec, q_code.data());
}
} else if (use_precomputed_table == 1) {
dis0 = coarse_dis;
fvec_madd_dispatch(
pq.M * pq.ksub,
ivfpq.precomputed_table.data() + key * pq.ksub * pq.M,
-2.0,
sim_table_2,
sim_table);
if (polysemous_ht != 0) {
ivfpq.quantizer->compute_residual(qi, residual_vec, key);
pq.compute_code(residual_vec, q_code.data());
}
} else if (use_precomputed_table == 2) {
dis0 = coarse_dis;
const MultiIndexQuantizer* miq =
dynamic_cast<const MultiIndexQuantizer*>(ivfpq.quantizer);
FAISS_THROW_IF_NOT(miq);
const ProductQuantizer& cpq = miq->pq;
size_t Mf = pq.M / cpq.M;
const float* qtab = sim_table_2; // query-specific table
float* ltab = sim_table; // (output) list-specific table
long k = key;
for (size_t cm = 0; cm < cpq.M; cm++) {
// compute PQ index
size_t ki = k & ((uint64_t(1) << cpq.nbits) - 1);
k >>= cpq.nbits;
// get corresponding table
const float* pc = ivfpq.precomputed_table.data() +
(ki * pq.M + cm * Mf) * pq.ksub;
if (polysemous_ht == 0) {
// sum up with query-specific table
fvec_madd_dispatch(Mf * pq.ksub, pc, -2.0, qtab, ltab);
ltab += Mf * pq.ksub;
qtab += Mf * pq.ksub;
} else {
for (size_t m = cm * Mf; m < (cm + 1) * Mf; m++) {
q_code[m] = fvec_madd_and_argmin_dispatch(
pq.ksub, pc, -2, qtab, ltab);
pc += pq.ksub;
ltab += pq.ksub;
qtab += pq.ksub;
}
}
}
}
return dis0;
}
float precompute_list_table_pointers_L2() {
float dis0 = 0;
if (use_precomputed_table == 1) {
dis0 = coarse_dis;
const float* s =
ivfpq.precomputed_table.data() + key * pq.ksub * pq.M;
for (size_t m = 0; m < pq.M; m++) {
sim_table_ptrs[m] = s;
s += pq.ksub;
}
} else if (use_precomputed_table == 2) {
dis0 = coarse_dis;
const MultiIndexQuantizer* miq =
dynamic_cast<const MultiIndexQuantizer*>(ivfpq.quantizer);
FAISS_THROW_IF_NOT(miq);
const ProductQuantizer& cpq = miq->pq;
size_t Mf = pq.M / cpq.M;
long k = key;
size_t m0 = 0;
for (size_t cm = 0; cm < cpq.M; cm++) {
size_t ki = k & ((uint64_t(1) << cpq.nbits) - 1);
k >>= cpq.nbits;
const float* pc = ivfpq.precomputed_table.data() +
(ki * pq.M + cm * Mf) * pq.ksub;
for (size_t m = m0; m < m0 + Mf; m++) {
sim_table_ptrs[m] = pc;
pc += pq.ksub;
}
m0 += Mf;
}
} else {
FAISS_THROW_MSG("need precomputed tables");
}
if (polysemous_ht) {
FAISS_THROW_MSG("not implemented");
// Not clear that it makes sense to implemente this,
// because it costs M * ksub, which is what we wanted to
// avoid with the tables pointers.
}
return dis0;
}
};
template <class C, bool use_sel>
struct WrappedSearchResult {
ResultHandler& res;
size_t nup = 0;
idx_t list_no;
const idx_t* ids;
const IDSelector* sel;
WrappedSearchResult(
idx_t list_no_in,
const idx_t* ids_in,
const IDSelector* sel_in,
ResultHandler& res_in)
: res(res_in), list_no(list_no_in), ids(ids_in), sel(sel_in) {}
inline bool skip_entry(idx_t j) {
return use_sel && !sel->is_member(ids[j]);
}
inline void add(idx_t j, float dis) {
// Reached only for codes that passed skip_entry — i.e. distance
// was actually computed for this code (post-filter).
res.stats.scan_cnt++;
if (C::cmp(res.threshold, dis)) {
idx_t id = ids ? ids[j] : lo_build(this->list_no, j);
if (res.add_result(dis, id)) {
res.stats.nheap_updates++;
nup++;
}
}
}
};
/*****************************************************
* Scaning the codes.
* The scanning functions call their favorite precompute_*
* function to precompute the tables they need.
*****************************************************/
template <typename IDType, MetricType METRIC_TYPE, class PQCodeDist>
struct IVFPQScannerT : QueryTables {
using PQDecoder = typename PQCodeDist::PQDecoder;
const uint8_t* list_codes = nullptr;
const IDType* list_ids;
size_t list_size = 0;
IVFPQScannerT(
const IndexIVFPQ& ivfpq_in,
const IVFSearchParameters* params_in)
: QueryTables(ivfpq_in, params_in) {
FAISS_THROW_IF_NOT(METRIC_TYPE == metric_type);
}
float dis0 = 0.0f;
void init_list(idx_t list_no, float coarse_dis_in, int mode) {
this->key = list_no;
this->coarse_dis = coarse_dis_in;
if (mode == 2) {
dis0 = precompute_list_tables();
} else if (mode == 1) {
dis0 = precompute_list_table_pointers();
}
}
/*****************************************************
* Scaning the codes: simple PQ scan.
*****************************************************/
// This is the baseline version of scan_list_with_tables().
// It demonstrates what this function actually does.
//
// /// version of the scan where we use precomputed tables.
// template <class SearchResultType>
// void scan_list_with_table(
// size_t ncode,
// const uint8_t* codes,
// SearchResultType& res) const {
//
// for (size_t j = 0; j < ncode; j++, codes += pq.code_size) {
// if (res.skip_entry(j)) {
// continue;
// }
// float dis = dis0 + PQCodeDist::distance_single_code(
// pq, sim_table, codes);
// res.add(j, dis);
// }
// }
// This is the modified version of scan_list_with_tables().
// It was observed that doing manual unrolling of the loop that
// utilizes distance_single_code() speeds up the computations.
/// version of the scan where we use precomputed tables.
template <class SearchResultType>
void scan_list_with_table(
size_t ncode,
const uint8_t* codes,
SearchResultType& res) const {
int counter = 0;
size_t saved_j[4] = {0, 0, 0, 0};
for (size_t j = 0; j < ncode; j++) {
if (res.skip_entry(j)) {
continue;
}
saved_j[0] = (counter == 0) ? j : saved_j[0];
saved_j[1] = (counter == 1) ? j : saved_j[1];
saved_j[2] = (counter == 2) ? j : saved_j[2];
saved_j[3] = (counter == 3) ? j : saved_j[3];
counter += 1;
if (counter == 4) {
float distance_0 = 0;
float distance_1 = 0;
float distance_2 = 0;
float distance_3 = 0;
PQCodeDist::distance_four_codes(
pq.M,
pq.nbits,
sim_table,
codes + saved_j[0] * pq.code_size,
codes + saved_j[1] * pq.code_size,
codes + saved_j[2] * pq.code_size,
codes + saved_j[3] * pq.code_size,
distance_0,
distance_1,
distance_2,
distance_3);
res.add(saved_j[0], dis0 + distance_0);
res.add(saved_j[1], dis0 + distance_1);
res.add(saved_j[2], dis0 + distance_2);
res.add(saved_j[3], dis0 + distance_3);
counter = 0;
}
}
if (counter >= 1) {
float dis = dis0 +
PQCodeDist::distance_single_code(
pq.M,
pq.nbits,
sim_table,
codes + saved_j[0] * pq.code_size);
res.add(saved_j[0], dis);
}
if (counter >= 2) {
float dis = dis0 +
PQCodeDist::distance_single_code(
pq.M,
pq.nbits,
sim_table,
codes + saved_j[1] * pq.code_size);
res.add(saved_j[1], dis);
}
if (counter >= 3) {
float dis = dis0 +
PQCodeDist::distance_single_code(
pq.M,
pq.nbits,
sim_table,
codes + saved_j[2] * pq.code_size);
res.add(saved_j[2], dis);
}
}
/// tables are not precomputed, but pointers are provided to the
/// relevant X_c|x_r tables
template <class SearchResultType>
void scan_list_with_pointer(
size_t ncode,
const uint8_t* codes,
SearchResultType& res) const {
for (size_t j = 0; j < ncode; j++, codes += pq.code_size) {
if (res.skip_entry(j)) {
continue;
}
PQDecoder decoder(codes, pq.nbits);
float dis = dis0;
const float* tab = sim_table_2;
for (size_t m = 0; m < pq.M; m++) {
int ci = decoder.decode();
dis += sim_table_ptrs[m][ci] - 2 * tab[ci];
tab += pq.ksub;
}
res.add(j, dis);
}
}
/// nothing is precomputed: access residuals on-the-fly
template <class SearchResultType>
void scan_on_the_fly_dist(
size_t ncode,
const uint8_t* codes,
SearchResultType& res) const {
const float* dvec;
float local_dis0 = 0;
if (by_residual) {
if (METRIC_TYPE == METRIC_INNER_PRODUCT) {
ivfpq.quantizer->reconstruct(key, residual_vec);
local_dis0 = fvec_inner_product_dispatch(residual_vec, qi, d);
} else {
ivfpq.quantizer->compute_residual(qi, residual_vec, key);
}
dvec = residual_vec;
} else {
dvec = qi;
local_dis0 = 0;
}
for (size_t j = 0; j < ncode; j++, codes += pq.code_size) {
if (res.skip_entry(j)) {
continue;
}
pq.decode(codes, decoded_vec);
float dis;
if (METRIC_TYPE == METRIC_INNER_PRODUCT) {
dis = local_dis0 +
fvec_inner_product_dispatch(decoded_vec, qi, d);
} else {
dis = fvec_L2sqr_dispatch(decoded_vec, dvec, d);
}
res.add(j, dis);
}
}
/*****************************************************
* Scanning codes with polysemous filtering
*****************************************************/
// This is the baseline version of scan_list_polysemous_hc().
// It demonstrates what this function actually does.
// template <class HammingComputer, class SearchResultType>
// void scan_list_polysemous_hc(
// size_t ncode,
// const uint8_t* codes,
// SearchResultType& res) const {
// int ht = ivfpq.polysemous_ht;
// size_t n_hamming_pass = 0, nup = 0;
//
// int code_size = pq.code_size;
//
// HammingComputer hc(q_code.data(), code_size);
//
// for (size_t j = 0; j < ncode; j++, codes += code_size) {
// if (res.skip_entry(j)) {
// continue;
// }
// const uint8_t* b_code = codes;
// int hd = hc.hamming(b_code);
// if (hd < ht) {
// n_hamming_pass++;
//
// float dis =
// dis0 +
// PQCodeDist::distance_single_code(
// pq, sim_table, codes);
//
// res.add(j, dis);
// }
// }
// #pragma omp critical
// { indexIVFPQ_stats.n_hamming_pass += n_hamming_pass; }
// }
// This is the modified version of scan_list_with_tables().
// It was observed that doing manual unrolling of the loop that
// utilizes distance_single_code() speeds up the computations.
template <class HammingComputer, class SearchResultType>
void scan_list_polysemous_hc(
size_t ncode,
const uint8_t* codes,
SearchResultType& res) const {
int ht = ivfpq.polysemous_ht;
size_t n_hamming_pass = 0;
int code_size = static_cast<int>(pq.code_size);
size_t saved_j[8];
int counter = 0;
HammingComputer hc(q_code.data(), code_size);
for (size_t j = 0; j < (ncode / 4) * 4; j += 4) {
const uint8_t* b_code = codes + j * code_size;
// Unrolling is a key. Basically, doing multiple popcount
// operations one after another speeds things up.
// 9999999 is just an arbitrary large number
int hd0 = (res.skip_entry(j + 0))
? 99999999
: hc.hamming(b_code + 0 * code_size);
int hd1 = (res.skip_entry(j + 1))
? 99999999
: hc.hamming(b_code + 1 * code_size);
int hd2 = (res.skip_entry(j + 2))
? 99999999
: hc.hamming(b_code + 2 * code_size);
int hd3 = (res.skip_entry(j + 3))
? 99999999
: hc.hamming(b_code + 3 * code_size);
saved_j[counter] = j + 0;
counter = (hd0 < ht) ? (counter + 1) : counter;
saved_j[counter] = j + 1;
counter = (hd1 < ht) ? (counter + 1) : counter;
saved_j[counter] = j + 2;
counter = (hd2 < ht) ? (counter + 1) : counter;
saved_j[counter] = j + 3;
counter = (hd3 < ht) ? (counter + 1) : counter;
if (counter >= 4) {
// process four codes at the same time
n_hamming_pass += 4;
float distance_0 = dis0;
float distance_1 = dis0;
float distance_2 = dis0;
float distance_3 = dis0;
PQCodeDist::distance_four_codes(
pq.M,
pq.nbits,
sim_table,
codes + saved_j[0] * pq.code_size,
codes + saved_j[1] * pq.code_size,
codes + saved_j[2] * pq.code_size,
codes + saved_j[3] * pq.code_size,
distance_0,
distance_1,
distance_2,
distance_3);
res.add(saved_j[0], dis0 + distance_0);
res.add(saved_j[1], dis0 + distance_1);
res.add(saved_j[2], dis0 + distance_2);
res.add(saved_j[3], dis0 + distance_3);
//
counter -= 4;
saved_j[0] = saved_j[4];
saved_j[1] = saved_j[5];
saved_j[2] = saved_j[6];
saved_j[3] = saved_j[7];
}
}
for (int kk = 0; kk < counter; kk++) {
n_hamming_pass++;
float dis = dis0 +
PQCodeDist::distance_single_code(
pq.M,
pq.nbits,
sim_table,
codes + saved_j[kk] * pq.code_size);
res.add(saved_j[kk], dis);
}
// process leftovers
for (size_t j = (ncode / 4) * 4; j < ncode; j++) {
if (res.skip_entry(j)) {
continue;
}
const uint8_t* b_code = codes + j * code_size;
int hd = hc.hamming(b_code);
if (hd < ht) {
n_hamming_pass++;
float dis = dis0 +
PQCodeDist::distance_single_code(
pq.M,
pq.nbits,
sim_table,
codes + j * code_size);
res.add(j, dis);
}
}
#pragma omp critical
{
indexIVFPQ_stats.n_hamming_pass += n_hamming_pass;
}
}
template <class SearchResultType>
void scan_list_polysemous(
size_t ncode,
const uint8_t* codes,
SearchResultType& res) const {
with_HammingComputer<
SIMDLevel::NONE>(pq.code_size, [&]<class HammingComputer>() {
this->scan_list_polysemous_hc<HammingComputer, SearchResultType>(
ncode, codes, res);
});
}
};
/* We put as many parameters as possible in template. Hopefully the
* gain in runtime is worth the code bloat.
*
* C is the comparator < or >, it is directly related to METRIC_TYPE.
*
* precompute_mode is how much we precompute (2 = precompute distance tables,
* 1 = precompute pointers to distances, 0 = compute distances one by one).
* Currently only 2 is supported
*
* use_sel: store or ignore the IDSelector
*/
template <MetricType METRIC_TYPE, class C, class PQCodeDist, bool use_sel>
struct IVFPQScanner : IVFPQScannerT<idx_t, METRIC_TYPE, PQCodeDist>,
InvertedListScanner {
using InvertedListScanner::scan_codes;
int precompute_mode;
const IDSelector* sel;
IVFPQScanner(
const IndexIVFPQ& ivfpq_in,
bool store_pairs_in,
int precompute_mode_in,
const IDSelector* sel_in)
: IVFPQScannerT<idx_t, METRIC_TYPE, PQCodeDist>(ivfpq_in, nullptr),
precompute_mode(precompute_mode_in),
sel(sel_in) {
this->store_pairs = store_pairs_in;
this->keep_max = is_similarity_metric(METRIC_TYPE);
this->code_size = this->pq.code_size;
}
void set_query(const float* query) override {
this->init_query(query);
}
void set_list(idx_t list_no_in, float coarse_dis_in) override {
this->list_no = list_no_in;
this->init_list(list_no_in, coarse_dis_in, precompute_mode);
}
float distance_to_code(const uint8_t* code) const override {
FAISS_THROW_IF_NOT(precompute_mode == 2);
float dis = this->dis0 +
PQCodeDist::distance_single_code(
this->pq.M, this->pq.nbits, this->sim_table, code);
return dis;
}
size_t scan_codes(
size_t ncode,
const uint8_t* codes,
const idx_t* ids,
ResultHandler& handler) const override {
WrappedSearchResult<C, use_sel> res(
this->key,
this->store_pairs ? nullptr : ids,
this->sel,
handler);
if (this->polysemous_ht > 0) {
FAISS_THROW_IF_NOT(precompute_mode == 2);
this->scan_list_polysemous(ncode, codes, res);
} else if (precompute_mode == 2) {
this->scan_list_with_table(ncode, codes, res);
} else if (precompute_mode == 1) {
this->scan_list_with_pointer(ncode, codes, res);
} else if (precompute_mode == 0) {
this->scan_on_the_fly_dist(ncode, codes, res);
} else {
FAISS_THROW_MSG("bad precomp mode");
}
return res.nup;
}
};
} // anonymous namespace
InvertedListScanner* IndexIVFPQ::get_InvertedListScanner(
bool store_pairs,
const IDSelector* sel,
const IVFSearchParameters*) const {
return with_simd_level([&]<SIMDLevel SL>() -> InvertedListScanner* {
auto make =
[&]<class PQCodeDist, bool use_sel>() -> InvertedListScanner* {
if (metric_type == METRIC_INNER_PRODUCT) {
return new IVFPQScanner<
METRIC_INNER_PRODUCT,
CMin<float, idx_t>,
PQCodeDist,
use_sel>(*this, store_pairs, 2, sel);
} else if (metric_type == METRIC_L2) {
return new IVFPQScanner<
METRIC_L2,
CMax<float, idx_t>,
PQCodeDist,
use_sel>(*this, store_pairs, 2, sel);
} else {
FAISS_THROW_MSG("unsupported metric type");
}
};
auto with_decoder = [&]<bool use_sel>() -> InvertedListScanner* {
if (pq.nbits == 8) {
return make.template
operator()<PQCodeDistance<PQDecoder8, SL>, use_sel>();
} else if (pq.nbits == 16) {
return make.template
operator()<PQCodeDistance<PQDecoder16, SL>, use_sel>();
} else {
return make.template
operator()<PQCodeDistance<PQDecoderGeneric, SL>, use_sel>();
}
};
if (sel) {
return with_decoder.template operator()<true>();
} else {
return with_decoder.template operator()<false>();
}
return pq_code_distance::make_IVFPQInvertedListScanner<SL>(
*this, store_pairs, sel);
});
}
faiss/IndexPQ.cpp
+4 -82
View File
@@ -19,12 +19,14 @@
#include <faiss/impl/FaissAssert.h>
#include <faiss/utils/hamming.h>
#include <faiss/impl/pq_code_distance/pq_code_distance-inl.h>
#include <faiss/impl/pq_code_distance/pq_code_distance-generic.h>
#include <faiss/impl/simd_dispatch.h>
// Scalar (NONE) fallback for dynamic dispatch
#define THE_SIMD_LEVEL SIMDLevel::NONE
// NOLINTNEXTLINE(facebook-hte-InlineHeader)
#include <faiss/impl/pq_code_distance/PQDistanceComputer_impl.h>
// NOLINTNEXTLINE(facebook-hte-InlineHeader)
#include <faiss/impl/binary_hamming/IndexPQ_impl.h>
#include <faiss/utils/hamming_distance/hamming_computer-generic.h>
#undef THE_SIMD_LEVEL
@@ -79,89 +81,9 @@ void IndexPQ::train(idx_t n, const float* x) {
is_trained = true;
}
namespace {
template <class PQCodeDist>
struct PQDistanceComputer : FlatCodesDistanceComputer {
using PQDecoder = typename PQCodeDist::PQDecoder;
size_t d;
MetricType metric;
idx_t nb;
const ProductQuantizer& pq;
const float* sdc;
std::vector<float> precomputed_table;
size_t ndis;
const float* q;
float distance_to_code(const uint8_t* code) final {
ndis++;
float dis = PQCodeDist::distance_single_code(
pq.M, pq.nbits, precomputed_table.data(), code);
return dis;
}
float symmetric_dis(idx_t i, idx_t j) override {
FAISS_THROW_IF_NOT(sdc);
const float* sdci = sdc;
float accu = 0;
PQDecoder codei(codes + i * code_size, pq.nbits);
PQDecoder codej(codes + j * code_size, pq.nbits);
for (size_t l = 0; l < pq.M; l++) {
accu += sdci[codei.decode() + (codej.decode() << codei.nbits)];
sdci += uint64_t(1) << (2 * codei.nbits);
}
ndis++;
return accu;
}
explicit PQDistanceComputer(const IndexPQ& storage)
: FlatCodesDistanceComputer(
storage.codes.data(),
storage.code_size),
pq(storage.pq),
q(nullptr) {
precomputed_table.resize(pq.M * pq.ksub);
nb = storage.ntotal;
d = storage.d;
metric = storage.metric_type;
if (pq.sdc_table.size() == pq.ksub * pq.ksub * pq.M) {
sdc = pq.sdc_table.data();
} else {
sdc = nullptr;
}
ndis = 0;
}
void set_query(const float* x) override {
q = x;
if (metric == METRIC_L2) {
pq.compute_distance_table(x, precomputed_table.data());
} else {
pq.compute_inner_prod_table(x, precomputed_table.data());
}
}
};
template <SIMDLevel SL>
FlatCodesDistanceComputer* get_FlatCodesDistanceComputer1(
const IndexPQ& index) {
if (index.pq.nbits == 8) {
return new PQDistanceComputer<PQCodeDistance<PQDecoder8, SL>>(index);
} else if (index.pq.nbits == 16) {
return new PQDistanceComputer<PQCodeDistance<PQDecoder16, SL>>(index);
} else {
return new PQDistanceComputer<PQCodeDistance<PQDecoderGeneric, SL>>(
index);
}
}
} // namespace
FlatCodesDistanceComputer* IndexPQ::get_FlatCodesDistanceComputer() const {
return with_simd_level([&]<SIMDLevel SL>() {
return get_FlatCodesDistanceComputer1<SL>(*this);
return pq_code_distance::get_PQFlatCodesDistanceComputer<SL>(*this);
});
}
faiss/impl/pq_code_distance/IVFPQScanner_impl.h
+549
View File
@@ -0,0 +1,549 @@
/*
* 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
#ifndef THE_SIMD_LEVEL
#error "THE_SIMD_LEVEL must be defined before including IVFPQScanner_impl.h"
#endif
#include <faiss/IndexIVFPQ.h>
#include <faiss/impl/IDSelector.h>
#include <faiss/impl/ResultHandler.h>
#include <faiss/impl/pq_code_distance/IVFPQ_QueryTables.h>
#include <faiss/impl/pq_code_distance/pq_code_distance-inl.h>
#include <faiss/impl/simd_dispatch.h>
#include <faiss/invlists/DirectMap.h>
#include <faiss/utils/distances_dispatch.h>
#include <faiss/utils/hamming.h>
namespace faiss {
namespace pq_code_distance {
template <class C, bool use_sel>
struct WrappedSearchResult {
ResultHandler& res;
size_t nup = 0;
idx_t list_no;
const idx_t* ids;
const IDSelector* sel;
WrappedSearchResult(
idx_t list_no_in,
const idx_t* ids_in,
const IDSelector* sel_in,
ResultHandler& res_in)
: res(res_in), list_no(list_no_in), ids(ids_in), sel(sel_in) {}
inline bool skip_entry(idx_t j) {
return use_sel && !sel->is_member(ids[j]);
}
inline void add(idx_t j, float dis) {
// Reached only for codes that passed skip_entry — i.e. distance
// was actually computed for this code (post-filter).
res.stats.scan_cnt++;
if (C::cmp(res.threshold, dis)) {
idx_t id = ids ? ids[j] : lo_build(this->list_no, j);
if (res.add_result(dis, id)) {
res.stats.nheap_updates++;
nup++;
}
}
}
};
/*****************************************************
* Scaning the codes.
* The scanning functions call their favorite precompute_*
* function to precompute the tables they need.
*****************************************************/
template <typename IDType, MetricType METRIC_TYPE, class PQCodeDist>
struct IVFPQScannerT : QueryTables {
using PQDecoder = typename PQCodeDist::PQDecoder;
const uint8_t* list_codes = nullptr;
const IDType* list_ids;
size_t list_size = 0;
IVFPQScannerT(
const IndexIVFPQ& ivfpq_in,
const IVFSearchParameters* params_in)
: QueryTables(ivfpq_in, params_in) {
FAISS_THROW_IF_NOT(METRIC_TYPE == metric_type);
}
/*****************************************************
* Scaning the codes: simple PQ scan.
*****************************************************/
// This is the baseline version of scan_list_with_tables().
// It demonstrates what this function actually does.
//
// /// version of the scan where we use precomputed tables.
// template <class SearchResultType>
// void scan_list_with_table(
// size_t ncode,
// const uint8_t* codes,
// SearchResultType& res) const {
//
// for (size_t j = 0; j < ncode; j++, codes += pq.code_size) {
// if (res.skip_entry(j)) {
// continue;
// }
// float dis = dis0 + PQCodeDist::distance_single_code(
// pq, sim_table, codes);
// res.add(j, dis);
// }
// }
// This is the modified version of scan_list_with_tables().
// It was observed that doing manual unrolling of the loop that
// utilizes distance_single_code() speeds up the computations.
/// version of the scan where we use precomputed tables.
template <class SearchResultType>
void scan_list_with_table(
size_t ncode,
const uint8_t* codes,
SearchResultType& res) const {
int counter = 0;
size_t saved_j[4] = {0, 0, 0, 0};
for (size_t j = 0; j < ncode; j++) {
if (res.skip_entry(j)) {
continue;
}
saved_j[0] = (counter == 0) ? j : saved_j[0];
saved_j[1] = (counter == 1) ? j : saved_j[1];
saved_j[2] = (counter == 2) ? j : saved_j[2];
saved_j[3] = (counter == 3) ? j : saved_j[3];
counter += 1;
if (counter == 4) {
float distance_0 = 0;
float distance_1 = 0;
float distance_2 = 0;
float distance_3 = 0;
PQCodeDist::distance_four_codes(
pq.M,
pq.nbits,
sim_table,
codes + saved_j[0] * pq.code_size,
codes + saved_j[1] * pq.code_size,
codes + saved_j[2] * pq.code_size,
codes + saved_j[3] * pq.code_size,
distance_0,
distance_1,
distance_2,
distance_3);
res.add(saved_j[0], dis0 + distance_0);
res.add(saved_j[1], dis0 + distance_1);
res.add(saved_j[2], dis0 + distance_2);
res.add(saved_j[3], dis0 + distance_3);
counter = 0;
}
}
if (counter >= 1) {
float dis = dis0 +
PQCodeDist::distance_single_code(
pq.M,
pq.nbits,
sim_table,
codes + saved_j[0] * pq.code_size);
res.add(saved_j[0], dis);
}
if (counter >= 2) {
float dis = dis0 +
PQCodeDist::distance_single_code(
pq.M,
pq.nbits,
sim_table,
codes + saved_j[1] * pq.code_size);
res.add(saved_j[1], dis);
}
if (counter >= 3) {
float dis = dis0 +
PQCodeDist::distance_single_code(
pq.M,
pq.nbits,
sim_table,
codes + saved_j[2] * pq.code_size);
res.add(saved_j[2], dis);
}
}
/// tables are not precomputed, but pointers are provided to the
/// relevant X_c|x_r tables
template <class SearchResultType>
void scan_list_with_pointer(
size_t ncode,
const uint8_t* codes,
SearchResultType& res) const {
for (size_t j = 0; j < ncode; j++, codes += pq.code_size) {
if (res.skip_entry(j)) {
continue;
}
PQDecoder decoder(codes, pq.nbits);
float dis = dis0;
const float* tab = sim_table_2;
for (size_t m = 0; m < pq.M; m++) {
int ci = decoder.decode();
dis += sim_table_ptrs[m][ci] - 2 * tab[ci];
tab += pq.ksub;
}
res.add(j, dis);
}
}
/// nothing is precomputed: access residuals on-the-fly
template <class SearchResultType>
void scan_on_the_fly_dist(
size_t ncode,
const uint8_t* codes,
SearchResultType& res) const {
const float* dvec;
float local_dis0 = 0;
if (by_residual) {
if (METRIC_TYPE == METRIC_INNER_PRODUCT) {
ivfpq.quantizer->reconstruct(key, residual_vec);
local_dis0 = fvec_inner_product_dispatch(residual_vec, qi, d);
} else {
ivfpq.quantizer->compute_residual(qi, residual_vec, key);
}
dvec = residual_vec;
} else {
dvec = qi;
local_dis0 = 0;
}
for (size_t j = 0; j < ncode; j++, codes += pq.code_size) {
if (res.skip_entry(j)) {
continue;
}
pq.decode(codes, decoded_vec);
float dis;
if (METRIC_TYPE == METRIC_INNER_PRODUCT) {
dis = local_dis0 +
fvec_inner_product_dispatch(decoded_vec, qi, d);
} else {
dis = fvec_L2sqr_dispatch(decoded_vec, dvec, d);
}
res.add(j, dis);
}
}
/*****************************************************
* Scanning codes with polysemous filtering
*****************************************************/
// This is the baseline version of scan_list_polysemous_hc().
// It demonstrates what this function actually does.
// template <class HammingComputer, class SearchResultType>
// void scan_list_polysemous_hc(
// size_t ncode,
// const uint8_t* codes,
// SearchResultType& res) const {
// int ht = ivfpq.polysemous_ht;
// size_t n_hamming_pass = 0, nup = 0;
//
// int code_size = pq.code_size;
//
// HammingComputer hc(q_code.data(), code_size);
//
// for (size_t j = 0; j < ncode; j++, codes += code_size) {
// if (res.skip_entry(j)) {
// continue;
// }
// const uint8_t* b_code = codes;
// int hd = hc.hamming(b_code);
// if (hd < ht) {
// n_hamming_pass++;
//
// float dis =
// dis0 +
// PQCodeDist::distance_single_code(
// pq, sim_table, codes);
//
// res.add(j, dis);
// }
// }
// #pragma omp critical
// { indexIVFPQ_stats.n_hamming_pass += n_hamming_pass; }
// }
// This is the modified version of scan_list_with_tables().
// It was observed that doing manual unrolling of the loop that
// utilizes distance_single_code() speeds up the computations.
template <class HammingComputer, class SearchResultType>
void scan_list_polysemous_hc(
size_t ncode,
const uint8_t* codes,
SearchResultType& res) const {
int ht = ivfpq.polysemous_ht;
size_t n_hamming_pass = 0;
int code_size = static_cast<int>(pq.code_size);
size_t saved_j[8];
int counter = 0;
HammingComputer hc(q_code.data(), code_size);
for (size_t j = 0; j < (ncode / 4) * 4; j += 4) {
const uint8_t* b_code = codes + j * code_size;
// Unrolling is a key. Basically, doing multiple popcount
// operations one after another speeds things up.
// 9999999 is just an arbitrary large number
int hd0 = (res.skip_entry(j + 0))
? 99999999
: hc.hamming(b_code + 0 * code_size);
int hd1 = (res.skip_entry(j + 1))
? 99999999
: hc.hamming(b_code + 1 * code_size);
int hd2 = (res.skip_entry(j + 2))
? 99999999
: hc.hamming(b_code + 2 * code_size);
int hd3 = (res.skip_entry(j + 3))
? 99999999
: hc.hamming(b_code + 3 * code_size);
saved_j[counter] = j + 0;
counter = (hd0 < ht) ? (counter + 1) : counter;
saved_j[counter] = j + 1;
counter = (hd1 < ht) ? (counter + 1) : counter;
saved_j[counter] = j + 2;
counter = (hd2 < ht) ? (counter + 1) : counter;
saved_j[counter] = j + 3;
counter = (hd3 < ht) ? (counter + 1) : counter;
if (counter >= 4) {
// process four codes at the same time
n_hamming_pass += 4;
float distance_0 = dis0;
float distance_1 = dis0;
float distance_2 = dis0;
float distance_3 = dis0;
PQCodeDist::distance_four_codes(
pq.M,
pq.nbits,
sim_table,
codes + saved_j[0] * pq.code_size,
codes + saved_j[1] * pq.code_size,
codes + saved_j[2] * pq.code_size,
codes + saved_j[3] * pq.code_size,
distance_0,
distance_1,
distance_2,
distance_3);
res.add(saved_j[0], dis0 + distance_0);
res.add(saved_j[1], dis0 + distance_1);
res.add(saved_j[2], dis0 + distance_2);
res.add(saved_j[3], dis0 + distance_3);
//
counter -= 4;
saved_j[0] = saved_j[4];
saved_j[1] = saved_j[5];
saved_j[2] = saved_j[6];
saved_j[3] = saved_j[7];
}
}
for (int kk = 0; kk < counter; kk++) {
n_hamming_pass++;
float dis = dis0 +
PQCodeDist::distance_single_code(
pq.M,
pq.nbits,
sim_table,
codes + saved_j[kk] * pq.code_size);
res.add(saved_j[kk], dis);
}
// process leftovers
for (size_t j = (ncode / 4) * 4; j < ncode; j++) {
if (res.skip_entry(j)) {
continue;
}
const uint8_t* b_code = codes + j * code_size;
int hd = hc.hamming(b_code);
if (hd < ht) {
n_hamming_pass++;
float dis = dis0 +
PQCodeDist::distance_single_code(
pq.M,
pq.nbits,
sim_table,
codes + j * code_size);
res.add(j, dis);
}
}
#pragma omp critical
{
indexIVFPQ_stats.n_hamming_pass += n_hamming_pass;
}
}
template <class SearchResultType>
void scan_list_polysemous(
size_t ncode,
const uint8_t* codes,
SearchResultType& res) const {
with_HammingComputer<PQCodeDist::simd_level>(
pq.code_size, [&]<class HammingComputer>() {
this->scan_list_polysemous_hc<
HammingComputer,
SearchResultType>(ncode, codes, res);
});
}
};
/* We put as many parameters as possible in template. Hopefully the
* gain in runtime is worth the code bloat.
*
* C is the comparator < or >, it is directly related to METRIC_TYPE.
*
* precompute_mode is how much we precompute (2 = precompute distance tables,
* 1 = precompute pointers to distances, 0 = compute distances one by one).
* Currently only 2 is supported
*
* use_sel: store or ignore the IDSelector
*/
template <MetricType METRIC_TYPE, class C, class PQCodeDist, bool use_sel>
struct IVFPQScanner : IVFPQScannerT<idx_t, METRIC_TYPE, PQCodeDist>,
InvertedListScanner {
using InvertedListScanner::scan_codes;
int precompute_mode;
const IDSelector* sel;
IVFPQScanner(
const IndexIVFPQ& ivfpq_in,
bool store_pairs_in,
int precompute_mode_in,
const IDSelector* sel_in)
: IVFPQScannerT<idx_t, METRIC_TYPE, PQCodeDist>(ivfpq_in, nullptr),
precompute_mode(precompute_mode_in),
sel(sel_in) {
this->store_pairs = store_pairs_in;
this->keep_max = is_similarity_metric(METRIC_TYPE);
this->code_size = this->pq.code_size;
}
void set_query(const float* query) override {
this->init_query(query);
}
void set_list(idx_t list_no_in, float coarse_dis_in) override {
this->list_no = list_no_in;
this->init_list(list_no_in, coarse_dis_in, precompute_mode);
}
float distance_to_code(const uint8_t* code) const override {
FAISS_THROW_IF_NOT(precompute_mode == 2);
float dis = this->dis0 +
PQCodeDist::distance_single_code(
this->pq.M, this->pq.nbits, this->sim_table, code);
return dis;
}
size_t scan_codes(
size_t ncode,
const uint8_t* codes,
const idx_t* ids,
ResultHandler& handler) const override {
WrappedSearchResult<C, use_sel> res(
this->key,
this->store_pairs ? nullptr : ids,
this->sel,
handler);
if (this->polysemous_ht > 0) {
FAISS_THROW_IF_NOT(precompute_mode == 2);
this->scan_list_polysemous(ncode, codes, res);
} else if (precompute_mode == 2) {
this->scan_list_with_table(ncode, codes, res);
} else if (precompute_mode == 1) {
this->scan_list_with_pointer(ncode, codes, res);
} else if (precompute_mode == 0) {
this->scan_on_the_fly_dist(ncode, codes, res);
} else {
FAISS_THROW_MSG("bad precomp mode");
}
return res.nup;
}
};
template <SIMDLevel SL>
InvertedListScanner* make_IVFPQInvertedListScanner(
const IndexIVFPQ& ivfpq,
bool store_pairs,
const IDSelector* sel);
// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization)
template <>
InvertedListScanner* make_IVFPQInvertedListScanner<THE_SIMD_LEVEL>(
const IndexIVFPQ& ivfpq,
bool store_pairs,
const IDSelector* sel) {
auto make = [&]<class PQCodeDist, bool use_sel>() -> InvertedListScanner* {
if (ivfpq.metric_type == METRIC_INNER_PRODUCT) {
return new IVFPQScanner<
METRIC_INNER_PRODUCT,
CMin<float, idx_t>,
PQCodeDist,
use_sel>(ivfpq, store_pairs, 2, sel);
} else if (ivfpq.metric_type == METRIC_L2) {
return new IVFPQScanner<
METRIC_L2,
CMax<float, idx_t>,
PQCodeDist,
use_sel>(ivfpq, store_pairs, 2, sel);
} else {
FAISS_THROW_MSG("unsupported metric type");
}
};
auto with_decoder = [&]<bool use_sel>() -> InvertedListScanner* {
if (ivfpq.pq.nbits == 8) {
return make.template
operator()<PQCodeDistance<PQDecoder8, THE_SIMD_LEVEL>, use_sel>();
} else if (ivfpq.pq.nbits == 16) {
return make.template
operator()<PQCodeDistance<PQDecoder16, THE_SIMD_LEVEL>, use_sel>();
} else {
return make.template operator()<
PQCodeDistance<PQDecoderGeneric, THE_SIMD_LEVEL>,
use_sel>();
}
};
if (sel) {
return with_decoder.template operator()<true>();
} else {
return with_decoder.template operator()<false>();
}
}
} // namespace pq_code_distance
} // namespace faiss
faiss/impl/pq_code_distance/IVFPQ_QueryTables.cpp
+245
View File
@@ -0,0 +1,245 @@
/*
* 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/impl/pq_code_distance/IVFPQ_QueryTables.h>
#include <faiss/Clustering.h>
#include <faiss/utils/distances_dispatch.h>
namespace faiss {
namespace pq_code_distance {
QueryTables::QueryTables(
const IndexIVFPQ& ivfpq_in,
const IVFSearchParameters* params_in)
: ivfpq(ivfpq_in),
params(params_in),
d(ivfpq_in.d),
pq(ivfpq_in.pq),
metric_type(ivfpq_in.metric_type),
by_residual(ivfpq_in.by_residual),
use_precomputed_table(ivfpq_in.use_precomputed_table) {
mem.resize(pq.ksub * pq.M * 2 + d * 2);
sim_table = mem.data();
sim_table_2 = sim_table + pq.ksub * pq.M;
residual_vec = sim_table_2 + pq.ksub * pq.M;
decoded_vec = residual_vec + d;
// for polysemous
polysemous_ht = ivfpq_in.polysemous_ht;
if (auto ivfpq_params =
dynamic_cast<const IVFPQSearchParameters*>(params_in)) {
polysemous_ht = ivfpq_params->polysemous_ht;
}
if (polysemous_ht != 0) {
q_code.resize(pq.code_size);
}
init_list_cycles = 0;
sim_table_ptrs.resize(pq.M);
}
void QueryTables::init_query(const float* qi_in) {
this->qi = qi_in;
if (metric_type == METRIC_INNER_PRODUCT) {
init_query_IP();
} else {
init_query_L2();
}
if (!by_residual && polysemous_ht != 0) {
pq.compute_code(qi_in, q_code.data());
}
}
void QueryTables::init_query_IP() {
// precompute some tables specific to the query qi
pq.compute_inner_prod_table(qi, sim_table);
}
void QueryTables::init_query_L2() {
if (!by_residual) {
pq.compute_distance_table(qi, sim_table);
} else if (use_precomputed_table) {
pq.compute_inner_prod_table(qi, sim_table_2);
}
}
float QueryTables::precompute_list_tables() {
float dis0 = 0;
uint64_t t0;
TIC;
if (by_residual) {
if (metric_type == METRIC_INNER_PRODUCT) {
dis0 = precompute_list_tables_IP();
} else {
dis0 = precompute_list_tables_L2();
}
}
init_list_cycles += TOC;
return dis0;
}
float QueryTables::precompute_list_table_pointers() {
float dis0 = 0;
uint64_t t0;
TIC;
if (by_residual) {
if (metric_type == METRIC_INNER_PRODUCT) {
FAISS_THROW_MSG("not implemented");
} else {
dis0 = precompute_list_table_pointers_L2();
}
}
init_list_cycles += TOC;
return dis0;
}
void QueryTables::init_list(idx_t list_no, float coarse_dis_in, int mode) {
this->key = list_no;
this->coarse_dis = coarse_dis_in;
if (mode == 2) {
dis0 = precompute_list_tables();
} else if (mode == 1) {
dis0 = precompute_list_table_pointers();
}
}
float QueryTables::precompute_list_tables_IP() {
// prepare the sim_table that will be used for accumulation
// and dis0, the initial value
ivfpq.quantizer->reconstruct(key, decoded_vec);
// decoded_vec = centroid
float dis0 = fvec_inner_product_dispatch(qi, decoded_vec, d);
if (polysemous_ht) {
for (int i = 0; i < d; i++) {
residual_vec[i] = qi[i] - decoded_vec[i];
}
pq.compute_code(residual_vec, q_code.data());
}
return dis0;
}
float QueryTables::precompute_list_tables_L2() {
float dis0 = 0;
if (use_precomputed_table == 0 || use_precomputed_table == -1) {
ivfpq.quantizer->compute_residual(qi, residual_vec, key);
pq.compute_distance_table(residual_vec, sim_table);
if (polysemous_ht != 0) {
pq.compute_code(residual_vec, q_code.data());
}
} else if (use_precomputed_table == 1) {
dis0 = coarse_dis;
fvec_madd_dispatch(
pq.M * pq.ksub,
ivfpq.precomputed_table.data() + key * pq.ksub * pq.M,
-2.0,
sim_table_2,
sim_table);
if (polysemous_ht != 0) {
ivfpq.quantizer->compute_residual(qi, residual_vec, key);
pq.compute_code(residual_vec, q_code.data());
}
} else if (use_precomputed_table == 2) {
dis0 = coarse_dis;
const MultiIndexQuantizer* miq =
dynamic_cast<const MultiIndexQuantizer*>(ivfpq.quantizer);
FAISS_THROW_IF_NOT(miq);
const ProductQuantizer& cpq = miq->pq;
size_t Mf = pq.M / cpq.M;
const float* qtab = sim_table_2; // query-specific table
float* ltab = sim_table; // (output) list-specific table
long k = key;
for (size_t cm = 0; cm < cpq.M; cm++) {
// compute PQ index
size_t ki = k & ((uint64_t(1) << cpq.nbits) - 1);
k >>= cpq.nbits;
// get corresponding table
const float* pc = ivfpq.precomputed_table.data() +
(ki * pq.M + cm * Mf) * pq.ksub;
if (polysemous_ht == 0) {
// sum up with query-specific table
fvec_madd_dispatch(Mf * pq.ksub, pc, -2.0, qtab, ltab);
ltab += Mf * pq.ksub;
qtab += Mf * pq.ksub;
} else {
for (size_t m = cm * Mf; m < (cm + 1) * Mf; m++) {
q_code[m] = fvec_madd_and_argmin_dispatch(
pq.ksub, pc, -2, qtab, ltab);
pc += pq.ksub;
ltab += pq.ksub;
qtab += pq.ksub;
}
}
}
}
return dis0;
}
float QueryTables::precompute_list_table_pointers_L2() {
float dis0 = 0;
if (use_precomputed_table == 1) {
dis0 = coarse_dis;
const float* s = ivfpq.precomputed_table.data() + key * pq.ksub * pq.M;
for (size_t m = 0; m < pq.M; m++) {
sim_table_ptrs[m] = s;
s += pq.ksub;
}
} else if (use_precomputed_table == 2) {
dis0 = coarse_dis;
const MultiIndexQuantizer* miq =
dynamic_cast<const MultiIndexQuantizer*>(ivfpq.quantizer);
FAISS_THROW_IF_NOT(miq);
const ProductQuantizer& cpq = miq->pq;
size_t Mf = pq.M / cpq.M;
long k = key;
size_t m0 = 0;
for (size_t cm = 0; cm < cpq.M; cm++) {
size_t ki = k & ((uint64_t(1) << cpq.nbits) - 1);
k >>= cpq.nbits;
const float* pc = ivfpq.precomputed_table.data() +
(ki * pq.M + cm * Mf) * pq.ksub;
for (size_t m = m0; m < m0 + Mf; m++) {
sim_table_ptrs[m] = pc;
pc += pq.ksub;
}
m0 += Mf;
}
} else {
FAISS_THROW_MSG("need precomputed tables");
}
if (polysemous_ht) {
FAISS_THROW_MSG("not implemented");
// Not clear that it makes sense to implemente this,
// because it costs M * ksub, which is what we wanted to
// avoid with the tables pointers.
}
return dis0;
}
} // namespace pq_code_distance
} // namespace faiss
faiss/impl/pq_code_distance/IVFPQ_QueryTables.h
+105
View File
@@ -0,0 +1,105 @@
/*
* 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 <cstdint>
#include <vector>
#include <faiss/IndexIVFPQ.h>
#include <faiss/impl/FaissAssert.h>
#include <faiss/impl/ProductQuantizer.h>
#include <faiss/utils/utils.h>
namespace faiss {
namespace pq_code_distance {
#define TIC t0 = get_cycles()
#define TOC get_cycles() - t0
/** QueryTables manages the various ways of searching an
* IndexIVFPQ. The code contains a lot of branches, depending on:
* - metric_type: are we computing L2 or Inner product similarity?
* - by_residual: do we encode raw vectors or residuals?
* - use_precomputed_table: are x_R|x_C tables precomputed?
* - polysemous_ht: are we filtering with polysemous codes?
*/
struct QueryTables {
/*****************************************************
* General data from the IVFPQ
*****************************************************/
const IndexIVFPQ& ivfpq;
const IVFSearchParameters* params;
// copied from IndexIVFPQ for easier access
int d;
const ProductQuantizer& pq;
MetricType metric_type;
bool by_residual;
int use_precomputed_table;
int polysemous_ht;
// pre-allocated data buffers
float *sim_table, *sim_table_2;
float *residual_vec, *decoded_vec;
// single data buffer
std::vector<float> mem;
// for table pointers
std::vector<const float*> sim_table_ptrs;
explicit QueryTables(
const IndexIVFPQ& ivfpq_in,
const IVFSearchParameters* params_in);
/*****************************************************
* What we do when query is known
*****************************************************/
// field specific to query
const float* qi = nullptr;
// query-specific initialization
void init_query(const float* qi_in);
void init_query_IP();
void init_query_L2();
/*****************************************************
* When inverted list is known: prepare computations
*****************************************************/
// fields specific to list
idx_t key = 0;
float coarse_dis = 0.0f;
std::vector<uint8_t> q_code;
uint64_t init_list_cycles;
/// once we know the query and the centroid, we can prepare the
/// sim_table that will be used for accumulation
/// and dis0, the initial value
float precompute_list_tables();
float precompute_list_table_pointers();
float dis0 = 0.0f;
void init_list(idx_t list_no, float coarse_dis_in, int mode);
float precompute_list_tables_IP();
float precompute_list_tables_L2();
float precompute_list_table_pointers_L2();
};
} // namespace pq_code_distance
} // namespace faiss
faiss/impl/pq_code_distance/PQDistanceComputer_impl.h
+106
View File
@@ -0,0 +1,106 @@
/*
* 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
#ifndef THE_SIMD_LEVEL
#error "THE_SIMD_LEVEL must be defined before including PQDistanceComputer_impl.h"
#endif
#include <faiss/IndexPQ.h>
#include <faiss/impl/DistanceComputer.h>
#include <faiss/impl/FaissAssert.h>
#include <faiss/impl/pq_code_distance/pq_code_distance-inl.h>
namespace faiss {
namespace pq_code_distance {
template <class PQCodeDist>
struct PQDistanceComputer : FlatCodesDistanceComputer {
using PQDecoder = typename PQCodeDist::PQDecoder;
size_t d;
MetricType metric;
idx_t nb;
const ProductQuantizer& pq;
const float* sdc;
std::vector<float> precomputed_table;
size_t ndis;
const float* q;
float distance_to_code(const uint8_t* code) final {
ndis++;
float dis = PQCodeDist::distance_single_code(
pq.M, pq.nbits, precomputed_table.data(), code);
return dis;
}
float symmetric_dis(idx_t i, idx_t j) override {
FAISS_THROW_IF_NOT(sdc);
const float* sdci = sdc;
float accu = 0;
PQDecoder codei(codes + i * code_size, pq.nbits);
PQDecoder codej(codes + j * code_size, pq.nbits);
for (size_t l = 0; l < pq.M; l++) {
accu += sdci[codei.decode() + (codej.decode() << codei.nbits)];
sdci += uint64_t(1) << (2 * codei.nbits);
}
ndis++;
return accu;
}
explicit PQDistanceComputer(const IndexPQ& storage)
: FlatCodesDistanceComputer(
storage.codes.data(),
storage.code_size),
pq(storage.pq),
q(nullptr) {
precomputed_table.resize(pq.M * pq.ksub);
nb = storage.ntotal;
d = storage.d;
metric = storage.metric_type;
if (pq.sdc_table.size() == pq.ksub * pq.ksub * pq.M) {
sdc = pq.sdc_table.data();
} else {
sdc = nullptr;
}
ndis = 0;
}
void set_query(const float* x) override {
q = x;
if (metric == METRIC_L2) {
pq.compute_distance_table(x, precomputed_table.data());
} else {
pq.compute_inner_prod_table(x, precomputed_table.data());
}
}
};
template <SIMDLevel SL>
FlatCodesDistanceComputer* get_PQFlatCodesDistanceComputer(
const IndexPQ& index);
// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization)
template <>
FlatCodesDistanceComputer* get_PQFlatCodesDistanceComputer<THE_SIMD_LEVEL>(
const IndexPQ& index) {
if (index.pq.nbits == 8) {
return new PQDistanceComputer<
PQCodeDistance<PQDecoder8, THE_SIMD_LEVEL>>(index);
} else if (index.pq.nbits == 16) {
return new PQDistanceComputer<
PQCodeDistance<PQDecoder16, THE_SIMD_LEVEL>>(index);
} else {
return new PQDistanceComputer<
PQCodeDistance<PQDecoderGeneric, THE_SIMD_LEVEL>>(index);
}
}
} // namespace pq_code_distance
} // namespace faiss
faiss/impl/pq_code_distance/avx2.cpp
+21
View File
@@ -0,0 +1,21 @@
/*
* 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.
*/
#ifdef COMPILE_SIMD_AVX2
#define THE_SIMD_LEVEL SIMDLevel::AVX2
// NOLINTNEXTLINE(facebook-hte-InlineHeader)
#include <faiss/impl/pq_code_distance/pq_code_distance-avx2.h>
// NOLINTNEXTLINE(facebook-hte-InlineHeader)
#include <faiss/utils/hamming_distance/hamming_computer-avx2.h>
// NOLINTNEXTLINE(facebook-hte-InlineHeader)
#include <faiss/impl/pq_code_distance/PQDistanceComputer_impl.h>
// NOLINTNEXTLINE(facebook-hte-InlineHeader)
#include <faiss/impl/pq_code_distance/IVFPQScanner_impl.h>
#endif // COMPILE_SIMD_AVX2
faiss/impl/pq_code_distance/avx512.cpp
+21
View File
@@ -0,0 +1,21 @@
/*
* 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.
*/
#ifdef COMPILE_SIMD_AVX512
#define THE_SIMD_LEVEL SIMDLevel::AVX512
// NOLINTNEXTLINE(facebook-hte-InlineHeader)
#include <faiss/impl/pq_code_distance/pq_code_distance-avx512.h>
// NOLINTNEXTLINE(facebook-hte-InlineHeader)
#include <faiss/utils/hamming_distance/hamming_computer-avx512.h>
// NOLINTNEXTLINE(facebook-hte-InlineHeader)
#include <faiss/impl/pq_code_distance/PQDistanceComputer_impl.h>
// NOLINTNEXTLINE(facebook-hte-InlineHeader)
#include <faiss/impl/pq_code_distance/IVFPQScanner_impl.h>
#endif // COMPILE_SIMD_AVX512
faiss/impl/pq_code_distance/neon.cpp
+21
View File
@@ -0,0 +1,21 @@
/*
* 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.
*/
#ifdef COMPILE_SIMD_ARM_NEON
#define THE_SIMD_LEVEL SIMDLevel::ARM_NEON
// NOLINTNEXTLINE(facebook-hte-InlineHeader)
#include <faiss/impl/pq_code_distance/pq_code_distance-generic.h>
// NOLINTNEXTLINE(facebook-hte-InlineHeader)
#include <faiss/utils/hamming_distance/hamming_computer-neon.h>
// NOLINTNEXTLINE(facebook-hte-InlineHeader)
#include <faiss/impl/pq_code_distance/PQDistanceComputer_impl.h>
// NOLINTNEXTLINE(facebook-hte-InlineHeader)
#include <faiss/impl/pq_code_distance/IVFPQScanner_impl.h>
#endif // COMPILE_SIMD_ARM_NEON
faiss/impl/pq_code_distance/pq_code_distance-avx2.cpp → faiss/impl/pq_code_distance/pq_code_distance-avx2.h
+9 -13
View File
@@ -5,7 +5,7 @@
* LICENSE file in the root directory of this source tree.
*/
#ifdef COMPILE_SIMD_AVX2
#pragma once
#include <immintrin.h>
@@ -162,9 +162,8 @@ namespace pq_code_distance {
// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization)
template <>
float pq_code_distance_single_impl<SIMDLevel::AVX2>(
float pq_code_distance_8bit_single_impl<SIMDLevel::AVX2>(
size_t M,
size_t nbits,
const float* sim_table,
const uint8_t* code) {
if (M == 4) {
@@ -223,7 +222,7 @@ float pq_code_distance_single_impl<SIMDLevel::AVX2>(
// Process leftovers.
if (m < M) {
PQDecoder8 decoder(code + m, nbits);
PQDecoder8 decoder(code + m, 8);
for (; m < M; m++) {
result += tab[decoder.decode()];
tab += ksub;
@@ -233,12 +232,11 @@ float pq_code_distance_single_impl<SIMDLevel::AVX2>(
return result;
}
// Combines 4 operations of pq_code_distance_single_impl().
// Combines 4 operations of pq_code_distance_8bit_single_impl().
// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization)
template <>
void pq_code_distance_four_impl<SIMDLevel::AVX2>(
void pq_code_distance_8bit_four_impl<SIMDLevel::AVX2>(
size_t M,
size_t nbits,
const float* sim_table,
const uint8_t* __restrict code0,
const uint8_t* __restrict code1,
@@ -341,10 +339,10 @@ void pq_code_distance_four_impl<SIMDLevel::AVX2>(
// Process leftovers.
if (m < M) {
PQDecoder8 decoder0(code0 + m, nbits);
PQDecoder8 decoder1(code1 + m, nbits);
PQDecoder8 decoder2(code2 + m, nbits);
PQDecoder8 decoder3(code3 + m, nbits);
PQDecoder8 decoder0(code0 + m, 8);
PQDecoder8 decoder1(code1 + m, 8);
PQDecoder8 decoder2(code2 + m, 8);
PQDecoder8 decoder3(code3 + m, 8);
for (; m < M; m++) {
result0 += tab[decoder0.decode()];
result1 += tab[decoder1.decode()];
@@ -357,5 +355,3 @@ void pq_code_distance_four_impl<SIMDLevel::AVX2>(
} // namespace pq_code_distance
} // namespace faiss
#endif // COMPILE_SIMD_AVX2
faiss/impl/pq_code_distance/pq_code_distance-avx512.cpp → faiss/impl/pq_code_distance/pq_code_distance-avx512.h
+9 -13
View File
@@ -5,7 +5,7 @@
* LICENSE file in the root directory of this source tree.
*/
#ifdef COMPILE_SIMD_AVX512
#pragma once
#include <immintrin.h>
@@ -22,9 +22,8 @@ namespace pq_code_distance {
// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization)
template <>
float pq_code_distance_single_impl<SIMDLevel::AVX512>(
float pq_code_distance_8bit_single_impl<SIMDLevel::AVX512>(
size_t M,
size_t nbits,
const float* sim_table,
const uint8_t* code0) {
float result0 = 0;
@@ -69,7 +68,7 @@ float pq_code_distance_single_impl<SIMDLevel::AVX512>(
// Process leftovers.
if (m < M) {
PQDecoder8 decoder0(code0 + m, nbits);
PQDecoder8 decoder0(code0 + m, 8);
for (; m < M; m++) {
result0 += tab[decoder0.decode()];
tab += ksub;
@@ -79,12 +78,11 @@ float pq_code_distance_single_impl<SIMDLevel::AVX512>(
return result0;
}
// Combines 4 operations of pq_code_distance_single_impl().
// Combines 4 operations of pq_code_distance_8bit_single_impl().
// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization)
template <>
void pq_code_distance_four_impl<SIMDLevel::AVX512>(
void pq_code_distance_8bit_four_impl<SIMDLevel::AVX512>(
size_t M,
size_t nbits,
const float* sim_table,
const uint8_t* __restrict code0,
const uint8_t* __restrict code1,
@@ -145,10 +143,10 @@ void pq_code_distance_four_impl<SIMDLevel::AVX512>(
// Process leftovers.
if (m < M) {
PQDecoder8 decoder0(code0 + m, nbits);
PQDecoder8 decoder1(code1 + m, nbits);
PQDecoder8 decoder2(code2 + m, nbits);
PQDecoder8 decoder3(code3 + m, nbits);
PQDecoder8 decoder0(code0 + m, 8);
PQDecoder8 decoder1(code1 + m, 8);
PQDecoder8 decoder2(code2 + m, 8);
PQDecoder8 decoder3(code3 + m, 8);
for (; m < M; m++) {
result0 += tab[decoder0.decode()];
result1 += tab[decoder1.decode()];
@@ -161,5 +159,3 @@ void pq_code_distance_four_impl<SIMDLevel::AVX512>(
} // namespace pq_code_distance
} // namespace faiss
#endif // COMPILE_SIMD_AVX512
faiss/impl/pq_code_distance/pq_code_distance-generic.cpp
+14 -99
View File
@@ -5,116 +5,32 @@
* LICENSE file in the root directory of this source tree.
*/
// This TU provides:
// 1. _impl specializations for NONE (and ARM_NEON), using scalar code.
// 2. Non-templated PQ code distance dispatch wrappers
// (pq_code_distance_single, pq_code_distance_four) declared in
// pq_code_distance.h. These use with_simd_level to route to the
// best available SIMD implementation via pq_code_distance_*_impl
// function template specializations defined in the per-SIMD .cpp files.
// This TU provides non-templated PQ code distance dispatch wrappers
// (pq_code_distance_8bit_single, pq_code_distance_8bit_four) declared
// in pq_code_distance-inl.h. These use with_simd_level to route to the
// best available SIMD implementation via pq_code_distance_8bit_*_impl
// function template specializations.
//
// The NONE and ARM_NEON _impl specializations are defined inline in
// pq_code_distance-generic.h (included transitively). The AVX2, AVX512,
// and ARM_SVE specializations are in their respective per-SIMD files.
#include <faiss/impl/pq_code_distance/pq_code_distance-inl.h>
#include <faiss/impl/pq_code_distance/pq_code_distance-generic.h>
namespace faiss {
namespace pq_code_distance {
// NONE: use scalar directly.
// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization)
template <>
float pq_code_distance_single_impl<SIMDLevel::NONE>(
float pq_code_distance_8bit_single(
size_t M,
size_t nbits,
const float* sim_table,
const uint8_t* code) {
return PQCodeDistanceScalar<PQDecoder8>::distance_single_code(
M, nbits, sim_table, code);
}
// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization)
template <>
void pq_code_distance_four_impl<SIMDLevel::NONE>(
size_t M,
size_t nbits,
const float* sim_table,
const uint8_t* __restrict code0,
const uint8_t* __restrict code1,
const uint8_t* __restrict code2,
const uint8_t* __restrict code3,
float& result0,
float& result1,
float& result2,
float& result3) {
PQCodeDistanceScalar<PQDecoder8>::distance_four_codes(
M,
nbits,
sim_table,
code0,
code1,
code2,
code3,
result0,
result1,
result2,
result3);
}
#ifdef COMPILE_SIMD_ARM_NEON
// ARM_NEON: No NEON-optimized PQ code distance exists. Use scalar.
// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization)
template <>
float pq_code_distance_single_impl<SIMDLevel::ARM_NEON>(
size_t M,
size_t nbits,
const float* sim_table,
const uint8_t* code) {
return PQCodeDistanceScalar<PQDecoder8>::distance_single_code(
M, nbits, sim_table, code);
}
// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization)
template <>
void pq_code_distance_four_impl<SIMDLevel::ARM_NEON>(
size_t M,
size_t nbits,
const float* sim_table,
const uint8_t* __restrict code0,
const uint8_t* __restrict code1,
const uint8_t* __restrict code2,
const uint8_t* __restrict code3,
float& result0,
float& result1,
float& result2,
float& result3) {
PQCodeDistanceScalar<PQDecoder8>::distance_four_codes(
M,
nbits,
sim_table,
code0,
code1,
code2,
code3,
result0,
result1,
result2,
result3);
}
#endif // COMPILE_SIMD_ARM_NEON
float pq_code_distance_single(
size_t M,
size_t nbits,
const float* sim_table,
const uint8_t* code) {
return with_simd_level([&]<SIMDLevel SL>() {
return pq_code_distance_single_impl<SL>(M, nbits, sim_table, code);
return pq_code_distance_8bit_single_impl<SL>(M, sim_table, code);
});
}
void pq_code_distance_four(
void pq_code_distance_8bit_four(
size_t M,
size_t nbits,
const float* sim_table,
const uint8_t* __restrict code0,
const uint8_t* __restrict code1,
@@ -125,9 +41,8 @@ void pq_code_distance_four(
float& result2,
float& result3) {
with_simd_level([&]<SIMDLevel SL>() {
pq_code_distance_four_impl<SL>(
pq_code_distance_8bit_four_impl<SL>(
M,
nbits,
sim_table,
code0,
code1,
faiss/impl/pq_code_distance/pq_code_distance-generic.h
+96
View File
@@ -0,0 +1,96 @@
/*
* 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 <faiss/impl/pq_code_distance/pq_code_distance-inl.h>
namespace faiss {
namespace pq_code_distance {
// NONE: use scalar directly.
// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization)
template <>
inline float pq_code_distance_8bit_single_impl<SIMDLevel::NONE>(
size_t M,
const float* sim_table,
const uint8_t* code) {
return PQCodeDistanceScalar<PQDecoder8>::distance_single_code(
M, 8, sim_table, code);
}
// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization)
template <>
inline void pq_code_distance_8bit_four_impl<SIMDLevel::NONE>(
size_t M,
const float* sim_table,
const uint8_t* __restrict code0,
const uint8_t* __restrict code1,
const uint8_t* __restrict code2,
const uint8_t* __restrict code3,
float& result0,
float& result1,
float& result2,
float& result3) {
PQCodeDistanceScalar<PQDecoder8>::distance_four_codes(
M,
8,
sim_table,
code0,
code1,
code2,
code3,
result0,
result1,
result2,
result3);
}
#ifdef COMPILE_SIMD_ARM_NEON
// ARM_NEON: No NEON-optimized PQ code distance exists. Use scalar.
// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization)
template <>
inline float pq_code_distance_8bit_single_impl<SIMDLevel::ARM_NEON>(
size_t M,
const float* sim_table,
const uint8_t* code) {
return PQCodeDistanceScalar<PQDecoder8>::distance_single_code(
M, 8, sim_table, code);
}
// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization)
template <>
inline void pq_code_distance_8bit_four_impl<SIMDLevel::ARM_NEON>(
size_t M,
const float* sim_table,
const uint8_t* __restrict code0,
const uint8_t* __restrict code1,
const uint8_t* __restrict code2,
const uint8_t* __restrict code3,
float& result0,
float& result1,
float& result2,
float& result3) {
PQCodeDistanceScalar<PQDecoder8>::distance_four_codes(
M,
8,
sim_table,
code0,
code1,
code2,
code3,
result0,
result1,
result2,
result3);
}
#endif // COMPILE_SIMD_ARM_NEON
} // namespace pq_code_distance
} // namespace faiss
faiss/impl/pq_code_distance/pq_code_distance-inl.h
+238 -5
View File
@@ -9,15 +9,248 @@
/**
* @file pq_code_distance-inl.h
* @brief Private header for PQ code distance SIMD implementations.
* @brief PQ code distance SIMD-dispatched implementations.
*
* This is a PRIVATE header — do not include in public APIs or user code.
* Only faiss internal .cpp files (the per-SIMD implementation files and
* pq_code_distance-generic.cpp) should include this header.
*
* This header re-exports the public API (pq_code_distance.h) plus the
* simd_dispatch.h machinery needed by the implementation files.
*/
#include <cstddef>
#include <cstdint>
#include <type_traits>
#include <faiss/impl/ProductQuantizer.h>
#include <faiss/impl/platform_macros.h>
#include <faiss/impl/simd_dispatch.h>
#include <faiss/utils/pq_code_distance.h>
namespace faiss {
namespace pq_code_distance {
/*********************************************************************
* PQCodeDistance — SIMD-dispatched PQ code distance
*
* Computes the distance from a PQ-encoded vector to a query vector,
* given a precomputed table of sub-distances (one per subquantizer
* per centroid). Originally extracted from IndexIVFPQ.cpp.
*
* DESIGN:
*
* PQCodeDistance<PQDecoderT, SL> computes PQ code distances at a given
* SIMD level. The dispatch site (IndexIVFPQ.cpp, IndexPQ.cpp) uses
* with_simd_level to select SL at runtime, which instantiates
* PQCodeDistance for ALL decoder types (PQDecoder8, PQDecoder16,
* PQDecoderGeneric) at the chosen level.
*
* Only PQDecoder8 has SIMD-optimized implementations (AVX2, AVX512,
* ARM_SVE). The other decoders always use scalar code — their decode()
* method is inherently sequential, so SIMD doesn't help.
*
* The primary template is always complete (no forward declarations
* needed). For PQDecoder8, it delegates to _impl dispatch bridge
* functions whose specializations are defined in per-SIMD .cpp files
* and resolved at link time. For other decoders, it uses scalar.
*
* ADDING A NEW SIMD LEVEL:
*
* 1. Add the level to SIMDLevel enum (simd_levels.h)
* 2. Add dispatch_config entry (simd_dispatch.bzl)
* 3. Define pq_code_distance_8bit_single_impl<NEW_LEVEL> and
* pq_code_distance_8bit_four_impl<NEW_LEVEL> specializations in a
* new .cpp file compiled with appropriate SIMD flags
* 4. Add the .cpp to the build (CMakeLists.txt, xplat.bzl)
*********************************************************************/
/// Scalar PQ code distance implementation.
/// Templated only on decoder type, independent of SIMD level.
/// Used directly by non-PQDecoder8 decoders (PQDecoder16,
/// PQDecoderGeneric) and as fallback for PQDecoder8 at NONE/NEON.
template <typename PQDecoderT>
struct PQCodeDistanceScalar {
using PQDecoder = PQDecoderT;
static float distance_single_code(
// number of subquantizers
size_t M,
size_t nbits,
// precomputed distances, layout (M, ksub)
const float* sim_table,
const uint8_t* code) {
PQDecoderT decoder(code, nbits);
const size_t ksub = 1 << nbits;
const float* tab = sim_table;
float result = 0;
for (size_t m = 0; m < M; m++) {
result += tab[decoder.decode()];
tab += ksub;
}
return result;
}
static void distance_four_codes(
size_t M,
size_t nbits,
const float* sim_table,
const uint8_t* __restrict code0,
const uint8_t* __restrict code1,
const uint8_t* __restrict code2,
const uint8_t* __restrict code3,
float& result0,
float& result1,
float& result2,
float& result3) {
PQDecoderT decoder0(code0, nbits);
PQDecoderT decoder1(code1, nbits);
PQDecoderT decoder2(code2, nbits);
PQDecoderT decoder3(code3, nbits);
const size_t ksub = 1 << nbits;
const float* tab = sim_table;
result0 = 0;
result1 = 0;
result2 = 0;
result3 = 0;
for (size_t m = 0; m < M; m++) {
result0 += tab[decoder0.decode()];
result1 += tab[decoder1.decode()];
result2 += tab[decoder2.decode()];
result3 += tab[decoder3.decode()];
tab += ksub;
}
}
};
/*********************************************************************
* Dispatch bridge — function templates for PQDecoder8 SIMD dispatch.
*
* Primary declarations only; specializations are defined in per-SIMD
* .cpp files (AVX2, AVX512, ARM_SVE) and pq_code_distance-generic.cpp
* (NONE, ARM_NEON). Same pattern as fvec_L2sqr et al. in distances.h.
*********************************************************************/
template <SIMDLevel SL>
float pq_code_distance_8bit_single_impl(
size_t M,
const float* sim_table,
const uint8_t* code);
template <SIMDLevel SL>
void pq_code_distance_8bit_four_impl(
size_t M,
const float* sim_table,
const uint8_t* __restrict code0,
const uint8_t* __restrict code1,
const uint8_t* __restrict code2,
const uint8_t* __restrict code3,
float& result0,
float& result1,
float& result2,
float& result3);
/// Primary template — always complete.
/// For PQDecoder8, delegates to _impl dispatch bridges (resolved at
/// link time to per-SIMD implementations). For other decoders, uses
/// scalar — their sequential decode() methods don't benefit from SIMD.
template <typename PQDecoderT, SIMDLevel SL>
struct PQCodeDistance {
using PQDecoder = PQDecoderT;
static constexpr SIMDLevel simd_level = SL;
static float distance_single_code(
size_t M,
size_t nbits,
const float* sim_table,
const uint8_t* code) {
if constexpr (std::is_same_v<PQDecoderT, PQDecoder8>) {
return pq_code_distance_8bit_single_impl<SL>(M, sim_table, code);
} else {
return PQCodeDistanceScalar<PQDecoderT>::distance_single_code(
M, nbits, sim_table, code);
}
}
static void distance_four_codes(
size_t M,
size_t nbits,
const float* sim_table,
const uint8_t* __restrict code0,
const uint8_t* __restrict code1,
const uint8_t* __restrict code2,
const uint8_t* __restrict code3,
float& result0,
float& result1,
float& result2,
float& result3) {
if constexpr (std::is_same_v<PQDecoderT, PQDecoder8>) {
pq_code_distance_8bit_four_impl<SL>(
M,
sim_table,
code0,
code1,
code2,
code3,
result0,
result1,
result2,
result3);
} else {
PQCodeDistanceScalar<PQDecoderT>::distance_four_codes(
M,
nbits,
sim_table,
code0,
code1,
code2,
code3,
result0,
result1,
result2,
result3);
}
}
};
/*********************************************************************
* Non-templated PQ code distance dispatch (PQDecoder8 only).
*
* These follow the same pattern as distances.h: the caller does not
* name a SIMDLevel. Internally they dispatch via with_simd_level
* to the best available SIMD implementation (DD: runtime detection,
* static: compile-time selection). Definitions are in
* pq_code_distance-generic.cpp.
*********************************************************************/
/// Compute PQ distance for a single code, dispatching to the best
/// available SIMD level.
FAISS_API float pq_code_distance_8bit_single(
size_t M,
const float* sim_table,
const uint8_t* code);
/// Compute PQ distances for four codes simultaneously, dispatching
/// to the best available SIMD level.
FAISS_API void pq_code_distance_8bit_four(
size_t M,
const float* sim_table,
const uint8_t* __restrict code0,
const uint8_t* __restrict code1,
const uint8_t* __restrict code2,
const uint8_t* __restrict code3,
float& result0,
float& result1,
float& result2,
float& result3);
} // namespace pq_code_distance
// Re-export public API into namespace faiss for convenience
using pq_code_distance::pq_code_distance_8bit_four;
using pq_code_distance::pq_code_distance_8bit_single;
using pq_code_distance::PQCodeDistance;
using pq_code_distance::PQCodeDistanceScalar;
} // namespace faiss
faiss/impl/pq_code_distance/pq_code_distance-sve.cpp
+5 -7
View File
@@ -83,16 +83,15 @@ namespace pq_code_distance {
// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization)
template <>
float pq_code_distance_single_impl<SIMDLevel::ARM_SVE>(
float pq_code_distance_8bit_single_impl<SIMDLevel::ARM_SVE>(
size_t M,
size_t nbits,
const float* sim_table,
const uint8_t* code) {
if (M <= svcntw())
return distance_single_code_sve_for_small_m(M, sim_table, code);
const float* tab = sim_table;
const size_t ksub = 1 << nbits;
constexpr size_t ksub = 1 << 8;
const auto offsets_0 = svindex_u32(0, static_cast<uint32_t>(ksub));
auto partialSum = svdup_n_f32(0.f);
@@ -159,12 +158,11 @@ float pq_code_distance_single_impl<SIMDLevel::ARM_SVE>(
return svaddv_f32(svptrue_b32(), partialSum);
}
// Combines 4 operations of pq_code_distance_single_impl().
// Combines 4 operations of pq_code_distance_8bit_single_impl().
// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization)
template <>
void pq_code_distance_four_impl<SIMDLevel::ARM_SVE>(
void pq_code_distance_8bit_four_impl<SIMDLevel::ARM_SVE>(
size_t M,
size_t nbits,
const float* sim_table,
const uint8_t* __restrict code0,
const uint8_t* __restrict code1,
@@ -190,7 +188,7 @@ void pq_code_distance_four_impl<SIMDLevel::ARM_SVE>(
}
const float* tab = sim_table;
const size_t ksub = 1 << nbits;
constexpr size_t ksub = 1 << 8;
const auto offsets_0 = svindex_u32(0, static_cast<uint32_t>(ksub));
faiss/python/__init__.pyi
+2
View File
@@ -87,6 +87,8 @@ def unpack_bitstrings(
M_or_nbits: int | npt.NDArray[np.int32],
nbit: int | None = None,
) -> npt.NDArray[np.int32]: ...
def popcount32(x: int) -> int: ...
def popcount64(x: int) -> int: ...
# Version information
FAISS_VERSION_MAJOR: int
faiss/utils/simd_impl/hamming_impl.h
+481
View File
@@ -0,0 +1,481 @@
/*
* 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.
*/
// Shared implementation header for Hamming distance dynamic dispatch.
// Included by per-ISA TUs (hamming_avx2.cpp, hamming_avx512.cpp,
// hamming_neon.cpp) and by hamming.cpp (for the NONE fallback TU).
//
// THE_SIMD_LEVEL must be defined before including this header.
//
// ODR CONTRACT: each TU that includes this header MUST define a UNIQUE
// THE_SIMD_LEVEL value across the linked binary. The template
// specializations at the bottom of this file (hammings_knn_hc_fixSL<SL>,
// etc.) have external linkage with SL-keyed mangled names; two TUs sharing
// THE_SIMD_LEVEL would emit conflicting definitions of the same symbol.
// In DD mode, the per-ISA TUs use distinct values (NONE/AVX2/AVX512/NEON);
// in static mode, only hamming.cpp includes this header so uniqueness is
// trivially satisfied.
//
// with_HammingComputer<SL> is templatized on SIMDLevel and is called
// with THE_SIMD_LEVEL from each _fixSL specialization below.
#pragma once
#ifndef THE_SIMD_LEVEL
#error "Define THE_SIMD_LEVEL before including hamming_impl.h"
#endif
// ISA-specific struct definitions — MUST come first (see comment above).
#include <faiss/utils/hamming_distance/hamming_computer.h>
#include <faiss/utils/hamming.h>
#include <algorithm>
#include <cstdio>
#include <limits>
#include <memory>
#include <vector>
#include <faiss/impl/AuxIndexStructures.h>
#include <faiss/impl/FaissAssert.h>
#include <faiss/impl/IDSelector.h>
#include <faiss/utils/Heap.h>
#include <faiss/utils/approx_topk_hamming/approx_topk_hamming.h>
#include <faiss/utils/utils.h>
namespace faiss {
// All HammingComputer-dependent implementation templates live in an anonymous
// namespace for ODR safety. Different TUs compile this header with different
// HammingComputer struct layouts (generic vs AVX2 vs NEON). The anonymous
// namespace ensures each TU gets its own copy with internal linkage, preventing
// the linker from merging incompatible instantiations.
namespace {
/******************************************************************
* HammingComputer-based search templates
******************************************************************/
template <class HammingComputer>
void hammings_knn_hc_impl(
int bytes_per_code,
int_maxheap_array_t* __restrict ha,
const uint8_t* __restrict bs1,
const uint8_t* __restrict bs2,
size_t n2,
bool order = true,
bool init_heap = true,
ApproxTopK_mode_t approx_topk_mode = ApproxTopK_mode_t::EXACT_TOPK,
const faiss::IDSelector* sel = nullptr) {
size_t k = ha->k;
if (init_heap) {
ha->heapify();
}
const size_t block_size = hamming_batch_size;
for (size_t j0 = 0; j0 < n2; j0 += block_size) {
const size_t j1 = std::min(j0 + block_size, n2);
#pragma omp parallel for
for (int64_t i = 0; i < static_cast<int64_t>(ha->nh); i++) {
HammingComputer hc(bs1 + i * bytes_per_code, bytes_per_code);
const uint8_t* __restrict bs2_ = bs2 + j0 * bytes_per_code;
hamdis_t dis;
hamdis_t* __restrict bh_val_ = ha->val + i * k;
int64_t* __restrict bh_ids_ = ha->ids + i * k;
#define HANDLE_APPROX(NB, BD) \
case ApproxTopK_mode_t::APPROX_TOPK_BUCKETS_B##NB##_D##BD: \
FAISS_THROW_IF_NOT_FMT( \
k <= NB * BD, \
"The chosen mode (%d) of approximate top-k supports " \
"up to %d values, but %zd is requested.", \
(int)(ApproxTopK_mode_t::APPROX_TOPK_BUCKETS_B##NB##_D##BD), \
NB * BD, \
k); \
HeapWithBucketsForHamming32< \
CMax<hamdis_t, int64_t>, \
NB, \
BD, \
HammingComputer>:: \
addn(j1 - j0, hc, bs2_, k, bh_val_, bh_ids_, sel); \
break;
switch (approx_topk_mode) {
HANDLE_APPROX(8, 3)
HANDLE_APPROX(8, 2)
HANDLE_APPROX(16, 2)
HANDLE_APPROX(32, 2)
default: {
for (size_t j = j0; j < j1; j++, bs2_ += bytes_per_code) {
if (sel && !sel->is_member(j)) {
continue;
}
dis = hc.hamming(bs2_);
if (dis < bh_val_[0]) {
faiss::maxheap_replace_top<hamdis_t>(
k, bh_val_, bh_ids_, dis, j);
}
}
} break;
}
}
}
if (order) {
ha->reorder();
}
}
#undef HANDLE_APPROX
template <class HammingComputer>
void hammings_knn_mc_impl(
int bytes_per_code,
const uint8_t* __restrict a,
const uint8_t* __restrict b,
size_t na,
size_t nb,
size_t k,
int32_t* __restrict distances,
int64_t* __restrict labels,
const faiss::IDSelector* sel) {
const int nBuckets = bytes_per_code * 8 + 1;
std::vector<int> all_counters(na * nBuckets, 0);
std::unique_ptr<int64_t[]> all_ids_per_dis(new int64_t[na * nBuckets * k]);
std::vector<HCounterState<HammingComputer>> cs;
for (size_t i = 0; i < na; ++i) {
cs.push_back(
HCounterState<HammingComputer>(
all_counters.data() + i * nBuckets,
all_ids_per_dis.get() + i * nBuckets * k,
a + i * bytes_per_code,
8 * bytes_per_code,
k));
}
const size_t block_size = hamming_batch_size;
for (size_t j0 = 0; j0 < nb; j0 += block_size) {
const size_t j1 = std::min(j0 + block_size, nb);
#pragma omp parallel for
for (int64_t i = 0; i < static_cast<int64_t>(na); ++i) {
for (size_t j = j0; j < j1; ++j) {
if (!sel || sel->is_member(j)) {
cs[i].update_counter(b + j * bytes_per_code, j);
}
}
}
}
for (size_t i = 0; i < na; ++i) {
HCounterState<HammingComputer>& csi = cs[i];
size_t nres = 0;
for (int b_2 = 0; b_2 < nBuckets && nres < k; b_2++) {
for (int l = 0; l < csi.counters[b_2] && nres < k; l++) {
labels[i * k + nres] = csi.ids_per_dis[b_2 * k + l];
distances[i * k + nres] = b_2;
nres++;
}
}
while (nres < k) {
labels[i * k + nres] = -1;
distances[i * k + nres] = std::numeric_limits<int32_t>::max();
++nres;
}
}
}
template <class HammingComputer>
void hamming_range_search_impl(
const uint8_t* a,
const uint8_t* b,
size_t na,
size_t nb,
int radius,
size_t code_size,
RangeSearchResult* res,
const faiss::IDSelector* sel) {
#pragma omp parallel
{
RangeSearchPartialResult pres(res);
#pragma omp for
for (int64_t i = 0; i < static_cast<int64_t>(na); i++) {
HammingComputer hc(a + i * code_size, code_size);
const uint8_t* yi = b;
RangeQueryResult& qres = pres.new_result(i);
for (size_t j = 0; j < nb; j++) {
if (!sel || sel->is_member(j)) {
int dis = hc.hamming(yi);
if (dis < radius) {
qres.add(dis, j);
}
}
yi += code_size;
}
}
pres.finalize();
}
}
/******************************************************************
* Generalized Hamming distances
******************************************************************/
template <class HammingComputer>
void hamming_dis_inner_loop(
const uint8_t* __restrict ca,
const uint8_t* __restrict cb,
size_t nb,
size_t code_size,
int k,
hamdis_t* __restrict bh_val_,
int64_t* __restrict bh_ids_) {
HammingComputer hc(ca, code_size);
for (size_t j = 0; j < nb; j++) {
int ndiff = hc.hamming(cb);
cb += code_size;
if (ndiff < bh_val_[0]) {
maxheap_replace_top<hamdis_t>(k, bh_val_, bh_ids_, ndiff, j);
}
}
}
void generalized_hammings_knn_hc_impl(
int_maxheap_array_t* __restrict ha,
const uint8_t* __restrict a,
const uint8_t* __restrict b,
size_t nb,
size_t code_size,
int ordered) {
int na = ha->nh;
int k = ha->k;
if (ordered) {
ha->heapify();
}
#pragma omp parallel for
for (int i = 0; i < na; i++) {
const uint8_t* __restrict ca = a + i * code_size;
const uint8_t* __restrict cb = b;
hamdis_t* __restrict bh_val_ = ha->val + i * k;
int64_t* __restrict bh_ids_ = ha->ids + i * k;
switch (code_size) {
case 8:
hamming_dis_inner_loop<GenHammingComputer8_tpl<THE_SIMD_LEVEL>>(
ca, cb, nb, 8, k, bh_val_, bh_ids_);
break;
case 16:
hamming_dis_inner_loop<
GenHammingComputer16_tpl<THE_SIMD_LEVEL>>(
ca, cb, nb, 16, k, bh_val_, bh_ids_);
break;
case 32:
hamming_dis_inner_loop<
GenHammingComputer32_tpl<THE_SIMD_LEVEL>>(
ca, cb, nb, 32, k, bh_val_, bh_ids_);
break;
default:
hamming_dis_inner_loop<
GenHammingComputerM8_tpl<THE_SIMD_LEVEL>>(
ca, cb, nb, code_size, k, bh_val_, bh_ids_);
break;
}
}
if (ordered) {
ha->reorder();
}
}
} // anonymous namespace
/******************************************************************
* Entry point template specializations at THE_SIMD_LEVEL
******************************************************************/
#define C64(x) ((uint64_t*)x)
template <>
void hammings_knn_hc_fixSL<THE_SIMD_LEVEL>(
int_maxheap_array_t* ha,
const uint8_t* a,
const uint8_t* b,
size_t nb,
size_t ncodes,
int ordered,
ApproxTopK_mode_t approx_topk_mode,
const IDSelector* sel) {
with_HammingComputer<THE_SIMD_LEVEL>(ncodes, [&]<class HammingComputer>() {
hammings_knn_hc_impl<HammingComputer>(
ncodes, ha, a, b, nb, ordered, true, approx_topk_mode, sel);
});
}
template <>
void hammings_knn_mc_fixSL<THE_SIMD_LEVEL>(
const uint8_t* a,
const uint8_t* b,
size_t na,
size_t nb,
size_t k,
size_t ncodes,
int32_t* distances,
int64_t* labels,
const IDSelector* sel) {
with_HammingComputer<THE_SIMD_LEVEL>(ncodes, [&]<class HammingComputer>() {
hammings_knn_mc_impl<HammingComputer>(
ncodes, a, b, na, nb, k, distances, labels, sel);
});
}
template <>
void hamming_range_search_fixSL<THE_SIMD_LEVEL>(
const uint8_t* a,
const uint8_t* b,
size_t na,
size_t nb,
int radius,
size_t code_size,
RangeSearchResult* result,
const IDSelector* sel) {
with_HammingComputer<THE_SIMD_LEVEL>(
code_size, [&]<class HammingComputer>() {
hamming_range_search_impl<HammingComputer>(
a, b, na, nb, radius, code_size, result, sel);
});
}
template <>
void hammings_fixSL<THE_SIMD_LEVEL>(
const uint8_t* a,
const uint8_t* b,
size_t na,
size_t nb,
size_t ncodes,
hamdis_t* dis) {
FAISS_THROW_IF_NOT(ncodes % 8 == 0);
switch (ncodes) {
case 8:
hammings_impl<64>(C64(a), C64(b), na, nb, dis);
return;
case 16:
hammings_impl<128>(C64(a), C64(b), na, nb, dis);
return;
case 32:
hammings_impl<256>(C64(a), C64(b), na, nb, dis);
return;
case 64:
hammings_impl<512>(C64(a), C64(b), na, nb, dis);
return;
default:
hammings_impl_runtime(C64(a), C64(b), na, nb, ncodes * 8, dis);
return;
}
}
template <>
void generalized_hammings_knn_hc_fixSL<THE_SIMD_LEVEL>(
int_maxheap_array_t* ha,
const uint8_t* a,
const uint8_t* b,
size_t nb,
size_t code_size,
int ordered) {
generalized_hammings_knn_hc_impl(ha, a, b, nb, code_size, ordered);
}
template <>
void hamming_count_thres_fixSL<THE_SIMD_LEVEL>(
const uint8_t* bs1,
const uint8_t* bs2,
size_t n1,
size_t n2,
hamdis_t ht,
size_t ncodes,
size_t* nptr) {
switch (ncodes) {
case 8:
hamming_count_thres_impl<64>(C64(bs1), C64(bs2), n1, n2, ht, nptr);
return;
case 16:
hamming_count_thres_impl<128>(C64(bs1), C64(bs2), n1, n2, ht, nptr);
return;
case 32:
hamming_count_thres_impl<256>(C64(bs1), C64(bs2), n1, n2, ht, nptr);
return;
case 64:
hamming_count_thres_impl<512>(C64(bs1), C64(bs2), n1, n2, ht, nptr);
return;
default:
FAISS_THROW_FMT("not implemented for %zu bits", ncodes);
}
}
template <>
void crosshamming_count_thres_fixSL<THE_SIMD_LEVEL>(
const uint8_t* dbs,
size_t n,
hamdis_t ht,
size_t ncodes,
size_t* nptr) {
switch (ncodes) {
case 8:
crosshamming_count_thres_impl<64>(C64(dbs), n, ht, nptr);
return;
case 16:
crosshamming_count_thres_impl<128>(C64(dbs), n, ht, nptr);
return;
case 32:
crosshamming_count_thres_impl<256>(C64(dbs), n, ht, nptr);
return;
case 64:
crosshamming_count_thres_impl<512>(C64(dbs), n, ht, nptr);
return;
default:
FAISS_THROW_FMT("not implemented for %zu bits", ncodes);
}
}
template <>
size_t match_hamming_thres_fixSL<THE_SIMD_LEVEL>(
const uint8_t* bs1,
const uint8_t* bs2,
size_t n1,
size_t n2,
hamdis_t ht,
size_t ncodes,
int64_t* idx,
hamdis_t* dis) {
switch (ncodes) {
case 8:
return match_hamming_thres_impl<64>(
C64(bs1), C64(bs2), n1, n2, ht, idx, dis);
case 16:
return match_hamming_thres_impl<128>(
C64(bs1), C64(bs2), n1, n2, ht, idx, dis);
case 32:
return match_hamming_thres_impl<256>(
C64(bs1), C64(bs2), n1, n2, ht, idx, dis);
case 64:
return match_hamming_thres_impl<512>(
C64(bs1), C64(bs2), n1, n2, ht, idx, dis);
default:
FAISS_THROW_FMT("not implemented for %zu bits", ncodes);
return 0;
}
}
#undef C64
} // namespace faiss
tests/test_pq_code_distance.cpp
+8 -5
View File
@@ -19,7 +19,11 @@
#include <faiss/impl/FaissAssert.h>
#include <faiss/impl/ProductQuantizer.h>
#include <faiss/utils/pq_code_distance.h>
#include <faiss/impl/pq_code_distance/pq_code_distance-inl.h>
// Pull in inline NONE specializations so the test can call
// PQCodeDistance<PQDecoder8, SIMDLevel::NONE> directly.
#include <faiss/impl/pq_code_distance/pq_code_distance-generic.h>
size_t nMismatches(
const std::vector<float>& ref,
@@ -152,8 +156,8 @@ void test(
for (size_t k = 0; k < 1000; k++) {
#pragma omp parallel for schedule(guided)
for (int64_t i = 0; i < (int64_t)n; i++) {
resultsDispatched1x[i] = faiss::pq_code_distance_single(
subq, 8, lookup.data(), codes.data() + subq * i);
resultsDispatched1x[i] = faiss::pq_code_distance_8bit_single(
subq, lookup.data(), codes.data() + subq * i);
}
}
const auto endingTimepoint = std::chrono::steady_clock::now();
@@ -171,9 +175,8 @@ void test(
for (size_t k = 0; k < 1000; k++) {
#pragma omp parallel for schedule(guided)
for (int64_t i = 0; i < (int64_t)n; i += 4) {
faiss::pq_code_distance_four(
faiss::pq_code_distance_8bit_four(
subq,
8,
lookup.data(),
codes.data() + subq * (i + 0),
codes.data() + subq * (i + 1),