From 41007232d6c941d834e1fec709c6669917f3751d Mon Sep 17 00:00:00 2001 From: Check Deng Date: Tue, 22 Feb 2022 15:04:11 -0800 Subject: [PATCH] AQ fastscan (#2169) Summary: Work in progress. Pull Request resolved: https://github.com/facebookresearch/faiss/pull/2169 Test Plan: buck test mode/opt //faiss/tests/:test_fast_scan buck test mode/opt //faiss/tests/:test_fast_scan_ivf Reviewed By: beauby Differential Revision: D34208813 Pulled By: mdouze fbshipit-source-id: 74b72e07dc537667a7def403c4e46d3d05408c27 --- .gitignore | 1 + benchs/bench_ivf_fastscan.py | 112 +++ faiss/CMakeLists.txt | 10 +- faiss/IndexAdditiveQuantizer.cpp | 4 +- faiss/IndexAdditiveQuantizerFastScan.cpp | 247 +++++ faiss/IndexAdditiveQuantizerFastScan.h | 123 +++ faiss/IndexFastScan.cpp | 557 +++++++++++ faiss/IndexFastScan.h | 137 +++ faiss/IndexIVFAdditiveQuantizer.cpp | 8 +- faiss/IndexIVFAdditiveQuantizerFastScan.cpp | 530 +++++++++++ faiss/IndexIVFAdditiveQuantizerFastScan.h | 138 +++ faiss/IndexIVFFastScan.cpp | 995 ++++++++++++++++++++ faiss/IndexIVFFastScan.h | 193 ++++ faiss/IndexIVFPQFastScan.cpp | 875 +---------------- faiss/IndexIVFPQFastScan.h | 119 +-- faiss/IndexPQFastScan.cpp | 465 +-------- faiss/IndexPQFastScan.h | 84 +- faiss/impl/AdditiveQuantizer.cpp | 87 +- faiss/impl/AdditiveQuantizer.h | 32 +- faiss/impl/LocalSearchQuantizer.cpp | 29 +- faiss/impl/LocalSearchQuantizer.h | 7 +- faiss/impl/LookupTableScaler.h | 79 ++ faiss/impl/ResidualQuantizer.cpp | 33 +- faiss/impl/ResidualQuantizer.h | 7 +- faiss/impl/index_read.cpp | 93 +- faiss/impl/index_write.cpp | 93 +- faiss/impl/pq4_fast_scan.h | 16 +- faiss/impl/pq4_fast_scan_search_1.cpp | 89 +- faiss/impl/pq4_fast_scan_search_qbs.cpp | 102 +- faiss/index_factory.cpp | 39 +- faiss/python/swigfaiss.swig | 24 + faiss/utils/quantize_lut.cpp | 62 ++ faiss/utils/quantize_lut.h | 20 + faiss/utils/simdlib_avx2.h | 4 + faiss/utils/simdlib_emulated.h | 5 + faiss/utils/simdlib_neon.h | 5 + tests/test_fast_scan.py | 160 +++- tests/test_fast_scan_ivf.py | 229 ++++- 38 files changed, 4181 insertions(+), 1632 deletions(-) create mode 100644 benchs/bench_ivf_fastscan.py create mode 100644 faiss/IndexAdditiveQuantizerFastScan.cpp create mode 100644 faiss/IndexAdditiveQuantizerFastScan.h create mode 100644 faiss/IndexFastScan.cpp create mode 100644 faiss/IndexFastScan.h create mode 100644 faiss/IndexIVFAdditiveQuantizerFastScan.cpp create mode 100644 faiss/IndexIVFAdditiveQuantizerFastScan.h create mode 100644 faiss/IndexIVFFastScan.cpp create mode 100644 faiss/IndexIVFFastScan.h create mode 100644 faiss/impl/LookupTableScaler.h diff --git a/.gitignore b/.gitignore index cbd1fa70e..183e2066e 100644 --- a/.gitignore +++ b/.gitignore @@ -16,3 +16,4 @@ /c_api/gpu/bin/ /tests/test /tests/gtest/ +faiss/python/swigfaiss_avx2.swig diff --git a/benchs/bench_ivf_fastscan.py b/benchs/bench_ivf_fastscan.py new file mode 100644 index 000000000..c2d1d5e83 --- /dev/null +++ b/benchs/bench_ivf_fastscan.py @@ -0,0 +1,112 @@ +# Copyright (c) Facebook, Inc. and its affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import faiss +import time +import os +import multiprocessing as mp +import numpy as np +import matplotlib.pyplot as plt + +try: + from faiss.contrib.datasets_fb import \ + DatasetSIFT1M, DatasetDeep1B, DatasetBigANN +except ImportError: + from faiss.contrib.datasets import \ + DatasetSIFT1M, DatasetDeep1B, DatasetBigANN + + +# ds = DatasetDeep1B(10**6) +# ds = DatasetBigANN(nb_M=1) +ds = DatasetSIFT1M() + +xq = ds.get_queries() +xb = ds.get_database() +gt = ds.get_groundtruth() + +xt = ds.get_train() + +nb, d = xb.shape +nq, d = xq.shape +nt, d = xt.shape + +k = 1 +AQ = faiss.AdditiveQuantizer + + +def eval_recall(index, name): + t0 = time.time() + D, I = index.search(xq, k=k) + t = time.time() - t0 + speed = t * 1000 / nq + qps = 1000 / speed + + corrects = (gt == I).sum() + recall = corrects / nq + print( + f'\tnprobe {index.nprobe:3d}, Recall@{k}: ' + f'{recall:.6f}, speed: {speed:.6f} ms/query' + ) + + return recall, qps + + +def eval_and_plot(name, rescale_norm=True, plot=True): + index = faiss.index_factory(d, name) + index_path = f"indices/{name}.faissindex" + + if os.path.exists(index_path): + index = faiss.read_index(index_path) + else: + faiss.omp_set_num_threads(mp.cpu_count()) + index.train(xt) + index.add(xb) + faiss.write_index(index, index_path) + + # search params + if hasattr(index, 'rescale_norm'): + index.rescale_norm = rescale_norm + name += f"(rescale_norm={rescale_norm})" + faiss.omp_set_num_threads(1) + + data = [] + print(f"======{name}") + for nprobe in 1, 2, 4, 6, 8, 12, 16, 24, 32, 48, 64, 128: + index.nprobe = nprobe + recall, qps = eval_recall(index, name) + data.append((recall, qps)) + + if plot: + data = np.array(data) + plt.plot(data[:, 0], data[:, 1], label=name) # x - recall, y - qps + + +M, nlist = 32, 1024 + +# just for warmup... +# eval_and_plot(f"IVF{nlist},PQ{M}x4fs", plot=False) + +# benchmark +plt.figure(figsize=(8, 6), dpi=80) + +# PQ +eval_and_plot(f"IVF{nlist},PQ{M}x4fs") +eval_and_plot(f"IVF{nlist},PQ{M}x4fsr") + +# AQ, by_residual +eval_and_plot(f"IVF{nlist},LSQ{M-2}x4fsr_Nlsq2x4") +eval_and_plot(f"IVF{nlist},RQ{M-2}x4fsr_Nrq2x4") +eval_and_plot(f"IVF{nlist},LSQ{M-2}x4fsr_Nlsq2x4", rescale_norm=False) +eval_and_plot(f"IVF{nlist},RQ{M-2}x4fsr_Nrq2x4", rescale_norm=False) + +# AQ, no by_residual +eval_and_plot(f"IVF{nlist},LSQ{M-2}x4fs_Nlsq2x4") +eval_and_plot(f"IVF{nlist},RQ{M-2}x4fs_Nrq2x4") + +plt.title("Indices on SIFT1M") +plt.xlabel("Recall@1") +plt.ylabel("QPS") +plt.legend(bbox_to_anchor=(1.02, 0.1), loc='upper left', borderaxespad=0) +plt.savefig("bench_ivf_fastscan.png", bbox_inches='tight') diff --git a/faiss/CMakeLists.txt b/faiss/CMakeLists.txt index 1fbc8492b..e8a4c8b8b 100644 --- a/faiss/CMakeLists.txt +++ b/faiss/CMakeLists.txt @@ -24,6 +24,8 @@ set(FAISS_SRC IndexIVFAdditiveQuantizer.cpp IndexIVFFlat.cpp IndexIVFPQ.cpp + IndexIVFFastScan.cpp + IndexIVFAdditiveQuantizerFastScan.cpp IndexIVFPQFastScan.cpp IndexIVFPQR.cpp IndexIVFSpectralHash.cpp @@ -32,6 +34,8 @@ set(FAISS_SRC IndexLattice.cpp IndexNSG.cpp IndexPQ.cpp + IndexFastScan.cpp + IndexAdditiveQuantizerFastScan.cpp IndexPQFastScan.cpp IndexPreTransform.cpp IndexRefine.cpp @@ -100,6 +104,8 @@ set(FAISS_HEADERS IndexIVFAdditiveQuantizer.h IndexIVFFlat.h IndexIVFPQ.h + IndexIVFFastScan.h + IndexIVFAdditiveQuantizerFastScan.h IndexIVFPQFastScan.h IndexIVFPQR.h IndexIVFSpectralHash.h @@ -108,11 +114,12 @@ set(FAISS_HEADERS IndexNNDescent.h IndexNSG.h IndexPQ.h + IndexFastScan.h + IndexAdditiveQuantizerFastScan.h IndexPQFastScan.h IndexPreTransform.h IndexRefine.h IndexReplicas.h - IndexAdditiveQuantizer.h IndexScalarQuantizer.h IndexShards.h MatrixStats.h @@ -128,6 +135,7 @@ set(FAISS_HEADERS impl/FaissException.h impl/HNSW.h impl/LocalSearchQuantizer.h + impl/LookupTableScaler.h impl/NNDescent.h impl/NSG.h impl/PolysemousTraining.h diff --git a/faiss/IndexAdditiveQuantizer.cpp b/faiss/IndexAdditiveQuantizer.cpp index b126c6e37..02a6edd7b 100644 --- a/faiss/IndexAdditiveQuantizer.cpp +++ b/faiss/IndexAdditiveQuantizer.cpp @@ -144,7 +144,9 @@ void IndexAdditiveQuantizer::search( search_with_LUT (*this, x, rh); } else if (aq->search_type == AdditiveQuantizer::ST_norm_qint4) { search_with_LUT (*this, x, rh); - } else if (aq->search_type == AdditiveQuantizer::ST_norm_cqint8) { + } else if (aq->search_type == AdditiveQuantizer::ST_norm_cqint8 || + aq->search_type == AdditiveQuantizer::ST_norm_lsq2x4 || + aq->search_type == AdditiveQuantizer::ST_norm_rq2x4) { search_with_LUT (*this, x, rh); } else if (aq->search_type == AdditiveQuantizer::ST_norm_cqint4) { search_with_LUT (*this, x, rh); diff --git a/faiss/IndexAdditiveQuantizerFastScan.cpp b/faiss/IndexAdditiveQuantizerFastScan.cpp new file mode 100644 index 000000000..bf25a2fd7 --- /dev/null +++ b/faiss/IndexAdditiveQuantizerFastScan.cpp @@ -0,0 +1,247 @@ +/** + * Copyright (c) Facebook, Inc. and its 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 +#include + +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace faiss { + +inline size_t roundup(size_t a, size_t b) { + return (a + b - 1) / b * b; +} + +IndexAdditiveQuantizerFastScan::IndexAdditiveQuantizerFastScan( + AdditiveQuantizer* aq, + MetricType metric, + int bbs) { + init(aq, metric, bbs); +} + +void IndexAdditiveQuantizerFastScan::init( + AdditiveQuantizer* aq, + MetricType metric, + int bbs) { + FAISS_THROW_IF_NOT(aq != nullptr); + FAISS_THROW_IF_NOT(!aq->nbits.empty()); + FAISS_THROW_IF_NOT(aq->nbits[0] == 4); + if (metric == METRIC_INNER_PRODUCT) { + FAISS_THROW_IF_NOT_MSG( + aq->search_type == AdditiveQuantizer::ST_LUT_nonorm, + "Search type must be ST_LUT_nonorm for IP metric"); + } else { + FAISS_THROW_IF_NOT_MSG( + aq->search_type == AdditiveQuantizer::ST_norm_lsq2x4 || + aq->search_type == AdditiveQuantizer::ST_norm_rq2x4, + "Search type must be lsq2x4 or rq2x4 for L2 metric"); + } + + this->aq = aq; + if (metric == METRIC_L2) { + M = aq->M + 2; // 2x4 bits AQ + } else { + M = aq->M; + } + init_fastscan(aq->d, M, 4, metric, bbs); + + max_train_points = 1024 * ksub * M; +} + +IndexAdditiveQuantizerFastScan::IndexAdditiveQuantizerFastScan() + : IndexFastScan() { + is_trained = false; + aq = nullptr; +} + +IndexAdditiveQuantizerFastScan::IndexAdditiveQuantizerFastScan( + const IndexAdditiveQuantizer& orig, + int bbs) { + init(orig.aq, orig.metric_type, bbs); + + ntotal = orig.ntotal; + is_trained = orig.is_trained; + orig_codes = orig.codes.data(); + + ntotal2 = roundup(ntotal, bbs); + codes.resize(ntotal2 * M2 / 2); + pq4_pack_codes(orig_codes, ntotal, M, ntotal2, bbs, M2, codes.get()); +} + +IndexAdditiveQuantizerFastScan::~IndexAdditiveQuantizerFastScan() {} + +void IndexAdditiveQuantizerFastScan::train(idx_t n, const float* x_in) { + if (is_trained) { + return; + } + + const int seed = 0x12345; + size_t nt = n; + const float* x = fvecs_maybe_subsample( + d, &nt, max_train_points, x_in, verbose, seed); + n = nt; + if (verbose) { + printf("training additive quantizer on %zd vectors\n", nt); + } + + aq->verbose = verbose; + aq->train(n, x); + if (metric_type == METRIC_L2) { + estimate_norm_scale(n, x); + } + + is_trained = true; +} + +void IndexAdditiveQuantizerFastScan::estimate_norm_scale( + idx_t n, + const float* x_in) { + FAISS_THROW_IF_NOT(metric_type == METRIC_L2); + + constexpr int seed = 0x980903; + constexpr size_t max_points_estimated = 65536; + size_t ns = n; + const float* x = fvecs_maybe_subsample( + d, &ns, max_points_estimated, x_in, verbose, seed); + n = ns; + std::unique_ptr del_x; + if (x != x_in) { + del_x.reset((float*)x); + } + + std::vector dis_tables(n * M * ksub); + compute_float_LUT(dis_tables.data(), n, x); + + // here we compute the mean of scales for each query + // TODO: try max of scales + double scale = 0; + +#pragma omp parallel for reduction(+ : scale) + for (idx_t i = 0; i < n; i++) { + const float* lut = dis_tables.data() + i * M * ksub; + scale += quantize_lut::aq_estimate_norm_scale(M, ksub, 2, lut); + } + scale /= n; + norm_scale = (int)std::roundf(std::max(scale, 1.0)); + + if (verbose) { + printf("estimated norm scale: %lf\n", scale); + printf("rounded norm scale: %d\n", norm_scale); + } +} + +void IndexAdditiveQuantizerFastScan::compute_codes( + uint8_t* tmp_codes, + idx_t n, + const float* x) const { + aq->compute_codes(x, tmp_codes, n); +} + +void IndexAdditiveQuantizerFastScan::compute_float_LUT( + float* lut, + idx_t n, + const float* x) const { + if (metric_type == METRIC_INNER_PRODUCT) { + aq->compute_LUT(n, x, lut, 1.0f); + } else { + // compute inner product look-up tables + const size_t ip_dim12 = aq->M * ksub; + const size_t norm_dim12 = 2 * ksub; + std::vector ip_lut(n * ip_dim12); + aq->compute_LUT(n, x, ip_lut.data(), -2.0f); + + // copy and rescale norm look-up tables + auto norm_tabs = aq->norm_tabs; + if (rescale_norm && norm_scale > 1 && metric_type == METRIC_L2) { + for (size_t i = 0; i < norm_tabs.size(); i++) { + norm_tabs[i] /= norm_scale; + } + } + const float* norm_lut = norm_tabs.data(); + FAISS_THROW_IF_NOT(norm_tabs.size() == norm_dim12); + + // combine them + for (idx_t i = 0; i < n; i++) { + memcpy(lut, ip_lut.data() + i * ip_dim12, ip_dim12 * sizeof(*lut)); + lut += ip_dim12; + memcpy(lut, norm_lut, norm_dim12 * sizeof(*lut)); + lut += norm_dim12; + } + } +} + +void IndexAdditiveQuantizerFastScan::search( + idx_t n, + const float* x, + idx_t k, + float* distances, + idx_t* labels) const { + FAISS_THROW_IF_NOT(k > 0); + bool rescale = (rescale_norm && norm_scale > 1 && metric_type == METRIC_L2); + if (!rescale) { + IndexFastScan::search(n, x, k, distances, labels); + return; + } + + NormTableScaler scaler(norm_scale); + if (metric_type == METRIC_L2) { + search_dispatch_implem(n, x, k, distances, labels, scaler); + } else { + search_dispatch_implem(n, x, k, distances, labels, scaler); + } +} + +/************************************************************************************** + * IndexResidualQuantizerFastScan + **************************************************************************************/ + +IndexResidualQuantizerFastScan::IndexResidualQuantizerFastScan( + int d, ///< dimensionality of the input vectors + size_t M, ///< number of subquantizers + size_t nbits, ///< number of bit per subvector index + MetricType metric, + Search_type_t search_type, + int bbs) + : rq(d, M, nbits, search_type) { + init(&rq, metric, bbs); +} + +IndexResidualQuantizerFastScan::IndexResidualQuantizerFastScan() { + aq = &rq; +} + +/************************************************************************************** + * IndexLocalSearchQuantizerFastScan + **************************************************************************************/ + +IndexLocalSearchQuantizerFastScan::IndexLocalSearchQuantizerFastScan( + int d, + size_t M, ///< number of subquantizers + size_t nbits, ///< number of bit per subvector index + MetricType metric, + Search_type_t search_type, + int bbs) + : lsq(d, M, nbits, search_type) { + init(&lsq, metric, bbs); +} + +IndexLocalSearchQuantizerFastScan::IndexLocalSearchQuantizerFastScan() { + aq = &lsq; +} + +} // namespace faiss diff --git a/faiss/IndexAdditiveQuantizerFastScan.h b/faiss/IndexAdditiveQuantizerFastScan.h new file mode 100644 index 000000000..fdc139b5a --- /dev/null +++ b/faiss/IndexAdditiveQuantizerFastScan.h @@ -0,0 +1,123 @@ +/** + * Copyright (c) Facebook, Inc. and its 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 + +namespace faiss { + +/** Fast scan version of IndexAQ. Works for 4-bit AQ for now. + * + * The codes are not stored sequentially but grouped in blocks of size bbs. + * This makes it possible to compute distances quickly with SIMD instructions. + * + * Implementations: + * 12: blocked loop with internal loop on Q with qbs + * 13: same with reservoir accumulator to store results + * 14: no qbs with heap accumulator + * 15: no qbs with reservoir accumulator + */ + +struct IndexAdditiveQuantizerFastScan : IndexFastScan { + AdditiveQuantizer* aq; + using Search_type_t = AdditiveQuantizer::Search_type_t; + + bool rescale_norm = true; + int norm_scale = 1; + + // max number of training vectors + size_t max_train_points = 0; + + explicit IndexAdditiveQuantizerFastScan( + AdditiveQuantizer* aq, + MetricType metric = METRIC_L2, + int bbs = 32); + + void init( + AdditiveQuantizer* aq, + MetricType metric = METRIC_L2, + int bbs = 32); + + IndexAdditiveQuantizerFastScan(); + + ~IndexAdditiveQuantizerFastScan() override; + + /// build from an existing IndexAQ + explicit IndexAdditiveQuantizerFastScan( + const IndexAdditiveQuantizer& orig, + int bbs = 32); + + void train(idx_t n, const float* x) override; + + void estimate_norm_scale(idx_t n, const float* x); + + void compute_codes(uint8_t* codes, idx_t n, const float* x) const override; + + void compute_float_LUT(float* lut, idx_t n, const float* x) const override; + + void search( + idx_t n, + const float* x, + idx_t k, + float* distances, + idx_t* labels) const override; +}; + +/** Index based on a residual quantizer. Stored vectors are + * approximated by residual quantization codes. + * Can also be used as a codec + */ +struct IndexResidualQuantizerFastScan : IndexAdditiveQuantizerFastScan { + /// The residual quantizer used to encode the vectors + ResidualQuantizer rq; + + /** Constructor. + * + * @param d dimensionality of the input vectors + * @param M number of subquantizers + * @param nbits number of bit per subvector index + * @param metric metric type + * @param search_type AQ search type + */ + IndexResidualQuantizerFastScan( + int d, ///< dimensionality of the input vectors + size_t M, ///< number of subquantizers + size_t nbits, ///< number of bit per subvector index + MetricType metric = METRIC_L2, + Search_type_t search_type = AdditiveQuantizer::ST_norm_rq2x4, + int bbs = 32); + + IndexResidualQuantizerFastScan(); +}; + +struct IndexLocalSearchQuantizerFastScan : IndexAdditiveQuantizerFastScan { + LocalSearchQuantizer lsq; + + /** Constructor. + * + * @param d dimensionality of the input vectors + * @param M number of subquantizers + * @param nbits number of bit per subvector index + * @param metric metric type + * @param search_type AQ search type + */ + IndexLocalSearchQuantizerFastScan( + int d, ///< dimensionality of the input vectors + size_t M, ///< number of subquantizers + size_t nbits, ///< number of bit per subvector index + MetricType metric = METRIC_L2, + Search_type_t search_type = AdditiveQuantizer::ST_norm_lsq2x4, + int bbs = 32); + + IndexLocalSearchQuantizerFastScan(); +}; + +} // namespace faiss diff --git a/faiss/IndexFastScan.cpp b/faiss/IndexFastScan.cpp new file mode 100644 index 000000000..272157b2c --- /dev/null +++ b/faiss/IndexFastScan.cpp @@ -0,0 +1,557 @@ +/** + * Copyright (c) Facebook, Inc. and its 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 +#include + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace faiss { + +using namespace simd_result_handlers; + +inline size_t roundup(size_t a, size_t b) { + return (a + b - 1) / b * b; +} + +void IndexFastScan::init_fastscan( + int d, + size_t M, + size_t nbits, + MetricType metric, + int bbs) { + FAISS_THROW_IF_NOT(nbits == 4); + FAISS_THROW_IF_NOT(bbs % 32 == 0); + this->d = d; + this->M = M; + this->nbits = nbits; + this->metric_type = metric; + this->bbs = bbs; + ksub = (1 << nbits); + + code_size = (M * nbits + 7) / 8; + ntotal = ntotal2 = 0; + M2 = roundup(M, 2); + is_trained = false; +} + +IndexFastScan::IndexFastScan() + : bbs(0), M(0), code_size(0), ntotal2(0), M2(0) {} + +void IndexFastScan::reset() { + codes.resize(0); + ntotal = 0; +} + +void IndexFastScan::add(idx_t n, const float* x) { + FAISS_THROW_IF_NOT(is_trained); + + // do some blocking to avoid excessive allocs + constexpr idx_t bs = 65536; + if (n > bs) { + for (idx_t i0 = 0; i0 < n; i0 += bs) { + idx_t i1 = std::min(n, i0 + bs); + if (verbose) { + printf("IndexFastScan::add %zd/%zd", size_t(i1), size_t(n)); + } + add(i1 - i0, x + i0 * d); + } + return; + } + InterruptCallback::check(); + + AlignedTable tmp_codes(n * code_size); + compute_codes(tmp_codes.get(), n, x); + + ntotal2 = roundup(ntotal + n, bbs); + size_t new_size = ntotal2 * M2 / 2; // assume nbits = 4 + size_t old_size = codes.size(); + if (new_size > old_size) { + codes.resize(new_size); + memset(codes.get() + old_size, 0, new_size - old_size); + } + + pq4_pack_codes_range( + tmp_codes.get(), M, ntotal, ntotal + n, bbs, M2, codes.get()); + + ntotal += n; +} + +namespace { + +template +void estimators_from_tables_generic( + const IndexFastScan& index, + const uint8_t* codes, + size_t ncodes, + const dis_t* dis_table, + size_t k, + typename C::T* heap_dis, + int64_t* heap_ids, + const Scaler& scaler) { + using accu_t = typename C::T; + + for (size_t j = 0; j < ncodes; ++j) { + BitstringReader bsr(codes + j * index.code_size, index.code_size); + accu_t dis = 0; + const dis_t* dt = dis_table; + for (size_t m = 0; m < index.M - scaler.nscale; m++) { + uint64_t c = bsr.read(index.nbits); + dis += dt[c]; + dt += index.ksub; + } + + for (size_t m = 0; m < scaler.nscale; m++) { + uint64_t c = bsr.read(index.nbits); + dis += scaler.scale_one(dt[c]); + dt += index.ksub; + } + + if (C::cmp(heap_dis[0], dis)) { + heap_pop(k, heap_dis, heap_ids); + heap_push(k, heap_dis, heap_ids, dis, j); + } + } +} + +} // anonymous namespace + +using namespace quantize_lut; + +void IndexFastScan::compute_quantized_LUT( + idx_t n, + const float* x, + uint8_t* lut, + float* normalizers) const { + size_t dim12 = ksub * M; + std::unique_ptr dis_tables(new float[n * dim12]); + compute_float_LUT(dis_tables.get(), n, x); + + for (uint64_t i = 0; i < n; i++) { + round_uint8_per_column( + dis_tables.get() + i * dim12, + M, + ksub, + &normalizers[2 * i], + &normalizers[2 * i + 1]); + } + + for (uint64_t i = 0; i < n; i++) { + const float* t_in = dis_tables.get() + i * dim12; + uint8_t* t_out = lut + i * M2 * ksub; + + for (int j = 0; j < dim12; j++) { + t_out[j] = int(t_in[j]); + } + memset(t_out + dim12, 0, (M2 - M) * ksub); + } +} + +/****************************************************************************** + * Search driver routine + ******************************************************************************/ + +void IndexFastScan::search( + idx_t n, + const float* x, + idx_t k, + float* distances, + idx_t* labels) const { + FAISS_THROW_IF_NOT(k > 0); + + DummyScaler scaler; + if (metric_type == METRIC_L2) { + search_dispatch_implem(n, x, k, distances, labels, scaler); + } else { + search_dispatch_implem(n, x, k, distances, labels, scaler); + } +} + +template +void IndexFastScan::search_dispatch_implem( + idx_t n, + const float* x, + idx_t k, + float* distances, + idx_t* labels, + const Scaler& scaler) const { + using Cfloat = typename std::conditional< + is_max, + CMax, + CMin>::type; + + using C = typename std:: + conditional, CMin>::type; + + if (n == 0) { + return; + } + + // actual implementation used + int impl = implem; + + if (impl == 0) { + if (bbs == 32) { + impl = 12; + } else { + impl = 14; + } + if (k > 20) { + impl++; + } + } + + if (implem == 1) { + FAISS_THROW_MSG("not implemented"); + } else if (implem == 2 || implem == 3 || implem == 4) { + FAISS_THROW_IF_NOT(orig_codes != nullptr); + search_implem_234(n, x, k, distances, labels, scaler); + } else if (impl >= 12 && impl <= 15) { + FAISS_THROW_IF_NOT(ntotal < INT_MAX); + int nt = std::min(omp_get_max_threads(), int(n)); + if (nt < 2) { + if (impl == 12 || impl == 13) { + search_implem_12(n, x, k, distances, labels, impl, scaler); + } else { + search_implem_14(n, x, k, distances, labels, impl, scaler); + } + } else { + // explicitly slice over threads +#pragma omp parallel for num_threads(nt) + for (int slice = 0; slice < nt; slice++) { + idx_t i0 = n * slice / nt; + idx_t i1 = n * (slice + 1) / nt; + float* dis_i = distances + i0 * k; + idx_t* lab_i = labels + i0 * k; + if (impl == 12 || impl == 13) { + search_implem_12( + i1 - i0, x + i0 * d, k, dis_i, lab_i, impl, scaler); + } else { + search_implem_14( + i1 - i0, x + i0 * d, k, dis_i, lab_i, impl, scaler); + } + } + } + } else { + FAISS_THROW_FMT("invalid implem %d impl=%d", implem, impl); + } +} + +template +void IndexFastScan::search_implem_234( + idx_t n, + const float* x, + idx_t k, + float* distances, + idx_t* labels, + const Scaler& scaler) const { + FAISS_THROW_IF_NOT(implem == 2 || implem == 3 || implem == 4); + + const size_t dim12 = ksub * M; + std::unique_ptr dis_tables(new float[n * dim12]); + compute_float_LUT(dis_tables.get(), n, x); + + std::vector normalizers(n * 2); + + if (implem == 2) { + // default float + } else if (implem == 3 || implem == 4) { + for (uint64_t i = 0; i < n; i++) { + round_uint8_per_column( + dis_tables.get() + i * dim12, + M, + ksub, + &normalizers[2 * i], + &normalizers[2 * i + 1]); + } + } + +#pragma omp parallel for if (n > 1000) + for (int64_t i = 0; i < n; i++) { + int64_t* heap_ids = labels + i * k; + float* heap_dis = distances + i * k; + + heap_heapify(k, heap_dis, heap_ids); + + estimators_from_tables_generic( + *this, + orig_codes, + ntotal, + dis_tables.get() + i * dim12, + k, + heap_dis, + heap_ids, + scaler); + + heap_reorder(k, heap_dis, heap_ids); + + if (implem == 4) { + float a = normalizers[2 * i]; + float b = normalizers[2 * i + 1]; + + for (int j = 0; j < k; j++) { + heap_dis[j] = heap_dis[j] / a + b; + } + } + } +} + +template +void IndexFastScan::search_implem_12( + idx_t n, + const float* x, + idx_t k, + float* distances, + idx_t* labels, + int impl, + const Scaler& scaler) const { + FAISS_THROW_IF_NOT(bbs == 32); + + // handle qbs2 blocking by recursive call + int64_t qbs2 = this->qbs == 0 ? 11 : pq4_qbs_to_nq(this->qbs); + if (n > qbs2) { + for (int64_t i0 = 0; i0 < n; i0 += qbs2) { + int64_t i1 = std::min(i0 + qbs2, n); + search_implem_12( + i1 - i0, + x + d * i0, + k, + distances + i0 * k, + labels + i0 * k, + impl, + scaler); + } + return; + } + + size_t dim12 = ksub * M2; + AlignedTable quantized_dis_tables(n * dim12); + std::unique_ptr normalizers(new float[2 * n]); + + if (skip & 1) { + quantized_dis_tables.clear(); + } else { + compute_quantized_LUT( + n, x, quantized_dis_tables.get(), normalizers.get()); + } + + AlignedTable LUT(n * dim12); + + // block sizes are encoded in qbs, 4 bits at a time + + // caution: we override an object field + int qbs = this->qbs; + + if (n != pq4_qbs_to_nq(qbs)) { + qbs = pq4_preferred_qbs(n); + } + + int LUT_nq = + pq4_pack_LUT_qbs(qbs, M2, quantized_dis_tables.get(), LUT.get()); + FAISS_THROW_IF_NOT(LUT_nq == n); + + if (k == 1) { + SingleResultHandler handler(n, ntotal); + if (skip & 4) { + // pass + } else { + handler.disable = bool(skip & 2); + pq4_accumulate_loop_qbs( + qbs, ntotal2, M2, codes.get(), LUT.get(), handler, scaler); + } + + handler.to_flat_arrays(distances, labels, normalizers.get()); + + } else if (impl == 12) { + std::vector tmp_dis(n * k); + std::vector tmp_ids(n * k); + + if (skip & 4) { + // skip + } else { + HeapHandler handler( + n, tmp_dis.data(), tmp_ids.data(), k, ntotal); + handler.disable = bool(skip & 2); + + pq4_accumulate_loop_qbs( + qbs, ntotal2, M2, codes.get(), LUT.get(), handler, scaler); + + if (!(skip & 8)) { + handler.to_flat_arrays(distances, labels, normalizers.get()); + } + } + + } else { // impl == 13 + + ReservoirHandler handler(n, ntotal, k, 2 * k); + handler.disable = bool(skip & 2); + + if (skip & 4) { + // skip + } else { + pq4_accumulate_loop_qbs( + qbs, ntotal2, M2, codes.get(), LUT.get(), handler, scaler); + } + + if (!(skip & 8)) { + handler.to_flat_arrays(distances, labels, normalizers.get()); + } + + FastScan_stats.t0 += handler.times[0]; + FastScan_stats.t1 += handler.times[1]; + FastScan_stats.t2 += handler.times[2]; + FastScan_stats.t3 += handler.times[3]; + } +} + +FastScanStats FastScan_stats; + +template +void IndexFastScan::search_implem_14( + idx_t n, + const float* x, + idx_t k, + float* distances, + idx_t* labels, + int impl, + const Scaler& scaler) const { + FAISS_THROW_IF_NOT(bbs % 32 == 0); + + int qbs2 = qbs == 0 ? 4 : qbs; + + // handle qbs2 blocking by recursive call + if (n > qbs2) { + for (int64_t i0 = 0; i0 < n; i0 += qbs2) { + int64_t i1 = std::min(i0 + qbs2, n); + search_implem_14( + i1 - i0, + x + d * i0, + k, + distances + i0 * k, + labels + i0 * k, + impl, + scaler); + } + return; + } + + size_t dim12 = ksub * M2; + AlignedTable quantized_dis_tables(n * dim12); + std::unique_ptr normalizers(new float[2 * n]); + + if (skip & 1) { + quantized_dis_tables.clear(); + } else { + compute_quantized_LUT( + n, x, quantized_dis_tables.get(), normalizers.get()); + } + + AlignedTable LUT(n * dim12); + pq4_pack_LUT(n, M2, quantized_dis_tables.get(), LUT.get()); + + if (k == 1) { + SingleResultHandler handler(n, ntotal); + if (skip & 4) { + // pass + } else { + handler.disable = bool(skip & 2); + pq4_accumulate_loop( + n, + ntotal2, + bbs, + M2, + codes.get(), + LUT.get(), + handler, + scaler); + } + handler.to_flat_arrays(distances, labels, normalizers.get()); + + } else if (impl == 14) { + std::vector tmp_dis(n * k); + std::vector tmp_ids(n * k); + + if (skip & 4) { + // skip + } else if (k > 1) { + HeapHandler handler( + n, tmp_dis.data(), tmp_ids.data(), k, ntotal); + handler.disable = bool(skip & 2); + + pq4_accumulate_loop( + n, + ntotal2, + bbs, + M2, + codes.get(), + LUT.get(), + handler, + scaler); + + if (!(skip & 8)) { + handler.to_flat_arrays(distances, labels, normalizers.get()); + } + } + + } else { // impl == 15 + + ReservoirHandler handler(n, ntotal, k, 2 * k); + handler.disable = bool(skip & 2); + + if (skip & 4) { + // skip + } else { + pq4_accumulate_loop( + n, + ntotal2, + bbs, + M2, + codes.get(), + LUT.get(), + handler, + scaler); + } + + if (!(skip & 8)) { + handler.to_flat_arrays(distances, labels, normalizers.get()); + } + } +} + +template void IndexFastScan::search_dispatch_implem( + idx_t n, + const float* x, + idx_t k, + float* distances, + idx_t* labels, + const NormTableScaler& scaler) const; + +template void IndexFastScan::search_dispatch_implem( + idx_t n, + const float* x, + idx_t k, + float* distances, + idx_t* labels, + const NormTableScaler& scaler) const; + +} // namespace faiss diff --git a/faiss/IndexFastScan.h b/faiss/IndexFastScan.h new file mode 100644 index 000000000..7c33f0f57 --- /dev/null +++ b/faiss/IndexFastScan.h @@ -0,0 +1,137 @@ +/** + * Copyright (c) Facebook, Inc. and its 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 + +namespace faiss { + +/** Fast scan version of IndexPQ and IndexAQ. Works for 4-bit PQ and AQ for now. + * + * The codes are not stored sequentially but grouped in blocks of size bbs. + * This makes it possible to compute distances quickly with SIMD instructions. + * + * Implementations: + * 12: blocked loop with internal loop on Q with qbs + * 13: same with reservoir accumulator to store results + * 14: no qbs with heap accumulator + * 15: no qbs with reservoir accumulator + */ + +struct IndexFastScan : Index { + // implementation to select + int implem = 0; + // skip some parts of the computation (for timing) + int skip = 0; + + // size of the kernel + int bbs; // set at build time + int qbs = 0; // query block size 0 = use default + + // vector quantizer + size_t M; + size_t nbits; + size_t ksub; + size_t code_size; + + // packed version of the codes + size_t ntotal2; + size_t M2; + + AlignedTable codes; + + // this is for testing purposes only + // (set when initialized by IndexPQ or IndexAQ) + const uint8_t* orig_codes = nullptr; + + void init_fastscan( + int d, + size_t M, + size_t nbits, + MetricType metric, + int bbs); + + IndexFastScan(); + + void reset() override; + + void search( + idx_t n, + const float* x, + idx_t k, + float* distances, + idx_t* labels) const override; + + void add(idx_t n, const float* x) override; + + virtual void compute_codes(uint8_t* codes, idx_t n, const float* x) + const = 0; + + virtual void compute_float_LUT(float* lut, idx_t n, const float* x) + const = 0; + + // called by search function + void compute_quantized_LUT( + idx_t n, + const float* x, + uint8_t* lut, + float* normalizers) const; + + template + void search_dispatch_implem( + idx_t n, + const float* x, + idx_t k, + float* distances, + idx_t* labels, + const Scaler& scaler) const; + + template + void search_implem_234( + idx_t n, + const float* x, + idx_t k, + float* distances, + idx_t* labels, + const Scaler& scaler) const; + + template + void search_implem_12( + idx_t n, + const float* x, + idx_t k, + float* distances, + idx_t* labels, + int impl, + const Scaler& scaler) const; + + template + void search_implem_14( + idx_t n, + const float* x, + idx_t k, + float* distances, + idx_t* labels, + int impl, + const Scaler& scaler) const; +}; + +struct FastScanStats { + uint64_t t0, t1, t2, t3; + FastScanStats() { + reset(); + } + void reset() { + memset(this, 0, sizeof(*this)); + } +}; + +FAISS_API extern FastScanStats FastScan_stats; + +} // namespace faiss diff --git a/faiss/IndexIVFAdditiveQuantizer.cpp b/faiss/IndexIVFAdditiveQuantizer.cpp index 3cdac091e..41fd2f6f0 100644 --- a/faiss/IndexIVFAdditiveQuantizer.cpp +++ b/faiss/IndexIVFAdditiveQuantizer.cpp @@ -44,6 +44,10 @@ void IndexIVFAdditiveQuantizer::train_residual(idx_t n, const float* x) { const float* x_in = x; size_t max_train_points = 1024 * ((size_t)1 << aq->nbits[0]); + // we need more data to train LSQ + if (dynamic_cast(aq)) { + max_train_points = 1024 * aq->M * ((size_t)1 << aq->nbits[0]); + } x = fvecs_maybe_subsample( d, (size_t*)&n, max_train_points, x, verbose, 1234); @@ -244,8 +248,10 @@ InvertedListScanner* IndexIVFAdditiveQuantizer::get_InvertedListScanner( A(ST_norm_float) A(ST_norm_qint8) A(ST_norm_qint4) - A(ST_norm_cqint8) A(ST_norm_cqint4) + case AdditiveQuantizer::ST_norm_lsq2x4: + case AdditiveQuantizer::ST_norm_rq2x4: + A(ST_norm_cqint8) #undef A default: FAISS_THROW_FMT( diff --git a/faiss/IndexIVFAdditiveQuantizerFastScan.cpp b/faiss/IndexIVFAdditiveQuantizerFastScan.cpp new file mode 100644 index 000000000..2be4c3559 --- /dev/null +++ b/faiss/IndexIVFAdditiveQuantizerFastScan.cpp @@ -0,0 +1,530 @@ +/** + * Copyright (c) Facebook, Inc. and its 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 +#include + +#include + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace faiss { + +inline size_t roundup(size_t a, size_t b) { + return (a + b - 1) / b * b; +} + +IndexIVFAdditiveQuantizerFastScan::IndexIVFAdditiveQuantizerFastScan( + Index* quantizer, + AdditiveQuantizer* aq, + size_t d, + size_t nlist, + MetricType metric, + int bbs) + : IndexIVFFastScan(quantizer, d, nlist, 0, metric) { + if (aq != nullptr) { + init(aq, nlist, metric, bbs); + } +} + +void IndexIVFAdditiveQuantizerFastScan::init( + AdditiveQuantizer* aq, + size_t nlist, + MetricType metric, + int bbs) { + FAISS_THROW_IF_NOT(aq != nullptr); + FAISS_THROW_IF_NOT(!aq->nbits.empty()); + FAISS_THROW_IF_NOT(aq->nbits[0] == 4); + if (metric == METRIC_INNER_PRODUCT) { + FAISS_THROW_IF_NOT_MSG( + aq->search_type == AdditiveQuantizer::ST_LUT_nonorm, + "Search type must be ST_LUT_nonorm for IP metric"); + } else { + FAISS_THROW_IF_NOT_MSG( + aq->search_type == AdditiveQuantizer::ST_norm_lsq2x4 || + aq->search_type == AdditiveQuantizer::ST_norm_rq2x4, + "Search type must be lsq2x4 or rq2x4 for L2 metric"); + } + + this->aq = aq; + if (metric_type == METRIC_L2) { + M = aq->M + 2; // 2x4 bits AQ + } else { + M = aq->M; + } + init_fastscan(M, 4, nlist, metric, bbs); + + max_train_points = 1024 * ksub * M; + by_residual = true; +} + +IndexIVFAdditiveQuantizerFastScan::IndexIVFAdditiveQuantizerFastScan( + const IndexIVFAdditiveQuantizer& orig, + int bbs) + : IndexIVFFastScan( + orig.quantizer, + orig.d, + orig.nlist, + 0, + orig.metric_type), + aq(orig.aq) { + FAISS_THROW_IF_NOT( + metric_type == METRIC_INNER_PRODUCT || !orig.by_residual); + + init(aq, nlist, metric_type, bbs); + + is_trained = orig.is_trained; + ntotal = orig.ntotal; + nprobe = orig.nprobe; + + for (size_t i = 0; i < nlist; i++) { + size_t nb = orig.invlists->list_size(i); + size_t nb2 = roundup(nb, bbs); + AlignedTable tmp(nb2 * M2 / 2); + pq4_pack_codes( + InvertedLists::ScopedCodes(orig.invlists, i).get(), + nb, + M, + nb2, + bbs, + M2, + tmp.get()); + invlists->add_entries( + i, + nb, + InvertedLists::ScopedIds(orig.invlists, i).get(), + tmp.get()); + } + + orig_invlists = orig.invlists; +} + +IndexIVFAdditiveQuantizerFastScan::IndexIVFAdditiveQuantizerFastScan() { + bbs = 0; + M2 = 0; + aq = nullptr; + + is_trained = false; +} + +IndexIVFAdditiveQuantizerFastScan::~IndexIVFAdditiveQuantizerFastScan() {} + +/********************************************************* + * Training + *********************************************************/ + +void IndexIVFAdditiveQuantizerFastScan::train_residual( + idx_t n, + const float* x_in) { + if (aq->is_trained) { + return; + } + + const int seed = 0x12345; + size_t nt = n; + const float* x = fvecs_maybe_subsample( + d, &nt, max_train_points, x_in, verbose, seed); + n = nt; + if (verbose) { + printf("training additive quantizer on %zd vectors\n", nt); + } + aq->verbose = verbose; + + std::unique_ptr del_x; + if (x != x_in) { + del_x.reset((float*)x); + } + + const float* trainset; + std::vector residuals(n * d); + std::vector assign(n); + + if (by_residual) { + if (verbose) { + printf("computing residuals\n"); + } + quantizer->assign(n, x, assign.data()); + residuals.resize(n * d); + for (idx_t i = 0; i < n; i++) { + quantizer->compute_residual( + x + i * d, residuals.data() + i * d, assign[i]); + } + trainset = residuals.data(); + } else { + trainset = x; + } + + if (verbose) { + printf("training %zdx%zd additive quantizer on " + "%" PRId64 " vectors in %dD\n", + aq->M, + ksub, + n, + d); + } + aq->verbose = verbose; + aq->train(n, trainset); + + // train norm quantizer + if (by_residual && metric_type == METRIC_L2) { + std::vector decoded_x(n * d); + std::vector x_codes(n * aq->code_size); + aq->compute_codes(residuals.data(), x_codes.data(), n); + aq->decode(x_codes.data(), decoded_x.data(), n); + + // add coarse centroids + FAISS_THROW_IF_NOT(assign.size() == n); + std::vector centroid(d); + for (idx_t i = 0; i < n; i++) { + auto xi = decoded_x.data() + i * d; + quantizer->reconstruct(assign[i], centroid.data()); + fvec_add(d, centroid.data(), xi, xi); + } + + std::vector norms(n, 0); + fvec_norms_L2sqr(norms.data(), decoded_x.data(), d, n); + + // re-train norm tables + aq->train_norm(n, norms.data()); + } + + if (metric_type == METRIC_L2) { + estimate_norm_scale(n, x); + } +} + +void IndexIVFAdditiveQuantizerFastScan::estimate_norm_scale( + idx_t n, + const float* x_in) { + FAISS_THROW_IF_NOT(metric_type == METRIC_L2); + + constexpr int seed = 0x980903; + constexpr size_t max_points_estimated = 65536; + size_t ns = n; + const float* x = fvecs_maybe_subsample( + d, &ns, max_points_estimated, x_in, verbose, seed); + n = ns; + std::unique_ptr del_x; + if (x != x_in) { + del_x.reset((float*)x); + } + + std::vector coarse_ids(n); + std::vector coarse_dis(n); + quantizer->search(n, x, 1, coarse_dis.data(), coarse_ids.data()); + + AlignedTable dis_tables; + AlignedTable biases; + + size_t index_nprobe = nprobe; + nprobe = 1; + compute_LUT(n, x, coarse_ids.data(), coarse_dis.data(), dis_tables, biases); + nprobe = index_nprobe; + + float scale = 0; + +#pragma omp parallel for reduction(+ : scale) + for (idx_t i = 0; i < n; i++) { + const float* lut = dis_tables.get() + i * M * ksub; + scale += quantize_lut::aq_estimate_norm_scale(M, ksub, 2, lut); + } + scale /= n; + norm_scale = (int)std::roundf(std::max(scale, 1.0f)); + + if (verbose) { + printf("estimated norm scale: %lf\n", scale); + printf("rounded norm scale: %d\n", norm_scale); + } +} + +/********************************************************* + * Code management functions + *********************************************************/ + +void IndexIVFAdditiveQuantizerFastScan::encode_vectors( + idx_t n, + const float* x, + const idx_t* list_nos, + uint8_t* codes, + bool include_listnos) const { + idx_t bs = 65536; + if (n > bs) { + for (idx_t i0 = 0; i0 < n; i0 += bs) { + idx_t i1 = std::min(n, i0 + bs); + encode_vectors( + i1 - i0, + x + i0 * d, + list_nos + i0, + codes + i0 * code_size, + include_listnos); + } + return; + } + + if (by_residual) { + std::vector residuals(n * d); + std::vector centroids(n * d); + +#pragma omp parallel for if (n > 1000) + for (idx_t i = 0; i < n; i++) { + if (list_nos[i] < 0) { + memset(residuals.data() + i * d, 0, sizeof(residuals[0]) * d); + } else { + quantizer->compute_residual( + x + i * d, residuals.data() + i * d, list_nos[i]); + } + } + +#pragma omp parallel for if (n > 1000) + for (idx_t i = 0; i < n; i++) { + auto c = centroids.data() + i * d; + quantizer->reconstruct(list_nos[i], c); + } + + aq->compute_codes(residuals.data(), codes, n, centroids.data()); + } else { + aq->compute_codes(x, codes, n); + } + + if (include_listnos) { + size_t coarse_size = coarse_code_size(); + for (idx_t i = n - 1; i >= 0; i--) { + uint8_t* code = codes + i * (coarse_size + code_size); + memmove(code + coarse_size, codes + i * code_size, code_size); + encode_listno(list_nos[i], code); + } + } +} + +/********************************************************* + * Search functions + *********************************************************/ + +void IndexIVFAdditiveQuantizerFastScan::search( + idx_t n, + const float* x, + idx_t k, + float* distances, + idx_t* labels) const { + FAISS_THROW_IF_NOT(k > 0); + bool rescale = (rescale_norm && norm_scale > 1 && metric_type == METRIC_L2); + if (!rescale) { + IndexIVFFastScan::search(n, x, k, distances, labels); + return; + } + + NormTableScaler scaler(norm_scale); + if (metric_type == METRIC_L2) { + search_dispatch_implem(n, x, k, distances, labels, scaler); + } else { + search_dispatch_implem(n, x, k, distances, labels, scaler); + } +} + +/********************************************************* + * Look-Up Table functions + *********************************************************/ + +/******************************************************** + +Let q denote the query vector, + x denote the quantized database vector, + c denote the corresponding IVF centroid, + r denote the residual (x - c). + +The L2 distance between q and x is: + + d(q, x) = (q - x)^2 + = (q - c - r)^2 + = q^2 - 2 - 2 + x^2 + +where q^2 is a constant for all x, is only relevant to c, +and x^2 is the quantized database vector norm. + +Different from IVFAdditiveQuantizer, we encode the quantized vector norm x^2 +instead of r^2. So that we only need to compute one LUT for each query vector: + + LUT[m][k] = -2 * + +`-2` could be precomputed in `compute_LUT` and store in `biases`. +if `by_residual=False`, `` is simply 0. + + + +About norm look-up tables: + +To take advantage of the fast SIMD table lookups, we encode the norm by a 2x4 +bits 1D additive quantizer (simply treat the scalar norm as a 1D vector). + +Let `cm` denote the codebooks of the trained 2x4 bits 1D additive quantizer, +size (2, 16); `bm` denote the encoding code of the norm, a 8-bit integer; `cb` +denote the codebooks of the additive quantizer to encode the database vector, +size (M, 16). + +The decoded norm is: + + decoded_norm = cm[0][bm & 15] + cm[1][bm >> 4] + +The decoding is actually doing a table look-up. + +We combine the norm LUTs and the IP LUTs together: + + LUT is a 2D table, size (M + 2, 16) + if m < M : + LUT[m][k] = -2 * + else: + LUT[m][k] = cm[m - M][k] + +********************************************************/ + +bool IndexIVFAdditiveQuantizerFastScan::lookup_table_is_3d() const { + return false; +} + +void IndexIVFAdditiveQuantizerFastScan::compute_LUT( + size_t n, + const float* x, + const idx_t* coarse_ids, + const float* , + AlignedTable& dis_tables, + AlignedTable& biases) const { + const size_t dim12 = ksub * M; + const size_t ip_dim12 = aq->M * ksub; + + dis_tables.resize(n * dim12); + + float coef = 1.0f; + if (metric_type == METRIC_L2) { + coef = -2.0f; + } + + if (by_residual) { + // bias = coef * + // NOTE: q^2 is not added to `biases` + biases.resize(n * nprobe); +#pragma omp parallel + { + std::vector centroid(d); + float* c = centroid.data(); + +#pragma omp for + for (idx_t ij = 0; ij < n * nprobe; ij++) { + int i = ij / nprobe; + quantizer->reconstruct(coarse_ids[ij], c); + biases[ij] = coef * fvec_inner_product(c, x + i * d, d); + } + } + } + + if (metric_type == METRIC_L2) { + const size_t norm_dim12 = 2 * ksub; + + // inner product look-up tables + aq->compute_LUT(n, x, dis_tables.data(), -2.0f, dim12); + + // copy and rescale norm look-up tables + auto norm_tabs = aq->norm_tabs; + if (rescale_norm && norm_scale > 1 && metric_type == METRIC_L2) { + for (size_t i = 0; i < norm_tabs.size(); i++) { + norm_tabs[i] /= norm_scale; + } + } + const float* norm_lut = norm_tabs.data(); + FAISS_THROW_IF_NOT(norm_tabs.size() == norm_dim12); + + // combine them +#pragma omp parallel for if (n > 100) + for (idx_t i = 0; i < n; i++) { + float* tab = dis_tables.data() + i * dim12 + ip_dim12; + memcpy(tab, norm_lut, norm_dim12 * sizeof(*tab)); + } + + } else if (metric_type == METRIC_INNER_PRODUCT) { + aq->compute_LUT(n, x, dis_tables.get()); + } else { + FAISS_THROW_FMT("metric %d not supported", metric_type); + } +} + +void IndexIVFAdditiveQuantizerFastScan::sa_decode( + idx_t n, + const uint8_t* bytes, + float* x) const { + aq->decode(bytes, x, n); +} + +/********** IndexIVFLocalSearchQuantizerFastScan ************/ +IndexIVFLocalSearchQuantizerFastScan::IndexIVFLocalSearchQuantizerFastScan( + Index* quantizer, + size_t d, + size_t nlist, + size_t M, + size_t nbits, + MetricType metric, + Search_type_t search_type, + int bbs) + : IndexIVFAdditiveQuantizerFastScan( + quantizer, + nullptr, + d, + nlist, + metric, + bbs), + lsq(d, M, nbits, search_type) { + FAISS_THROW_IF_NOT(nbits == 4); + init(&lsq, nlist, metric, bbs); +} + +IndexIVFLocalSearchQuantizerFastScan::IndexIVFLocalSearchQuantizerFastScan() { + aq = &lsq; +} + +IndexIVFLocalSearchQuantizerFastScan::~IndexIVFLocalSearchQuantizerFastScan() {} + +/********** IndexIVFResidualQuantizerFastScan ************/ +IndexIVFResidualQuantizerFastScan::IndexIVFResidualQuantizerFastScan( + Index* quantizer, + size_t d, + size_t nlist, + size_t M, + size_t nbits, + MetricType metric, + Search_type_t search_type, + int bbs) + : IndexIVFAdditiveQuantizerFastScan( + quantizer, + nullptr, + d, + nlist, + metric, + bbs), + rq(d, M, nbits, search_type) { + FAISS_THROW_IF_NOT(nbits == 4); + init(&rq, nlist, metric, bbs); +} + +IndexIVFResidualQuantizerFastScan::IndexIVFResidualQuantizerFastScan() { + aq = &rq; +} + +IndexIVFResidualQuantizerFastScan::~IndexIVFResidualQuantizerFastScan() {} + +} // namespace faiss diff --git a/faiss/IndexIVFAdditiveQuantizerFastScan.h b/faiss/IndexIVFAdditiveQuantizerFastScan.h new file mode 100644 index 000000000..87c066189 --- /dev/null +++ b/faiss/IndexIVFAdditiveQuantizerFastScan.h @@ -0,0 +1,138 @@ +/** + * Copyright (c) Facebook, Inc. and its 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 { + +/** Fast scan version of IVFAQ. Works for 4-bit AQ for now. + * + * The codes in the inverted lists are not stored sequentially but + * grouped in blocks of size bbs. This makes it possible to very quickly + * compute distances with SIMD instructions. + * + * Implementations (implem): + * 0: auto-select implementation (default) + * 1: orig's search, re-implemented + * 2: orig's search, re-ordered by invlist + * 10: optimizer int16 search, collect results in heap, no qbs + * 11: idem, collect results in reservoir + * 12: optimizer int16 search, collect results in heap, uses qbs + * 13: idem, collect results in reservoir + */ + +struct IndexIVFAdditiveQuantizerFastScan : IndexIVFFastScan { + using Search_type_t = AdditiveQuantizer::Search_type_t; + + AdditiveQuantizer* aq; + + bool rescale_norm = false; + int norm_scale = 1; + + // max number of training vectors + size_t max_train_points; + + IndexIVFAdditiveQuantizerFastScan( + Index* quantizer, + AdditiveQuantizer* aq, + size_t d, + size_t nlist, + MetricType metric = METRIC_L2, + int bbs = 32); + + void init(AdditiveQuantizer* aq, size_t nlist, MetricType metric, int bbs); + + IndexIVFAdditiveQuantizerFastScan(); + + ~IndexIVFAdditiveQuantizerFastScan() override; + + // built from an IndexIVFAQ + explicit IndexIVFAdditiveQuantizerFastScan( + const IndexIVFAdditiveQuantizer& orig, + int bbs = 32); + + void train_residual(idx_t n, const float* x) override; + + void estimate_norm_scale(idx_t n, const float* x); + + /// same as the regular IVFAQ encoder. The codes are not reorganized by + /// blocks a that point + void encode_vectors( + idx_t n, + const float* x, + const idx_t* list_nos, + uint8_t* codes, + bool include_listno = false) const override; + + void search( + idx_t n, + const float* x, + idx_t k, + float* distances, + idx_t* labels) const override; + + // prepare look-up tables + + bool lookup_table_is_3d() const override; + + void compute_LUT( + size_t n, + const float* x, + const idx_t* coarse_ids, + const float* coarse_dis, + AlignedTable& dis_tables, + AlignedTable& biases) const override; + + void sa_decode(idx_t n, const uint8_t* bytes, float* x) const override; +}; + +struct IndexIVFLocalSearchQuantizerFastScan + : IndexIVFAdditiveQuantizerFastScan { + LocalSearchQuantizer lsq; + + IndexIVFLocalSearchQuantizerFastScan( + Index* quantizer, + size_t d, + size_t nlist, + size_t M, + size_t nbits, + MetricType metric = METRIC_L2, + Search_type_t search_type = AdditiveQuantizer::ST_norm_lsq2x4, + int bbs = 32); + + IndexIVFLocalSearchQuantizerFastScan(); + + ~IndexIVFLocalSearchQuantizerFastScan(); +}; + +struct IndexIVFResidualQuantizerFastScan : IndexIVFAdditiveQuantizerFastScan { + ResidualQuantizer rq; + + IndexIVFResidualQuantizerFastScan( + Index* quantizer, + size_t d, + size_t nlist, + size_t M, + size_t nbits, + MetricType metric = METRIC_L2, + Search_type_t search_type = AdditiveQuantizer::ST_norm_lsq2x4, + int bbs = 32); + + IndexIVFResidualQuantizerFastScan(); + + ~IndexIVFResidualQuantizerFastScan(); +}; + +} // namespace faiss diff --git a/faiss/IndexIVFFastScan.cpp b/faiss/IndexIVFFastScan.cpp new file mode 100644 index 000000000..9f76da719 --- /dev/null +++ b/faiss/IndexIVFFastScan.cpp @@ -0,0 +1,995 @@ +/** + * Copyright (c) Facebook, Inc. and its 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 +#include + +#include + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace faiss { + +using namespace simd_result_handlers; + +inline size_t roundup(size_t a, size_t b) { + return (a + b - 1) / b * b; +} + +IndexIVFFastScan::IndexIVFFastScan( + Index* quantizer, + size_t d, + size_t nlist, + size_t code_size, + MetricType metric) + : IndexIVF(quantizer, d, nlist, code_size, metric) {} + +IndexIVFFastScan::IndexIVFFastScan() { + bbs = 0; + M2 = 0; + is_trained = false; +} + +void IndexIVFFastScan::init_fastscan( + size_t M, + size_t nbits, + size_t nlist, + MetricType /* metric */, + int bbs) { + FAISS_THROW_IF_NOT(bbs % 32 == 0); + FAISS_THROW_IF_NOT(nbits == 4); + + this->M = M; + this->nbits = nbits; + this->bbs = bbs; + ksub = (1 << nbits); + M2 = roundup(M, 2); + code_size = M2 / 2; + + is_trained = false; + replace_invlists(new BlockInvertedLists(nlist, bbs, bbs * M2 / 2), true); +} + +IndexIVFFastScan::~IndexIVFFastScan() {} + +/********************************************************* + * Code management functions + *********************************************************/ + +void IndexIVFFastScan::add_with_ids( + idx_t n, + const float* x, + const idx_t* xids) { + FAISS_THROW_IF_NOT(is_trained); + + // do some blocking to avoid excessive allocs + constexpr idx_t bs = 65536; + if (n > bs) { + double t0 = getmillisecs(); + for (idx_t i0 = 0; i0 < n; i0 += bs) { + idx_t i1 = std::min(n, i0 + bs); + if (verbose) { + double t1 = getmillisecs(); + double elapsed_time = (t1 - t0) / 1000; + double total_time = 0; + if (i0 != 0) { + total_time = elapsed_time / i0 * n; + } + size_t mem = get_mem_usage_kb() / (1 << 10); + + printf("IndexIVFFastScan::add_with_ids %zd/%zd, time %.2f/%.2f, RSS %zdMB", + size_t(i1), + size_t(n), + elapsed_time, + total_time, + mem); + } + add_with_ids(i1 - i0, x + i0 * d, xids ? xids + i0 : nullptr); + } + return; + } + InterruptCallback::check(); + + AlignedTable codes(n * code_size); + direct_map.check_can_add(xids); + std::unique_ptr idx(new idx_t[n]); + quantizer->assign(n, x, idx.get()); + size_t nadd = 0, nminus1 = 0; + + for (size_t i = 0; i < n; i++) { + if (idx[i] < 0) { + nminus1++; + } + } + + AlignedTable flat_codes(n * code_size); + encode_vectors(n, x, idx.get(), flat_codes.get()); + + DirectMapAdd dm_adder(direct_map, n, xids); + BlockInvertedLists* bil = dynamic_cast(invlists); + FAISS_THROW_IF_NOT_MSG(bil, "only block inverted lists supported"); + + // prepare batches + std::vector order(n); + for (idx_t i = 0; i < n; i++) { + order[i] = i; + } + + // TODO should not need stable + std::stable_sort(order.begin(), order.end(), [&idx](idx_t a, idx_t b) { + return idx[a] < idx[b]; + }); + + // TODO parallelize + idx_t i0 = 0; + while (i0 < n) { + idx_t list_no = idx[order[i0]]; + idx_t i1 = i0 + 1; + while (i1 < n && idx[order[i1]] == list_no) { + i1++; + } + + if (list_no == -1) { + i0 = i1; + continue; + } + + // make linear array + AlignedTable list_codes((i1 - i0) * code_size); + size_t list_size = bil->list_size(list_no); + + bil->resize(list_no, list_size + i1 - i0); + + for (idx_t i = i0; i < i1; i++) { + size_t ofs = list_size + i - i0; + idx_t id = xids ? xids[order[i]] : ntotal + order[i]; + dm_adder.add(order[i], list_no, ofs); + bil->ids[list_no][ofs] = id; + memcpy(list_codes.data() + (i - i0) * code_size, + flat_codes.data() + order[i] * code_size, + code_size); + nadd++; + } + pq4_pack_codes_range( + list_codes.data(), + M, + list_size, + list_size + i1 - i0, + bbs, + M2, + bil->codes[list_no].data()); + + i0 = i1; + } + + ntotal += n; +} + +/********************************************************* + * search + *********************************************************/ + +namespace { + +template +void estimators_from_tables_generic( + const IndexIVFFastScan& index, + const uint8_t* codes, + size_t ncodes, + const dis_t* dis_table, + const int64_t* ids, + float bias, + size_t k, + typename C::T* heap_dis, + int64_t* heap_ids, + const Scaler& scaler) { + using accu_t = typename C::T; + for (size_t j = 0; j < ncodes; ++j) { + BitstringReader bsr(codes + j * index.code_size, index.code_size); + accu_t dis = bias; + const dis_t* __restrict dt = dis_table; + for (size_t m = 0; m < index.M - scaler.nscale; m++) { + uint64_t c = bsr.read(index.nbits); + dis += dt[c]; + dt += index.ksub; + } + + for (size_t m = 0; m < scaler.nscale; m++) { + uint64_t c = bsr.read(index.nbits); + dis += scaler.scale_one(dt[c]); + dt += index.ksub; + } + + if (C::cmp(heap_dis[0], dis)) { + heap_pop(k, heap_dis, heap_ids); + heap_push(k, heap_dis, heap_ids, dis, ids[j]); + } + } +} + +using idx_t = Index::idx_t; +using namespace quantize_lut; + +} // anonymous namespace + +/********************************************************* + * Look-Up Table functions + *********************************************************/ + +void IndexIVFFastScan::compute_LUT_uint8( + size_t n, + const float* x, + const idx_t* coarse_ids, + const float* coarse_dis, + AlignedTable& dis_tables, + AlignedTable& biases, + float* normalizers) const { + AlignedTable dis_tables_float; + AlignedTable biases_float; + + uint64_t t0 = get_cy(); + compute_LUT(n, x, coarse_ids, coarse_dis, dis_tables_float, biases_float); + IVFFastScan_stats.t_compute_distance_tables += get_cy() - t0; + + bool lut_is_3d = lookup_table_is_3d(); + size_t dim123 = ksub * M; + size_t dim123_2 = ksub * M2; + if (lut_is_3d) { + dim123 *= nprobe; + dim123_2 *= nprobe; + } + dis_tables.resize(n * dim123_2); + if (biases_float.get()) { + biases.resize(n * nprobe); + } + uint64_t t1 = get_cy(); + +#pragma omp parallel for if (n > 100) + for (int64_t i = 0; i < n; i++) { + const float* t_in = dis_tables_float.get() + i * dim123; + const float* b_in = nullptr; + uint8_t* t_out = dis_tables.get() + i * dim123_2; + uint16_t* b_out = nullptr; + if (biases_float.get()) { + b_in = biases_float.get() + i * nprobe; + b_out = biases.get() + i * nprobe; + } + + quantize_LUT_and_bias( + nprobe, + M, + ksub, + lut_is_3d, + t_in, + b_in, + t_out, + M2, + b_out, + normalizers + 2 * i, + normalizers + 2 * i + 1); + } + IVFFastScan_stats.t_round += get_cy() - t1; +} + +/********************************************************* + * Search functions + *********************************************************/ + +void IndexIVFFastScan::search( + idx_t n, + const float* x, + idx_t k, + float* distances, + idx_t* labels) const { + FAISS_THROW_IF_NOT(k > 0); + + DummyScaler scaler; + if (metric_type == METRIC_L2) { + search_dispatch_implem(n, x, k, distances, labels, scaler); + } else { + search_dispatch_implem(n, x, k, distances, labels, scaler); + } +} + +template +void IndexIVFFastScan::search_dispatch_implem( + idx_t n, + const float* x, + idx_t k, + float* distances, + idx_t* labels, + const Scaler& scaler) const { + using Cfloat = typename std::conditional< + is_max, + CMax, + CMin>::type; + + using C = typename std::conditional< + is_max, + CMax, + CMin>::type; + + if (n == 0) { + return; + } + + // actual implementation used + int impl = implem; + + if (impl == 0) { + if (bbs == 32) { + impl = 12; + } else { + impl = 10; + } + if (k > 20) { + impl++; + } + } + + if (impl == 1) { + search_implem_1(n, x, k, distances, labels, scaler); + } else if (impl == 2) { + search_implem_2(n, x, k, distances, labels, scaler); + + } else if (impl >= 10 && impl <= 13) { + size_t ndis = 0, nlist_visited = 0; + + if (n < 2) { + if (impl == 12 || impl == 13) { + search_implem_12( + n, + x, + k, + distances, + labels, + impl, + &ndis, + &nlist_visited, + scaler); + } else { + search_implem_10( + n, + x, + k, + distances, + labels, + impl, + &ndis, + &nlist_visited, + scaler); + } + } else { + // explicitly slice over threads + int nslice; + if (n <= omp_get_max_threads()) { + nslice = n; + } else if (lookup_table_is_3d()) { + // make sure we don't make too big LUT tables + size_t lut_size_per_query = + M * ksub * nprobe * (sizeof(float) + sizeof(uint8_t)); + + size_t max_lut_size = precomputed_table_max_bytes; + // how many queries we can handle within mem budget + size_t nq_ok = + std::max(max_lut_size / lut_size_per_query, size_t(1)); + nslice = + roundup(std::max(size_t(n / nq_ok), size_t(1)), + omp_get_max_threads()); + } else { + // LUTs unlikely to be a limiting factor + nslice = omp_get_max_threads(); + } + +#pragma omp parallel for reduction(+ : ndis, nlist_visited) + for (int slice = 0; slice < nslice; slice++) { + idx_t i0 = n * slice / nslice; + idx_t i1 = n * (slice + 1) / nslice; + float* dis_i = distances + i0 * k; + idx_t* lab_i = labels + i0 * k; + if (impl == 12 || impl == 13) { + search_implem_12( + i1 - i0, + x + i0 * d, + k, + dis_i, + lab_i, + impl, + &ndis, + &nlist_visited, + scaler); + } else { + search_implem_10( + i1 - i0, + x + i0 * d, + k, + dis_i, + lab_i, + impl, + &ndis, + &nlist_visited, + scaler); + } + } + } + indexIVF_stats.nq += n; + indexIVF_stats.ndis += ndis; + indexIVF_stats.nlist += nlist_visited; + } else { + FAISS_THROW_FMT("implem %d does not exist", implem); + } +} + +template +void IndexIVFFastScan::search_implem_1( + idx_t n, + const float* x, + idx_t k, + float* distances, + idx_t* labels, + const Scaler& scaler) const { + FAISS_THROW_IF_NOT(orig_invlists); + + std::unique_ptr coarse_ids(new idx_t[n * nprobe]); + std::unique_ptr coarse_dis(new float[n * nprobe]); + + quantizer->search(n, x, nprobe, coarse_dis.get(), coarse_ids.get()); + + size_t dim12 = ksub * M; + AlignedTable dis_tables; + AlignedTable biases; + + compute_LUT(n, x, coarse_ids.get(), coarse_dis.get(), dis_tables, biases); + + bool single_LUT = !lookup_table_is_3d(); + + size_t ndis = 0, nlist_visited = 0; + +#pragma omp parallel for reduction(+ : ndis, nlist_visited) + for (idx_t i = 0; i < n; i++) { + int64_t* heap_ids = labels + i * k; + float* heap_dis = distances + i * k; + heap_heapify(k, heap_dis, heap_ids); + float* LUT = nullptr; + + if (single_LUT) { + LUT = dis_tables.get() + i * dim12; + } + for (idx_t j = 0; j < nprobe; j++) { + if (!single_LUT) { + LUT = dis_tables.get() + (i * nprobe + j) * dim12; + } + idx_t list_no = coarse_ids[i * nprobe + j]; + if (list_no < 0) + continue; + size_t ls = orig_invlists->list_size(list_no); + if (ls == 0) + continue; + InvertedLists::ScopedCodes codes(orig_invlists, list_no); + InvertedLists::ScopedIds ids(orig_invlists, list_no); + + float bias = biases.get() ? biases[i * nprobe + j] : 0; + + estimators_from_tables_generic( + *this, + codes.get(), + ls, + LUT, + ids.get(), + bias, + k, + heap_dis, + heap_ids, + scaler); + nlist_visited++; + ndis++; + } + heap_reorder(k, heap_dis, heap_ids); + } + indexIVF_stats.nq += n; + indexIVF_stats.ndis += ndis; + indexIVF_stats.nlist += nlist_visited; +} + +template +void IndexIVFFastScan::search_implem_2( + idx_t n, + const float* x, + idx_t k, + float* distances, + idx_t* labels, + const Scaler& scaler) const { + FAISS_THROW_IF_NOT(orig_invlists); + + std::unique_ptr coarse_ids(new idx_t[n * nprobe]); + std::unique_ptr coarse_dis(new float[n * nprobe]); + + quantizer->search(n, x, nprobe, coarse_dis.get(), coarse_ids.get()); + + size_t dim12 = ksub * M2; + AlignedTable dis_tables; + AlignedTable biases; + std::unique_ptr normalizers(new float[2 * n]); + + compute_LUT_uint8( + n, + x, + coarse_ids.get(), + coarse_dis.get(), + dis_tables, + biases, + normalizers.get()); + + bool single_LUT = !lookup_table_is_3d(); + + size_t ndis = 0, nlist_visited = 0; + +#pragma omp parallel for reduction(+ : ndis, nlist_visited) + for (idx_t i = 0; i < n; i++) { + std::vector tmp_dis(k); + int64_t* heap_ids = labels + i * k; + uint16_t* heap_dis = tmp_dis.data(); + heap_heapify(k, heap_dis, heap_ids); + const uint8_t* LUT = nullptr; + + if (single_LUT) { + LUT = dis_tables.get() + i * dim12; + } + for (idx_t j = 0; j < nprobe; j++) { + if (!single_LUT) { + LUT = dis_tables.get() + (i * nprobe + j) * dim12; + } + idx_t list_no = coarse_ids[i * nprobe + j]; + if (list_no < 0) + continue; + size_t ls = orig_invlists->list_size(list_no); + if (ls == 0) + continue; + InvertedLists::ScopedCodes codes(orig_invlists, list_no); + InvertedLists::ScopedIds ids(orig_invlists, list_no); + + uint16_t bias = biases.get() ? biases[i * nprobe + j] : 0; + + estimators_from_tables_generic( + *this, + codes.get(), + ls, + LUT, + ids.get(), + bias, + k, + heap_dis, + heap_ids, + scaler); + + nlist_visited++; + ndis += ls; + } + heap_reorder(k, heap_dis, heap_ids); + // convert distances to float + { + float one_a = 1 / normalizers[2 * i], b = normalizers[2 * i + 1]; + if (skip & 16) { + one_a = 1; + b = 0; + } + float* heap_dis_float = distances + i * k; + for (int j = 0; j < k; j++) { + heap_dis_float[j] = b + heap_dis[j] * one_a; + } + } + } + indexIVF_stats.nq += n; + indexIVF_stats.ndis += ndis; + indexIVF_stats.nlist += nlist_visited; +} + +template +void IndexIVFFastScan::search_implem_10( + idx_t n, + const float* x, + idx_t k, + float* distances, + idx_t* labels, + int impl, + size_t* ndis_out, + size_t* nlist_out, + const Scaler& scaler) const { + memset(distances, -1, sizeof(float) * k * n); + memset(labels, -1, sizeof(idx_t) * k * n); + + using HeapHC = HeapHandler; + using ReservoirHC = ReservoirHandler; + using SingleResultHC = SingleResultHandler; + + std::unique_ptr coarse_ids(new idx_t[n * nprobe]); + std::unique_ptr coarse_dis(new float[n * nprobe]); + + uint64_t times[10]; + memset(times, 0, sizeof(times)); + int ti = 0; +#define TIC times[ti++] = get_cy() + TIC; + + quantizer->search(n, x, nprobe, coarse_dis.get(), coarse_ids.get()); + + TIC; + + size_t dim12 = ksub * M2; + AlignedTable dis_tables; + AlignedTable biases; + std::unique_ptr normalizers(new float[2 * n]); + + compute_LUT_uint8( + n, + x, + coarse_ids.get(), + coarse_dis.get(), + dis_tables, + biases, + normalizers.get()); + + TIC; + + bool single_LUT = !lookup_table_is_3d(); + + TIC; + size_t ndis = 0, nlist_visited = 0; + + { + AlignedTable tmp_distances(k); + for (idx_t i = 0; i < n; i++) { + const uint8_t* LUT = nullptr; + int qmap1[1] = {0}; + std::unique_ptr> handler; + + if (k == 1) { + handler.reset(new SingleResultHC(1, 0)); + } else if (impl == 10) { + handler.reset(new HeapHC( + 1, tmp_distances.get(), labels + i * k, k, 0)); + } else if (impl == 11) { + handler.reset(new ReservoirHC(1, 0, k, 2 * k)); + } else { + FAISS_THROW_MSG("invalid"); + } + + handler->q_map = qmap1; + + if (single_LUT) { + LUT = dis_tables.get() + i * dim12; + } + for (idx_t j = 0; j < nprobe; j++) { + size_t ij = i * nprobe + j; + if (!single_LUT) { + LUT = dis_tables.get() + ij * dim12; + } + if (biases.get()) { + handler->dbias = biases.get() + ij; + } + + idx_t list_no = coarse_ids[ij]; + if (list_no < 0) + continue; + size_t ls = invlists->list_size(list_no); + if (ls == 0) + continue; + + InvertedLists::ScopedCodes codes(invlists, list_no); + InvertedLists::ScopedIds ids(invlists, list_no); + + handler->ntotal = ls; + handler->id_map = ids.get(); + +#define DISPATCH(classHC) \ + if (dynamic_cast(handler.get())) { \ + auto* res = static_cast(handler.get()); \ + pq4_accumulate_loop( \ + 1, roundup(ls, bbs), bbs, M2, codes.get(), LUT, *res, scaler); \ + } + DISPATCH(HeapHC) + else DISPATCH(ReservoirHC) else DISPATCH(SingleResultHC) +#undef DISPATCH + + nlist_visited++; + ndis++; + } + + handler->to_flat_arrays( + distances + i * k, + labels + i * k, + skip & 16 ? nullptr : normalizers.get() + i * 2); + } + } + *ndis_out = ndis; + *nlist_out = nlist; +} + +template +void IndexIVFFastScan::search_implem_12( + idx_t n, + const float* x, + idx_t k, + float* distances, + idx_t* labels, + int impl, + size_t* ndis_out, + size_t* nlist_out, + const Scaler& scaler) const { + if (n == 0) { // does not work well with reservoir + return; + } + FAISS_THROW_IF_NOT(bbs == 32); + + std::unique_ptr coarse_ids(new idx_t[n * nprobe]); + std::unique_ptr coarse_dis(new float[n * nprobe]); + + uint64_t times[10]; + memset(times, 0, sizeof(times)); + int ti = 0; +#define TIC times[ti++] = get_cy() + TIC; + + quantizer->search(n, x, nprobe, coarse_dis.get(), coarse_ids.get()); + + TIC; + + size_t dim12 = ksub * M2; + AlignedTable dis_tables; + AlignedTable biases; + std::unique_ptr normalizers(new float[2 * n]); + + compute_LUT_uint8( + n, + x, + coarse_ids.get(), + coarse_dis.get(), + dis_tables, + biases, + normalizers.get()); + + TIC; + + struct QC { + int qno; // sequence number of the query + int list_no; // list to visit + int rank; // this is the rank'th result of the coarse quantizer + }; + bool single_LUT = !lookup_table_is_3d(); + + std::vector qcs; + { + int ij = 0; + for (int i = 0; i < n; i++) { + for (int j = 0; j < nprobe; j++) { + if (coarse_ids[ij] >= 0) { + qcs.push_back(QC{i, int(coarse_ids[ij]), int(j)}); + } + ij++; + } + } + std::sort(qcs.begin(), qcs.end(), [](const QC& a, const QC& b) { + return a.list_no < b.list_no; + }); + } + TIC; + + // prepare the result handlers + + std::unique_ptr> handler; + AlignedTable tmp_distances; + + using HeapHC = HeapHandler; + using ReservoirHC = ReservoirHandler; + using SingleResultHC = SingleResultHandler; + + if (k == 1) { + handler.reset(new SingleResultHC(n, 0)); + } else if (impl == 12) { + tmp_distances.resize(n * k); + handler.reset(new HeapHC(n, tmp_distances.get(), labels, k, 0)); + } else if (impl == 13) { + handler.reset(new ReservoirHC(n, 0, k, 2 * k)); + } + + int qbs2 = this->qbs2 ? this->qbs2 : 11; + + std::vector tmp_bias; + if (biases.get()) { + tmp_bias.resize(qbs2); + handler->dbias = tmp_bias.data(); + } + TIC; + + size_t ndis = 0; + + size_t i0 = 0; + uint64_t t_copy_pack = 0, t_scan = 0; + while (i0 < qcs.size()) { + uint64_t tt0 = get_cy(); + + // find all queries that access this inverted list + int list_no = qcs[i0].list_no; + size_t i1 = i0 + 1; + + while (i1 < qcs.size() && i1 < i0 + qbs2) { + if (qcs[i1].list_no != list_no) { + break; + } + i1++; + } + + size_t list_size = invlists->list_size(list_no); + + if (list_size == 0) { + i0 = i1; + continue; + } + + // re-organize LUTs and biases into the right order + int nc = i1 - i0; + + std::vector q_map(nc), lut_entries(nc); + AlignedTable LUT(nc * dim12); + memset(LUT.get(), -1, nc * dim12); + int qbs = pq4_preferred_qbs(nc); + + for (size_t i = i0; i < i1; i++) { + const QC& qc = qcs[i]; + q_map[i - i0] = qc.qno; + int ij = qc.qno * nprobe + qc.rank; + lut_entries[i - i0] = single_LUT ? qc.qno : ij; + if (biases.get()) { + tmp_bias[i - i0] = biases[ij]; + } + } + pq4_pack_LUT_qbs_q_map( + qbs, M2, dis_tables.get(), lut_entries.data(), LUT.get()); + + // access the inverted list + + ndis += (i1 - i0) * list_size; + + InvertedLists::ScopedCodes codes(invlists, list_no); + InvertedLists::ScopedIds ids(invlists, list_no); + + // prepare the handler + + handler->ntotal = list_size; + handler->q_map = q_map.data(); + handler->id_map = ids.get(); + uint64_t tt1 = get_cy(); + +#define DISPATCH(classHC) \ + if (dynamic_cast(handler.get())) { \ + auto* res = static_cast(handler.get()); \ + pq4_accumulate_loop_qbs( \ + qbs, list_size, M2, codes.get(), LUT.get(), *res, scaler); \ + } + DISPATCH(HeapHC) + else DISPATCH(ReservoirHC) else DISPATCH(SingleResultHC) + + // prepare for next loop + i0 = i1; + + uint64_t tt2 = get_cy(); + t_copy_pack += tt1 - tt0; + t_scan += tt2 - tt1; + } + TIC; + + // labels is in-place for HeapHC + handler->to_flat_arrays( + distances, labels, skip & 16 ? nullptr : normalizers.get()); + + TIC; + + // these stats are not thread-safe + + for (int i = 1; i < ti; i++) { + IVFFastScan_stats.times[i] += times[i] - times[i - 1]; + } + IVFFastScan_stats.t_copy_pack += t_copy_pack; + IVFFastScan_stats.t_scan += t_scan; + + if (auto* rh = dynamic_cast(handler.get())) { + for (int i = 0; i < 4; i++) { + IVFFastScan_stats.reservoir_times[i] += rh->times[i]; + } + } + + *ndis_out = ndis; + *nlist_out = nlist; +} + +void IndexIVFFastScan::reconstruct_from_offset( + int64_t list_no, + int64_t offset, + float* recons) const { + // unpack codes + InvertedLists::ScopedCodes list_codes(invlists, list_no); + std::vector code(code_size, 0); + BitstringWriter bsw(code.data(), code_size); + for (size_t m = 0; m < M; m++) { + uint8_t c = + pq4_get_packed_element(list_codes.get(), bbs, M2, offset, m); + bsw.write(c, nbits); + } + sa_decode(1, code.data(), recons); + + // add centroid to it + if (by_residual) { + std::vector centroid(d); + quantizer->reconstruct(list_no, centroid.data()); + for (int i = 0; i < d; ++i) { + recons[i] += centroid[i]; + } + } +} + +void IndexIVFFastScan::reconstruct_orig_invlists() { + FAISS_THROW_IF_NOT(orig_invlists != nullptr); + FAISS_THROW_IF_NOT(orig_invlists->list_size(0) == 0); + + for (size_t list_no = 0; list_no < nlist; list_no++) { + InvertedLists::ScopedCodes codes(invlists, list_no); + InvertedLists::ScopedIds ids(invlists, list_no); + size_t list_size = orig_invlists->list_size(list_no); + std::vector code(code_size, 0); + + for (size_t offset = 0; offset < list_size; offset++) { + // unpack codes + BitstringWriter bsw(code.data(), code_size); + for (size_t m = 0; m < M; m++) { + uint8_t c = + pq4_get_packed_element(codes.get(), bbs, M2, offset, m); + bsw.write(c, nbits); + } + + // get id + idx_t id = ids.get()[offset]; + + orig_invlists->add_entry(list_no, id, code.data()); + } + } +} + +IVFFastScanStats IVFFastScan_stats; + +template void IndexIVFFastScan::search_dispatch_implem( + idx_t n, + const float* x, + idx_t k, + float* distances, + idx_t* labels, + const NormTableScaler& scaler) const; + +template void IndexIVFFastScan::search_dispatch_implem( + idx_t n, + const float* x, + idx_t k, + float* distances, + idx_t* labels, + const NormTableScaler& scaler) const; + +} // namespace faiss diff --git a/faiss/IndexIVFFastScan.h b/faiss/IndexIVFFastScan.h new file mode 100644 index 000000000..00fdfd349 --- /dev/null +++ b/faiss/IndexIVFFastScan.h @@ -0,0 +1,193 @@ +/** + * Copyright (c) Facebook, Inc. and its 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 + +namespace faiss { + +/** Fast scan version of IVFPQ and IVFAQ. Works for 4-bit PQ/AQ for now. + * + * The codes in the inverted lists are not stored sequentially but + * grouped in blocks of size bbs. This makes it possible to very quickly + * compute distances with SIMD instructions. + * + * Implementations (implem): + * 0: auto-select implementation (default) + * 1: orig's search, re-implemented + * 2: orig's search, re-ordered by invlist + * 10: optimizer int16 search, collect results in heap, no qbs + * 11: idem, collect results in reservoir + * 12: optimizer int16 search, collect results in heap, uses qbs + * 13: idem, collect results in reservoir + */ + +struct IndexIVFFastScan : IndexIVF { + // size of the kernel + int bbs; // set at build time + + size_t M; + size_t nbits; + size_t ksub; + + // M rounded up to a multiple of 2 + size_t M2; + + // search-time implementation + int implem = 0; + // skip some parts of the computation (for timing) + int skip = 0; + bool by_residual = false; + + // batching factors at search time (0 = default) + int qbs = 0; + size_t qbs2 = 0; + + IndexIVFFastScan( + Index* quantizer, + size_t d, + size_t nlist, + size_t code_size, + MetricType metric = METRIC_L2); + + IndexIVFFastScan(); + + void init_fastscan( + size_t M, + size_t nbits, + size_t nlist, + MetricType metric, + int bbs); + + ~IndexIVFFastScan() override; + + /// orig's inverted lists (for debugging) + InvertedLists* orig_invlists = nullptr; + + void add_with_ids(idx_t n, const float* x, const idx_t* xids) override; + + // prepare look-up tables + + virtual bool lookup_table_is_3d() const = 0; + + virtual void compute_LUT( + size_t n, + const float* x, + const idx_t* coarse_ids, + const float* coarse_dis, + AlignedTable& dis_tables, + AlignedTable& biases) const = 0; + + void compute_LUT_uint8( + size_t n, + const float* x, + const idx_t* coarse_ids, + const float* coarse_dis, + AlignedTable& dis_tables, + AlignedTable& biases, + float* normalizers) const; + + void search( + idx_t n, + const float* x, + idx_t k, + float* distances, + idx_t* labels) const override; + + // internal search funcs + + template + void search_dispatch_implem( + idx_t n, + const float* x, + idx_t k, + float* distances, + idx_t* labels, + const Scaler& scaler) const; + + template + void search_implem_1( + idx_t n, + const float* x, + idx_t k, + float* distances, + idx_t* labels, + const Scaler& scaler) const; + + template + void search_implem_2( + idx_t n, + const float* x, + idx_t k, + float* distances, + idx_t* labels, + const Scaler& scaler) const; + + // implem 10 and 12 are not multithreaded internally, so + // export search stats + template + void search_implem_10( + idx_t n, + const float* x, + idx_t k, + float* distances, + idx_t* labels, + int impl, + size_t* ndis_out, + size_t* nlist_out, + const Scaler& scaler) const; + + template + void search_implem_12( + idx_t n, + const float* x, + idx_t k, + float* distances, + idx_t* labels, + int impl, + size_t* ndis_out, + size_t* nlist_out, + const Scaler& scaler) const; + + // reconstruct vectors from packed invlists + void reconstruct_from_offset(int64_t list_no, int64_t offset, float* recons) + const override; + + // reconstruct orig invlists (for debugging) + void reconstruct_orig_invlists(); +}; + +struct IVFFastScanStats { + uint64_t times[10]; + uint64_t t_compute_distance_tables, t_round; + uint64_t t_copy_pack, t_scan, t_to_flat; + uint64_t reservoir_times[4]; + double t_aq_encode; + double t_aq_norm_encode; + + double Mcy_at(int i) { + return times[i] / (1000 * 1000.0); + } + + double Mcy_reservoir_at(int i) { + return reservoir_times[i] / (1000 * 1000.0); + } + IVFFastScanStats() { + reset(); + } + void reset() { + memset(this, 0, sizeof(*this)); + } +}; + +FAISS_API extern IVFFastScanStats IVFFastScan_stats; + +} // namespace faiss diff --git a/faiss/IndexIVFPQFastScan.cpp b/faiss/IndexIVFPQFastScan.cpp index a3dd457fd..91a054bc4 100644 --- a/faiss/IndexIVFPQFastScan.cpp +++ b/faiss/IndexIVFPQFastScan.cpp @@ -40,19 +40,13 @@ IndexIVFPQFastScan::IndexIVFPQFastScan( size_t d, size_t nlist, size_t M, - size_t nbits_per_idx, + size_t nbits, MetricType metric, int bbs) - : IndexIVF(quantizer, d, nlist, 0, metric), - pq(d, M, nbits_per_idx), - bbs(bbs) { - FAISS_THROW_IF_NOT(nbits_per_idx == 4); - M2 = roundup(pq.M, 2); + : IndexIVFFastScan(quantizer, d, nlist, 0, metric), pq(d, M, nbits) { by_residual = false; // set to false by default because it's much faster - is_trained = false; - code_size = pq.code_size; - replace_invlists(new BlockInvertedLists(nlist, bbs, bbs * M2 / 2), true); + init_fastscan(M, nbits, nlist, metric, bbs); } IndexIVFPQFastScan::IndexIVFPQFastScan() { @@ -62,26 +56,21 @@ IndexIVFPQFastScan::IndexIVFPQFastScan() { } IndexIVFPQFastScan::IndexIVFPQFastScan(const IndexIVFPQ& orig, int bbs) - : IndexIVF( + : IndexIVFFastScan( orig.quantizer, orig.d, orig.nlist, orig.pq.code_size, orig.metric_type), - pq(orig.pq), - bbs(bbs) { + pq(orig.pq) { FAISS_THROW_IF_NOT(orig.pq.nbits == 4); + init_fastscan(orig.pq.M, orig.pq.nbits, orig.nlist, orig.metric_type, bbs); + by_residual = orig.by_residual; ntotal = orig.ntotal; is_trained = orig.is_trained; nprobe = orig.nprobe; - size_t M = pq.M; - - M2 = roundup(M, 2); - - replace_invlists( - new BlockInvertedLists(orig.nlist, bbs, bbs * M2 / 2), true); precomputed_table.resize(orig.precomputed_table.size()); @@ -205,150 +194,10 @@ void IndexIVFPQFastScan::encode_vectors( } } -void IndexIVFPQFastScan::add_with_ids( - idx_t n, - const float* x, - const idx_t* xids) { - // copied from IndexIVF::add_with_ids ---> - - // do some blocking to avoid excessive allocs - idx_t bs = 65536; - if (n > bs) { - for (idx_t i0 = 0; i0 < n; i0 += bs) { - idx_t i1 = std::min(n, i0 + bs); - if (verbose) { - printf(" IndexIVFPQFastScan::add_with_ids %zd: %zd", - size_t(i0), - size_t(i1)); - } - add_with_ids(i1 - i0, x + i0 * d, xids ? xids + i0 : nullptr); - } - return; - } - InterruptCallback::check(); - - AlignedTable codes(n * code_size); - - FAISS_THROW_IF_NOT(is_trained); - direct_map.check_can_add(xids); - - std::unique_ptr idx(new idx_t[n]); - quantizer->assign(n, x, idx.get()); - size_t nadd = 0, nminus1 = 0; - - for (size_t i = 0; i < n; i++) { - if (idx[i] < 0) - nminus1++; - } - - AlignedTable flat_codes(n * code_size); - encode_vectors(n, x, idx.get(), flat_codes.get()); - - DirectMapAdd dm_adder(direct_map, n, xids); - - // <--- - - BlockInvertedLists* bil = dynamic_cast(invlists); - FAISS_THROW_IF_NOT_MSG(bil, "only block inverted lists supported"); - - // prepare batches - std::vector order(n); - for (idx_t i = 0; i < n; i++) { - order[i] = i; - } - - // TODO should not need stable - std::stable_sort(order.begin(), order.end(), [&idx](idx_t a, idx_t b) { - return idx[a] < idx[b]; - }); - - // TODO parallelize - idx_t i0 = 0; - while (i0 < n) { - idx_t list_no = idx[order[i0]]; - idx_t i1 = i0 + 1; - while (i1 < n && idx[order[i1]] == list_no) { - i1++; - } - - if (list_no == -1) { - i0 = i1; - continue; - } - - // make linear array - AlignedTable list_codes((i1 - i0) * code_size); - size_t list_size = bil->list_size(list_no); - - bil->resize(list_no, list_size + i1 - i0); - - for (idx_t i = i0; i < i1; i++) { - size_t ofs = list_size + i - i0; - idx_t id = xids ? xids[order[i]] : ntotal + order[i]; - dm_adder.add(order[i], list_no, ofs); - bil->ids[list_no][ofs] = id; - memcpy(list_codes.data() + (i - i0) * code_size, - flat_codes.data() + order[i] * code_size, - code_size); - nadd++; - } - pq4_pack_codes_range( - list_codes.data(), - pq.M, - list_size, - list_size + i1 - i0, - bbs, - M2, - bil->codes[list_no].data()); - - i0 = i1; - } - - ntotal += n; -} - /********************************************************* - * search + * Look-Up Table functions *********************************************************/ -namespace { - -// from impl/ProductQuantizer.cpp -template -void pq_estimators_from_tables_generic( - const ProductQuantizer& pq, - size_t nbits, - const uint8_t* codes, - size_t ncodes, - const dis_t* dis_table, - const int64_t* ids, - float dis0, - size_t k, - typename C::T* heap_dis, - int64_t* heap_ids) { - using accu_t = typename C::T; - const size_t M = pq.M; - const size_t ksub = pq.ksub; - for (size_t j = 0; j < ncodes; ++j) { - PQDecoderGeneric decoder(codes + j * pq.code_size, nbits); - accu_t dis = dis0; - const dis_t* dt = dis_table; - for (size_t m = 0; m < M; m++) { - uint64_t c = decoder.decode(); - dis += dt[c]; - dt += ksub; - } - - if (C::cmp(heap_dis[0], dis)) { - heap_pop(k, heap_dis, heap_ids); - heap_push(k, heap_dis, heap_ids, dis, ids[j]); - } - } -} - -using idx_t = Index::idx_t; -using namespace quantize_lut; - void fvec_madd_avx( size_t n, const float* a, @@ -373,11 +222,9 @@ void fvec_madd_avx( } } -} // anonymous namespace - -/********************************************************* - * Look-Up Table functions - *********************************************************/ +bool IndexIVFPQFastScan::lookup_table_is_3d() const { + return by_residual && metric_type == METRIC_L2; +} void IndexIVFPQFastScan::compute_LUT( size_t n, @@ -386,16 +233,14 @@ void IndexIVFPQFastScan::compute_LUT( const float* coarse_dis, AlignedTable& dis_tables, AlignedTable& biases) const { - const IndexIVFPQFastScan& ivfpq = *this; size_t dim12 = pq.ksub * pq.M; size_t d = pq.d; - size_t nprobe = ivfpq.nprobe; - if (ivfpq.by_residual) { - if (ivfpq.metric_type == METRIC_L2) { + if (by_residual) { + if (metric_type == METRIC_L2) { dis_tables.resize(n * nprobe * dim12); - if (ivfpq.use_precomputed_table == 1) { + if (use_precomputed_table == 1) { biases.resize(n * nprobe); memcpy(biases.get(), coarse_dis, sizeof(float) * n * nprobe); @@ -434,7 +279,7 @@ void IndexIVFPQFastScan::compute_LUT( idx_t cij = coarse_ids[ij]; if (cij >= 0) { - ivfpq.quantizer->compute_residual(x + i * d, xij, cij); + quantizer->compute_residual(x + i * d, xij, cij); } else { // will fill with NaNs memset(xij, -1, sizeof(float) * d); @@ -445,7 +290,7 @@ void IndexIVFPQFastScan::compute_LUT( n * nprobe, xrel.get(), dis_tables.get()); } - } else if (ivfpq.metric_type == METRIC_INNER_PRODUCT) { + } else if (metric_type == METRIC_INNER_PRODUCT) { dis_tables.resize(n * dim12); pq.compute_inner_prod_tables(n, x, dis_tables.get()); // compute_inner_prod_tables(pq, n, x, dis_tables.get()); @@ -453,698 +298,24 @@ void IndexIVFPQFastScan::compute_LUT( biases.resize(n * nprobe); memcpy(biases.get(), coarse_dis, sizeof(float) * n * nprobe); } else { - FAISS_THROW_FMT("metric %d not supported", ivfpq.metric_type); + FAISS_THROW_FMT("metric %d not supported", metric_type); } } else { dis_tables.resize(n * dim12); - if (ivfpq.metric_type == METRIC_L2) { + if (metric_type == METRIC_L2) { pq.compute_distance_tables(n, x, dis_tables.get()); - } else if (ivfpq.metric_type == METRIC_INNER_PRODUCT) { + } else if (metric_type == METRIC_INNER_PRODUCT) { pq.compute_inner_prod_tables(n, x, dis_tables.get()); } else { - FAISS_THROW_FMT("metric %d not supported", ivfpq.metric_type); + FAISS_THROW_FMT("metric %d not supported", metric_type); } } } -void IndexIVFPQFastScan::compute_LUT_uint8( - size_t n, - const float* x, - const idx_t* coarse_ids, - const float* coarse_dis, - AlignedTable& dis_tables, - AlignedTable& biases, - float* normalizers) const { - const IndexIVFPQFastScan& ivfpq = *this; - AlignedTable dis_tables_float; - AlignedTable biases_float; - - uint64_t t0 = get_cy(); - compute_LUT(n, x, coarse_ids, coarse_dis, dis_tables_float, biases_float); - IVFFastScan_stats.t_compute_distance_tables += get_cy() - t0; - - bool lut_is_3d = ivfpq.by_residual && ivfpq.metric_type == METRIC_L2; - size_t dim123 = pq.ksub * pq.M; - size_t dim123_2 = pq.ksub * M2; - if (lut_is_3d) { - dim123 *= nprobe; - dim123_2 *= nprobe; - } - dis_tables.resize(n * dim123_2); - if (biases_float.get()) { - biases.resize(n * nprobe); - } - uint64_t t1 = get_cy(); - -#pragma omp parallel for if (n > 100) - for (int64_t i = 0; i < n; i++) { - const float* t_in = dis_tables_float.get() + i * dim123; - const float* b_in = nullptr; - uint8_t* t_out = dis_tables.get() + i * dim123_2; - uint16_t* b_out = nullptr; - if (biases_float.get()) { - b_in = biases_float.get() + i * nprobe; - b_out = biases.get() + i * nprobe; - } - - quantize_LUT_and_bias( - nprobe, - pq.M, - pq.ksub, - lut_is_3d, - t_in, - b_in, - t_out, - M2, - b_out, - normalizers + 2 * i, - normalizers + 2 * i + 1); - } - IVFFastScan_stats.t_round += get_cy() - t1; +void IndexIVFPQFastScan::sa_decode(idx_t n, const uint8_t* bytes, float* x) + const { + pq.decode(bytes, x, n); } -/********************************************************* - * Search functions - *********************************************************/ - -template -void IndexIVFPQFastScan::search_dispatch_implem( - idx_t n, - const float* x, - idx_t k, - float* distances, - idx_t* labels) const { - using Cfloat = typename std::conditional< - is_max, - CMax, - CMin>::type; - - using C = typename std::conditional< - is_max, - CMax, - CMin>::type; - - if (n == 0) { - return; - } - - // actual implementation used - int impl = implem; - - if (impl == 0) { - if (bbs == 32) { - impl = 12; - } else { - impl = 10; - } - if (k > 20) { - impl++; - } - } - - if (impl == 1) { - search_implem_1(n, x, k, distances, labels); - } else if (impl == 2) { - search_implem_2(n, x, k, distances, labels); - - } else if (impl >= 10 && impl <= 13) { - size_t ndis = 0, nlist_visited = 0; - - if (n < 2) { - if (impl == 12 || impl == 13) { - search_implem_12( - n, - x, - k, - distances, - labels, - impl, - &ndis, - &nlist_visited); - } else { - search_implem_10( - n, - x, - k, - distances, - labels, - impl, - &ndis, - &nlist_visited); - } - } else { - // explicitly slice over threads - int nslice; - if (n <= omp_get_max_threads()) { - nslice = n; - } else if (by_residual && metric_type == METRIC_L2) { - // make sure we don't make too big LUT tables - size_t lut_size_per_query = pq.M * pq.ksub * nprobe * - (sizeof(float) + sizeof(uint8_t)); - - size_t max_lut_size = precomputed_table_max_bytes; - // how many queries we can handle within mem budget - size_t nq_ok = - std::max(max_lut_size / lut_size_per_query, size_t(1)); - nslice = - roundup(std::max(size_t(n / nq_ok), size_t(1)), - omp_get_max_threads()); - } else { - // LUTs unlikely to be a limiting factor - nslice = omp_get_max_threads(); - } - -#pragma omp parallel for reduction(+ : ndis, nlist_visited) - for (int slice = 0; slice < nslice; slice++) { - idx_t i0 = n * slice / nslice; - idx_t i1 = n * (slice + 1) / nslice; - float* dis_i = distances + i0 * k; - idx_t* lab_i = labels + i0 * k; - if (impl == 12 || impl == 13) { - search_implem_12( - i1 - i0, - x + i0 * d, - k, - dis_i, - lab_i, - impl, - &ndis, - &nlist_visited); - } else { - search_implem_10( - i1 - i0, - x + i0 * d, - k, - dis_i, - lab_i, - impl, - &ndis, - &nlist_visited); - } - } - } - indexIVF_stats.nq += n; - indexIVF_stats.ndis += ndis; - indexIVF_stats.nlist += nlist_visited; - } else { - FAISS_THROW_FMT("implem %d does not exist", implem); - } -} - -void IndexIVFPQFastScan::search( - idx_t n, - const float* x, - idx_t k, - float* distances, - idx_t* labels) const { - FAISS_THROW_IF_NOT(k > 0); - - if (metric_type == METRIC_L2) { - search_dispatch_implem(n, x, k, distances, labels); - } else { - search_dispatch_implem(n, x, k, distances, labels); - } -} - -template -void IndexIVFPQFastScan::search_implem_1( - idx_t n, - const float* x, - idx_t k, - float* distances, - idx_t* labels) const { - FAISS_THROW_IF_NOT(orig_invlists); - - std::unique_ptr coarse_ids(new idx_t[n * nprobe]); - std::unique_ptr coarse_dis(new float[n * nprobe]); - - quantizer->search(n, x, nprobe, coarse_dis.get(), coarse_ids.get()); - - size_t dim12 = pq.ksub * pq.M; - AlignedTable dis_tables; - AlignedTable biases; - - compute_LUT(n, x, coarse_ids.get(), coarse_dis.get(), dis_tables, biases); - - bool single_LUT = !(by_residual && metric_type == METRIC_L2); - - size_t ndis = 0, nlist_visited = 0; - -#pragma omp parallel for reduction(+ : ndis, nlist_visited) - for (idx_t i = 0; i < n; i++) { - int64_t* heap_ids = labels + i * k; - float* heap_dis = distances + i * k; - heap_heapify(k, heap_dis, heap_ids); - float* LUT = nullptr; - - if (single_LUT) { - LUT = dis_tables.get() + i * dim12; - } - for (idx_t j = 0; j < nprobe; j++) { - if (!single_LUT) { - LUT = dis_tables.get() + (i * nprobe + j) * dim12; - } - idx_t list_no = coarse_ids[i * nprobe + j]; - if (list_no < 0) - continue; - size_t ls = orig_invlists->list_size(list_no); - if (ls == 0) - continue; - InvertedLists::ScopedCodes codes(orig_invlists, list_no); - InvertedLists::ScopedIds ids(orig_invlists, list_no); - - float bias = biases.get() ? biases[i * nprobe + j] : 0; - - pq_estimators_from_tables_generic( - pq, - pq.nbits, - codes.get(), - ls, - LUT, - ids.get(), - bias, - k, - heap_dis, - heap_ids); - nlist_visited++; - ndis++; - } - heap_reorder(k, heap_dis, heap_ids); - } - indexIVF_stats.nq += n; - indexIVF_stats.ndis += ndis; - indexIVF_stats.nlist += nlist_visited; -} - -template -void IndexIVFPQFastScan::search_implem_2( - idx_t n, - const float* x, - idx_t k, - float* distances, - idx_t* labels) const { - FAISS_THROW_IF_NOT(orig_invlists); - - std::unique_ptr coarse_ids(new idx_t[n * nprobe]); - std::unique_ptr coarse_dis(new float[n * nprobe]); - - quantizer->search(n, x, nprobe, coarse_dis.get(), coarse_ids.get()); - - size_t dim12 = pq.ksub * M2; - AlignedTable dis_tables; - AlignedTable biases; - std::unique_ptr normalizers(new float[2 * n]); - - compute_LUT_uint8( - n, - x, - coarse_ids.get(), - coarse_dis.get(), - dis_tables, - biases, - normalizers.get()); - - bool single_LUT = !(by_residual && metric_type == METRIC_L2); - - size_t ndis = 0, nlist_visited = 0; - -#pragma omp parallel for reduction(+ : ndis, nlist_visited) - for (idx_t i = 0; i < n; i++) { - std::vector tmp_dis(k); - int64_t* heap_ids = labels + i * k; - uint16_t* heap_dis = tmp_dis.data(); - heap_heapify(k, heap_dis, heap_ids); - const uint8_t* LUT = nullptr; - - if (single_LUT) { - LUT = dis_tables.get() + i * dim12; - } - for (idx_t j = 0; j < nprobe; j++) { - if (!single_LUT) { - LUT = dis_tables.get() + (i * nprobe + j) * dim12; - } - idx_t list_no = coarse_ids[i * nprobe + j]; - if (list_no < 0) - continue; - size_t ls = orig_invlists->list_size(list_no); - if (ls == 0) - continue; - InvertedLists::ScopedCodes codes(orig_invlists, list_no); - InvertedLists::ScopedIds ids(orig_invlists, list_no); - - uint16_t bias = biases.get() ? biases[i * nprobe + j] : 0; - - pq_estimators_from_tables_generic( - pq, - pq.nbits, - codes.get(), - ls, - LUT, - ids.get(), - bias, - k, - heap_dis, - heap_ids); - - nlist_visited++; - ndis += ls; - } - heap_reorder(k, heap_dis, heap_ids); - // convert distances to float - { - float one_a = 1 / normalizers[2 * i], b = normalizers[2 * i + 1]; - if (skip & 16) { - one_a = 1; - b = 0; - } - float* heap_dis_float = distances + i * k; - for (int j = 0; j < k; j++) { - heap_dis_float[j] = b + heap_dis[j] * one_a; - } - } - } - indexIVF_stats.nq += n; - indexIVF_stats.ndis += ndis; - indexIVF_stats.nlist += nlist_visited; -} - -template -void IndexIVFPQFastScan::search_implem_10( - idx_t n, - const float* x, - idx_t k, - float* distances, - idx_t* labels, - int impl, - size_t* ndis_out, - size_t* nlist_out) const { - memset(distances, -1, sizeof(float) * k * n); - memset(labels, -1, sizeof(idx_t) * k * n); - - using HeapHC = HeapHandler; - using ReservoirHC = ReservoirHandler; - using SingleResultHC = SingleResultHandler; - - std::unique_ptr coarse_ids(new idx_t[n * nprobe]); - std::unique_ptr coarse_dis(new float[n * nprobe]); - - uint64_t times[10]; - memset(times, 0, sizeof(times)); - int ti = 0; -#define TIC times[ti++] = get_cy() - TIC; - - quantizer->search(n, x, nprobe, coarse_dis.get(), coarse_ids.get()); - - TIC; - - size_t dim12 = pq.ksub * M2; - AlignedTable dis_tables; - AlignedTable biases; - std::unique_ptr normalizers(new float[2 * n]); - - compute_LUT_uint8( - n, - x, - coarse_ids.get(), - coarse_dis.get(), - dis_tables, - biases, - normalizers.get()); - - TIC; - - bool single_LUT = !(by_residual && metric_type == METRIC_L2); - - TIC; - size_t ndis = 0, nlist_visited = 0; - - { - AlignedTable tmp_distances(k); - for (idx_t i = 0; i < n; i++) { - const uint8_t* LUT = nullptr; - int qmap1[1] = {0}; - std::unique_ptr> handler; - - if (k == 1) { - handler.reset(new SingleResultHC(1, 0)); - } else if (impl == 10) { - handler.reset(new HeapHC( - 1, tmp_distances.get(), labels + i * k, k, 0)); - } else if (impl == 11) { - handler.reset(new ReservoirHC(1, 0, k, 2 * k)); - } else { - FAISS_THROW_MSG("invalid"); - } - - handler->q_map = qmap1; - - if (single_LUT) { - LUT = dis_tables.get() + i * dim12; - } - for (idx_t j = 0; j < nprobe; j++) { - size_t ij = i * nprobe + j; - if (!single_LUT) { - LUT = dis_tables.get() + ij * dim12; - } - if (biases.get()) { - handler->dbias = biases.get() + ij; - } - - idx_t list_no = coarse_ids[ij]; - if (list_no < 0) - continue; - size_t ls = invlists->list_size(list_no); - if (ls == 0) - continue; - - InvertedLists::ScopedCodes codes(invlists, list_no); - InvertedLists::ScopedIds ids(invlists, list_no); - - handler->ntotal = ls; - handler->id_map = ids.get(); - -#define DISPATCH(classHC) \ - if (dynamic_cast(handler.get())) { \ - auto* res = static_cast(handler.get()); \ - pq4_accumulate_loop( \ - 1, roundup(ls, bbs), bbs, M2, codes.get(), LUT, *res); \ - } - DISPATCH(HeapHC) - else DISPATCH(ReservoirHC) else DISPATCH(SingleResultHC) -#undef DISPATCH - - nlist_visited++; - ndis++; - } - - handler->to_flat_arrays( - distances + i * k, - labels + i * k, - skip & 16 ? nullptr : normalizers.get() + i * 2); - } - } - *ndis_out = ndis; - *nlist_out = nlist; -} - -template -void IndexIVFPQFastScan::search_implem_12( - idx_t n, - const float* x, - idx_t k, - float* distances, - idx_t* labels, - int impl, - size_t* ndis_out, - size_t* nlist_out) const { - if (n == 0) { // does not work well with reservoir - return; - } - FAISS_THROW_IF_NOT(bbs == 32); - - std::unique_ptr coarse_ids(new idx_t[n * nprobe]); - std::unique_ptr coarse_dis(new float[n * nprobe]); - - uint64_t times[10]; - memset(times, 0, sizeof(times)); - int ti = 0; -#define TIC times[ti++] = get_cy() - TIC; - - quantizer->search(n, x, nprobe, coarse_dis.get(), coarse_ids.get()); - - TIC; - - size_t dim12 = pq.ksub * M2; - AlignedTable dis_tables; - AlignedTable biases; - std::unique_ptr normalizers(new float[2 * n]); - - compute_LUT_uint8( - n, - x, - coarse_ids.get(), - coarse_dis.get(), - dis_tables, - biases, - normalizers.get()); - - TIC; - - struct QC { - int qno; // sequence number of the query - int list_no; // list to visit - int rank; // this is the rank'th result of the coarse quantizer - }; - bool single_LUT = !(by_residual && metric_type == METRIC_L2); - - std::vector qcs; - { - int ij = 0; - for (int i = 0; i < n; i++) { - for (int j = 0; j < nprobe; j++) { - if (coarse_ids[ij] >= 0) { - qcs.push_back(QC{i, int(coarse_ids[ij]), int(j)}); - } - ij++; - } - } - std::sort(qcs.begin(), qcs.end(), [](const QC& a, const QC& b) { - return a.list_no < b.list_no; - }); - } - TIC; - - // prepare the result handlers - - std::unique_ptr> handler; - AlignedTable tmp_distances; - - using HeapHC = HeapHandler; - using ReservoirHC = ReservoirHandler; - using SingleResultHC = SingleResultHandler; - - if (k == 1) { - handler.reset(new SingleResultHC(n, 0)); - } else if (impl == 12) { - tmp_distances.resize(n * k); - handler.reset(new HeapHC(n, tmp_distances.get(), labels, k, 0)); - } else if (impl == 13) { - handler.reset(new ReservoirHC(n, 0, k, 2 * k)); - } - - int qbs2 = this->qbs2 ? this->qbs2 : 11; - - std::vector tmp_bias; - if (biases.get()) { - tmp_bias.resize(qbs2); - handler->dbias = tmp_bias.data(); - } - TIC; - - size_t ndis = 0; - - size_t i0 = 0; - uint64_t t_copy_pack = 0, t_scan = 0; - while (i0 < qcs.size()) { - uint64_t tt0 = get_cy(); - - // find all queries that access this inverted list - int list_no = qcs[i0].list_no; - size_t i1 = i0 + 1; - - while (i1 < qcs.size() && i1 < i0 + qbs2) { - if (qcs[i1].list_no != list_no) { - break; - } - i1++; - } - - size_t list_size = invlists->list_size(list_no); - - if (list_size == 0) { - i0 = i1; - continue; - } - - // re-organize LUTs and biases into the right order - int nc = i1 - i0; - - std::vector q_map(nc), lut_entries(nc); - AlignedTable LUT(nc * dim12); - memset(LUT.get(), -1, nc * dim12); - int qbs = pq4_preferred_qbs(nc); - - for (size_t i = i0; i < i1; i++) { - const QC& qc = qcs[i]; - q_map[i - i0] = qc.qno; - int ij = qc.qno * nprobe + qc.rank; - lut_entries[i - i0] = single_LUT ? qc.qno : ij; - if (biases.get()) { - tmp_bias[i - i0] = biases[ij]; - } - } - pq4_pack_LUT_qbs_q_map( - qbs, M2, dis_tables.get(), lut_entries.data(), LUT.get()); - - // access the inverted list - - ndis += (i1 - i0) * list_size; - - InvertedLists::ScopedCodes codes(invlists, list_no); - InvertedLists::ScopedIds ids(invlists, list_no); - - // prepare the handler - - handler->ntotal = list_size; - handler->q_map = q_map.data(); - handler->id_map = ids.get(); - uint64_t tt1 = get_cy(); - -#define DISPATCH(classHC) \ - if (dynamic_cast(handler.get())) { \ - auto* res = static_cast(handler.get()); \ - pq4_accumulate_loop_qbs( \ - qbs, list_size, M2, codes.get(), LUT.get(), *res); \ - } - DISPATCH(HeapHC) - else DISPATCH(ReservoirHC) else DISPATCH(SingleResultHC) - - // prepare for next loop - i0 = i1; - - uint64_t tt2 = get_cy(); - t_copy_pack += tt1 - tt0; - t_scan += tt2 - tt1; - } - TIC; - - // labels is in-place for HeapHC - handler->to_flat_arrays( - distances, labels, skip & 16 ? nullptr : normalizers.get()); - - TIC; - - // these stats are not thread-safe - - for (int i = 1; i < ti; i++) { - IVFFastScan_stats.times[i] += times[i] - times[i - 1]; - } - IVFFastScan_stats.t_copy_pack += t_copy_pack; - IVFFastScan_stats.t_scan += t_scan; - - if (auto* rh = dynamic_cast(handler.get())) { - for (int i = 0; i < 4; i++) { - IVFFastScan_stats.reservoir_times[i] += rh->times[i]; - } - } - - *ndis_out = ndis; - *nlist_out = nlist; -} - -IVFFastScanStats IVFFastScan_stats; - } // namespace faiss diff --git a/faiss/IndexIVFPQFastScan.h b/faiss/IndexIVFPQFastScan.h index 24277dc12..55c5430b6 100644 --- a/faiss/IndexIVFPQFastScan.h +++ b/faiss/IndexIVFPQFastScan.h @@ -9,6 +9,7 @@ #include +#include #include #include #include @@ -31,36 +32,20 @@ namespace faiss { * 13: idem, collect results in reservoir */ -struct IndexIVFPQFastScan : IndexIVF { - bool by_residual; ///< Encode residual or plain vector? +struct IndexIVFPQFastScan : IndexIVFFastScan { ProductQuantizer pq; ///< produces the codes - // size of the kernel - int bbs; // set at build time - - // M rounded up to a multiple of 2 - size_t M2; - /// precomputed tables management int use_precomputed_table = 0; /// if use_precompute_table size (nlist, pq.M, pq.ksub) AlignedTable precomputed_table; - // search-time implementation - int implem = 0; - // skip some parts of the computation (for timing) - int skip = 0; - - // batching factors at search time (0 = default) - int qbs = 0; - size_t qbs2 = 0; - IndexIVFPQFastScan( Index* quantizer, size_t d, size_t nlist, size_t M, - size_t nbits_per_idx, + size_t nbits, MetricType metric = METRIC_L2, int bbs = 32); @@ -69,9 +54,6 @@ struct IndexIVFPQFastScan : IndexIVF { // built from an IndexIVFPQ explicit IndexIVFPQFastScan(const IndexIVFPQ& orig, int bbs = 32); - /// orig's inverted lists (for debugging) - InvertedLists* orig_invlists = nullptr; - void train_residual(idx_t n, const float* x) override; /// build precomputed table, possibly updating use_precomputed_table @@ -86,106 +68,19 @@ struct IndexIVFPQFastScan : IndexIVF { uint8_t* codes, bool include_listno = false) const override; - void add_with_ids(idx_t n, const float* x, const idx_t* xids) override; - - void search( - idx_t n, - const float* x, - idx_t k, - float* distances, - idx_t* labels) const override; - // prepare look-up tables + bool lookup_table_is_3d() const override; + void compute_LUT( size_t n, const float* x, const idx_t* coarse_ids, const float* coarse_dis, AlignedTable& dis_tables, - AlignedTable& biases) const; + AlignedTable& biases) const override; - void compute_LUT_uint8( - size_t n, - const float* x, - const idx_t* coarse_ids, - const float* coarse_dis, - AlignedTable& dis_tables, - AlignedTable& biases, - float* normalizers) const; - - // internal search funcs - - template - void search_dispatch_implem( - idx_t n, - const float* x, - idx_t k, - float* distances, - idx_t* labels) const; - - template - void search_implem_1( - idx_t n, - const float* x, - idx_t k, - float* distances, - idx_t* labels) const; - - template - void search_implem_2( - idx_t n, - const float* x, - idx_t k, - float* distances, - idx_t* labels) const; - - // implem 10 and 12 are not multithreaded internally, so - // export search stats - template - void search_implem_10( - idx_t n, - const float* x, - idx_t k, - float* distances, - idx_t* labels, - int impl, - size_t* ndis_out, - size_t* nlist_out) const; - - template - void search_implem_12( - idx_t n, - const float* x, - idx_t k, - float* distances, - idx_t* labels, - int impl, - size_t* ndis_out, - size_t* nlist_out) const; + void sa_decode(idx_t n, const uint8_t* bytes, float* x) const override; }; -struct IVFFastScanStats { - uint64_t times[10]; - uint64_t t_compute_distance_tables, t_round; - uint64_t t_copy_pack, t_scan, t_to_flat; - uint64_t reservoir_times[4]; - - double Mcy_at(int i) { - return times[i] / (1000 * 1000.0); - } - - double Mcy_reservoir_at(int i) { - return reservoir_times[i] / (1000 * 1000.0); - } - IVFFastScanStats() { - reset(); - } - void reset() { - memset(this, 0, sizeof(*this)); - } -}; - -FAISS_API extern IVFFastScanStats IVFFastScan_stats; - } // namespace faiss diff --git a/faiss/IndexPQFastScan.cpp b/faiss/IndexPQFastScan.cpp index ed8467fcb..d754e620f 100644 --- a/faiss/IndexPQFastScan.cpp +++ b/faiss/IndexPQFastScan.cpp @@ -14,17 +14,11 @@ #include #include -#include +#include #include -#include -#include -#include - namespace faiss { -using namespace simd_result_handlers; - inline size_t roundup(size_t a, size_t b) { return (a + b - 1) / b * b; } @@ -35,37 +29,19 @@ IndexPQFastScan::IndexPQFastScan( size_t nbits, MetricType metric, int bbs) - : Index(d, metric), - pq(d, M, nbits), - bbs(bbs), - ntotal2(0), - M2(roundup(M, 2)) { - FAISS_THROW_IF_NOT(nbits == 4); - is_trained = false; + : pq(d, M, nbits) { + init_fastscan(d, M, nbits, metric, bbs); } -IndexPQFastScan::IndexPQFastScan() : bbs(0), ntotal2(0), M2(0) {} - -IndexPQFastScan::IndexPQFastScan(const IndexPQ& orig, int bbs) - : Index(orig.d, orig.metric_type), pq(orig.pq), bbs(bbs) { - FAISS_THROW_IF_NOT(orig.pq.nbits == 4); +IndexPQFastScan::IndexPQFastScan(const IndexPQ& orig, int bbs) : pq(orig.pq) { + init_fastscan(orig.d, pq.M, pq.nbits, orig.metric_type, bbs); ntotal = orig.ntotal; + ntotal2 = roundup(ntotal, bbs); is_trained = orig.is_trained; orig_codes = orig.codes.data(); - qbs = 0; // means use default - // pack the codes - - size_t M = pq.M; - - FAISS_THROW_IF_NOT(bbs % 32 == 0); - M2 = roundup(M, 2); - ntotal2 = roundup(ntotal, bbs); - codes.resize(ntotal2 * M2 / 2); - - // printf("M=%d M2=%d code_size=%d\n", M, M2, pq.code_size); pq4_pack_codes(orig.codes.data(), ntotal, M, ntotal2, bbs, M2, codes.get()); } @@ -77,432 +53,17 @@ void IndexPQFastScan::train(idx_t n, const float* x) { is_trained = true; } -void IndexPQFastScan::add(idx_t n, const float* x) { - FAISS_THROW_IF_NOT(is_trained); - AlignedTable tmp_codes(n * pq.code_size); - pq.compute_codes(x, tmp_codes.get(), n); - ntotal2 = roundup(ntotal + n, bbs); - size_t new_size = ntotal2 * M2 / 2; - size_t old_size = codes.size(); - if (new_size > old_size) { - codes.resize(new_size); - memset(codes.get() + old_size, 0, new_size - old_size); - } - pq4_pack_codes_range( - tmp_codes.get(), pq.M, ntotal, ntotal + n, bbs, M2, codes.get()); - ntotal += n; +void IndexPQFastScan::compute_codes(uint8_t* codes, idx_t n, const float* x) + const { + pq.compute_codes(x, codes, n); } -void IndexPQFastScan::reset() { - codes.resize(0); - ntotal = 0; -} - -namespace { - -// from impl/ProductQuantizer.cpp -template -void pq_estimators_from_tables_generic( - const ProductQuantizer& pq, - size_t nbits, - const uint8_t* codes, - size_t ncodes, - const dis_t* dis_table, - size_t k, - typename C::T* heap_dis, - int64_t* heap_ids) { - using accu_t = typename C::T; - const size_t M = pq.M; - const size_t ksub = pq.ksub; - for (size_t j = 0; j < ncodes; ++j) { - PQDecoderGeneric decoder(codes + j * pq.code_size, nbits); - accu_t dis = 0; - const dis_t* __restrict dt = dis_table; - for (size_t m = 0; m < M; m++) { - uint64_t c = decoder.decode(); - dis += dt[c]; - dt += ksub; - } - - if (C::cmp(heap_dis[0], dis)) { - heap_pop(k, heap_dis, heap_ids); - heap_push(k, heap_dis, heap_ids, dis, j); - } - } -} - -} // anonymous namespace - -using namespace quantize_lut; - -void IndexPQFastScan::compute_quantized_LUT( - idx_t n, - const float* x, - uint8_t* lut, - float* normalizers) const { - size_t dim12 = pq.ksub * pq.M; - std::unique_ptr dis_tables(new float[n * dim12]); +void IndexPQFastScan::compute_float_LUT(float* lut, idx_t n, const float* x) + const { if (metric_type == METRIC_L2) { - pq.compute_distance_tables(n, x, dis_tables.get()); + pq.compute_distance_tables(n, x, lut); } else { - pq.compute_inner_prod_tables(n, x, dis_tables.get()); - } - - for (uint64_t i = 0; i < n; i++) { - round_uint8_per_column( - dis_tables.get() + i * dim12, - pq.M, - pq.ksub, - &normalizers[2 * i], - &normalizers[2 * i + 1]); - } - - for (uint64_t i = 0; i < n; i++) { - const float* t_in = dis_tables.get() + i * dim12; - uint8_t* t_out = lut + i * M2 * pq.ksub; - - for (int j = 0; j < dim12; j++) { - t_out[j] = int(t_in[j]); - } - memset(t_out + dim12, 0, (M2 - pq.M) * pq.ksub); - } -} - -/****************************************************************************** - * Search driver routine - ******************************************************************************/ - -void IndexPQFastScan::search( - idx_t n, - const float* x, - idx_t k, - float* distances, - idx_t* labels) const { - FAISS_THROW_IF_NOT(k > 0); - - if (metric_type == METRIC_L2) { - search_dispatch_implem(n, x, k, distances, labels); - } else { - search_dispatch_implem(n, x, k, distances, labels); - } -} - -template -void IndexPQFastScan::search_dispatch_implem( - idx_t n, - const float* x, - idx_t k, - float* distances, - idx_t* labels) const { - using Cfloat = typename std::conditional< - is_max, - CMax, - CMin>::type; - - using C = typename std:: - conditional, CMin>::type; - - if (n == 0) { - return; - } - - // actual implementation used - int impl = implem; - - if (impl == 0) { - if (bbs == 32) { - impl = 12; - } else { - impl = 14; - } - if (k > 20) { - impl++; - } - } - - if (implem == 1) { - FAISS_THROW_IF_NOT(orig_codes); - FAISS_THROW_IF_NOT(is_max); - float_maxheap_array_t res = {size_t(n), size_t(k), labels, distances}; - pq.search(x, n, orig_codes, ntotal, &res, true); - } else if (implem == 2 || implem == 3 || implem == 4) { - FAISS_THROW_IF_NOT(orig_codes); - - size_t dim12 = pq.ksub * pq.M; - std::unique_ptr dis_tables(new float[n * dim12]); - if (is_max) { - pq.compute_distance_tables(n, x, dis_tables.get()); - } else { - pq.compute_inner_prod_tables(n, x, dis_tables.get()); - } - - std::vector normalizers(n * 2); - - if (implem == 2) { - // default float - } else if (implem == 3 || implem == 4) { - for (uint64_t i = 0; i < n; i++) { - round_uint8_per_column( - dis_tables.get() + i * dim12, - pq.M, - pq.ksub, - &normalizers[2 * i], - &normalizers[2 * i + 1]); - } - } - - for (int64_t i = 0; i < n; i++) { - int64_t* heap_ids = labels + i * k; - float* heap_dis = distances + i * k; - - heap_heapify(k, heap_dis, heap_ids); - - pq_estimators_from_tables_generic( - pq, - pq.nbits, - orig_codes, - ntotal, - dis_tables.get() + i * dim12, - k, - heap_dis, - heap_ids); - - heap_reorder(k, heap_dis, heap_ids); - - if (implem == 4) { - float a = normalizers[2 * i]; - float b = normalizers[2 * i + 1]; - - for (int j = 0; j < k; j++) { - heap_dis[j] = heap_dis[j] / a + b; - } - } - } - } else if (impl >= 12 && impl <= 15) { - FAISS_THROW_IF_NOT(ntotal < INT_MAX); - int nt = std::min(omp_get_max_threads(), int(n)); - if (nt < 2) { - if (impl == 12 || impl == 13) { - search_implem_12(n, x, k, distances, labels, impl); - } else { - search_implem_14(n, x, k, distances, labels, impl); - } - } else { - // explicitly slice over threads -#pragma omp parallel for num_threads(nt) - for (int slice = 0; slice < nt; slice++) { - idx_t i0 = n * slice / nt; - idx_t i1 = n * (slice + 1) / nt; - float* dis_i = distances + i0 * k; - idx_t* lab_i = labels + i0 * k; - if (impl == 12 || impl == 13) { - search_implem_12( - i1 - i0, x + i0 * d, k, dis_i, lab_i, impl); - } else { - search_implem_14( - i1 - i0, x + i0 * d, k, dis_i, lab_i, impl); - } - } - } - } else { - FAISS_THROW_FMT("invalid implem %d impl=%d", implem, impl); - } -} - -template -void IndexPQFastScan::search_implem_12( - idx_t n, - const float* x, - idx_t k, - float* distances, - idx_t* labels, - int impl) const { - FAISS_THROW_IF_NOT(bbs == 32); - - // handle qbs2 blocking by recursive call - int64_t qbs2 = this->qbs == 0 ? 11 : pq4_qbs_to_nq(this->qbs); - if (n > qbs2) { - for (int64_t i0 = 0; i0 < n; i0 += qbs2) { - int64_t i1 = std::min(i0 + qbs2, n); - search_implem_12( - i1 - i0, - x + d * i0, - k, - distances + i0 * k, - labels + i0 * k, - impl); - } - return; - } - - size_t dim12 = pq.ksub * M2; - AlignedTable quantized_dis_tables(n * dim12); - std::unique_ptr normalizers(new float[2 * n]); - - if (skip & 1) { - quantized_dis_tables.clear(); - } else { - compute_quantized_LUT( - n, x, quantized_dis_tables.get(), normalizers.get()); - } - - AlignedTable LUT(n * dim12); - - // block sizes are encoded in qbs, 4 bits at a time - - // caution: we override an object field - int qbs = this->qbs; - - if (n != pq4_qbs_to_nq(qbs)) { - qbs = pq4_preferred_qbs(n); - } - - int LUT_nq = - pq4_pack_LUT_qbs(qbs, M2, quantized_dis_tables.get(), LUT.get()); - FAISS_THROW_IF_NOT(LUT_nq == n); - - if (k == 1) { - SingleResultHandler handler(n, ntotal); - if (skip & 4) { - // pass - } else { - handler.disable = bool(skip & 2); - pq4_accumulate_loop_qbs( - qbs, ntotal2, M2, codes.get(), LUT.get(), handler); - } - - handler.to_flat_arrays(distances, labels, normalizers.get()); - - } else if (impl == 12) { - std::vector tmp_dis(n * k); - std::vector tmp_ids(n * k); - - if (skip & 4) { - // skip - } else { - HeapHandler handler( - n, tmp_dis.data(), tmp_ids.data(), k, ntotal); - handler.disable = bool(skip & 2); - - pq4_accumulate_loop_qbs( - qbs, ntotal2, M2, codes.get(), LUT.get(), handler); - - if (!(skip & 8)) { - handler.to_flat_arrays(distances, labels, normalizers.get()); - } - } - - } else { // impl == 13 - - ReservoirHandler handler(n, ntotal, k, 2 * k); - handler.disable = bool(skip & 2); - - if (skip & 4) { - // skip - } else { - pq4_accumulate_loop_qbs( - qbs, ntotal2, M2, codes.get(), LUT.get(), handler); - } - - if (!(skip & 8)) { - handler.to_flat_arrays(distances, labels, normalizers.get()); - } - - FastScan_stats.t0 += handler.times[0]; - FastScan_stats.t1 += handler.times[1]; - FastScan_stats.t2 += handler.times[2]; - FastScan_stats.t3 += handler.times[3]; - } -} - -FastScanStats FastScan_stats; - -template -void IndexPQFastScan::search_implem_14( - idx_t n, - const float* x, - idx_t k, - float* distances, - idx_t* labels, - int impl) const { - FAISS_THROW_IF_NOT(bbs % 32 == 0); - - int qbs2 = qbs == 0 ? 4 : qbs; - - // handle qbs2 blocking by recursive call - if (n > qbs2) { - for (int64_t i0 = 0; i0 < n; i0 += qbs2) { - int64_t i1 = std::min(i0 + qbs2, n); - search_implem_14( - i1 - i0, - x + d * i0, - k, - distances + i0 * k, - labels + i0 * k, - impl); - } - return; - } - - size_t dim12 = pq.ksub * M2; - AlignedTable quantized_dis_tables(n * dim12); - std::unique_ptr normalizers(new float[2 * n]); - - if (skip & 1) { - quantized_dis_tables.clear(); - } else { - compute_quantized_LUT( - n, x, quantized_dis_tables.get(), normalizers.get()); - } - - AlignedTable LUT(n * dim12); - pq4_pack_LUT(n, M2, quantized_dis_tables.get(), LUT.get()); - - if (k == 1) { - SingleResultHandler handler(n, ntotal); - if (skip & 4) { - // pass - } else { - handler.disable = bool(skip & 2); - pq4_accumulate_loop( - n, ntotal2, bbs, M2, codes.get(), LUT.get(), handler); - } - handler.to_flat_arrays(distances, labels, normalizers.get()); - - } else if (impl == 14) { - std::vector tmp_dis(n * k); - std::vector tmp_ids(n * k); - - if (skip & 4) { - // skip - } else if (k > 1) { - HeapHandler handler( - n, tmp_dis.data(), tmp_ids.data(), k, ntotal); - handler.disable = bool(skip & 2); - - pq4_accumulate_loop( - n, ntotal2, bbs, M2, codes.get(), LUT.get(), handler); - - if (!(skip & 8)) { - handler.to_flat_arrays(distances, labels, normalizers.get()); - } - } - - } else { // impl == 15 - - ReservoirHandler handler(n, ntotal, k, 2 * k); - handler.disable = bool(skip & 2); - - if (skip & 4) { - // skip - } else { - pq4_accumulate_loop( - n, ntotal2, bbs, M2, codes.get(), LUT.get(), handler); - } - - if (!(skip & 8)) { - handler.to_flat_arrays(distances, labels, normalizers.get()); - } + pq.compute_inner_prod_tables(n, x, lut); } } diff --git a/faiss/IndexPQFastScan.h b/faiss/IndexPQFastScan.h index c2309e31e..ab620bb92 100644 --- a/faiss/IndexPQFastScan.h +++ b/faiss/IndexPQFastScan.h @@ -7,6 +7,7 @@ #pragma once +#include #include #include #include @@ -25,27 +26,9 @@ namespace faiss { * 15: no qbs with reservoir accumulator */ -struct IndexPQFastScan : Index { +struct IndexPQFastScan : IndexFastScan { ProductQuantizer pq; - // implementation to select - int implem = 0; - // skip some parts of the computation (for timing) - int skip = 0; - - // size of the kernel - int bbs; // set at build time - int qbs = 0; // query block size 0 = use default - - // packed version of the codes - size_t ntotal2; - size_t M2; - - AlignedTable codes; - - // this is for testing purposes only (set when initialized by IndexPQ) - const uint8_t* orig_codes = nullptr; - IndexPQFastScan( int d, size_t M, @@ -53,73 +36,16 @@ struct IndexPQFastScan : Index { MetricType metric = METRIC_L2, int bbs = 32); - IndexPQFastScan(); + IndexPQFastScan() = default; /// build from an existing IndexPQ explicit IndexPQFastScan(const IndexPQ& orig, int bbs = 32); void train(idx_t n, const float* x) override; - void add(idx_t n, const float* x) override; - void reset() override; - void search( - idx_t n, - const float* x, - idx_t k, - float* distances, - idx_t* labels) const override; - // called by search function - void compute_quantized_LUT( - idx_t n, - const float* x, - uint8_t* lut, - float* normalizers) const; + void compute_codes(uint8_t* codes, idx_t n, const float* x) const override; - template - void search_dispatch_implem( - idx_t n, - const float* x, - idx_t k, - float* distances, - idx_t* labels) const; - - template - void search_implem_2( - idx_t n, - const float* x, - idx_t k, - float* distances, - idx_t* labels) const; - - template - void search_implem_12( - idx_t n, - const float* x, - idx_t k, - float* distances, - idx_t* labels, - int impl) const; - - template - void search_implem_14( - idx_t n, - const float* x, - idx_t k, - float* distances, - idx_t* labels, - int impl) const; + void compute_float_LUT(float* lut, idx_t n, const float* x) const override; }; -struct FastScanStats { - uint64_t t0, t1, t2, t3; - FastScanStats() { - reset(); - } - void reset() { - memset(this, 0, sizeof(*this)); - } -}; - -FAISS_API extern FastScanStats FastScan_stats; - } // namespace faiss diff --git a/faiss/impl/AdditiveQuantizer.cpp b/faiss/impl/AdditiveQuantizer.cpp index d08a30044..d3a7bccd0 100644 --- a/faiss/impl/AdditiveQuantizer.cpp +++ b/faiss/impl/AdditiveQuantizer.cpp @@ -17,7 +17,10 @@ #include +#include #include +#include +#include #include #include #include @@ -89,6 +92,8 @@ void AdditiveQuantizer::set_derived_values() { break; case ST_norm_qint8: case ST_norm_cqint8: + case ST_norm_lsq2x4: + case ST_norm_rq2x4: tot_bits += 8; break; case ST_norm_qint4: @@ -101,6 +106,56 @@ void AdditiveQuantizer::set_derived_values() { code_size = (tot_bits + 7) / 8; } +void AdditiveQuantizer::train_norm(size_t n, const float* norms) { + norm_min = HUGE_VALF; + norm_max = -HUGE_VALF; + for (idx_t i = 0; i < n; i++) { + if (norms[i] < norm_min) { + norm_min = norms[i]; + } + if (norms[i] > norm_max) { + norm_max = norms[i]; + } + } + + if (search_type == ST_norm_cqint8 || search_type == ST_norm_cqint4) { + size_t k = (1 << 8); + if (search_type == ST_norm_cqint4) { + k = (1 << 4); + } + Clustering1D clus(k); + clus.train_exact(n, norms); + qnorm.add(clus.k, clus.centroids.data()); + } else if (search_type == ST_norm_lsq2x4 || search_type == ST_norm_rq2x4) { + std::unique_ptr aq; + if (search_type == ST_norm_lsq2x4) { + aq.reset(new LocalSearchQuantizer(1, 2, 4)); + } else { + aq.reset(new ResidualQuantizer(1, 2, 4)); + } + + aq->train(n, norms); + // flatten aq codebooks + std::vector flat_codebooks(1 << 8); + FAISS_THROW_IF_NOT(aq->codebooks.size() == 32); + + // save norm tables for 4-bit fastscan search + norm_tabs = aq->codebooks; + + // assume big endian + const float* c = norm_tabs.data(); + for (size_t i = 0; i < 16; i++) { + for (size_t j = 0; j < 16; j++) { + flat_codebooks[i * 16 + j] = c[j] + c[16 + i]; + } + } + + qnorm.reset(); + qnorm.add(1 << 8, flat_codebooks.data()); + FAISS_THROW_IF_NOT(qnorm.ntotal == (1 << 8)); + } +} + namespace { // TODO @@ -132,7 +187,7 @@ float decode_qint4(uint8_t i, float amin, float amax) { uint32_t AdditiveQuantizer::encode_qcint(float x) const { idx_t id; - qnorm.assign(idx_t(1), &x, &id, idx_t(1)); + qnorm.assign(1, &x, &id, 1); return uint32_t(id); } @@ -145,18 +200,25 @@ void AdditiveQuantizer::pack_codes( const int32_t* codes, uint8_t* packed_codes, int64_t ld_codes, - const float* norms) const { + const float* norms, + const float* centroids) const { if (ld_codes == -1) { ld_codes = M; } std::vector norm_buf; if (search_type == ST_norm_float || search_type == ST_norm_qint4 || search_type == ST_norm_qint8 || search_type == ST_norm_cqint8 || - search_type == ST_norm_cqint4) { - if (!norms) { + search_type == ST_norm_cqint4 || search_type == ST_norm_lsq2x4 || + search_type == ST_norm_rq2x4) { + if (centroids != nullptr || !norms) { norm_buf.resize(n); std::vector x_recons(n * d); decode_unpacked(codes, x_recons.data(), n, ld_codes); + + if (centroids != nullptr) { + // x = x + c + fvec_add(n * d, x_recons.data(), centroids, x_recons.data()); + } fvec_norms_L2sqr(norm_buf.data(), x_recons.data(), d, n); norms = norm_buf.data(); } @@ -186,6 +248,8 @@ void AdditiveQuantizer::pack_codes( bsw.write(b, 4); break; } + case ST_norm_lsq2x4: + case ST_norm_rq2x4: case ST_norm_cqint8: { uint32_t b = encode_qcint(norms[i]); bsw.write(b, 8); @@ -283,28 +347,33 @@ void AdditiveQuantizer::decode_64bit(idx_t bits, float* xi) const { } } -void AdditiveQuantizer::compute_LUT(size_t n, const float* xq, float* LUT) - const { +void AdditiveQuantizer::compute_LUT( + size_t n, + const float* xq, + float* LUT, + float alpha, + long ld_lut) const { // in all cases, it is large matrix multiplication FINTEGER ncenti = total_codebook_size; FINTEGER di = d; FINTEGER nqi = n; - float one = 1, zero = 0; + FINTEGER ldc = ld_lut > 0 ? ld_lut : ncenti; + float zero = 0; sgemm_("Transposed", "Not transposed", &ncenti, &nqi, &di, - &one, + &alpha, codebooks.data(), &di, xq, &di, &zero, LUT, - &ncenti); + &ldc); } namespace { diff --git a/faiss/impl/AdditiveQuantizer.h b/faiss/impl/AdditiveQuantizer.h index 4bcc1b9ce..986e4c54a 100644 --- a/faiss/impl/AdditiveQuantizer.h +++ b/faiss/impl/AdditiveQuantizer.h @@ -37,7 +37,9 @@ struct AdditiveQuantizer { bool verbose; ///< verbose during training? bool is_trained; ///< is trained or not - IndexFlat1D qnorm; ///< store and search norms + IndexFlat1D qnorm; ///< store and search norms + std::vector norm_tabs; ///< store norms of codebook entries for 4-bit + ///< fastscan search uint32_t encode_qcint( float x) const; ///< encode norm by non-uniform scalar quantization @@ -57,6 +59,10 @@ struct AdditiveQuantizer { ST_norm_qint4, ST_norm_cqint8, ///< use a LUT, and store non-uniform quantized norm ST_norm_cqint4, + + ST_norm_lsq2x4, ///< use a 2x4 bits lsq as norm quantizer (for fast + ///< scan) + ST_norm_rq2x4, ///< use a 2x4 bits rq as norm quantizer (for fast scan) }; AdditiveQuantizer( @@ -72,13 +78,20 @@ struct AdditiveQuantizer { ///< Train the additive quantizer virtual void train(size_t n, const float* x) = 0; + ///< Train the norm quantizer + void train_norm(size_t n, const float* norms); + /** Encode a set of vectors * * @param x vectors to encode, size n * d * @param codes output codes, size n * code_size + * @param centroids centroids to be added to x, size n * d */ - virtual void compute_codes(const float* x, uint8_t* codes, size_t n) - const = 0; + virtual void compute_codes( + const float* x, + uint8_t* codes, + size_t n, + const float* centroids = nullptr) const = 0; /** pack a series of code to bit-compact format * @@ -87,13 +100,15 @@ struct AdditiveQuantizer { * @param ld_codes leading dimension of codes * @param norms norms of the vectors (size n). Will be computed if * needed but not provided + * @param centroids centroids to be added to x, size n * d */ void pack_codes( size_t n, const int32_t* codes, uint8_t* packed_codes, int64_t ld_codes = -1, - const float* norms = nullptr) const; + const float* norms = nullptr, + const float* centroids = nullptr) const; /** Decode a set of vectors * @@ -143,8 +158,15 @@ struct AdditiveQuantizer { * * @param xq query vector, size (n, d) * @param LUT look-up table, size (n, total_codebook_size) + * @param alpha compute alpha * inner-product + * @param ld leading dimension of LUT */ - void compute_LUT(size_t n, const float* xq, float* LUT) const; + void compute_LUT( + size_t n, + const float* xq, + float* LUT, + float alpha = 1.0f, + long ld_lut = -1) const; /// exact IP search void knn_centroids_inner_product( diff --git a/faiss/impl/LocalSearchQuantizer.cpp b/faiss/impl/LocalSearchQuantizer.cpp index 2bb4aea0a..96f6d7de0 100644 --- a/faiss/impl/LocalSearchQuantizer.cpp +++ b/faiss/impl/LocalSearchQuantizer.cpp @@ -15,7 +15,6 @@ #include -#include #include #include #include @@ -182,7 +181,7 @@ LocalSearchQuantizer::LocalSearchQuantizer() : LocalSearchQuantizer(0, 0, 0) {} void LocalSearchQuantizer::train(size_t n, const float* x) { FAISS_THROW_IF_NOT(K == (1 << nbits[0])); - FAISS_THROW_IF_NOT(nperts <= M); + nperts = std::min(nperts, M); lsq_timer.reset(); LSQTimerScope scope(&lsq_timer, "train"); @@ -264,26 +263,7 @@ void LocalSearchQuantizer::train(size_t n, const float* x) { decode_unpacked(codes.data(), x_recons.data(), n); fvec_norms_L2sqr(norms.data(), x_recons.data(), d, n); - norm_min = HUGE_VALF; - norm_max = -HUGE_VALF; - for (idx_t i = 0; i < n; i++) { - if (norms[i] < norm_min) { - norm_min = norms[i]; - } - if (norms[i] > norm_max) { - norm_max = norms[i]; - } - } - - if (search_type == ST_norm_cqint8 || search_type == ST_norm_cqint4) { - size_t k = (1 << 8); - if (search_type == ST_norm_cqint4) { - k = (1 << 4); - } - Clustering1D clus(k); - clus.train_exact(n, norms.data()); - qnorm.add(clus.k, clus.centroids.data()); - } + train_norm(n, norms.data()); } if (verbose) { @@ -321,7 +301,8 @@ void LocalSearchQuantizer::perturb_codebooks( void LocalSearchQuantizer::compute_codes( const float* x, uint8_t* codes_out, - size_t n) const { + size_t n, + const float* centroids) const { FAISS_THROW_IF_NOT_MSG(is_trained, "LSQ is not trained yet."); lsq_timer.reset(); @@ -335,7 +316,7 @@ void LocalSearchQuantizer::compute_codes( random_int32(codes, 0, K - 1, gen); icm_encode(codes.data(), x, n, encode_ils_iters, gen); - pack_codes(n, codes.data(), codes_out); + pack_codes(n, codes.data(), codes_out, -1, nullptr, centroids); if (verbose) { scope.finish(); diff --git a/faiss/impl/LocalSearchQuantizer.h b/faiss/impl/LocalSearchQuantizer.h index 162a70a0b..e90106421 100644 --- a/faiss/impl/LocalSearchQuantizer.h +++ b/faiss/impl/LocalSearchQuantizer.h @@ -83,8 +83,13 @@ struct LocalSearchQuantizer : AdditiveQuantizer { * @param x vectors to encode, size n * d * @param codes output codes, size n * code_size * @param n number of vectors + * @param centroids centroids to be added to x, size n * d */ - void compute_codes(const float* x, uint8_t* codes, size_t n) const override; + void compute_codes( + const float* x, + uint8_t* codes, + size_t n, + const float* centroids = nullptr) const override; /** Update codebooks given encodings * diff --git a/faiss/impl/LookupTableScaler.h b/faiss/impl/LookupTableScaler.h new file mode 100644 index 000000000..cf41ad880 --- /dev/null +++ b/faiss/impl/LookupTableScaler.h @@ -0,0 +1,79 @@ +/** + * Copyright (c) Facebook, Inc. and its 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 + +/******************************************* + * The Scaler objects are used to specialize the handling of the + * norm components in Additive quantizer fast-scan. + ********************************************/ + +namespace faiss { + +/// no-op handler +struct DummyScaler { + static constexpr int nscale = 0; + + inline simd32uint8 lookup(const simd32uint8&, const simd32uint8&) + const { + FAISS_THROW_MSG("DummyScaler::lookup should not be called."); + return simd32uint8(0); + } + + inline simd16uint16 scale_lo(const simd32uint8&) const { + FAISS_THROW_MSG("DummyScaler::scale_lo should not be called."); + return simd16uint16(0); + } + + inline simd16uint16 scale_hi(const simd32uint8&) const { + FAISS_THROW_MSG("DummyScaler::scale_hi should not be called."); + return simd16uint16(0); + } + + template + inline dist_t scale_one(const dist_t&) const { + FAISS_THROW_MSG("DummyScaler::scale_one should not be called."); + return 0; + } +}; + +/// consumes 2x4 bits to encode a norm as a scalar additive quantizer +/// the norm is scaled because its range if larger than other components +struct NormTableScaler { + static constexpr int nscale = 2; + int scale_int; + simd16uint16 scale_simd; + + explicit NormTableScaler(int scale) : + scale_int(scale), scale_simd(scale) {} + + inline simd32uint8 lookup(const simd32uint8& lut, const simd32uint8& c) + const { + return lut.lookup_2_lanes(c); + } + + inline simd16uint16 scale_lo(const simd32uint8& res) const { + return simd16uint16(res) * scale_simd; + } + + inline simd16uint16 scale_hi(const simd32uint8& res) const { + return (simd16uint16(res) >> 8) * scale_simd; + } + + // for non-SIMD implem 2, 3, 4 + template + inline dist_t scale_one(const dist_t& x) const { + return x * scale_int; + } +}; + +} // namespace faiss diff --git a/faiss/impl/ResidualQuantizer.cpp b/faiss/impl/ResidualQuantizer.cpp index 3954d653e..bd3871c6e 100644 --- a/faiss/impl/ResidualQuantizer.cpp +++ b/faiss/impl/ResidualQuantizer.cpp @@ -16,11 +16,6 @@ #include #include -#include -#include -#include - -#include #include #include #include @@ -388,26 +383,8 @@ void ResidualQuantizer::train(size_t n, const float* x) { x + i * d, residuals.data() + i * cur_beam_size * d, d); } - norm_min = HUGE_VALF; - norm_max = -HUGE_VALF; - for (idx_t i = 0; i < n; i++) { - if (norms[i] < norm_min) { - norm_min = norms[i]; - } - if (norms[i] > norm_max) { - norm_max = norms[i]; - } - } - - if (search_type == ST_norm_cqint8 || search_type == ST_norm_cqint4) { - size_t k = (1 << 8); - if (search_type == ST_norm_cqint4) { - k = (1 << 4); - } - Clustering1D clus(k); - clus.train_exact(n, norms.data()); - qnorm.add(clus.k, clus.centroids.data()); - } + // fvec_norms_L2sqr(norms.data(), x, d, n); + train_norm(n, norms.data()); if (!(train_type & Skip_codebook_tables)) { compute_codebook_tables(); @@ -546,7 +523,8 @@ size_t ResidualQuantizer::memory_per_point(int beam_size) const { void ResidualQuantizer::compute_codes( const float* x, uint8_t* codes_out, - size_t n) const { + size_t n, + const float* centroids) const { FAISS_THROW_IF_NOT_MSG(is_trained, "RQ is not trained yet."); size_t mem = memory_per_point(); @@ -632,7 +610,8 @@ void ResidualQuantizer::compute_codes( codes.data(), codes_out, M * max_beam_size, - norms.size() > 0 ? norms.data() : nullptr); + norms.size() > 0 ? norms.data() : nullptr, + centroids); } void ResidualQuantizer::refine_beam( diff --git a/faiss/impl/ResidualQuantizer.h b/faiss/impl/ResidualQuantizer.h index 18c30e865..3f7f29fc2 100644 --- a/faiss/impl/ResidualQuantizer.h +++ b/faiss/impl/ResidualQuantizer.h @@ -97,8 +97,13 @@ struct ResidualQuantizer : AdditiveQuantizer { * * @param x vectors to encode, size n * d * @param codes output codes, size n * code_size + * @param centroids centroids to be added to x, size n * d */ - void compute_codes(const float* x, uint8_t* codes, size_t n) const override; + void compute_codes( + const float* x, + uint8_t* codes, + size_t n, + const float* centroids = nullptr) const override; /** lower-level encode function * diff --git a/faiss/impl/index_read.cpp b/faiss/impl/index_read.cpp index 53f6485c0..a3f112e3e 100644 --- a/faiss/impl/index_read.cpp +++ b/faiss/impl/index_read.cpp @@ -26,10 +26,12 @@ #include #include +#include #include #include #include #include +#include #include #include #include @@ -271,9 +273,17 @@ static void read_AdditiveQuantizer(AdditiveQuantizer* aq, IOReader* f) { READ1(aq->norm_min); READ1(aq->norm_max); if (aq->search_type == AdditiveQuantizer::ST_norm_cqint8 || - aq->search_type == AdditiveQuantizer::ST_norm_cqint4) { + aq->search_type == AdditiveQuantizer::ST_norm_cqint4 || + aq->search_type == AdditiveQuantizer::ST_norm_lsq2x4 || + aq->search_type == AdditiveQuantizer::ST_norm_rq2x4) { READXBVECTOR(aq->qnorm.codes); } + + if (aq->search_type == AdditiveQuantizer::ST_norm_lsq2x4 || + aq->search_type == AdditiveQuantizer::ST_norm_rq2x4) { + READVECTOR(aq->norm_tabs); + } + aq->set_derived_values(); } @@ -556,6 +566,73 @@ Index* read_index(IOReader* f, int io_flags) { READ1(idxr->beam_factor); idxr->set_beam_factor(idxr->beam_factor); idx = idxr; + } else if (h == fourcc("ILfs") || h == fourcc("IRfs")) { + bool is_LSQ = h == fourcc("ILfs"); + IndexAdditiveQuantizerFastScan* idxaqfs; + if (is_LSQ) { + idxaqfs = new IndexLocalSearchQuantizerFastScan(); + } else { + idxaqfs = new IndexResidualQuantizerFastScan(); + } + read_index_header(idxaqfs, f); + + if (is_LSQ) { + read_LocalSearchQuantizer((LocalSearchQuantizer*)idxaqfs->aq, f); + } else { + read_ResidualQuantizer((ResidualQuantizer*)idxaqfs->aq, f); + } + + READ1(idxaqfs->implem); + READ1(idxaqfs->bbs); + READ1(idxaqfs->qbs); + + READ1(idxaqfs->M); + READ1(idxaqfs->nbits); + READ1(idxaqfs->ksub); + READ1(idxaqfs->code_size); + READ1(idxaqfs->ntotal2); + READ1(idxaqfs->M2); + + READ1(idxaqfs->rescale_norm); + READ1(idxaqfs->norm_scale); + READ1(idxaqfs->max_train_points); + + READVECTOR(idxaqfs->codes); + idx = idxaqfs; + } else if (h == fourcc("IVLf") || h == fourcc("IVRf")) { + bool is_LSQ = h == fourcc("IVLf"); + IndexIVFAdditiveQuantizerFastScan* ivaqfs; + if (is_LSQ) { + ivaqfs = new IndexIVFLocalSearchQuantizerFastScan(); + } else { + ivaqfs = new IndexIVFResidualQuantizerFastScan(); + } + read_ivf_header(ivaqfs, f); + + if (is_LSQ) { + read_LocalSearchQuantizer((LocalSearchQuantizer*)ivaqfs->aq, f); + } else { + read_ResidualQuantizer((ResidualQuantizer*)ivaqfs->aq, f); + } + + READ1(ivaqfs->by_residual); + READ1(ivaqfs->implem); + READ1(ivaqfs->bbs); + READ1(ivaqfs->qbs); + + READ1(ivaqfs->M); + READ1(ivaqfs->nbits); + READ1(ivaqfs->ksub); + READ1(ivaqfs->code_size); + READ1(ivaqfs->qbs2); + READ1(ivaqfs->M2); + + READ1(ivaqfs->rescale_norm); + READ1(ivaqfs->norm_scale); + READ1(ivaqfs->max_train_points); + + read_InvertedLists(ivaqfs, f, io_flags); + idx = ivaqfs; } else if (h == fourcc("IvFl") || h == fourcc("IvFL")) { // legacy IndexIVFFlat* ivfl = new IndexIVFFlat(); std::vector> ids; @@ -781,6 +858,13 @@ Index* read_index(IOReader* f, int io_flags) { READ1(idxpqfs->ntotal2); READ1(idxpqfs->M2); READVECTOR(idxpqfs->codes); + + const auto& pq = idxpqfs->pq; + idxpqfs->M = pq.M; + idxpqfs->nbits = pq.nbits; + idxpqfs->ksub = (1 << pq.nbits); + idxpqfs->code_size = pq.code_size; + idx = idxpqfs; } else if (h == fourcc("IwPf")) { @@ -795,6 +879,13 @@ Index* read_index(IOReader* f, int io_flags) { read_ProductQuantizer(&ivpq->pq, f); read_InvertedLists(ivpq, f, io_flags); ivpq->precompute_table(); + + const auto& pq = ivpq->pq; + ivpq->M = pq.M; + ivpq->nbits = pq.nbits; + ivpq->ksub = (1 << pq.nbits); + ivpq->code_size = pq.code_size; + idx = ivpq; } else { FAISS_THROW_FMT( diff --git a/faiss/impl/index_write.cpp b/faiss/impl/index_write.cpp index aa21c3752..1c1fd5c60 100644 --- a/faiss/impl/index_write.cpp +++ b/faiss/impl/index_write.cpp @@ -27,10 +27,12 @@ #include #include +#include #include #include #include #include +#include #include #include #include @@ -170,9 +172,16 @@ static void write_AdditiveQuantizer(const AdditiveQuantizer* aq, IOWriter* f) { WRITE1(aq->norm_min); WRITE1(aq->norm_max); if (aq->search_type == AdditiveQuantizer::ST_norm_cqint8 || - aq->search_type == AdditiveQuantizer::ST_norm_cqint4) { + aq->search_type == AdditiveQuantizer::ST_norm_cqint4 || + aq->search_type == AdditiveQuantizer::ST_norm_lsq2x4 || + aq->search_type == AdditiveQuantizer::ST_norm_rq2x4) { WRITEXBVECTOR(aq->qnorm.codes); } + + if (aq->search_type == AdditiveQuantizer::ST_norm_lsq2x4 || + aq->search_type == AdditiveQuantizer::ST_norm_rq2x4) { + WRITEVECTOR(aq->norm_tabs); + } } static void write_ResidualQuantizer(const ResidualQuantizer* rq, IOWriter* f) { @@ -385,6 +394,88 @@ void write_index(const Index* idx, IOWriter* f) { write_LocalSearchQuantizer(&idxr->lsq, f); WRITE1(idxr->code_size); WRITEVECTOR(idxr->codes); + } else if ( + auto* idxaqfs = + dynamic_cast(idx)) { + auto idxlsqfs = + dynamic_cast(idx); + auto idxrqfs = dynamic_cast(idx); + FAISS_THROW_IF_NOT(idxlsqfs || idxrqfs); + + if (idxlsqfs) { + uint32_t h = fourcc("ILfs"); + WRITE1(h); + } else { + uint32_t h = fourcc("IRfs"); + WRITE1(h); + } + + write_index_header(idxaqfs, f); + + if (idxlsqfs) { + write_LocalSearchQuantizer(&idxlsqfs->lsq, f); + } else { + write_ResidualQuantizer(&idxrqfs->rq, f); + } + WRITE1(idxaqfs->implem); + WRITE1(idxaqfs->bbs); + WRITE1(idxaqfs->qbs); + + WRITE1(idxaqfs->M); + WRITE1(idxaqfs->nbits); + WRITE1(idxaqfs->ksub); + WRITE1(idxaqfs->code_size); + WRITE1(idxaqfs->ntotal2); + WRITE1(idxaqfs->M2); + + WRITE1(idxaqfs->rescale_norm); + WRITE1(idxaqfs->norm_scale); + WRITE1(idxaqfs->max_train_points); + + WRITEVECTOR(idxaqfs->codes); + } else if ( + auto* ivaqfs = + dynamic_cast( + idx)) { + auto ivlsqfs = + dynamic_cast(idx); + auto ivrqfs = + dynamic_cast(idx); + FAISS_THROW_IF_NOT(ivlsqfs || ivrqfs); + + if (ivlsqfs) { + uint32_t h = fourcc("IVLf"); + WRITE1(h); + } else { + uint32_t h = fourcc("IVRf"); + WRITE1(h); + } + + write_ivf_header(ivaqfs, f); + + if (ivlsqfs) { + write_LocalSearchQuantizer(&ivlsqfs->lsq, f); + } else { + write_ResidualQuantizer(&ivrqfs->rq, f); + } + + WRITE1(ivaqfs->by_residual); + WRITE1(ivaqfs->implem); + WRITE1(ivaqfs->bbs); + WRITE1(ivaqfs->qbs); + + WRITE1(ivaqfs->M); + WRITE1(ivaqfs->nbits); + WRITE1(ivaqfs->ksub); + WRITE1(ivaqfs->code_size); + WRITE1(ivaqfs->qbs2); + WRITE1(ivaqfs->M2); + + WRITE1(ivaqfs->rescale_norm); + WRITE1(ivaqfs->norm_scale); + WRITE1(ivaqfs->max_train_points); + + write_InvertedLists(ivaqfs->invlists, f); } else if ( const ResidualCoarseQuantizer* idxr = dynamic_cast(idx)) { diff --git a/faiss/impl/pq4_fast_scan.h b/faiss/impl/pq4_fast_scan.h index 5a6a69a32..c082bed43 100644 --- a/faiss/impl/pq4_fast_scan.h +++ b/faiss/impl/pq4_fast_scan.h @@ -26,7 +26,7 @@ namespace faiss { * The unused bytes are set to 0. * * @param codes input codes, size (ntotal, ceil(M / 2)) - * @param nototal number of input codes + * @param ntotal number of input codes * @param nb output number of codes (ntotal rounded up to a multiple of * bbs) * @param M2 number of sub-quantizers (=M rounded up to a muliple of 2) @@ -88,8 +88,9 @@ void pq4_pack_LUT(int nq, int nsq, const uint8_t* src, uint8_t* dest); * @param nsq number of sub-quantizers (muliple of 2) * @param codes packed codes array * @param LUT packed look-up table + * @param scaler scaler to scale the encoded norm */ -template +template void pq4_accumulate_loop( int nq, size_t nb, @@ -97,7 +98,8 @@ void pq4_accumulate_loop( int nsq, const uint8_t* codes, const uint8_t* LUT, - ResultHandler& res); + ResultHandler& res, + const Scaler& scaler); /* qbs versions, supported only for bbs=32. * @@ -141,20 +143,22 @@ int pq4_pack_LUT_qbs_q_map( /** Run accumulation loop. * - * @param qbs 4-bit encded number of queries + * @param qbs 4-bit encoded number of queries * @param nb number of database codes (mutliple of bbs) * @param nsq number of sub-quantizers * @param codes encoded database vectors (packed) * @param LUT look-up table (packed) * @param res call-back for the resutls + * @param scaler scaler to scale the encoded norm */ -template +template void pq4_accumulate_loop_qbs( int qbs, size_t nb, int nsq, const uint8_t* codes, const uint8_t* LUT, - ResultHandler& res); + ResultHandler& res, + const Scaler& scaler); } // namespace faiss diff --git a/faiss/impl/pq4_fast_scan_search_1.cpp b/faiss/impl/pq4_fast_scan_search_1.cpp index b835bb648..672c7d6db 100644 --- a/faiss/impl/pq4_fast_scan_search_1.cpp +++ b/faiss/impl/pq4_fast_scan_search_1.cpp @@ -8,6 +8,7 @@ #include #include +#include #include namespace faiss { @@ -26,12 +27,13 @@ namespace { * writes results in a ResultHandler */ -template +template void kernel_accumulate_block( int nsq, const uint8_t* codes, const uint8_t* LUT, - ResultHandler& res) { + ResultHandler& res, + const Scaler& scaler) { // distance accumulators simd16uint16 accu[NQ][BB][4]; @@ -44,7 +46,7 @@ void kernel_accumulate_block( } } - for (int sq = 0; sq < nsq; sq += 2) { + for (int sq = 0; sq < nsq - scaler.nscale; sq += 2) { simd32uint8 lut_cache[NQ]; for (int q = 0; q < NQ; q++) { lut_cache[q] = simd32uint8(LUT); @@ -72,6 +74,35 @@ void kernel_accumulate_block( } } + for (int sq = 0; sq < scaler.nscale; sq += 2) { + simd32uint8 lut_cache[NQ]; + for (int q = 0; q < NQ; q++) { + lut_cache[q] = simd32uint8(LUT); + LUT += 32; + } + + for (int b = 0; b < BB; b++) { + simd32uint8 c = simd32uint8(codes); + codes += 32; + simd32uint8 mask(15); + simd32uint8 chi = simd32uint8(simd16uint16(c) >> 4) & mask; + simd32uint8 clo = c & mask; + + for (int q = 0; q < NQ; q++) { + simd32uint8 lut = lut_cache[q]; + + simd32uint8 res0 = scaler.lookup(lut, clo); + accu[q][b][0] += scaler.scale_lo(res0); // handle vectors 0..7 + accu[q][b][1] += scaler.scale_hi(res0); // handle vectors 8..15 + + simd32uint8 res1 = scaler.lookup(lut, chi); + accu[q][b][2] += scaler.scale_lo(res1); // handle vectors 16..23 + accu[q][b][3] += + scaler.scale_hi(res1); // handle vectors 24..31 + } + } + } + for (int q = 0; q < NQ; q++) { for (int b = 0; b < BB; b++) { accu[q][b][0] -= accu[q][b][1] << 8; @@ -85,17 +116,18 @@ void kernel_accumulate_block( } } -template +template void accumulate_fixed_blocks( size_t nb, int nsq, const uint8_t* codes, const uint8_t* LUT, - ResultHandler& res) { + ResultHandler& res, + const Scaler& scaler) { constexpr int bbs = 32 * BB; for (int64_t j0 = 0; j0 < nb; j0 += bbs) { FixedStorageHandler res2; - kernel_accumulate_block(nsq, codes, LUT, res2); + kernel_accumulate_block(nsq, codes, LUT, res2, scaler); res.set_block_origin(0, j0); res2.to_other_handler(res); codes += bbs * nsq / 2; @@ -104,7 +136,7 @@ void accumulate_fixed_blocks( } // anonymous namespace -template +template void pq4_accumulate_loop( int nq, size_t nb, @@ -112,15 +144,16 @@ void pq4_accumulate_loop( int nsq, const uint8_t* codes, const uint8_t* LUT, - ResultHandler& res) { + ResultHandler& res, + const Scaler& scaler) { FAISS_THROW_IF_NOT(is_aligned_pointer(codes)); FAISS_THROW_IF_NOT(is_aligned_pointer(LUT)); FAISS_THROW_IF_NOT(bbs % 32 == 0); FAISS_THROW_IF_NOT(nb % bbs == 0); -#define DISPATCH(NQ, BB) \ - case NQ * 1000 + BB: \ - accumulate_fixed_blocks(nb, nsq, codes, LUT, res); \ +#define DISPATCH(NQ, BB) \ + case NQ * 1000 + BB: \ + accumulate_fixed_blocks(nb, nsq, codes, LUT, res, scaler); \ break switch (nq * 1000 + bbs / 32) { @@ -141,20 +174,28 @@ void pq4_accumulate_loop( // explicit template instantiations -#define INSTANTIATE_ACCUMULATE(TH, C, with_id_map) \ - template void pq4_accumulate_loop>( \ - int, \ - size_t, \ - int, \ - int, \ - const uint8_t*, \ - const uint8_t*, \ - TH&); +#define INSTANTIATE_ACCUMULATE(TH, C, with_id_map, S) \ + template void pq4_accumulate_loop, S>( \ + int, \ + size_t, \ + int, \ + int, \ + const uint8_t*, \ + const uint8_t*, \ + TH&, \ + const S&); -#define INSTANTIATE_3(C, with_id_map) \ - INSTANTIATE_ACCUMULATE(SingleResultHandler, C, with_id_map) \ - INSTANTIATE_ACCUMULATE(HeapHandler, C, with_id_map) \ - INSTANTIATE_ACCUMULATE(ReservoirHandler, C, with_id_map) +using DS = DummyScaler; +using NS = NormTableScaler; + +#define INSTANTIATE_3(C, with_id_map) \ + INSTANTIATE_ACCUMULATE(SingleResultHandler, C, with_id_map, DS) \ + INSTANTIATE_ACCUMULATE(HeapHandler, C, with_id_map, DS) \ + INSTANTIATE_ACCUMULATE(ReservoirHandler, C, with_id_map, DS) \ + \ + INSTANTIATE_ACCUMULATE(SingleResultHandler, C, with_id_map, NS) \ + INSTANTIATE_ACCUMULATE(HeapHandler, C, with_id_map, NS) \ + INSTANTIATE_ACCUMULATE(ReservoirHandler, C, with_id_map, NS) using Csi = CMax; INSTANTIATE_3(Csi, false); diff --git a/faiss/impl/pq4_fast_scan_search_qbs.cpp b/faiss/impl/pq4_fast_scan_search_qbs.cpp index b06d494a7..06525987a 100644 --- a/faiss/impl/pq4_fast_scan_search_qbs.cpp +++ b/faiss/impl/pq4_fast_scan_search_qbs.cpp @@ -8,6 +8,7 @@ #include #include +#include #include #include @@ -27,15 +28,17 @@ namespace { * writes results in a ResultHandler */ -template +template void kernel_accumulate_block( int nsq, const uint8_t* codes, const uint8_t* LUT, - ResultHandler& res) { + ResultHandler& res, + const Scaler& scaler) { // dummy alloc to keep the windows compiler happy constexpr int NQA = NQ > 0 ? NQ : 1; // distance accumulators + // layout: accu[q][b]: distance accumulator for vectors 8*b..8*b+7 simd16uint16 accu[NQA][4]; for (int q = 0; q < NQ; q++) { @@ -45,7 +48,7 @@ void kernel_accumulate_block( } // _mm_prefetch(codes + 768, 0); - for (int sq = 0; sq < nsq; sq += 2) { + for (int sq = 0; sq < nsq - scaler.nscale; sq += 2) { // prefetch simd32uint8 c(codes); codes += 32; @@ -71,6 +74,31 @@ void kernel_accumulate_block( } } + for (int sq = 0; sq < scaler.nscale; sq += 2) { + // prefetch + simd32uint8 c(codes); + codes += 32; + + simd32uint8 mask(0xf); + // shift op does not exist for int8... + simd32uint8 chi = simd32uint8(simd16uint16(c) >> 4) & mask; + simd32uint8 clo = c & mask; + + for (int q = 0; q < NQ; q++) { + // load LUTs for 2 quantizers + simd32uint8 lut(LUT); + LUT += 32; + + simd32uint8 res0 = scaler.lookup(lut, clo); + accu[q][0] += scaler.scale_lo(res0); // handle vectors 0..7 + accu[q][1] += scaler.scale_hi(res0); // handle vectors 8..15 + + simd32uint8 res1 = scaler.lookup(lut, chi); + accu[q][2] += scaler.scale_lo(res1); // handle vectors 16..23 + accu[q][3] += scaler.scale_hi(res1); // handle vectors 24..31 + } + } + for (int q = 0; q < NQ; q++) { accu[q][0] -= accu[q][1] << 8; simd16uint16 dis0 = combine2x2(accu[q][0], accu[q][1]); @@ -81,13 +109,14 @@ void kernel_accumulate_block( } // handle at most 4 blocks of queries -template +template void accumulate_q_4step( size_t ntotal2, int nsq, const uint8_t* codes, const uint8_t* LUT0, - ResultHandler& res) { + ResultHandler& res, + const Scaler& scaler) { constexpr int Q1 = QBS & 15; constexpr int Q2 = (QBS >> 4) & 15; constexpr int Q3 = (QBS >> 8) & 15; @@ -97,21 +126,21 @@ void accumulate_q_4step( for (int64_t j0 = 0; j0 < ntotal2; j0 += 32) { FixedStorageHandler res2; const uint8_t* LUT = LUT0; - kernel_accumulate_block(nsq, codes, LUT, res2); + kernel_accumulate_block(nsq, codes, LUT, res2, scaler); LUT += Q1 * nsq * 16; if (Q2 > 0) { res2.set_block_origin(Q1, 0); - kernel_accumulate_block(nsq, codes, LUT, res2); + kernel_accumulate_block(nsq, codes, LUT, res2, scaler); LUT += Q2 * nsq * 16; } if (Q3 > 0) { res2.set_block_origin(Q1 + Q2, 0); - kernel_accumulate_block(nsq, codes, LUT, res2); + kernel_accumulate_block(nsq, codes, LUT, res2, scaler); LUT += Q3 * nsq * 16; } if (Q4 > 0) { res2.set_block_origin(Q1 + Q2 + Q3, 0); - kernel_accumulate_block(nsq, codes, LUT, res2); + kernel_accumulate_block(nsq, codes, LUT, res2, scaler); } res.set_block_origin(0, j0); res2.to_other_handler(res); @@ -119,29 +148,31 @@ void accumulate_q_4step( } } -template +template void kernel_accumulate_block_loop( size_t ntotal2, int nsq, const uint8_t* codes, const uint8_t* LUT, - ResultHandler& res) { + ResultHandler& res, + const Scaler& scaler) { for (int64_t j0 = 0; j0 < ntotal2; j0 += 32) { res.set_block_origin(0, j0); kernel_accumulate_block( - nsq, codes + j0 * nsq / 2, LUT, res); + nsq, codes + j0 * nsq / 2, LUT, res, scaler); } } // non-template version of accumulate kernel -- dispatches dynamically -template +template void accumulate( int nq, size_t ntotal2, int nsq, const uint8_t* codes, const uint8_t* LUT, - ResultHandler& res) { + ResultHandler& res, + const Scaler& scaler) { assert(nsq % 2 == 0); assert(is_aligned_pointer(codes)); assert(is_aligned_pointer(LUT)); @@ -149,7 +180,7 @@ void accumulate( #define DISPATCH(NQ) \ case NQ: \ kernel_accumulate_block_loop( \ - ntotal2, nsq, codes, LUT, res); \ + ntotal2, nsq, codes, LUT, res, scaler); \ return switch (nq) { @@ -165,23 +196,24 @@ void accumulate( } // namespace -template +template void pq4_accumulate_loop_qbs( int qbs, size_t ntotal2, int nsq, const uint8_t* codes, const uint8_t* LUT0, - ResultHandler& res) { + ResultHandler& res, + const Scaler& scaler) { assert(nsq % 2 == 0); assert(is_aligned_pointer(codes)); assert(is_aligned_pointer(LUT0)); // try out optimized versions switch (qbs) { -#define DISPATCH(QBS) \ - case QBS: \ - accumulate_q_4step(ntotal2, nsq, codes, LUT0, res); \ +#define DISPATCH(QBS) \ + case QBS: \ + accumulate_q_4step(ntotal2, nsq, codes, LUT0, res, scaler); \ return; DISPATCH(0x3333); // 12 DISPATCH(0x2333); // 11 @@ -219,9 +251,10 @@ void pq4_accumulate_loop_qbs( int nq = qi & 15; qi >>= 4; res.set_block_origin(i0, j0); -#define DISPATCH(NQ) \ - case NQ: \ - kernel_accumulate_block(nsq, codes, LUT, res); \ +#define DISPATCH(NQ) \ + case NQ: \ + kernel_accumulate_block( \ + nsq, codes, LUT, res, scaler); \ break switch (nq) { DISPATCH(1); @@ -241,9 +274,23 @@ void pq4_accumulate_loop_qbs( // explicit template instantiations -#define INSTANTIATE_ACCUMULATE_Q(RH) \ - template void pq4_accumulate_loop_qbs( \ - int, size_t, int, const uint8_t*, const uint8_t*, RH&); +#define INSTANTIATE_ACCUMULATE_Q(RH) \ + template void pq4_accumulate_loop_qbs( \ + int, \ + size_t, \ + int, \ + const uint8_t*, \ + const uint8_t*, \ + RH&, \ + const DummyScaler&); \ + template void pq4_accumulate_loop_qbs( \ + int, \ + size_t, \ + int, \ + const uint8_t*, \ + const uint8_t*, \ + RH&, \ + const NormTableScaler&); using Csi = CMax; INSTANTIATE_ACCUMULATE_Q(SingleResultHandler) @@ -293,7 +340,8 @@ void accumulate_to_mem( uint16_t* accu) { FAISS_THROW_IF_NOT(ntotal2 % 32 == 0); StoreResultHandler handler(accu, ntotal2); - accumulate(nq, ntotal2, nsq, codes, LUT, handler); + DummyScaler scaler; + accumulate(nq, ntotal2, nsq, codes, LUT, handler, scaler); } int pq4_preferred_qbs(int n) { diff --git a/faiss/index_factory.cpp b/faiss/index_factory.cpp index 0ae6900c0..7b808fd5d 100644 --- a/faiss/index_factory.cpp +++ b/faiss/index_factory.cpp @@ -26,10 +26,12 @@ #include #include +#include #include #include #include #include +#include #include #include #include @@ -149,11 +151,13 @@ std::map aq_search_type = { {"_Nqint4", AdditiveQuantizer::ST_norm_qint4}, {"_Ncqint8", AdditiveQuantizer::ST_norm_cqint8}, {"_Ncqint4", AdditiveQuantizer::ST_norm_cqint4}, + {"_Nlsq2x4", AdditiveQuantizer::ST_norm_lsq2x4}, + {"_Nrq2x4", AdditiveQuantizer::ST_norm_rq2x4}, }; const std::string aq_def_pattern = "[0-9]+x[0-9]+(_[0-9]+x[0-9]+)*"; const std::string aq_norm_pattern = - "(|_Nnone|_Nfloat|_Nqint8|_Nqint4|_Ncqint8|_Ncqint4)"; + "(|_Nnone|_Nfloat|_Nqint8|_Nqint4|_Ncqint8|_Ncqint4|_Nlsq2x4|_Nrq2x4)"; AdditiveQuantizer::Search_type_t aq_parse_search_type( std::string stok, @@ -169,7 +173,7 @@ AdditiveQuantizer::Search_type_t aq_parse_search_type( std::vector aq_parse_nbits(std::string stok) { std::vector nbits; std::smatch sm; - while (std::regex_search(stok, sm, std::regex("([0-9]+)x([0-9]+)"))) { + while (std::regex_search(stok, sm, std::regex("[^q]([0-9]+)x([0-9]+)"))) { int M = std::stoi(sm[1].str()); int nbit = std::stoi(sm[2].str()); nbits.resize(nbits.size() + M, nbit); @@ -341,6 +345,21 @@ IndexIVF* parse_IndexIVF( } return index_ivf; } + if (match("(RQ|LSQ)([0-9]+)x4fs(r?)(_[0-9]+)?" + aq_norm_pattern)) { + int M = std::stoi(sm[2].str()); + int bbs = mres_to_int(sm[4], 32, 1); + auto st = aq_parse_search_type(sm[sm.size() - 1].str(), mt); + IndexIVFAdditiveQuantizerFastScan* index_ivf; + if (sm[1].str() == "RQ") { + index_ivf = new IndexIVFResidualQuantizerFastScan( + get_q(), d, nlist, M, 4, mt, st, bbs); + } else { + index_ivf = new IndexIVFLocalSearchQuantizerFastScan( + get_q(), d, nlist, M, 4, mt, st, bbs); + } + index_ivf->by_residual = (sm[3].str() == "r"); + return index_ivf; + } if (match("(ITQ|PCA|PCAR)([0-9]+)?,SH([-0-9.e]+)?([gcm])?")) { int outdim = mres_to_int(sm[2], d); // is also the number of bits std::unique_ptr vt; @@ -531,6 +550,22 @@ Index* parse_other_indexes( return new IndexLocalSearchQuantizer(d, M, nbit, metric, st); } + // IndexAdditiveQuantizerFastScan + // RQ{M}x4fs_{bbs}_{search_type} + pattern = "(LSQ|RQ)([0-9]+)x4fs(_[0-9]+)?" + aq_norm_pattern; + if (match(pattern)) { + int M = std::stoi(sm[2].str()); + int bbs = mres_to_int(sm[3], 32, 1); + auto st = aq_parse_search_type(sm[sm.size() - 1].str(), metric); + + if (sm[1].str() == "RQ") { + return new IndexResidualQuantizerFastScan(d, M, 4, metric, st, bbs); + } else if (sm[1].str() == "LSQ") { + return new IndexLocalSearchQuantizerFastScan( + d, M, 4, metric, st, bbs); + } + } + return nullptr; } diff --git a/faiss/python/swigfaiss.swig b/faiss/python/swigfaiss.swig index c995a25c5..2d05d0c2a 100644 --- a/faiss/python/swigfaiss.swig +++ b/faiss/python/swigfaiss.swig @@ -77,7 +77,11 @@ typedef uint64_t size_t; #include #include +#include +#include #include +#include +#include #include #include @@ -438,7 +442,11 @@ void gpu_sync_all_devices() %include %include +%include +%include %include +%include +%include %include %include @@ -554,6 +562,8 @@ void gpu_sync_all_devices() DOWNCAST ( IndexIVFScalarQuantizer ) DOWNCAST ( IndexIVFResidualQuantizer ) DOWNCAST ( IndexIVFLocalSearchQuantizer ) + DOWNCAST ( IndexIVFResidualQuantizerFastScan ) + DOWNCAST ( IndexIVFLocalSearchQuantizerFastScan ) DOWNCAST ( IndexIVFFlatDedup ) DOWNCAST ( IndexIVFFlat ) DOWNCAST ( IndexIVF ) @@ -564,6 +574,8 @@ void gpu_sync_all_devices() DOWNCAST ( IndexPQ ) DOWNCAST ( IndexResidualQuantizer ) DOWNCAST ( IndexLocalSearchQuantizer ) + DOWNCAST ( IndexResidualQuantizerFastScan ) + DOWNCAST ( IndexLocalSearchQuantizerFastScan ) DOWNCAST ( ResidualCoarseQuantizer ) DOWNCAST ( LocalSearchCoarseQuantizer ) DOWNCAST ( IndexScalarQuantizer ) @@ -654,6 +666,14 @@ void gpu_sync_all_devices() } } +%typemap(out) faiss::AdditiveQuantizer * { + DOWNCAST (LocalSearchQuantizer) + DOWNCAST (ResidualQuantizer) + { + assert(false); + } +} + // just to downcast pointers that come from elsewhere (eg. direct // access to object fields) %inline %{ @@ -673,6 +693,10 @@ faiss::InvertedLists * downcast_InvertedLists (faiss::InvertedLists *il) { return il; } +faiss::AdditiveQuantizer * downcast_AdditiveQuantizer (faiss::AdditiveQuantizer *aq) +{ + return aq; +} %} %include diff --git a/faiss/utils/quantize_lut.cpp b/faiss/utils/quantize_lut.cpp index 12ca579b6..642f601d7 100644 --- a/faiss/utils/quantize_lut.cpp +++ b/faiss/utils/quantize_lut.cpp @@ -284,6 +284,68 @@ void quantize_LUT_and_bias( *b_out = b; } +void aq_quantize_LUT_and_bias( + size_t nprobe, + size_t M, + size_t ksub, + const float* LUT, + const float* bias, + size_t M_norm, + int norm_scale, + uint8_t* LUTq, + size_t M2, + uint16_t* biasq, + float* a_out, + float* b_out) { + float a, b; + std::vector mins(M); + float max_span_LUT = -HUGE_VAL, max_span_dis; + float bias_min = tab_min(bias, nprobe); + float bias_max = tab_max(bias, nprobe); + max_span_dis = bias_max - bias_min; + b = 0; + for (int i = 0; i < M; i++) { + mins[i] = tab_min(LUT + i * ksub, ksub); + float span = tab_max(LUT + i * ksub, ksub) - mins[i]; + max_span_LUT = std::max(max_span_LUT, span); + max_span_dis += (i >= M - M_norm ? span * norm_scale : span); + b += mins[i]; + } + a = std::min(255 / max_span_LUT, 65535 / max_span_dis); + b += bias_min; + + for (int i = 0; i < M; i++) { + round_tab(LUT + i * ksub, ksub, a, mins[i], LUTq + i * ksub); + } + memset(LUTq + M * ksub, 0, ksub * (M2 - M)); + round_tab(bias, nprobe, a, bias_min, biasq); + + *a_out = a; + *b_out = b; +} + +float aq_estimate_norm_scale( + size_t M, + size_t ksub, + size_t M_norm, + const float* LUT) { + float max_span_LUT = -HUGE_VAL; + for (int i = 0; i < M - M_norm; i++) { + float min = tab_min(LUT + i * ksub, ksub); + float span = tab_max(LUT + i * ksub, ksub) - min; + max_span_LUT = std::max(max_span_LUT, span); + } + + float max_span_LUT_norm = -HUGE_VAL; + for (int i = M - M_norm; i < M; i++) { + float min = tab_min(LUT + i * ksub, ksub); + float span = tab_max(LUT + i * ksub, ksub) - min; + max_span_LUT_norm = std::max(max_span_LUT_norm, span); + } + + return max_span_LUT_norm / max_span_LUT; +} + } // namespace quantize_lut } // namespace faiss diff --git a/faiss/utils/quantize_lut.h b/faiss/utils/quantize_lut.h index 332ebdbfa..b7d4fc42f 100644 --- a/faiss/utils/quantize_lut.h +++ b/faiss/utils/quantize_lut.h @@ -77,6 +77,26 @@ void quantize_LUT_and_bias( float* a_out = nullptr, float* b_out = nullptr); +void aq_quantize_LUT_and_bias( + size_t nprobe, + size_t M, + size_t ksub, + const float* LUT, + const float* bias, + size_t M_norm, + int norm_scale, + uint8_t* LUTq, + size_t M2, + uint16_t* biasq, + float* a_out, + float* b_out); + +float aq_estimate_norm_scale( + size_t M, + size_t ksub, + size_t M_norm, + const float* LUT); + } // namespace quantize_lut } // namespace faiss diff --git a/faiss/utils/simdlib_avx2.h b/faiss/utils/simdlib_avx2.h index c053334e3..6eee3fa0e 100644 --- a/faiss/utils/simdlib_avx2.h +++ b/faiss/utils/simdlib_avx2.h @@ -111,6 +111,10 @@ struct simd16uint16 : simd256bit { i = _mm256_set1_epi16((short)x); } + simd16uint16 operator*(const simd16uint16& other) const { + return simd16uint16(_mm256_mullo_epi16(i, other.i)); + } + // shift must be known at compile time simd16uint16 operator>>(const int shift) const { return simd16uint16(_mm256_srli_epi16(i, shift)); diff --git a/faiss/utils/simdlib_emulated.h b/faiss/utils/simdlib_emulated.h index 7ab52a5b7..c6166f626 100644 --- a/faiss/utils/simdlib_emulated.h +++ b/faiss/utils/simdlib_emulated.h @@ -120,6 +120,11 @@ struct simd16uint16 : simd256bit { } } + simd16uint16 operator*(const simd16uint16& other) const { + return binary_func( + *this, other, [](uint16_t a, uint16_t b) { return a * b; }); + } + // shift must be known at compile time simd16uint16 operator>>(const int shift) const { return unary_func(*this, [shift](uint16_t a) { return a >> shift; }); diff --git a/faiss/utils/simdlib_neon.h b/faiss/utils/simdlib_neon.h index a8f2b84cf..6f2e38f26 100644 --- a/faiss/utils/simdlib_neon.h +++ b/faiss/utils/simdlib_neon.h @@ -260,6 +260,11 @@ struct simd16uint16 { detail::simdlib::set1(data, &vdupq_n_u16, x); } + simd16uint16 operator*(const simd16uint16& other) const { + return simd16uint16{ + detail::simdlib::binary_func(data, other.data, &vmulq_u16)}; + } + // shift must be known at compile time simd16uint16 operator>>(const int shift) const { switch (shift) { diff --git a/tests/test_fast_scan.py b/tests/test_fast_scan.py index 8b18b62b8..e0f8ab3fc 100644 --- a/tests/test_fast_scan.py +++ b/tests/test_fast_scan.py @@ -6,6 +6,8 @@ import unittest import time +import os +import tempfile import numpy as np import faiss @@ -141,7 +143,6 @@ class TestRounding(unittest.TestCase): ######################################################### - def reference_accu(codes, LUT): nq, nsp, is_16 = LUT.shape nb, nsp_2 = codes.shape @@ -211,9 +212,6 @@ class ThisIsNotATestLoop5: # (unittest.TestCase): self.do_loop5_kernel(2, 2) - - - ########################################################## # Tests for various IndexPQFastScan implementations ########################################################## @@ -435,3 +433,157 @@ class TestAdd(unittest.TestCase): D3, I3 = index3.search(ds.get_queries(), 10) np.testing.assert_array_equal(D3, Dnew) np.testing.assert_array_equal(I3, Inew) + + +class TestAQFastScan(unittest.TestCase): + + def subtest_accuracy(self, aq, st, implem, metric_type='L2'): + """ + Compare IndexAdditiveQuantizerFastScan with IndexAQ (qint8) + """ + d = 16 + ds = datasets.SyntheticDataset(d, 1000, 2000, 1000, metric_type) + gt = ds.get_groundtruth(k=1) + + if metric_type == 'L2': + metric = faiss.METRIC_L2 + postfix1 = '_Nqint8' + postfix2 = f'_N{st}2x4' + else: + metric = faiss.METRIC_INNER_PRODUCT + postfix1 = postfix2 = '' + + index = faiss.index_factory(d, f'{aq}3x4{postfix1}', metric) + index.train(ds.get_train()) + index.add(ds.get_database()) + Dref, Iref = index.search(ds.get_queries(), 1) + + indexfs = faiss.index_factory(d, f'{aq}3x4fs_32{postfix2}', metric) + indexfs.train(ds.get_train()) + indexfs.add(ds.get_database()) + indexfs.implem = implem + Da, Ia = indexfs.search(ds.get_queries(), 1) + + nq = Iref.shape[0] + recall_ref = (Iref == gt).sum() / nq + recall = (Ia == gt).sum() / nq + + print(aq, st, implem, metric_type, recall_ref, recall) + assert abs(recall_ref - recall) < 0.05 + + def xx_test_accuracy(self): + for metric in 'L2', 'IP': + for implem in 0, 12, 13, 14, 15: + self.subtest_accuracy('RQ', 'rq', implem, metric) + self.subtest_accuracy('LSQ', 'lsq', implem, metric) + + def subtest_from_idxaq(self, implem, metric): + if metric == 'L2': + metric_type = faiss.METRIC_L2 + st = '_Nrq2x4' + else: + metric_type = faiss.METRIC_INNER_PRODUCT + st = '' + + d = 16 + ds = datasets.SyntheticDataset(d, 1000, 2000, 1000, metric=metric) + gt = ds.get_groundtruth(k=1) + index = faiss.index_factory(d, 'RQ8x4' + st, metric_type) + index.train(ds.get_train()) + index.add(ds.get_database()) + index.nprobe = 16 + Dref, Iref = index.search(ds.get_queries(), 1) + + indexfs = faiss.IndexAdditiveQuantizerFastScan(index) + indexfs.implem = implem + D1, I1 = indexfs.search(ds.get_queries(), 1) + + nq = Iref.shape[0] + recall_ref = (Iref == gt).sum() / nq + recall1 = (I1 == gt).sum() / nq + print(recall_ref, recall1) + assert abs(recall_ref - recall1) < 0.05 + + def test_from_idxaq(self): + for implem in 2, 3, 4: + self.subtest_from_idxaq(implem, 'L2') + self.subtest_from_idxaq(implem, 'IP') + + def subtest_factory(self, aq, M, bbs, st): + """ + Format: {AQ}{M}x4fs_{bbs}_N{st} + + AQ (str): `LSQ` or `RQ` + M (int): number of subquantizers + bbs (int): build block size + st (str): search type, `lsq2x4` or `rq2x4` + """ + AQ = faiss.AdditiveQuantizer + d = 16 + + if bbs > 0: + index = faiss.index_factory(d, f'{aq}{M}x4fs_{bbs}_N{st}2x4') + else: + index = faiss.index_factory(d, f'{aq}{M}x4fs_N{st}2x4') + bbs = 32 + + assert index.bbs == bbs + aq = faiss.downcast_AdditiveQuantizer(index.aq) + assert aq.M == M + + if aq == 'LSQ': + assert isinstance(aq, faiss.LocalSearchQuantizer) + if aq == 'RQ': + assert isinstance(aq, faiss.ResidualQuantizer) + + if st == 'lsq': + assert aq.search_type == AQ.ST_norm_lsq2x4 + if st == 'rq': + assert aq.search_type == AQ.ST_norm_rq2x4 + + def test_factory(self): + self.subtest_factory('LSQ', 16, 64, 'lsq') + self.subtest_factory('LSQ', 16, 64, 'rq') + self.subtest_factory('RQ', 16, 64, 'rq') + self.subtest_factory('RQ', 16, 64, 'lsq') + self.subtest_factory('LSQ', 64, 0, 'lsq') + + def subtest_io(self, factory_str): + d = 8 + ds = datasets.SyntheticDataset(d, 1000, 2000, 1000) + + index = faiss.index_factory(d, factory_str) + index.train(ds.get_train()) + index.add(ds.get_database()) + D1, I1 = index.search(ds.get_queries(), 1) + + fd, fname = tempfile.mkstemp() + os.close(fd) + try: + faiss.write_index(index, fname) + index2 = faiss.read_index(fname) + D2, I2 = index2.search(ds.get_queries(), 1) + np.testing.assert_array_equal(I1, I2) + finally: + if os.path.exists(fname): + os.unlink(fname) + + def test_io(self): + self.subtest_io('LSQ4x4fs_Nlsq2x4') + self.subtest_io('LSQ4x4fs_Nrq2x4') + self.subtest_io('RQ4x4fs_Nrq2x4') + self.subtest_io('RQ4x4fs_Nlsq2x4') + + +for metric in 'L2', 'IP': + for implem in 0, 12, 13, 14, 15: + setattr( + TestAQFastScan, + f"test_accuracy_{metric}_LSQ_implem{implem}", + lambda self: self.subtest_accuracy('LSQ', 'lsq', implem, metric) + ) + setattr( + TestAQFastScan, + f"test_accuracy_{metric}_RQ_implem{implem}", + lambda self: self.subtest_accuracy('RQ', 'rq', implem, metric) + ) diff --git a/tests/test_fast_scan_ivf.py b/tests/test_fast_scan_ivf.py index f903f9b6e..2c5e2f168 100644 --- a/tests/test_fast_scan_ivf.py +++ b/tests/test_fast_scan_ivf.py @@ -4,8 +4,9 @@ # LICENSE file in the root directory of this source tree. +import os import unittest -import platform +import tempfile import numpy as np import faiss @@ -351,7 +352,6 @@ class TestIVFImplem12(unittest.TestCase): class TestIVFImplem10(TestIVFImplem12): IMPLEM = 10 - class TestIVFImplem11(TestIVFImplem12): IMPLEM = 11 @@ -484,3 +484,228 @@ class TestIsTrained(unittest.TestCase): ) des = faiss.rand((1000, 32)) index.train(des) + + +class TestIVFAQFastScan(unittest.TestCase): + + def subtest_accuracy(self, aq, st, by_residual, implem, metric_type='L2'): + """ + Compare IndexIVFAdditiveQuantizerFastScan with + IndexIVFAdditiveQuantizer + """ + nlist, d = 16, 8 + ds = datasets.SyntheticDataset(d, 1000, 1000, 500, metric_type) + gt = ds.get_groundtruth(k=1) + + if metric_type == 'L2': + metric = faiss.METRIC_L2 + postfix1 = '_Nqint8' + postfix2 = f'_N{st}2x4' + else: + metric = faiss.METRIC_INNER_PRODUCT + postfix1 = postfix2 = '' + + index = faiss.index_factory(d, f'IVF{nlist},{aq}3x4{postfix1}', metric) + index.by_residual = by_residual + index.train(ds.get_train()) + index.add(ds.get_database()) + index.nprobe = 16 + Dref, Iref = index.search(ds.get_queries(), 1) + + indexfs = faiss.index_factory( + d, f'IVF{nlist},{aq}3x4fs_32{postfix2}', metric) + indexfs.by_residual = by_residual + indexfs.train(ds.get_train()) + indexfs.add(ds.get_database()) + indexfs.nprobe = 16 + indexfs.implem = implem + D1, I1 = indexfs.search(ds.get_queries(), 1) + + nq = Iref.shape[0] + recall_ref = (Iref == gt).sum() / nq + recall1 = (I1 == gt).sum() / nq + + print(aq, st, by_residual, implem, metric_type, recall_ref, recall1) + assert abs(recall_ref - recall1) < 0.05 + + def xx_test_accuracy(self): + # generated programatically below + for metric in 'L2', 'IP': + for byr in True, False: + for implem in 0, 10, 11, 12, 13: + self.subtest_accuracy('RQ', 'rq', byr, implem, metric) + self.subtest_accuracy('LSQ', 'lsq', byr, implem, metric) + + def subtest_rescale_accuracy(self, aq, st, by_residual, implem): + """ + we set norm_scale to 2 and compare it with IndexIVFAQ + """ + nlist, d = 16, 8 + ds = datasets.SyntheticDataset(d, 1000, 1000, 500) + gt = ds.get_groundtruth(k=1) + + # if metric_type == 'L2': + metric = faiss.METRIC_L2 + postfix1 = '_Nqint8' + postfix2 = f'_N{st}2x4' + + index = faiss.index_factory( + d, f'IVF{nlist},{aq}3x4{postfix1}', metric) + index.by_residual = by_residual + index.train(ds.get_train()) + index.add(ds.get_database()) + index.nprobe = 16 + Dref, Iref = index.search(ds.get_queries(), 1) + + indexfs = faiss.index_factory( + d, f'IVF{nlist},{aq}3x4fs_32{postfix2}', metric) + indexfs.by_residual = by_residual + indexfs.norm_scale = 2 + indexfs.train(ds.get_train()) + indexfs.add(ds.get_database()) + indexfs.nprobe = 16 + indexfs.implem = implem + D1, I1 = indexfs.search(ds.get_queries(), 1) + + nq = Iref.shape[0] + recall_ref = (Iref == gt).sum() / nq + recall1 = (I1 == gt).sum() / nq + + print(aq, st, by_residual, implem, recall_ref, recall1) + assert abs(recall_ref - recall1) < 0.05 + + def xx_test_rescale_accuracy(self): + for byr in True, False: + for implem in 0, 10, 11, 12, 13: + self.subtest_accuracy('RQ', 'rq', byr, implem, 'L2') + self.subtest_accuracy('LSQ', 'lsq', byr, implem, 'L2') + + def subtest_from_ivfaq(self, implem): + d = 8 + ds = datasets.SyntheticDataset(d, 1000, 2000, 1000, metric='IP') + gt = ds.get_groundtruth(k=1) + index = faiss.index_factory(d, 'IVF16,RQ8x4', faiss.METRIC_INNER_PRODUCT) + index.train(ds.get_train()) + index.add(ds.get_database()) + index.nprobe = 16 + Dref, Iref = index.search(ds.get_queries(), 1) + + indexfs = faiss.IndexIVFAdditiveQuantizerFastScan(index) + D1, I1 = indexfs.search(ds.get_queries(), 1) + + nq = Iref.shape[0] + recall_ref = (Iref == gt).sum() / nq + recall1 = (I1 == gt).sum() / nq + print(recall_ref, recall1) + assert abs(recall_ref - recall1) < 0.02 + + def test_from_ivfaq(self): + for implem in 0, 1, 2: + self.subtest_from_ivfaq(implem) + + def subtest_factory(self, aq, M, bbs, st, r='r'): + """ + Format: IVF{nlist},{AQ}{M}x4fs{r}_{bbs}_N{st} + + nlist (int): number of inverted lists + AQ (str): `LSQ` or `RQ` + M (int): number of sub-quantizers + bbs (int): build block size + st (str): search type, `lsq2x4` or `rq2x4` + r (str): `r` or ``, by_residual or not + """ + AQ = faiss.AdditiveQuantizer + nlist, d = 128, 16 + + if bbs > 0: + index = faiss.index_factory( + d, f'IVF{nlist},{aq}{M}x4fs{r}_{bbs}_N{st}2x4') + else: + index = faiss.index_factory( + d, f'IVF{nlist},{aq}{M}x4fs{r}_N{st}2x4') + bbs = 32 + + assert index.nlist == nlist + assert index.bbs == bbs + aq = faiss.downcast_AdditiveQuantizer(index.aq) + assert aq.M == M + + if aq == 'LSQ': + assert isinstance(aq, faiss.LocalSearchQuantizer) + if aq == 'RQ': + assert isinstance(aq, faiss.ResidualQuantizer) + + if st == 'lsq': + assert aq.search_type == AQ.ST_norm_lsq2x4 + if st == 'rq': + assert aq.search_type == AQ.ST_norm_rq2x4 + + assert index.by_residual == (r == 'r') + + def test_factory(self): + self.subtest_factory('LSQ', 16, 64, 'lsq') + self.subtest_factory('LSQ', 16, 64, 'rq') + self.subtest_factory('RQ', 16, 64, 'rq') + self.subtest_factory('RQ', 16, 64, 'lsq') + self.subtest_factory('LSQ', 64, 0, 'lsq') + + self.subtest_factory('LSQ', 64, 0, 'lsq', r='') + + def subtest_io(self, factory_str): + d = 8 + ds = datasets.SyntheticDataset(d, 1000, 2000, 1000) + + index = faiss.index_factory(d, factory_str) + index.train(ds.get_train()) + index.add(ds.get_database()) + D1, I1 = index.search(ds.get_queries(), 1) + + fd, fname = tempfile.mkstemp() + os.close(fd) + try: + faiss.write_index(index, fname) + index2 = faiss.read_index(fname) + D2, I2 = index2.search(ds.get_queries(), 1) + np.testing.assert_array_equal(I1, I2) + finally: + if os.path.exists(fname): + os.unlink(fname) + + def test_io(self): + self.subtest_io('IVF16,LSQ4x4fs_Nlsq2x4') + self.subtest_io('IVF16,LSQ4x4fs_Nrq2x4') + self.subtest_io('IVF16,RQ4x4fs_Nrq2x4') + self.subtest_io('IVF16,RQ4x4fs_Nlsq2x4') + + +# add more tests programatically + +for byr in True, False: + for implem in 0, 10, 11, 12, 13: + for metric in 'L2', 'IP': + setattr( + TestIVFAQFastScan, + f"test_accuracy_{metric}_RQ_implem{implem}_residual{byr}", + lambda self: + self.subtest_accuracy('RQ', 'rq', byr, implem, metric) + ) + setattr( + TestIVFAQFastScan, + f"test_accuracy_{metric}_LSQ_implem{implem}_residual{byr}", + lambda self: + self.subtest_accuracy('LSQ', 'lsq', byr, implem, metric) + ) + + setattr( + TestIVFAQFastScan, + f"test_rescale_accuracy_LSQ_implem{implem}_residual{byr}", + lambda self: + self.subtest_rescale_accuracy('LSQ', 'lsq', byr, implem) + ) + + setattr( + TestIVFAQFastScan, + f"test_rescale_accuracy_RQ_implem{implem}_residual{byr}", + lambda self: + self.subtest_rescale_accuracy('RQ', 'rq', byr, implem) + )