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

@@ -16,3 +16,4 @@
 /c_api/gpu/bin/
 /tests/test
 /tests/gtest/
+faiss/python/swigfaiss_avx2.swig

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')

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

faiss/IndexAdditiveQuantizer.cpp

@@ -144,7 +144,9 @@ void IndexAdditiveQuantizer::search(
                 search_with_LUT<false, AdditiveQuantizer::ST_norm_qint8> (*this, x, rh);
             } else if (aq->search_type == AdditiveQuantizer::ST_norm_qint4) {
                 search_with_LUT<false, AdditiveQuantizer::ST_norm_qint4> (*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<false, AdditiveQuantizer::ST_norm_cqint8> (*this, x, rh);
             } else if (aq->search_type == AdditiveQuantizer::ST_norm_cqint4) {
                 search_with_LUT<false, AdditiveQuantizer::ST_norm_cqint4> (*this, x, rh);

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 <faiss/IndexAdditiveQuantizerFastScan.h>
+
+#include <limits.h>
+#include <cassert>
+#include <memory>
+
+#include <omp.h>
+
+#include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/LocalSearchQuantizer.h>
+#include <faiss/impl/LookupTableScaler.h>
+#include <faiss/impl/ResidualQuantizer.h>
+#include <faiss/impl/pq4_fast_scan.h>
+#include <faiss/utils/quantize_lut.h>
+#include <faiss/utils/utils.h>
+
+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<float[]> del_x;
+    if (x != x_in) {
+        del_x.reset((float*)x);
+    }
+
+    std::vector<float> 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<float> 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<true>(n, x, k, distances, labels, scaler);
+    } else {
+        search_dispatch_implem<false>(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

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 <faiss/IndexAdditiveQuantizer.h>
+#include <faiss/IndexFastScan.h>
+#include <faiss/impl/AdditiveQuantizer.h>
+#include <faiss/utils/AlignedTable.h>
+
+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

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 <faiss/IndexFastScan.h>
+
+#include <limits.h>
+#include <cassert>
+#include <memory>
+
+#include <omp.h>
+
+#include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/LookupTableScaler.h>
+#include <faiss/impl/ResultHandler.h>
+#include <faiss/utils/distances.h>
+#include <faiss/utils/extra_distances.h>
+#include <faiss/utils/hamming.h>
+#include <faiss/utils/random.h>
+#include <faiss/utils/utils.h>
+
+#include <faiss/impl/pq4_fast_scan.h>
+#include <faiss/impl/simd_result_handlers.h>
+#include <faiss/utils/quantize_lut.h>
+
+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<uint8_t> 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 <class C, typename dis_t, class Scaler>
+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<C>(k, heap_dis, heap_ids);
+            heap_push<C>(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<float[]> 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<true>(n, x, k, distances, labels, scaler);
+    } else {
+        search_dispatch_implem<false>(n, x, k, distances, labels, scaler);
+    }
+}
+
+template <bool is_max, class Scaler>
+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<float, int64_t>,
+            CMin<float, int64_t>>::type;
+
+    using C = typename std::
+            conditional<is_max, CMax<uint16_t, int>, CMin<uint16_t, int>>::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<Cfloat>(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<C>(n, x, k, distances, labels, impl, scaler);
+            } else {
+                search_implem_14<C>(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<C>(
+                            i1 - i0, x + i0 * d, k, dis_i, lab_i, impl, scaler);
+                } else {
+                    search_implem_14<C>(
+                            i1 - i0, x + i0 * d, k, dis_i, lab_i, impl, scaler);
+                }
+            }
+        }
+    } else {
+        FAISS_THROW_FMT("invalid implem %d impl=%d", implem, impl);
+    }
+}
+
+template <class Cfloat, class Scaler>
+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<float[]> dis_tables(new float[n * dim12]);
+    compute_float_LUT(dis_tables.get(), n, x);
+
+    std::vector<float> 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<Cfloat>(k, heap_dis, heap_ids);
+
+        estimators_from_tables_generic<Cfloat>(
+                *this,
+                orig_codes,
+                ntotal,
+                dis_tables.get() + i * dim12,
+                k,
+                heap_dis,
+                heap_ids,
+                scaler);
+
+        heap_reorder<Cfloat>(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 <class C, class Scaler>
+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<C>(
+                    i1 - i0,
+                    x + d * i0,
+                    k,
+                    distances + i0 * k,
+                    labels + i0 * k,
+                    impl,
+                    scaler);
+        }
+        return;
+    }
+
+    size_t dim12 = ksub * M2;
+    AlignedTable<uint8_t> quantized_dis_tables(n * dim12);
+    std::unique_ptr<float[]> 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<uint8_t> 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<C> 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<uint16_t> tmp_dis(n * k);
+        std::vector<int32_t> tmp_ids(n * k);
+
+        if (skip & 4) {
+            // skip
+        } else {
+            HeapHandler<C> 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<C> 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 <class C, class Scaler>
+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<C>(
+                    i1 - i0,
+                    x + d * i0,
+                    k,
+                    distances + i0 * k,
+                    labels + i0 * k,
+                    impl,
+                    scaler);
+        }
+        return;
+    }
+
+    size_t dim12 = ksub * M2;
+    AlignedTable<uint8_t> quantized_dis_tables(n * dim12);
+    std::unique_ptr<float[]> 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<uint8_t> LUT(n * dim12);
+    pq4_pack_LUT(n, M2, quantized_dis_tables.get(), LUT.get());
+
+    if (k == 1) {
+        SingleResultHandler<C> 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<uint16_t> tmp_dis(n * k);
+        std::vector<int32_t> tmp_ids(n * k);
+
+        if (skip & 4) {
+            // skip
+        } else if (k > 1) {
+            HeapHandler<C> 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<C> 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<true, NormTableScaler>(
+        idx_t n,
+        const float* x,
+        idx_t k,
+        float* distances,
+        idx_t* labels,
+        const NormTableScaler& scaler) const;
+
+template void IndexFastScan::search_dispatch_implem<false, NormTableScaler>(
+        idx_t n,
+        const float* x,
+        idx_t k,
+        float* distances,
+        idx_t* labels,
+        const NormTableScaler& scaler) const;
+
+} // namespace faiss

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 <faiss/Index.h>
+#include <faiss/utils/AlignedTable.h>
+
+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<uint8_t> 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 <bool is_max, class Scaler>
+    void search_dispatch_implem(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            idx_t* labels,
+            const Scaler& scaler) const;
+
+    template <class Cfloat, class Scaler>
+    void search_implem_234(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            idx_t* labels,
+            const Scaler& scaler) const;
+
+    template <class C, class Scaler>
+    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 <class C, class Scaler>
+    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

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<LocalSearchQuantizer*>(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(

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 <faiss/IndexIVFAdditiveQuantizerFastScan.h>
+
+#include <cassert>
+#include <cinttypes>
+#include <cstdio>
+
+#include <omp.h>
+
+#include <memory>
+
+#include <faiss/impl/AuxIndexStructures.h>
+#include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/LookupTableScaler.h>
+#include <faiss/impl/pq4_fast_scan.h>
+#include <faiss/invlists/BlockInvertedLists.h>
+#include <faiss/utils/distances.h>
+#include <faiss/utils/hamming.h>
+#include <faiss/utils/quantize_lut.h>
+#include <faiss/utils/simdlib.h>
+#include <faiss/utils/utils.h>
+
+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<uint8_t> 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<float[]> del_x;
+    if (x != x_in) {
+        del_x.reset((float*)x);
+    }
+
+    const float* trainset;
+    std::vector<float> residuals(n * d);
+    std::vector<idx_t> 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<float> decoded_x(n * d);
+        std::vector<uint8_t> 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<float> 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<float> 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<float[]> del_x;
+    if (x != x_in) {
+        del_x.reset((float*)x);
+    }
+
+    std::vector<idx_t> coarse_ids(n);
+    std::vector<float> coarse_dis(n);
+    quantizer->search(n, x, 1, coarse_dis.data(), coarse_ids.data());
+
+    AlignedTable<float> dis_tables;
+    AlignedTable<float> 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<float> residuals(n * d);
+        std::vector<float> 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<true>(n, x, k, distances, labels, scaler);
+    } else {
+        search_dispatch_implem<false>(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<q, c> - 2<q, r> + x^2
+
+where q^2 is a constant for all x, <q,c> 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 * <q, codebooks[m][k]>
+
+`-2<q,c>` could be precomputed in `compute_LUT` and store in `biases`.
+if `by_residual=False`, `<q,c>` 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 * <q, cb[m][k]>
+    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<float>& dis_tables,
+        AlignedTable<float>& 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 * <q, c>
+        // NOTE: q^2 is not added to `biases`
+        biases.resize(n * nprobe);
+#pragma omp parallel
+        {
+            std::vector<float> 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

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 <memory>
+
+#include <faiss/IndexIVFAdditiveQuantizer.h>
+#include <faiss/IndexIVFFastScan.h>
+#include <faiss/impl/AdditiveQuantizer.h>
+#include <faiss/impl/LocalSearchQuantizer.h>
+#include <faiss/utils/AlignedTable.h>
+
+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<float>& dis_tables,
+            AlignedTable<float>& 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

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 <faiss/IndexIVFFastScan.h>
+
+#include <cassert>
+#include <cinttypes>
+#include <cstdio>
+
+#include <omp.h>
+
+#include <memory>
+
+#include <faiss/IndexIVFPQ.h>
+#include <faiss/impl/AuxIndexStructures.h>
+#include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/LookupTableScaler.h>
+#include <faiss/impl/pq4_fast_scan.h>
+#include <faiss/impl/simd_result_handlers.h>
+#include <faiss/invlists/BlockInvertedLists.h>
+#include <faiss/utils/distances.h>
+#include <faiss/utils/hamming.h>
+#include <faiss/utils/quantize_lut.h>
+#include <faiss/utils/simdlib.h>
+#include <faiss/utils/utils.h>
+
+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<uint8_t> codes(n * code_size);
+    direct_map.check_can_add(xids);
+    std::unique_ptr<idx_t[]> 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<uint8_t> 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<BlockInvertedLists*>(invlists);
+    FAISS_THROW_IF_NOT_MSG(bil, "only block inverted lists supported");
+
+    // prepare batches
+    std::vector<idx_t> 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<uint8_t> 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 <class C, typename dis_t, class Scaler>
+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<C>(k, heap_dis, heap_ids);
+            heap_push<C>(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<uint8_t>& dis_tables,
+        AlignedTable<uint16_t>& biases,
+        float* normalizers) const {
+    AlignedTable<float> dis_tables_float;
+    AlignedTable<float> 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<true>(n, x, k, distances, labels, scaler);
+    } else {
+        search_dispatch_implem<false>(n, x, k, distances, labels, scaler);
+    }
+}
+
+template <bool is_max, class Scaler>
+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<float, int64_t>,
+            CMin<float, int64_t>>::type;
+
+    using C = typename std::conditional<
+            is_max,
+            CMax<uint16_t, int64_t>,
+            CMin<uint16_t, int64_t>>::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<Cfloat>(n, x, k, distances, labels, scaler);
+    } else if (impl == 2) {
+        search_implem_2<C>(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<C>(
+                        n,
+                        x,
+                        k,
+                        distances,
+                        labels,
+                        impl,
+                        &ndis,
+                        &nlist_visited,
+                        scaler);
+            } else {
+                search_implem_10<C>(
+                        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<C>(
+                            i1 - i0,
+                            x + i0 * d,
+                            k,
+                            dis_i,
+                            lab_i,
+                            impl,
+                            &ndis,
+                            &nlist_visited,
+                            scaler);
+                } else {
+                    search_implem_10<C>(
+                            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 <class C, class Scaler>
+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<idx_t[]> coarse_ids(new idx_t[n * nprobe]);
+    std::unique_ptr<float[]> coarse_dis(new float[n * nprobe]);
+
+    quantizer->search(n, x, nprobe, coarse_dis.get(), coarse_ids.get());
+
+    size_t dim12 = ksub * M;
+    AlignedTable<float> dis_tables;
+    AlignedTable<float> 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<C>(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<C>(
+                    *this,
+                    codes.get(),
+                    ls,
+                    LUT,
+                    ids.get(),
+                    bias,
+                    k,
+                    heap_dis,
+                    heap_ids,
+                    scaler);
+            nlist_visited++;
+            ndis++;
+        }
+        heap_reorder<C>(k, heap_dis, heap_ids);
+    }
+    indexIVF_stats.nq += n;
+    indexIVF_stats.ndis += ndis;
+    indexIVF_stats.nlist += nlist_visited;
+}
+
+template <class C, class Scaler>
+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<idx_t[]> coarse_ids(new idx_t[n * nprobe]);
+    std::unique_ptr<float[]> coarse_dis(new float[n * nprobe]);
+
+    quantizer->search(n, x, nprobe, coarse_dis.get(), coarse_ids.get());
+
+    size_t dim12 = ksub * M2;
+    AlignedTable<uint8_t> dis_tables;
+    AlignedTable<uint16_t> biases;
+    std::unique_ptr<float[]> 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<uint16_t> tmp_dis(k);
+        int64_t* heap_ids = labels + i * k;
+        uint16_t* heap_dis = tmp_dis.data();
+        heap_heapify<C>(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<C>(
+                    *this,
+                    codes.get(),
+                    ls,
+                    LUT,
+                    ids.get(),
+                    bias,
+                    k,
+                    heap_dis,
+                    heap_ids,
+                    scaler);
+
+            nlist_visited++;
+            ndis += ls;
+        }
+        heap_reorder<C>(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 <class C, class Scaler>
+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<C, true>;
+    using ReservoirHC = ReservoirHandler<C, true>;
+    using SingleResultHC = SingleResultHandler<C, true>;
+
+    std::unique_ptr<idx_t[]> coarse_ids(new idx_t[n * nprobe]);
+    std::unique_ptr<float[]> 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<uint8_t> dis_tables;
+    AlignedTable<uint16_t> biases;
+    std::unique_ptr<float[]> 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<uint16_t> tmp_distances(k);
+        for (idx_t i = 0; i < n; i++) {
+            const uint8_t* LUT = nullptr;
+            int qmap1[1] = {0};
+            std::unique_ptr<SIMDResultHandler<C, true>> 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<classHC*>(handler.get())) {                               \
+        auto* res = static_cast<classHC*>(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 <class C, class Scaler>
+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<idx_t[]> coarse_ids(new idx_t[n * nprobe]);
+    std::unique_ptr<float[]> 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<uint8_t> dis_tables;
+    AlignedTable<uint16_t> biases;
+    std::unique_ptr<float[]> 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<QC> 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<SIMDResultHandler<C, true>> handler;
+    AlignedTable<uint16_t> tmp_distances;
+
+    using HeapHC = HeapHandler<C, true>;
+    using ReservoirHC = ReservoirHandler<C, true>;
+    using SingleResultHC = SingleResultHandler<C, true>;
+
+    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<uint16_t> 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<int> q_map(nc), lut_entries(nc);
+        AlignedTable<uint8_t> 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<classHC*>(handler.get())) {                           \
+        auto* res = static_cast<classHC*>(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<ReservoirHC*>(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<uint8_t> 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<float> 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<uint8_t> 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<true, NormTableScaler>(
+        idx_t n,
+        const float* x,
+        idx_t k,
+        float* distances,
+        idx_t* labels,
+        const NormTableScaler& scaler) const;
+
+template void IndexIVFFastScan::search_dispatch_implem<false, NormTableScaler>(
+        idx_t n,
+        const float* x,
+        idx_t k,
+        float* distances,
+        idx_t* labels,
+        const NormTableScaler& scaler) const;
+
+} // namespace faiss

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 <memory>
+
+#include <faiss/IndexIVF.h>
+#include <faiss/utils/AlignedTable.h>
+
+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<float>& dis_tables,
+            AlignedTable<float>& biases) const = 0;
+
+    void compute_LUT_uint8(
+            size_t n,
+            const float* x,
+            const idx_t* coarse_ids,
+            const float* coarse_dis,
+            AlignedTable<uint8_t>& dis_tables,
+            AlignedTable<uint16_t>& 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 <bool is_max, class Scaler>
+    void search_dispatch_implem(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            idx_t* labels,
+            const Scaler& scaler) const;
+
+    template <class C, class Scaler>
+    void search_implem_1(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            idx_t* labels,
+            const Scaler& scaler) const;
+
+    template <class C, class Scaler>
+    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 <class C, class Scaler>
+    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 <class C, class Scaler>
+    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

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<uint8_t> codes(n * code_size);
-
-    FAISS_THROW_IF_NOT(is_trained);
-    direct_map.check_can_add(xids);
-
-    std::unique_ptr<idx_t[]> 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<uint8_t> 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<BlockInvertedLists*>(invlists);
-    FAISS_THROW_IF_NOT_MSG(bil, "only block inverted lists supported");
-
-    // prepare batches
-    std::vector<idx_t> 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<uint8_t> 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 <class C, typename dis_t>
-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<C>(k, heap_dis, heap_ids);
-            heap_push<C>(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<float>& dis_tables,
         AlignedTable<float>& 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<uint8_t>& dis_tables,
-        AlignedTable<uint16_t>& biases,
-        float* normalizers) const {
-    const IndexIVFPQFastScan& ivfpq = *this;
-    AlignedTable<float> dis_tables_float;
-    AlignedTable<float> 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 <bool is_max>
-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<float, int64_t>,
-            CMin<float, int64_t>>::type;
-
-    using C = typename std::conditional<
-            is_max,
-            CMax<uint16_t, int64_t>,
-            CMin<uint16_t, int64_t>>::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<Cfloat>(n, x, k, distances, labels);
-    } else if (impl == 2) {
-        search_implem_2<C>(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<C>(
-                        n,
-                        x,
-                        k,
-                        distances,
-                        labels,
-                        impl,
-                        &ndis,
-                        &nlist_visited);
-            } else {
-                search_implem_10<C>(
-                        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<C>(
-                            i1 - i0,
-                            x + i0 * d,
-                            k,
-                            dis_i,
-                            lab_i,
-                            impl,
-                            &ndis,
-                            &nlist_visited);
-                } else {
-                    search_implem_10<C>(
-                            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<true>(n, x, k, distances, labels);
-    } else {
-        search_dispatch_implem<false>(n, x, k, distances, labels);
-    }
-}
-
-template <class C>
-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<idx_t[]> coarse_ids(new idx_t[n * nprobe]);
-    std::unique_ptr<float[]> 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<float> dis_tables;
-    AlignedTable<float> 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<C>(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<C>(
-                    pq,
-                    pq.nbits,
-                    codes.get(),
-                    ls,
-                    LUT,
-                    ids.get(),
-                    bias,
-                    k,
-                    heap_dis,
-                    heap_ids);
-            nlist_visited++;
-            ndis++;
-        }
-        heap_reorder<C>(k, heap_dis, heap_ids);
-    }
-    indexIVF_stats.nq += n;
-    indexIVF_stats.ndis += ndis;
-    indexIVF_stats.nlist += nlist_visited;
-}
-
-template <class C>
-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<idx_t[]> coarse_ids(new idx_t[n * nprobe]);
-    std::unique_ptr<float[]> coarse_dis(new float[n * nprobe]);
-
-    quantizer->search(n, x, nprobe, coarse_dis.get(), coarse_ids.get());
-
-    size_t dim12 = pq.ksub * M2;
-    AlignedTable<uint8_t> dis_tables;
-    AlignedTable<uint16_t> biases;
-    std::unique_ptr<float[]> 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<uint16_t> tmp_dis(k);
-        int64_t* heap_ids = labels + i * k;
-        uint16_t* heap_dis = tmp_dis.data();
-        heap_heapify<C>(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<C>(
-                    pq,
-                    pq.nbits,
-                    codes.get(),
-                    ls,
-                    LUT,
-                    ids.get(),
-                    bias,
-                    k,
-                    heap_dis,
-                    heap_ids);
-
-            nlist_visited++;
-            ndis += ls;
-        }
-        heap_reorder<C>(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 <class C>
-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<C, true>;
-    using ReservoirHC = ReservoirHandler<C, true>;
-    using SingleResultHC = SingleResultHandler<C, true>;
-
-    std::unique_ptr<idx_t[]> coarse_ids(new idx_t[n * nprobe]);
-    std::unique_ptr<float[]> 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<uint8_t> dis_tables;
-    AlignedTable<uint16_t> biases;
-    std::unique_ptr<float[]> 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<uint16_t> tmp_distances(k);
-        for (idx_t i = 0; i < n; i++) {
-            const uint8_t* LUT = nullptr;
-            int qmap1[1] = {0};
-            std::unique_ptr<SIMDResultHandler<C, true>> 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<classHC*>(handler.get())) {                       \
-        auto* res = static_cast<classHC*>(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 <class C>
-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<idx_t[]> coarse_ids(new idx_t[n * nprobe]);
-    std::unique_ptr<float[]> 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<uint8_t> dis_tables;
-    AlignedTable<uint16_t> biases;
-    std::unique_ptr<float[]> 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<QC> 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<SIMDResultHandler<C, true>> handler;
-    AlignedTable<uint16_t> tmp_distances;
-
-    using HeapHC = HeapHandler<C, true>;
-    using ReservoirHC = ReservoirHandler<C, true>;
-    using SingleResultHC = SingleResultHandler<C, true>;
-
-    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<uint16_t> 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<int> q_map(nc), lut_entries(nc);
-        AlignedTable<uint8_t> 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<classHC*>(handler.get())) {                   \
-        auto* res = static_cast<classHC*>(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<ReservoirHC*>(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

faiss/IndexIVFPQFastScan.h

@@ -9,6 +9,7 @@
 
 #include <memory>
 
+#include <faiss/IndexIVFFastScan.h>
 #include <faiss/IndexIVFPQ.h>
 #include <faiss/impl/ProductQuantizer.h>
 #include <faiss/utils/AlignedTable.h>
@@ -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<float> 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<float>& dis_tables,
-            AlignedTable<float>& biases) const;
+            AlignedTable<float>& biases) const override;
 
-    void compute_LUT_uint8(
-            size_t n,
-            const float* x,
-            const idx_t* coarse_ids,
-            const float* coarse_dis,
-            AlignedTable<uint8_t>& dis_tables,
-            AlignedTable<uint16_t>& biases,
-            float* normalizers) const;
-
-    // internal search funcs
-
-    template <bool is_max>
-    void search_dispatch_implem(
-            idx_t n,
-            const float* x,
-            idx_t k,
-            float* distances,
-            idx_t* labels) const;
-
-    template <class C>
-    void search_implem_1(
-            idx_t n,
-            const float* x,
-            idx_t k,
-            float* distances,
-            idx_t* labels) const;
-
-    template <class C>
-    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 <class C>
-    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 <class C>
-    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

faiss/IndexPQFastScan.cpp

@@ -14,17 +14,11 @@
 #include <omp.h>
 
 #include <faiss/impl/FaissAssert.h>
-#include <faiss/utils/random.h>
+#include <faiss/impl/pq4_fast_scan.h>
 #include <faiss/utils/utils.h>
 
-#include <faiss/impl/pq4_fast_scan.h>
-#include <faiss/impl/simd_result_handlers.h>
-#include <faiss/utils/quantize_lut.h>
-
 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<uint8_t> 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 <class C, typename dis_t>
-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<C>(k, heap_dis, heap_ids);
-            heap_push<C>(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<float[]> 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<true>(n, x, k, distances, labels);
-    } else {
-        search_dispatch_implem<false>(n, x, k, distances, labels);
-    }
-}
-
-template <bool is_max>
-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<float, int64_t>,
-            CMin<float, int64_t>>::type;
-
-    using C = typename std::
-            conditional<is_max, CMax<uint16_t, int>, CMin<uint16_t, int>>::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<float[]> 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<float> 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<Cfloat>(k, heap_dis, heap_ids);
-
-            pq_estimators_from_tables_generic<Cfloat>(
-                    pq,
-                    pq.nbits,
-                    orig_codes,
-                    ntotal,
-                    dis_tables.get() + i * dim12,
-                    k,
-                    heap_dis,
-                    heap_ids);
-
-            heap_reorder<Cfloat>(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<C>(n, x, k, distances, labels, impl);
-            } else {
-                search_implem_14<C>(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<C>(
-                            i1 - i0, x + i0 * d, k, dis_i, lab_i, impl);
-                } else {
-                    search_implem_14<C>(
-                            i1 - i0, x + i0 * d, k, dis_i, lab_i, impl);
-                }
-            }
-        }
-    } else {
-        FAISS_THROW_FMT("invalid implem %d impl=%d", implem, impl);
-    }
-}
-
-template <class C>
-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<C>(
-                    i1 - i0,
-                    x + d * i0,
-                    k,
-                    distances + i0 * k,
-                    labels + i0 * k,
-                    impl);
-        }
-        return;
-    }
-
-    size_t dim12 = pq.ksub * M2;
-    AlignedTable<uint8_t> quantized_dis_tables(n * dim12);
-    std::unique_ptr<float[]> 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<uint8_t> 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<C> 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<uint16_t> tmp_dis(n * k);
-        std::vector<int32_t> tmp_ids(n * k);
-
-        if (skip & 4) {
-            // skip
-        } else {
-            HeapHandler<C> 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<C> 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 <class C>
-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<C>(
-                    i1 - i0,
-                    x + d * i0,
-                    k,
-                    distances + i0 * k,
-                    labels + i0 * k,
-                    impl);
-        }
-        return;
-    }
-
-    size_t dim12 = pq.ksub * M2;
-    AlignedTable<uint8_t> quantized_dis_tables(n * dim12);
-    std::unique_ptr<float[]> 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<uint8_t> LUT(n * dim12);
-    pq4_pack_LUT(n, M2, quantized_dis_tables.get(), LUT.get());
-
-    if (k == 1) {
-        SingleResultHandler<C> 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<uint16_t> tmp_dis(n * k);
-        std::vector<int32_t> tmp_ids(n * k);
-
-        if (skip & 4) {
-            // skip
-        } else if (k > 1) {
-            HeapHandler<C> 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<C> 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);
     }
 }
 

faiss/IndexPQFastScan.h

@@ -7,6 +7,7 @@
 
 #pragma once
 
+#include <faiss/IndexFastScan.h>
 #include <faiss/IndexPQ.h>
 #include <faiss/impl/ProductQuantizer.h>
 #include <faiss/utils/AlignedTable.h>
@@ -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<uint8_t> 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 <bool is_max>
-    void search_dispatch_implem(
-            idx_t n,
-            const float* x,
-            idx_t k,
-            float* distances,
-            idx_t* labels) const;
-
-    template <class C>
-    void search_implem_2(
-            idx_t n,
-            const float* x,
-            idx_t k,
-            float* distances,
-            idx_t* labels) const;
-
-    template <class C>
-    void search_implem_12(
-            idx_t n,
-            const float* x,
-            idx_t k,
-            float* distances,
-            idx_t* labels,
-            int impl) const;
-
-    template <class C>
-    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

faiss/impl/AdditiveQuantizer.cpp

@@ -17,7 +17,10 @@
 
 #include <algorithm>
 
+#include <faiss/Clustering.h>
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/LocalSearchQuantizer.h>
+#include <faiss/impl/ResidualQuantizer.h>
 #include <faiss/utils/Heap.h>
 #include <faiss/utils/distances.h>
 #include <faiss/utils/hamming.h>
@@ -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<AdditiveQuantizer> 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<float> 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<float> 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<float> 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 {

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<float> 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(

faiss/impl/LocalSearchQuantizer.cpp

@@ -15,7 +15,6 @@
 
 #include <algorithm>
 
-#include <faiss/Clustering.h>
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/utils/distances.h>
@@ -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();

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
      *

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 <cstdint>
+#include <cstdlib>
+
+#include <faiss/utils/simdlib.h>
+
+/*******************************************
+ * 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 <class dist_t>
+    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 <class dist_t>
+    inline dist_t scale_one(const dist_t& x) const {
+        return x * scale_int;
+    }
+};
+
+} // namespace faiss

faiss/impl/ResidualQuantizer.cpp

@@ -16,11 +16,6 @@
 #include <cstring>
 #include <memory>
 
-#include <faiss/impl/FaissAssert.h>
-#include <faiss/impl/ResidualQuantizer.h>
-#include <faiss/utils/utils.h>
-
-#include <faiss/Clustering.h>
 #include <faiss/IndexFlat.h>
 #include <faiss/VectorTransform.h>
 #include <faiss/impl/AuxIndexStructures.h>
@@ -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(

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
      *

faiss/impl/index_read.cpp

@@ -26,10 +26,12 @@
 
 #include <faiss/Index2Layer.h>
 #include <faiss/IndexAdditiveQuantizer.h>
+#include <faiss/IndexAdditiveQuantizerFastScan.h>
 #include <faiss/IndexFlat.h>
 #include <faiss/IndexHNSW.h>
 #include <faiss/IndexIVF.h>
 #include <faiss/IndexIVFAdditiveQuantizer.h>
+#include <faiss/IndexIVFAdditiveQuantizerFastScan.h>
 #include <faiss/IndexIVFFlat.h>
 #include <faiss/IndexIVFPQ.h>
 #include <faiss/IndexIVFPQFastScan.h>
@@ -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<std::vector<Index::idx_t>> 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(

faiss/impl/index_write.cpp

@@ -27,10 +27,12 @@
 
 #include <faiss/Index2Layer.h>
 #include <faiss/IndexAdditiveQuantizer.h>
+#include <faiss/IndexAdditiveQuantizerFastScan.h>
 #include <faiss/IndexFlat.h>
 #include <faiss/IndexHNSW.h>
 #include <faiss/IndexIVF.h>
 #include <faiss/IndexIVFAdditiveQuantizer.h>
+#include <faiss/IndexIVFAdditiveQuantizerFastScan.h>
 #include <faiss/IndexIVFFlat.h>
 #include <faiss/IndexIVFPQ.h>
 #include <faiss/IndexIVFPQFastScan.h>
@@ -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<const IndexAdditiveQuantizerFastScan*>(idx)) {
+        auto idxlsqfs =
+                dynamic_cast<const IndexLocalSearchQuantizerFastScan*>(idx);
+        auto idxrqfs = dynamic_cast<const IndexResidualQuantizerFastScan*>(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<const IndexIVFAdditiveQuantizerFastScan*>(
+                            idx)) {
+        auto ivlsqfs =
+                dynamic_cast<const IndexIVFLocalSearchQuantizerFastScan*>(idx);
+        auto ivrqfs =
+                dynamic_cast<const IndexIVFResidualQuantizerFastScan*>(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<const ResidualCoarseQuantizer*>(idx)) {

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 <class ResultHandler>
+template <class ResultHandler, class Scaler>
 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 <class ResultHandler>
+template <class ResultHandler, class Scaler>
 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

faiss/impl/pq4_fast_scan_search_1.cpp

@@ -8,6 +8,7 @@
 #include <faiss/impl/pq4_fast_scan.h>
 
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/LookupTableScaler.h>
 #include <faiss/impl/simd_result_handlers.h>
 
 namespace faiss {
@@ -26,12 +27,13 @@ namespace {
  * writes results in a ResultHandler
  */
 
-template <int NQ, int BB, class ResultHandler>
+template <int NQ, int BB, class ResultHandler, class Scaler>
 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 <int NQ, int BB, class ResultHandler>
+template <int NQ, int BB, class ResultHandler, class Scaler>
 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<NQ, 2 * BB> res2;
-        kernel_accumulate_block<NQ, BB>(nsq, codes, LUT, res2);
+        kernel_accumulate_block<NQ, BB>(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 <class ResultHandler>
+template <class ResultHandler, class Scaler>
 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<NQ, BB>(nb, nsq, codes, LUT, res); \
+#define DISPATCH(NQ, BB)                                                   \
+    case NQ * 1000 + BB:                                                   \
+        accumulate_fixed_blocks<NQ, BB>(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<TH<C, with_id_map>>( \
-            int,                                           \
-            size_t,                                        \
-            int,                                           \
-            int,                                           \
-            const uint8_t*,                                \
-            const uint8_t*,                                \
-            TH<C, with_id_map>&);
+#define INSTANTIATE_ACCUMULATE(TH, C, with_id_map, S)         \
+    template void pq4_accumulate_loop<TH<C, with_id_map>, S>( \
+            int,                                              \
+            size_t,                                           \
+            int,                                              \
+            int,                                              \
+            const uint8_t*,                                   \
+            const uint8_t*,                                   \
+            TH<C, with_id_map>&,                              \
+            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<uint16_t, int>;
 INSTANTIATE_3(Csi, false);

faiss/impl/pq4_fast_scan_search_qbs.cpp

@@ -8,6 +8,7 @@
 #include <faiss/impl/pq4_fast_scan.h>
 
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/LookupTableScaler.h>
 #include <faiss/impl/simd_result_handlers.h>
 #include <faiss/utils/simdlib.h>
 
@@ -27,15 +28,17 @@ namespace {
  * writes results in a ResultHandler
  */
 
-template <int NQ, class ResultHandler>
+template <int NQ, class ResultHandler, class Scaler>
 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 <int QBS, class ResultHandler>
+template <int QBS, class ResultHandler, class Scaler>
 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<SQ, 2> res2;
         const uint8_t* LUT = LUT0;
-        kernel_accumulate_block<Q1>(nsq, codes, LUT, res2);
+        kernel_accumulate_block<Q1>(nsq, codes, LUT, res2, scaler);
         LUT += Q1 * nsq * 16;
         if (Q2 > 0) {
             res2.set_block_origin(Q1, 0);
-            kernel_accumulate_block<Q2>(nsq, codes, LUT, res2);
+            kernel_accumulate_block<Q2>(nsq, codes, LUT, res2, scaler);
             LUT += Q2 * nsq * 16;
         }
         if (Q3 > 0) {
             res2.set_block_origin(Q1 + Q2, 0);
-            kernel_accumulate_block<Q3>(nsq, codes, LUT, res2);
+            kernel_accumulate_block<Q3>(nsq, codes, LUT, res2, scaler);
             LUT += Q3 * nsq * 16;
         }
         if (Q4 > 0) {
             res2.set_block_origin(Q1 + Q2 + Q3, 0);
-            kernel_accumulate_block<Q4>(nsq, codes, LUT, res2);
+            kernel_accumulate_block<Q4>(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 <int NQ, class ResultHandler>
+template <int NQ, class ResultHandler, class Scaler>
 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<NQ, ResultHandler>(
-                nsq, codes + j0 * nsq / 2, LUT, res);
+                nsq, codes + j0 * nsq / 2, LUT, res, scaler);
     }
 }
 
 // non-template version of accumulate kernel -- dispatches dynamically
-template <class ResultHandler>
+template <class ResultHandler, class Scaler>
 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<NQ, ResultHandler>( \
-                ntotal2, nsq, codes, LUT, res);          \
+                ntotal2, nsq, codes, LUT, res, scaler);  \
         return
 
     switch (nq) {
@@ -165,23 +196,24 @@ void accumulate(
 
 } // namespace
 
-template <class ResultHandler>
+template <class ResultHandler, class Scaler>
 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<QBS>(ntotal2, nsq, codes, LUT0, res); \
+#define DISPATCH(QBS)                                                    \
+    case QBS:                                                            \
+        accumulate_q_4step<QBS>(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<NQ, ResultHandler>(nsq, codes, LUT, res); \
+#define DISPATCH(NQ)                                \
+    case NQ:                                        \
+        kernel_accumulate_block<NQ, ResultHandler>( \
+                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<RH>( \
-            int, size_t, int, const uint8_t*, const uint8_t*, RH&);
+#define INSTANTIATE_ACCUMULATE_Q(RH)                            \
+    template void pq4_accumulate_loop_qbs<RH, DummyScaler>(     \
+            int,                                                \
+            size_t,                                             \
+            int,                                                \
+            const uint8_t*,                                     \
+            const uint8_t*,                                     \
+            RH&,                                                \
+            const DummyScaler&);                                \
+    template void pq4_accumulate_loop_qbs<RH, NormTableScaler>( \
+            int,                                                \
+            size_t,                                             \
+            int,                                                \
+            const uint8_t*,                                     \
+            const uint8_t*,                                     \
+            RH&,                                                \
+            const NormTableScaler&);
 
 using Csi = CMax<uint16_t, int>;
 INSTANTIATE_ACCUMULATE_Q(SingleResultHandler<Csi>)
@@ -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) {

faiss/index_factory.cpp

@@ -26,10 +26,12 @@
 
 #include <faiss/Index2Layer.h>
 #include <faiss/IndexAdditiveQuantizer.h>
+#include <faiss/IndexAdditiveQuantizerFastScan.h>
 #include <faiss/IndexFlat.h>
 #include <faiss/IndexHNSW.h>
 #include <faiss/IndexIVF.h>
 #include <faiss/IndexIVFAdditiveQuantizer.h>
+#include <faiss/IndexIVFAdditiveQuantizerFastScan.h>
 #include <faiss/IndexIVFFlat.h>
 #include <faiss/IndexIVFPQ.h>
 #include <faiss/IndexIVFPQFastScan.h>
@@ -149,11 +151,13 @@ std::map<std::string, AdditiveQuantizer::Search_type_t> 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<size_t> aq_parse_nbits(std::string stok) {
     std::vector<size_t> 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<VectorTransform> 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;
 }
 

faiss/python/swigfaiss.swig

@@ -77,7 +77,11 @@ typedef uint64_t size_t;
 #include <faiss/IndexIVFPQR.h>
 #include <faiss/IndexIVFFlat.h>
 
+#include <faiss/IndexFastScan.h>
+#include <faiss/IndexAdditiveQuantizerFastScan.h>
 #include <faiss/IndexPQFastScan.h>
+#include <faiss/IndexIVFFastScan.h>
+#include <faiss/IndexIVFAdditiveQuantizerFastScan.h>
 #include <faiss/IndexIVFPQFastScan.h>
 #include <faiss/utils/quantize_lut.h>
 
@@ -438,7 +442,11 @@ void gpu_sync_all_devices()
 %include  <faiss/IndexIVFPQR.h>
 %include  <faiss/Index2Layer.h>
 
+%include  <faiss/IndexFastScan.h>
+%include  <faiss/IndexAdditiveQuantizerFastScan.h>
 %include  <faiss/IndexPQFastScan.h>
+%include  <faiss/IndexIVFFastScan.h>
+%include  <faiss/IndexIVFAdditiveQuantizerFastScan.h>
 %include  <faiss/IndexIVFPQFastScan.h>
 %include  <faiss/utils/quantize_lut.h>
 
@@ -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  <faiss/index_io.h>

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<float> 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

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

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));

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; });

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) {

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)
+        )

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)
+        )