diff --git a/faiss/CMakeLists.txt b/faiss/CMakeLists.txt index 92be9d07a..c62c0ab36 100644 --- a/faiss/CMakeLists.txt +++ b/faiss/CMakeLists.txt @@ -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 diff --git a/faiss/IndexIVFPQ.cpp b/faiss/IndexIVFPQ.cpp index 157e161da..7b98f3ab1 100644 --- a/faiss/IndexIVFPQ.cpp +++ b/faiss/IndexIVFPQ.cpp @@ -28,9 +28,15 @@ #include #include #include -#include +#include #include +// Scalar (NONE) fallback for dynamic dispatch +#define THE_SIMD_LEVEL SIMDLevel::NONE +// NOLINTNEXTLINE(facebook-hte-InlineHeader) +#include +#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 mem; - - // for table pointers - std::vector 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(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 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(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(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 -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 -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 - // 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 - 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 - 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 - 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 - // 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 - 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(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 - void scan_list_polysemous( - size_t ncode, - const uint8_t* codes, - SearchResultType& res) const { - with_HammingComputer< - SIMDLevel::NONE>(pq.code_size, [&]() { - this->scan_list_polysemous_hc( - 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 -struct IVFPQScanner : IVFPQScannerT, - 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(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 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([&]() -> InvertedListScanner* { - auto make = - [&]() -> InvertedListScanner* { - if (metric_type == METRIC_INNER_PRODUCT) { - return new IVFPQScanner< - METRIC_INNER_PRODUCT, - CMin, - PQCodeDist, - use_sel>(*this, store_pairs, 2, sel); - } else if (metric_type == METRIC_L2) { - return new IVFPQScanner< - METRIC_L2, - CMax, - PQCodeDist, - use_sel>(*this, store_pairs, 2, sel); - } else { - FAISS_THROW_MSG("unsupported metric type"); - } - }; - - auto with_decoder = [&]() -> InvertedListScanner* { - if (pq.nbits == 8) { - return make.template - operator(), use_sel>(); - } else if (pq.nbits == 16) { - return make.template - operator(), use_sel>(); - } else { - return make.template - operator(), use_sel>(); - } - }; - - if (sel) { - return with_decoder.template operator()(); - } else { - return with_decoder.template operator()(); - } + return pq_code_distance::make_IVFPQInvertedListScanner( + *this, store_pairs, sel); }); } diff --git a/faiss/IndexPQ.cpp b/faiss/IndexPQ.cpp index c02bda28d..c14de28ba 100644 --- a/faiss/IndexPQ.cpp +++ b/faiss/IndexPQ.cpp @@ -19,12 +19,14 @@ #include #include -#include +#include #include // Scalar (NONE) fallback for dynamic dispatch #define THE_SIMD_LEVEL SIMDLevel::NONE // NOLINTNEXTLINE(facebook-hte-InlineHeader) +#include +// NOLINTNEXTLINE(facebook-hte-InlineHeader) #include #include #undef THE_SIMD_LEVEL @@ -79,89 +81,9 @@ void IndexPQ::train(idx_t n, const float* x) { is_trained = true; } -namespace { - -template -struct PQDistanceComputer : FlatCodesDistanceComputer { - using PQDecoder = typename PQCodeDist::PQDecoder; - size_t d; - MetricType metric; - idx_t nb; - const ProductQuantizer& pq; - const float* sdc; - std::vector 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 -FlatCodesDistanceComputer* get_FlatCodesDistanceComputer1( - const IndexPQ& index) { - if (index.pq.nbits == 8) { - return new PQDistanceComputer>(index); - } else if (index.pq.nbits == 16) { - return new PQDistanceComputer>(index); - } else { - return new PQDistanceComputer>( - index); - } -} - -} // namespace - FlatCodesDistanceComputer* IndexPQ::get_FlatCodesDistanceComputer() const { return with_simd_level([&]() { - return get_FlatCodesDistanceComputer1(*this); + return pq_code_distance::get_PQFlatCodesDistanceComputer(*this); }); } diff --git a/faiss/impl/pq_code_distance/IVFPQScanner_impl.h b/faiss/impl/pq_code_distance/IVFPQScanner_impl.h new file mode 100644 index 000000000..1ce43d155 --- /dev/null +++ b/faiss/impl/pq_code_distance/IVFPQScanner_impl.h @@ -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 +#include +#include +#include +#include +#include +#include +#include +#include + +namespace faiss { +namespace pq_code_distance { + +template +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 +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 + // 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 + 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 + 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 + 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 + // 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 + 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(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 + void scan_list_polysemous( + size_t ncode, + const uint8_t* codes, + SearchResultType& res) const { + with_HammingComputer( + pq.code_size, [&]() { + 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 +struct IVFPQScanner : IVFPQScannerT, + 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(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 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 +InvertedListScanner* make_IVFPQInvertedListScanner( + const IndexIVFPQ& ivfpq, + bool store_pairs, + const IDSelector* sel); + +// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization) +template <> +InvertedListScanner* make_IVFPQInvertedListScanner( + const IndexIVFPQ& ivfpq, + bool store_pairs, + const IDSelector* sel) { + auto make = [&]() -> InvertedListScanner* { + if (ivfpq.metric_type == METRIC_INNER_PRODUCT) { + return new IVFPQScanner< + METRIC_INNER_PRODUCT, + CMin, + PQCodeDist, + use_sel>(ivfpq, store_pairs, 2, sel); + } else if (ivfpq.metric_type == METRIC_L2) { + return new IVFPQScanner< + METRIC_L2, + CMax, + PQCodeDist, + use_sel>(ivfpq, store_pairs, 2, sel); + } else { + FAISS_THROW_MSG("unsupported metric type"); + } + }; + + auto with_decoder = [&]() -> InvertedListScanner* { + if (ivfpq.pq.nbits == 8) { + return make.template + operator(), use_sel>(); + } else if (ivfpq.pq.nbits == 16) { + return make.template + operator(), use_sel>(); + } else { + return make.template operator()< + PQCodeDistance, + use_sel>(); + } + }; + + if (sel) { + return with_decoder.template operator()(); + } else { + return with_decoder.template operator()(); + } +} + +} // namespace pq_code_distance +} // namespace faiss diff --git a/faiss/impl/pq_code_distance/IVFPQ_QueryTables.cpp b/faiss/impl/pq_code_distance/IVFPQ_QueryTables.cpp new file mode 100644 index 000000000..920075df8 --- /dev/null +++ b/faiss/impl/pq_code_distance/IVFPQ_QueryTables.cpp @@ -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 + +#include +#include + +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(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(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(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 diff --git a/faiss/impl/pq_code_distance/IVFPQ_QueryTables.h b/faiss/impl/pq_code_distance/IVFPQ_QueryTables.h new file mode 100644 index 000000000..3d253bf76 --- /dev/null +++ b/faiss/impl/pq_code_distance/IVFPQ_QueryTables.h @@ -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 +#include + +#include +#include +#include +#include + +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 mem; + + // for table pointers + std::vector 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 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 diff --git a/faiss/impl/pq_code_distance/PQDistanceComputer_impl.h b/faiss/impl/pq_code_distance/PQDistanceComputer_impl.h new file mode 100644 index 000000000..6af74c1f7 --- /dev/null +++ b/faiss/impl/pq_code_distance/PQDistanceComputer_impl.h @@ -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 +#include +#include +#include + +namespace faiss { +namespace pq_code_distance { + +template +struct PQDistanceComputer : FlatCodesDistanceComputer { + using PQDecoder = typename PQCodeDist::PQDecoder; + size_t d; + MetricType metric; + idx_t nb; + const ProductQuantizer& pq; + const float* sdc; + std::vector 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 +FlatCodesDistanceComputer* get_PQFlatCodesDistanceComputer( + const IndexPQ& index); + +// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization) +template <> +FlatCodesDistanceComputer* get_PQFlatCodesDistanceComputer( + const IndexPQ& index) { + if (index.pq.nbits == 8) { + return new PQDistanceComputer< + PQCodeDistance>(index); + } else if (index.pq.nbits == 16) { + return new PQDistanceComputer< + PQCodeDistance>(index); + } else { + return new PQDistanceComputer< + PQCodeDistance>(index); + } +} + +} // namespace pq_code_distance +} // namespace faiss diff --git a/faiss/impl/pq_code_distance/avx2.cpp b/faiss/impl/pq_code_distance/avx2.cpp new file mode 100644 index 000000000..5538c7206 --- /dev/null +++ b/faiss/impl/pq_code_distance/avx2.cpp @@ -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 +// NOLINTNEXTLINE(facebook-hte-InlineHeader) +#include +// NOLINTNEXTLINE(facebook-hte-InlineHeader) +#include +// NOLINTNEXTLINE(facebook-hte-InlineHeader) +#include + +#endif // COMPILE_SIMD_AVX2 diff --git a/faiss/impl/pq_code_distance/avx512.cpp b/faiss/impl/pq_code_distance/avx512.cpp new file mode 100644 index 000000000..09dca42e2 --- /dev/null +++ b/faiss/impl/pq_code_distance/avx512.cpp @@ -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 +// NOLINTNEXTLINE(facebook-hte-InlineHeader) +#include +// NOLINTNEXTLINE(facebook-hte-InlineHeader) +#include +// NOLINTNEXTLINE(facebook-hte-InlineHeader) +#include + +#endif // COMPILE_SIMD_AVX512 diff --git a/faiss/impl/pq_code_distance/neon.cpp b/faiss/impl/pq_code_distance/neon.cpp new file mode 100644 index 000000000..8a5453d5a --- /dev/null +++ b/faiss/impl/pq_code_distance/neon.cpp @@ -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 +// NOLINTNEXTLINE(facebook-hte-InlineHeader) +#include +// NOLINTNEXTLINE(facebook-hte-InlineHeader) +#include +// NOLINTNEXTLINE(facebook-hte-InlineHeader) +#include + +#endif // COMPILE_SIMD_ARM_NEON diff --git a/faiss/impl/pq_code_distance/pq_code_distance-avx2.cpp b/faiss/impl/pq_code_distance/pq_code_distance-avx2.h similarity index 95% rename from faiss/impl/pq_code_distance/pq_code_distance-avx2.cpp rename to faiss/impl/pq_code_distance/pq_code_distance-avx2.h index 1d06dafad..377504bb8 100644 --- a/faiss/impl/pq_code_distance/pq_code_distance-avx2.cpp +++ b/faiss/impl/pq_code_distance/pq_code_distance-avx2.h @@ -5,7 +5,7 @@ * LICENSE file in the root directory of this source tree. */ -#ifdef COMPILE_SIMD_AVX2 +#pragma once #include @@ -162,9 +162,8 @@ namespace pq_code_distance { // NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization) template <> -float pq_code_distance_single_impl( +float pq_code_distance_8bit_single_impl( 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( // 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( 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( +void pq_code_distance_8bit_four_impl( 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( // 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( } // namespace pq_code_distance } // namespace faiss - -#endif // COMPILE_SIMD_AVX2 diff --git a/faiss/impl/pq_code_distance/pq_code_distance-avx512.cpp b/faiss/impl/pq_code_distance/pq_code_distance-avx512.h similarity index 90% rename from faiss/impl/pq_code_distance/pq_code_distance-avx512.cpp rename to faiss/impl/pq_code_distance/pq_code_distance-avx512.h index 157521385..1adb998f6 100644 --- a/faiss/impl/pq_code_distance/pq_code_distance-avx512.cpp +++ b/faiss/impl/pq_code_distance/pq_code_distance-avx512.h @@ -5,7 +5,7 @@ * LICENSE file in the root directory of this source tree. */ -#ifdef COMPILE_SIMD_AVX512 +#pragma once #include @@ -22,9 +22,8 @@ namespace pq_code_distance { // NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization) template <> -float pq_code_distance_single_impl( +float pq_code_distance_8bit_single_impl( 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( // 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( 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( +void pq_code_distance_8bit_four_impl( 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( // 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( } // namespace pq_code_distance } // namespace faiss - -#endif // COMPILE_SIMD_AVX512 diff --git a/faiss/impl/pq_code_distance/pq_code_distance-generic.cpp b/faiss/impl/pq_code_distance/pq_code_distance-generic.cpp index 8b8cc654d..ac40c58c1 100644 --- a/faiss/impl/pq_code_distance/pq_code_distance-generic.cpp +++ b/faiss/impl/pq_code_distance/pq_code_distance-generic.cpp @@ -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 +#include namespace faiss { namespace pq_code_distance { -// NONE: use scalar directly. - -// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization) -template <> -float pq_code_distance_single_impl( +float pq_code_distance_8bit_single( size_t M, - size_t nbits, - const float* sim_table, - const uint8_t* code) { - return PQCodeDistanceScalar::distance_single_code( - M, nbits, sim_table, code); -} - -// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization) -template <> -void pq_code_distance_four_impl( - 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::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( - size_t M, - size_t nbits, - const float* sim_table, - const uint8_t* code) { - return PQCodeDistanceScalar::distance_single_code( - M, nbits, sim_table, code); -} - -// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization) -template <> -void pq_code_distance_four_impl( - 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::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([&]() { - return pq_code_distance_single_impl(M, nbits, sim_table, code); + return pq_code_distance_8bit_single_impl(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([&]() { - pq_code_distance_four_impl( + pq_code_distance_8bit_four_impl( M, - nbits, sim_table, code0, code1, diff --git a/faiss/impl/pq_code_distance/pq_code_distance-generic.h b/faiss/impl/pq_code_distance/pq_code_distance-generic.h new file mode 100644 index 000000000..cc8171811 --- /dev/null +++ b/faiss/impl/pq_code_distance/pq_code_distance-generic.h @@ -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 + +namespace faiss { +namespace pq_code_distance { + +// NONE: use scalar directly. + +// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization) +template <> +inline float pq_code_distance_8bit_single_impl( + size_t M, + const float* sim_table, + const uint8_t* code) { + return PQCodeDistanceScalar::distance_single_code( + M, 8, sim_table, code); +} + +// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization) +template <> +inline 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) { + PQCodeDistanceScalar::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( + size_t M, + const float* sim_table, + const uint8_t* code) { + return PQCodeDistanceScalar::distance_single_code( + M, 8, sim_table, code); +} + +// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization) +template <> +inline 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) { + PQCodeDistanceScalar::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 diff --git a/faiss/impl/pq_code_distance/pq_code_distance-inl.h b/faiss/impl/pq_code_distance/pq_code_distance-inl.h index 4a4cb88a0..18d740d19 100644 --- a/faiss/impl/pq_code_distance/pq_code_distance-inl.h +++ b/faiss/impl/pq_code_distance/pq_code_distance-inl.h @@ -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 +#include +#include + +#include +#include #include -#include + +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 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 and + * pq_code_distance_8bit_four_impl 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 +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 +float pq_code_distance_8bit_single_impl( + size_t M, + const float* sim_table, + const uint8_t* code); + +template +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 +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) { + return pq_code_distance_8bit_single_impl(M, sim_table, code); + } else { + return PQCodeDistanceScalar::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) { + pq_code_distance_8bit_four_impl( + M, + sim_table, + code0, + code1, + code2, + code3, + result0, + result1, + result2, + result3); + } else { + PQCodeDistanceScalar::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 diff --git a/faiss/impl/pq_code_distance/pq_code_distance-sve.cpp b/faiss/impl/pq_code_distance/pq_code_distance-sve.cpp index d9caa2fcd..fa00cd2b9 100644 --- a/faiss/impl/pq_code_distance/pq_code_distance-sve.cpp +++ b/faiss/impl/pq_code_distance/pq_code_distance-sve.cpp @@ -83,16 +83,15 @@ namespace pq_code_distance { // NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization) template <> -float pq_code_distance_single_impl( +float pq_code_distance_8bit_single_impl( 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(ksub)); auto partialSum = svdup_n_f32(0.f); @@ -159,12 +158,11 @@ float pq_code_distance_single_impl( 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( +void pq_code_distance_8bit_four_impl( 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( } 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(ksub)); diff --git a/faiss/python/__init__.pyi b/faiss/python/__init__.pyi index 99d8d4f5d..a37b82fa8 100644 --- a/faiss/python/__init__.pyi +++ b/faiss/python/__init__.pyi @@ -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 diff --git a/faiss/utils/simd_impl/hamming_impl.h b/faiss/utils/simd_impl/hamming_impl.h new file mode 100644 index 000000000..a93c2143f --- /dev/null +++ b/faiss/utils/simd_impl/hamming_impl.h @@ -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, +// 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 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 + +#include + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +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 +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(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, \ + 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( + k, bh_val_, bh_ids_, dis, j); + } + } + } break; + } + } + } + if (order) { + ha->reorder(); + } +} + +#undef HANDLE_APPROX + +template +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 all_counters(na * nBuckets, 0); + std::unique_ptr all_ids_per_dis(new int64_t[na * nBuckets * k]); + + std::vector> cs; + for (size_t i = 0; i < na; ++i) { + cs.push_back( + HCounterState( + 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(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& 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::max(); + ++nres; + } + } +} + +template +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(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 +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(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>( + ca, cb, nb, 8, k, bh_val_, bh_ids_); + break; + case 16: + hamming_dis_inner_loop< + GenHammingComputer16_tpl>( + ca, cb, nb, 16, k, bh_val_, bh_ids_); + break; + case 32: + hamming_dis_inner_loop< + GenHammingComputer32_tpl>( + ca, cb, nb, 32, k, bh_val_, bh_ids_); + break; + default: + hamming_dis_inner_loop< + GenHammingComputerM8_tpl>( + 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( + 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(ncodes, [&]() { + hammings_knn_hc_impl( + ncodes, ha, a, b, nb, ordered, true, approx_topk_mode, sel); + }); +} + +template <> +void hammings_knn_mc_fixSL( + 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(ncodes, [&]() { + hammings_knn_mc_impl( + ncodes, a, b, na, nb, k, distances, labels, sel); + }); +} + +template <> +void hamming_range_search_fixSL( + 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( + code_size, [&]() { + hamming_range_search_impl( + a, b, na, nb, radius, code_size, result, sel); + }); +} + +template <> +void hammings_fixSL( + 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( + 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( + 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( + 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( + 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 diff --git a/tests/test_pq_code_distance.cpp b/tests/test_pq_code_distance.cpp index b6b2fc04e..45d4e8132 100644 --- a/tests/test_pq_code_distance.cpp +++ b/tests/test_pq_code_distance.cpp @@ -19,7 +19,11 @@ #include #include -#include +#include + +// Pull in inline NONE specializations so the test can call +// PQCodeDistance directly. +#include size_t nMismatches( const std::vector& 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),