Add reranked types to IndexSQFastScan and fix flat RaBitQ filtering (#5661)

Summary:
Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5661

Two gaps in the flat fast-scan indexes, both found while benchmarking
IDSelector filtering.

`IndexSQFastScan` accepted only the native 4-bit quantizers. It now accepts
the reranked types as well, the same set `IndexIVFSQFastScan` handles:
QT_6bit, QT_8bit, QT_8bit_uniform, QT_8bit_direct and QT_8bit_direct_signed.
Each code is SPLIT: the SIMD scan reads the top 4 bits, and `lo_codes` keeps
the rest for rerank. The total is the quantizer's own width, so `SQ8fs` holds
8 bits per dimension, the same as plain `SQ8`. A flat index needs no direct
map for this, unlike the IVF form, because the position is the id. `SQ8fs` and
`SQ6fs` now resolve in `index_factory`.

`IndexRaBitQFastScan` forwarded an `IDSelector` through every layer but its
kernel never applied it, so a filtered search returned ids outside the
admitted set. `RaBitQHeapHandler::handle` now tests the selector per code,
after the SIMD block, which is where the PQ and SQ handlers test theirs. The
entry point no longer refuses a parameter object.

The type predicates and the nibble packing move to `impl/sq_fastscan_utils.h`.
The flat and IVF classes then cannot disagree about which quantizers rerank,
or about how a nibble is laid out.

The `ISfs` payload gains the rerank factor and the original codes. Both are
written for every quantizer, and the codes are empty for a native 4-bit index.
`ISfs` has not been released, so no stored index is affected.

Reviewed By: junjieqi

Differential Revision: D121036639

fbshipit-source-id: 40e51fc348a483cf5da947fa7db56ff10f1af773

faiss/IndexIVFSQFastScan.cpp

@@ -6,6 +6,7 @@
  */
 
 #include <faiss/IndexIVFSQFastScan.h>
+#include <faiss/impl/sq_fastscan_utils.h>
 
 #include <omp.h>
 
@@ -29,107 +30,7 @@
 
 namespace faiss {
 
-namespace {
-
-size_t roundup(size_t a, size_t b) {
-    return (a + b - 1) / b * b;
-}
-
-bool is_native_4bit(ScalarQuantizer::QuantizerType qtype) {
-    return qtype == ScalarQuantizer::QT_4bit ||
-            qtype == ScalarQuantizer::QT_4bit_uniform;
-}
-
-bool needs_rerank(ScalarQuantizer::QuantizerType qtype) {
-    return qtype == ScalarQuantizer::QT_6bit ||
-            qtype == ScalarQuantizer::QT_8bit ||
-            qtype == ScalarQuantizer::QT_8bit_uniform ||
-            qtype == ScalarQuantizer::QT_8bit_direct ||
-            qtype == ScalarQuantizer::QT_8bit_direct_signed;
-}
-
-bool is_fallback(ScalarQuantizer::QuantizerType qtype) {
-    return !is_native_4bit(qtype) && !needs_rerank(qtype);
-}
-
-bool is_uniform_range(ScalarQuantizer::QuantizerType qtype) {
-    return qtype == ScalarQuantizer::QT_8bit_uniform ||
-            qtype == ScalarQuantizer::QT_8bit_direct ||
-            qtype == ScalarQuantizer::QT_8bit_direct_signed;
-}
-
-void get_uniform_range(const ScalarQuantizer& sq, float& vmin, float& vdiff) {
-    if (sq.qtype == ScalarQuantizer::QT_8bit_direct) {
-        vmin = 0;
-        vdiff = 255;
-    } else if (sq.qtype == ScalarQuantizer::QT_8bit_direct_signed) {
-        vmin = -128;
-        vdiff = 255;
-    } else {
-        vmin = sq.trained[0];
-        vdiff = sq.trained[1];
-    }
-}
-
-void float_to_4bit_nibbles(
-        const float* x,
-        uint8_t* nibbles,
-        idx_t n,
-        int d,
-        int M2,
-        const ScalarQuantizer& sq) {
-    const bool is_uniform = is_uniform_range(sq.qtype) ||
-            sq.qtype == ScalarQuantizer::QT_4bit_uniform;
-    const float* vmin_arr = nullptr;
-    const float* vdiff_arr = nullptr;
-    float vmin_s = 0, inv_vdiff_s = 0;
-
-    if (is_uniform) {
-        float vdiff;
-        get_uniform_range(sq, vmin_s, vdiff);
-        inv_vdiff_s = (vdiff > 0) ? (1.0f / vdiff) : 0;
-    } else {
-        vmin_arr = sq.trained.data();
-        vdiff_arr = sq.trained.data() + d;
-    }
-
-    const int half = M2 / 2;
-    for (idx_t i = 0; i < n; i++) {
-        const float* xi = x + i * d;
-        uint8_t* dst = nibbles + i * half;
-        memset(dst, 0, half);
-        for (int m = 0; m + 1 < d; m += 2) {
-            float f0, f1;
-            if (is_uniform) {
-                f0 = (xi[m] - vmin_s) * inv_vdiff_s;
-                f1 = (xi[m + 1] - vmin_s) * inv_vdiff_s;
-            } else {
-                f0 = (vdiff_arr[m] > 0) ? (xi[m] - vmin_arr[m]) / vdiff_arr[m]
-                                        : 0;
-                f1 = (vdiff_arr[m + 1] > 0)
-                        ? (xi[m + 1] - vmin_arr[m + 1]) / vdiff_arr[m + 1]
-                        : 0;
-            }
-            uint8_t lo = (uint8_t)std::min(15, std::max(0, (int)(f0 * 15.0f)));
-            uint8_t hi = (uint8_t)std::min(15, std::max(0, (int)(f1 * 15.0f)));
-            dst[m / 2] = lo | (hi << 4);
-        }
-        if (d & 1) {
-            float f;
-            if (is_uniform) {
-                f = (xi[d - 1] - vmin_s) * inv_vdiff_s;
-            } else {
-                f = (vdiff_arr[d - 1] > 0)
-                        ? (xi[d - 1] - vmin_arr[d - 1]) / vdiff_arr[d - 1]
-                        : 0;
-            }
-            dst[(d - 1) / 2] =
-                    (uint8_t)std::min(15, std::max(0, (int)(f * 15.0f)));
-        }
-    }
-}
-
-} // anonymous namespace
+using namespace faiss::sq_fastscan;
 
 // -----------------------------------------------------------------------
 // Constructors
@@ -500,10 +401,10 @@ void IndexIVFSQFastScan::search(
     if (is_fallback(sq.qtype)) {
         IndexIVF::search(n, x, k, distances, labels, params);
     } else if (needs_rerank(sq.qtype)) {
-        idx_t k_coarse = idx_t(k * rerank_factor);
-        if (k_coarse < k) {
-            k_coarse = k;
-        }
+        idx_t k_coarse = idx_t(double(k) * rerank_factor);
+        FAISS_THROW_IF_NOT(k_coarse >= k);
+        FAISS_THROW_IF_NOT_MSG(
+                n <= INT64_MAX / k_coarse, "n * k_coarse would overflow int64");
 
         std::vector<float> coarse_dis(n * k_coarse);
         std::vector<idx_t> coarse_ids(n * k_coarse);
@@ -592,10 +493,10 @@ void IndexIVFSQFastScan::search_preassigned(
                 params,
                 stats);
     } else if (needs_rerank(sq.qtype)) {
-        idx_t k_coarse = idx_t(k * rerank_factor);
-        if (k_coarse < k) {
-            k_coarse = k;
-        }
+        idx_t k_coarse = idx_t(double(k) * rerank_factor);
+        FAISS_THROW_IF_NOT(k_coarse >= k);
+        FAISS_THROW_IF_NOT_MSG(
+                n <= INT64_MAX / k_coarse, "n * k_coarse would overflow int64");
 
         std::vector<float> coarse_dis(n * k_coarse);
         std::vector<idx_t> coarse_ids(n * k_coarse);

faiss/IndexRaBitQFastScan.cpp

@@ -512,8 +512,8 @@ void IndexRaBitQFastScan::search(
         float* distances,
         idx_t* labels,
         const SearchParameters* params) const {
-    // The kernel gains its selector test in the next diff.
-    FAISS_THROW_IF_MSG(params, "search params not supported for this index");
+    // The kernel applies the selector after each SIMD block.
+    const IDSelector* rabitq_sel = params ? params->sel : nullptr;
 
     // Create query factors array on stack - memory managed by caller
     std::vector<rabitq_utils::QueryFactorsData> query_factors_storage(n);
@@ -524,9 +524,11 @@ void IndexRaBitQFastScan::search(
     FastScanDistancePostProcessing context;
     context.query_factors = query_factors_storage.data();
     if (metric_type == METRIC_L2) {
-        search_dispatch_implem<true>(n, x, k, distances, labels, context);
+        search_dispatch_implem<true>(
+                n, x, k, distances, labels, context, rabitq_sel);
     } else {
-        search_dispatch_implem<false>(n, x, k, distances, labels, context);
+        search_dispatch_implem<false>(
+                n, x, k, distances, labels, context, rabitq_sel);
     }
 }
 

faiss/IndexRaBitQFastScan.h

@@ -218,6 +218,10 @@ struct RaBitQHeapHandler
 
         for (size_t i = 0; i < max_vectors; i++) {
             const size_t db_idx = base_db_idx + i;
+            if (this->sel != nullptr &&
+                !this->sel->is_member(static_cast<int64_t>(db_idx))) {
+                continue;
+            }
             const float normalized_distance = d32tab[i] * one_a + bias;
             const uint8_t* base_ptr = aux_base + i * storage_size;
 

faiss/IndexSQFastScan.cpp

@@ -6,6 +6,7 @@
  */
 
 #include <faiss/IndexSQFastScan.h>
+#include <faiss/impl/sq_fastscan_utils.h>
 
 #include <omp.h>
 
@@ -28,18 +29,7 @@
 
 namespace faiss {
 
-namespace {
-
-size_t roundup(size_t a, size_t b) {
-    return (a + b - 1) / b * b;
-}
-
-bool is_native_4bit(ScalarQuantizer::QuantizerType qtype) {
-    return qtype == ScalarQuantizer::QT_4bit ||
-            qtype == ScalarQuantizer::QT_4bit_uniform;
-}
-
-} // anonymous namespace
+using namespace faiss::sq_fastscan;
 
 // -----------------------------------------------------------------------
 // Constructors
@@ -52,10 +42,10 @@ IndexSQFastScan::IndexSQFastScan(
         int bbs_in)
         : sq(d_in, qtype) {
     FAISS_THROW_IF_NOT_MSG(
-            is_native_4bit(qtype),
-            "IndexSQFastScan only supports QT_4bit and QT_4bit_uniform. "
-            "For higher-precision types, use "
-            "IndexRefine(IndexSQFastScan(...), IndexScalarQuantizer(...)).");
+            !is_fallback(qtype),
+            "IndexSQFastScan supports the native 4-bit and the reranked "
+            "types. QT_fp16 and QT_bf16 have no fast-scan path; use "
+            "IndexScalarQuantizer for those.");
     init_fastscan(d_in, d_in, 4, metric, bbs_in);
 }
 
@@ -64,17 +54,39 @@ IndexSQFastScan::IndexSQFastScan() = default;
 IndexSQFastScan::IndexSQFastScan(const IndexScalarQuantizer& orig, int bbs_in)
         : sq(orig.sq) {
     FAISS_THROW_IF_NOT_MSG(
-            is_native_4bit(sq.qtype),
-            "IndexSQFastScan conversion constructor only supports "
-            "QT_4bit and QT_4bit_uniform.");
+            !is_fallback(sq.qtype),
+            "IndexSQFastScan conversion constructor supports the native "
+            "4-bit and the reranked types only.");
     init_fastscan(orig.d, orig.d, 4, orig.metric_type, bbs_in);
     ntotal = orig.ntotal;
     is_trained = orig.is_trained;
 
     ntotal2 = roundup(ntotal, bbs_in);
     codes.resize(ntotal2 * M2 / 2);
-    pq4_pack_codes(
-            orig.codes.data(), ntotal, M, ntotal2, bbs_in, M2, codes.get());
+    if (is_native_4bit(sq.qtype)) {
+        pq4_pack_codes(
+                orig.codes.data(), ntotal, M, ntotal2, bbs_in, M2, codes.get());
+    } else {
+        // High nibbles feed the scan, low nibbles the rerank.
+        std::vector<uint8_t> nibbles(ntotal * (M2 / 2));
+        const int bits = sq_bits(sq.qtype), lo_bits = bits - 4;
+        const size_t half = sq_rerank_size(d, sq.qtype);
+        lo_codes.assign(size_t(ntotal) * half, 0);
+        for (idx_t i = 0; i < ntotal; i++) {
+            const uint8_t* src = orig.codes.data() + size_t(i) * sq.code_size;
+            uint8_t* hi = nibbles.data() + size_t(i) * (M2 / 2);
+            BitstringReader bsr(src, sq.code_size);
+            BitstringWriter bsw(lo_codes.data() + size_t(i) * half, half);
+            for (int m = 0; m < d; m++) {
+                const int code = int(bsr.read(bits));
+                const int h = code >> lo_bits;
+                hi[m / 2] |= (m & 1) ? uint8_t(h << 4) : uint8_t(h);
+                bsw.write(code & ((1 << lo_bits) - 1), lo_bits);
+            }
+        }
+        pq4_pack_codes(
+                nibbles.data(), ntotal, M, ntotal2, bbs_in, M2, codes.get());
+    }
 }
 
 // -----------------------------------------------------------------------
@@ -91,10 +103,31 @@ void IndexSQFastScan::train(idx_t n, const float* x) {
 
 void IndexSQFastScan::add(idx_t n, const float* x) {
     FAISS_THROW_IF_NOT(is_trained);
+    const size_t before = lo_codes.size();
     IndexFastScan::add(n, x);
+    if (!needs_rerank(sq.qtype) || lo_codes.size() != before) {
+        // IndexFastScan::add blocks a large batch and re-enters this override
+        // once per block, so the blocks have already appended the low bits.
+        return;
+    }
+    const int bits = sq_bits(sq.qtype), lo_bits = bits - 4;
+    const size_t stride = sq_rerank_size(d, sq.qtype);
+    lo_codes.resize(before + size_t(n) * stride, 0);
+    const SQRanges ranges(sq, d);
+    for (idx_t i = 0; i < n; i++) {
+        const float* xi = x + i * d;
+        uint8_t* dst = lo_codes.data() + before + size_t(i) * stride;
+        BitstringWriter bsw(dst, stride);
+        for (int m = 0; m < d; m++) {
+            const int code =
+                    sq_code_of(xi[m], ranges.vmin(m), ranges.vdiff(m), bits);
+            bsw.write(code & ((1 << lo_bits) - 1), lo_bits);
+        }
+    }
 }
 
 void IndexSQFastScan::reset() {
+    lo_codes.clear();
     IndexFastScan::reset();
 }
 
@@ -112,7 +145,87 @@ void IndexSQFastScan::search(
     FAISS_THROW_IF_NOT(k > 0);
     FAISS_THROW_IF_NOT(is_trained);
     // The kernel applies the selector itself, so no generic-scanner fallback.
-    IndexFastScan::search(n, x, k, distances, labels, params);
+    if (!needs_rerank(sq.qtype)) {
+        IndexFastScan::search(n, x, k, distances, labels, params);
+        return;
+    }
+
+    // The 4-bit scan is a shortlist. rerank_factor bounds the recall.
+    idx_t k_coarse = idx_t(double(k) * rerank_factor);
+    FAISS_THROW_IF_NOT(k_coarse >= k);
+    FAISS_THROW_IF_NOT_MSG(
+            n <= INT64_MAX / k_coarse, "n * k_coarse would overflow int64");
+    std::vector<float> coarse_dis(n * k_coarse);
+    std::vector<idx_t> coarse_ids(n * k_coarse);
+    IndexFastScan::search(
+            n, x, k_coarse, coarse_dis.data(), coarse_ids.data(), params);
+
+    const bool is_max = metric_type == METRIC_L2;
+#pragma omp parallel
+    {
+        std::unique_ptr<CodePacker> packer(get_CodePacker());
+        std::vector<uint8_t> hi_buf(M2 / 2);
+        std::vector<int> code_buf(d);
+        const size_t lo_stride = sq_rerank_size(d, sq.qtype);
+        const int maxv = (1 << sq_bits(sq.qtype)) - 1;
+        const SQRanges ranges(sq, d);
+#pragma omp for
+        for (idx_t i = 0; i < n; i++) {
+            float* D = distances + i * k;
+            idx_t* I = labels + i * k;
+            if (is_max) {
+                maxheap_heapify(k, D, I);
+            } else {
+                minheap_heapify(k, D, I);
+            }
+            for (idx_t j = 0; j < k_coarse; j++) {
+                idx_t id = coarse_ids[i * k_coarse + j];
+                if (id < 0) {
+                    continue;
+                }
+                float dis;
+                {
+                    // Recombine the two halves into the original code.
+                    packer->unpack_1(codes.get(), id, hi_buf.data());
+                    split_to_codes(
+                            hi_buf.data(),
+                            lo_codes.data() + size_t(id) * lo_stride,
+                            d,
+                            sq.qtype,
+                            code_buf.data());
+                    const float* xi = x + i * d;
+                    float acc = 0;
+                    for (int m = 0; m < d; m++) {
+                        // Matches Codec*bit::decode_component: (c+0.5)/maxv.
+                        const float v = ranges.vmin(m) +
+                                ((code_buf[m] + 0.5f) / float(maxv)) *
+                                        ranges.vdiff(m);
+                        if (is_max) {
+                            const float df = xi[m] - v;
+                            acc += df * df;
+                        } else {
+                            acc += xi[m] * v;
+                        }
+                    }
+                    dis = acc;
+                }
+                if (is_max) {
+                    if (dis < D[0]) {
+                        maxheap_replace_top(k, D, I, dis, id);
+                    }
+                } else {
+                    if (dis > D[0]) {
+                        minheap_replace_top(k, D, I, dis, id);
+                    }
+                }
+            }
+            if (is_max) {
+                maxheap_reorder(k, D, I);
+            } else {
+                minheap_reorder(k, D, I);
+            }
+        }
+    }
 }
 
 // -----------------------------------------------------------------------
@@ -121,7 +234,14 @@ void IndexSQFastScan::search(
 
 void IndexSQFastScan::compute_codes(uint8_t* out_codes, idx_t n, const float* x)
         const {
-    sq.compute_codes(x, out_codes, n);
+    if (is_native_4bit(sq.qtype)) {
+        sq.compute_codes(x, out_codes, n);
+    } else {
+        // The scanned nibble is the top half. lo_codes holds the bottom.
+        std::vector<uint8_t> lo_discard(
+                size_t(n) * sq_rerank_size(d, sq.qtype));
+        float_to_split_codes(x, out_codes, lo_discard.data(), n, d, M2, sq);
+    }
 }
 
 // -----------------------------------------------------------------------
@@ -136,15 +256,21 @@ void IndexSQFastScan::compute_float_LUT(
     const size_t dim = d;
     const size_t ksub_val = 16;
 
-    bool is_uniform = (sq.qtype == ScalarQuantizer::QT_4bit_uniform);
+    bool is_uniform = (sq.qtype == ScalarQuantizer::QT_4bit_uniform) ||
+            is_uniform_range(sq.qtype);
 
     std::vector<float> recon_table(dim * ksub_val);
 
+    const float lut_step = float(1 << sq_rerank_bits(sq.qtype));
+    const float lut_offset = lut_step / 2.0f;
+    const float lut_scale = float((1 << sq_bits(sq.qtype)) - 1);
+
     if (is_uniform) {
-        float vmin = sq.trained[0];
-        float vdiff = sq.trained[1];
+        float vmin, vdiff;
+        get_uniform_range(sq, vmin, vdiff);
         for (size_t c = 0; c < ksub_val; c++) {
-            float recon = vmin + ((c + 0.5f) / 15.0f) * vdiff;
+            float recon =
+                    vmin + ((c * lut_step + lut_offset) / lut_scale) * vdiff;
             for (size_t m = 0; m < dim; m++) {
                 recon_table[m * ksub_val + c] = recon;
             }
@@ -155,8 +281,8 @@ void IndexSQFastScan::compute_float_LUT(
         const float* vdiff = sq.trained.data() + dim;
         for (size_t m = 0; m < dim; m++) {
             for (size_t c = 0; c < ksub_val; c++) {
-                recon_table[m * ksub_val + c] =
-                        vmin[m] + ((c + 0.5f) / 15.0f) * vdiff[m];
+                recon_table[m * ksub_val + c] = vmin[m] +
+                        ((c * lut_step + lut_offset) / lut_scale) * vdiff[m];
             }
         }
     }
@@ -195,11 +321,37 @@ void IndexSQFastScan::sa_decode(idx_t n, const uint8_t* bytes, float* x) const {
     sq.decode(bytes, x, n);
 }
 
+void IndexSQFastScan::fill_sa_code(
+        const CodePacker& packer,
+        idx_t id,
+        uint8_t* scratch,
+        int* values,
+        uint8_t* code_out) const {
+    packer.unpack_1(codes.data(), id, scratch);
+    if (lo_codes.empty()) {
+        memcpy(code_out, scratch, sq.code_size);
+        return;
+    }
+    split_to_codes(
+            scratch,
+            lo_codes.data() + size_t(id) * sq_rerank_size(d, sq.qtype),
+            d,
+            sq.qtype,
+            values);
+    memset(code_out, 0, sq.code_size);
+    BitstringWriter out(code_out, sq.code_size);
+    const int bits = sq_bits(sq.qtype);
+    for (int m = 0; m < d; m++) {
+        out.write(values[m], bits);
+    }
+}
+
 void IndexSQFastScan::reconstruct(idx_t key, float* recons) const {
     FAISS_THROW_IF_NOT(key >= 0 && key < ntotal);
     std::unique_ptr<CodePacker> packer(get_CodePacker());
-    std::vector<uint8_t> flat_code(sq.code_size);
-    packer->unpack_1(codes.data(), key, flat_code.data());
+    std::vector<uint8_t> scratch(M2 / 2), flat_code(sq.code_size);
+    std::vector<int> values(d);
+    fill_sa_code(*packer, key, scratch.data(), values.data(), flat_code.data());
     sq.decode(flat_code.data(), recons, 1);
 }
 
@@ -227,9 +379,15 @@ void IndexSQFastScan::sa_encode(idx_t n, const float* x, uint8_t* bytes) const {
 void IndexSQFastScan::reconstruct_n(idx_t i0, idx_t ni, float* recons) const {
     FAISS_THROW_IF_NOT(i0 >= 0 && i0 + ni <= ntotal);
     std::unique_ptr<CodePacker> packer(get_CodePacker());
-    std::vector<uint8_t> flat_code(sq.code_size);
+    std::vector<uint8_t> scratch(M2 / 2), flat_code(sq.code_size);
+    std::vector<int> values(d);
     for (idx_t i = 0; i < ni; i++) {
-        packer->unpack_1(codes.data(), i0 + i, flat_code.data());
+        fill_sa_code(
+                *packer,
+                i0 + i,
+                scratch.data(),
+                values.data(),
+                flat_code.data());
         sq.decode(flat_code.data(), recons + i * d, 1);
     }
 }
@@ -246,7 +404,7 @@ void IndexSQFastScan::add_sa_codes(
     // Decode to float then re-add via IndexFastScan for SIMD packing
     std::vector<float> recon(n * d);
     sq.decode(code, recon.data(), n);
-    IndexFastScan::add(n, recon.data());
+    add(n, recon.data());
 }
 
 // -----------------------------------------------------------------------
@@ -272,6 +430,17 @@ void IndexSQFastScan::permute_entries(const idx_t* perm) {
                code_sz);
     }
     pq4_pack_codes(flat_new.data(), ntotal, M, ntotal2, bbs, M2, codes.get());
+
+    if (!lo_codes.empty()) {
+        const size_t lo_sz = sq_rerank_size(d, sq.qtype);
+        std::vector<uint8_t> lo_new(size_t(ntotal) * lo_sz);
+        for (idx_t i = 0; i < ntotal; i++) {
+            memcpy(lo_new.data() + size_t(i) * lo_sz,
+                   lo_codes.data() + size_t(perm[i]) * lo_sz,
+                   lo_sz);
+        }
+        lo_codes.swap(lo_new);
+    }
 }
 
 // -----------------------------------------------------------------------
@@ -327,8 +496,15 @@ FlatCodesDistanceComputer* IndexSQFastScan::get_FlatCodesDistanceComputer()
     std::unique_ptr<CodePacker> packer(get_CodePacker());
     size_t cs = sq.code_size;
     std::vector<uint8_t> flat(ntotal * cs);
+    std::vector<uint8_t> scratch(M2 / 2);
+    std::vector<int> values(d);
     for (idx_t i = 0; i < ntotal; i++) {
-        packer->unpack_1(codes.data(), i, flat.data() + i * cs);
+        fill_sa_code(
+                *packer,
+                i,
+                scratch.data(),
+                values.data(),
+                flat.data() + i * cs);
     }
 
     ScalarQuantizer::SQDistanceComputer* dc =
@@ -359,9 +535,15 @@ void IndexSQFastScan::range_search(
     std::vector<uint8_t> flat_codes(ntotal * sq.code_size);
     {
         std::unique_ptr<CodePacker> packer(get_CodePacker());
+        std::vector<uint8_t> scratch(M2 / 2);
+        std::vector<int> values(d);
         for (idx_t i = 0; i < ntotal; i++) {
-            packer->unpack_1(
-                    codes.data(), i, flat_codes.data() + i * sq.code_size);
+            fill_sa_code(
+                    *packer,
+                    i,
+                    scratch.data(),
+                    values.data(),
+                    flat_codes.data() + i * sq.code_size);
         }
     }
 
@@ -436,6 +618,10 @@ void IndexSQFastScan::check_compatible_for_merge(
 
 void IndexSQFastScan::merge_from(Index& otherIndex, idx_t add_id) {
     check_compatible_for_merge(otherIndex);
+    const IndexSQFastScan* other =
+            static_cast<const IndexSQFastScan*>(&otherIndex);
+    lo_codes.insert(
+            lo_codes.end(), other->lo_codes.begin(), other->lo_codes.end());
     IndexFastScan::merge_from(otherIndex, add_id);
 }
 

faiss/IndexSQFastScan.h

@@ -16,25 +16,33 @@
 
 namespace faiss {
 
-/** Fast scan version of IndexScalarQuantizer for native 4-bit types.
+/** Fast scan version of IndexScalarQuantizer.
  *
  * Supported quantizer types:
  *   - QT_4bit, QT_4bit_uniform: native 4-bit codes mapped onto the
- *     PQ4 FastScan SIMD infrastructure (vpshufb / equivalent).
- *     No precision loss.
+ *     PQ4 FastScan SIMD (vpshufb / equivalent).
  *
- * For higher-precision types (QT_8bit, QT_6bit, etc.), use
- * IndexRefine(IndexSQFastScan(QT_4bit), IndexScalarQuantizer(QT_8bit))
- * to get fast-scan with reranking.  For fallback types (QT_fp16,
- * QT_bf16), use IndexScalarQuantizer directly.
+ *   - QT_6bit, QT_8bit, QT_8bit_uniform, QT_8bit_direct,
+ *     QT_8bit_direct_signed: re-quantised to 4-bit for the scan, reranked
+ *     against the originals in orig_codes.
+ *
+ *   - All other types (QT_fp16, QT_bf16): use IndexScalarQuantizer directly.
  */
 struct IndexSQFastScan : IndexFastScan {
     ScalarQuantizer sq;
 
+    /// Overselection ratio for reranking.
+    float rerank_factor = 2;
+
+    /// Low nibbles of the 8-bit codes, two dimensions per byte, indexed by
+    /// id. The high nibbles live in the packed scan codes. Empty for the
+    /// native 4-bit types, which need no rerank.
+    std::vector<uint8_t> lo_codes;
+
     /** Constructor.
      *
      * @param d       dimensionality of input vectors
-     * @param qtype   QT_4bit or QT_4bit_uniform only
+     * @param qtype   any native 4-bit or reranked type (see above)
      * @param metric  distance metric (METRIC_L2 or METRIC_INNER_PRODUCT)
      * @param bbs     block size for SIMD processing (multiple of 32)
      */
@@ -46,7 +54,7 @@ struct IndexSQFastScan : IndexFastScan {
 
     IndexSQFastScan();
 
-    /// Build from an existing IndexScalarQuantizer (must be QT_4bit*)
+    /// Build from an existing IndexScalarQuantizer
     explicit IndexSQFastScan(const IndexScalarQuantizer& orig, int bbs = 32);
 
     void train(idx_t n, const float* x) override;
@@ -108,6 +116,13 @@ struct IndexSQFastScan : IndexFastScan {
 
     size_t sa_code_size() const override;
 
+    void fill_sa_code(
+            const CodePacker& packer,
+            idx_t id,
+            uint8_t* scratch,
+            int* values,
+            uint8_t* code_out) const;
+
     size_t fast_scan_code_size() const override;
 };
 

faiss/impl/index_read.cpp

@@ -2773,6 +2773,8 @@ std::unique_ptr<Index> read_index_up(IOReader* f, int io_flags) {
         READ1(idxsqfs->ntotal2);
         READ1(idxsqfs->M2);
         READVECTOR(idxsqfs->codes);
+        READ1(idxsqfs->rerank_factor);
+        READVECTOR(idxsqfs->lo_codes);
 
         // Restore FastScan base-class fields from the SQ
         idxsqfs->M = idxsqfs->sq.d;

faiss/impl/index_write.cpp

@@ -966,6 +966,8 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) {
         WRITE1(idxsqfs->ntotal2);
         WRITE1(idxsqfs->M2);
         WRITEVECTOR(idxsqfs->codes);
+        WRITE1(idxsqfs->rerank_factor);
+        WRITEVECTOR(idxsqfs->lo_codes);
     } else if (
             const IndexPQFastScan* idxpqfs =
                     dynamic_cast<const IndexPQFastScan*>(idx)) {

faiss/impl/sq_fastscan_utils.h

@@ -0,0 +1,232 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <algorithm>
+#include <cstring>
+
+#include <faiss/impl/ScalarQuantizer.h>
+#include <faiss/utils/hamming.h>
+
+/** Helpers shared by IndexSQFastScan and IndexIVFSQFastScan.
+ *
+ * Both re-quantise a higher-precision scalar quantizer down to 4-bit nibbles
+ * for the SIMD scan and keep the original codes to rerank with, so the
+ * classification of quantizer types and the nibble packing have to agree
+ * between them.
+ */
+
+namespace faiss {
+namespace sq_fastscan {
+
+inline size_t roundup(size_t a, size_t b) {
+    return (a + b - 1) / b * b;
+}
+
+inline bool is_native_4bit(ScalarQuantizer::QuantizerType qtype) {
+    return qtype == ScalarQuantizer::QT_4bit ||
+            qtype == ScalarQuantizer::QT_4bit_uniform;
+}
+
+inline bool needs_rerank(ScalarQuantizer::QuantizerType qtype) {
+    return qtype == ScalarQuantizer::QT_6bit ||
+            qtype == ScalarQuantizer::QT_8bit ||
+            qtype == ScalarQuantizer::QT_8bit_uniform ||
+            qtype == ScalarQuantizer::QT_8bit_direct ||
+            qtype == ScalarQuantizer::QT_8bit_direct_signed;
+}
+
+inline bool is_fallback(ScalarQuantizer::QuantizerType qtype) {
+    return !is_native_4bit(qtype) && !needs_rerank(qtype);
+}
+
+inline bool is_uniform_range(ScalarQuantizer::QuantizerType qtype) {
+    return qtype == ScalarQuantizer::QT_8bit_uniform ||
+            qtype == ScalarQuantizer::QT_8bit_direct ||
+            qtype == ScalarQuantizer::QT_8bit_direct_signed;
+}
+
+inline void get_uniform_range(
+        const ScalarQuantizer& sq,
+        float& vmin,
+        float& vdiff) {
+    if (sq.qtype == ScalarQuantizer::QT_8bit_direct) {
+        vmin = 0;
+        vdiff = 255;
+    } else if (sq.qtype == ScalarQuantizer::QT_8bit_direct_signed) {
+        vmin = -128;
+        vdiff = 255;
+    } else {
+        vmin = sq.trained[0];
+        vdiff = sq.trained[1];
+    }
+}
+
+/// Bits per dimension the quantizer stores.
+inline int sq_bits(ScalarQuantizer::QuantizerType qtype) {
+    if (qtype == ScalarQuantizer::QT_6bit) {
+        return 6;
+    }
+    return is_native_4bit(qtype) ? 4 : 8;
+}
+
+/// Bits left over once the scan has taken its 4. Zero for a native 4-bit
+/// code, which needs no rerank and therefore stores nothing extra.
+inline int sq_rerank_bits(ScalarQuantizer::QuantizerType qtype) {
+    return sq_bits(qtype) - 4;
+}
+
+/// Bytes per vector needed to hold the leftover bits of every dimension.
+inline size_t sq_rerank_size(int d, ScalarQuantizer::QuantizerType qtype) {
+    const int bits = sq_rerank_bits(qtype);
+    return bits > 0 ? (size_t(d) * bits + 7) / 8 : 0;
+}
+
+/// The code for one value at the quantizer's own width.
+inline int sq_code_of(float x, float vmin, float vdiff, int bits) {
+    const int maxv = (1 << bits) - 1;
+    const float f = (vdiff > 0) ? (x - vmin) / vdiff : 0.0f;
+    return std::min(maxv, std::max(0, int(f * float(maxv))));
+}
+
+/// Per-dimension range of the trained quantizer, resolved once for a batch.
+struct SQRanges {
+    SQRanges(const ScalarQuantizer& sq, int d)
+            : uniform(is_uniform_range(sq.qtype) ||
+                      sq.qtype == ScalarQuantizer::QT_4bit_uniform) {
+        if (uniform) {
+            get_uniform_range(sq, vmin_s, vdiff_s);
+        } else {
+            vmin_arr = sq.trained.data();
+            vdiff_arr = sq.trained.data() + d;
+        }
+    }
+
+    float vmin(int m) const {
+        return uniform ? vmin_s : vmin_arr[m];
+    }
+    float vdiff(int m) const {
+        return uniform ? vdiff_s : vdiff_arr[m];
+    }
+
+    const bool uniform;
+    const float* vmin_arr = nullptr;
+    const float* vdiff_arr = nullptr;
+    float vmin_s = 0, vdiff_s = 0;
+};
+
+/** Split each code into the top 4 bits the SIMD kernel scans and the rest,
+ *  which rerank keeps. `hi` is packed two dimensions per byte as
+ *  pq4_pack_codes expects. `lo` is bit-packed at `bits - 4` per dimension.
+ *  Total width is the quantizer's own, not that width plus 4.
+ */
+inline void float_to_split_codes(
+        const float* x,
+        uint8_t* hi,
+        uint8_t* lo,
+        idx_t n,
+        int d,
+        size_t M2,
+        const ScalarQuantizer& sq) {
+    const int bits = sq_bits(sq.qtype);
+    const int lo_bits = bits - 4;
+    const size_t hi_stride = M2 / 2;
+    const size_t lo_stride = sq_rerank_size(d, sq.qtype);
+    const SQRanges ranges(sq, d);
+    for (idx_t i = 0; i < n; i++) {
+        const float* xi = x + i * d;
+        uint8_t* hdst = hi + i * hi_stride;
+        memset(hdst, 0, hi_stride);
+        uint8_t* ldst = lo + i * lo_stride;
+        memset(ldst, 0, lo_stride);
+        BitstringWriter bsw(ldst, lo_stride);
+        for (int m = 0; m < d; m++) {
+            const int code =
+                    sq_code_of(xi[m], ranges.vmin(m), ranges.vdiff(m), bits);
+            const int h = code >> lo_bits;
+            hdst[m / 2] |= (m & 1) ? uint8_t(h << 4) : uint8_t(h);
+            bsw.write(code & ((1 << lo_bits) - 1), lo_bits);
+        }
+    }
+}
+
+/// Rebuild the code values of one vector from its two halves.
+inline void split_to_codes(
+        const uint8_t* hi,
+        const uint8_t* lo,
+        int d,
+        ScalarQuantizer::QuantizerType qtype,
+        int* codes_out) {
+    const int lo_bits = sq_rerank_bits(qtype);
+    BitstringReader bsr(lo, sq_rerank_size(d, qtype));
+    for (int m = 0; m < d; m++) {
+        const int h = (m & 1) ? (hi[m / 2] >> 4) : (hi[m / 2] & 0xF);
+        codes_out[m] = (h << lo_bits) | int(bsr.read(lo_bits));
+    }
+}
+
+inline void float_to_4bit_nibbles(
+        const float* x,
+        uint8_t* nibbles,
+        idx_t n,
+        int d,
+        int M2,
+        const ScalarQuantizer& sq) {
+    const bool is_uniform = is_uniform_range(sq.qtype) ||
+            sq.qtype == ScalarQuantizer::QT_4bit_uniform;
+    const float* vmin_arr = nullptr;
+    const float* vdiff_arr = nullptr;
+    float vmin_s = 0, inv_vdiff_s = 0;
+
+    if (is_uniform) {
+        float vdiff;
+        get_uniform_range(sq, vmin_s, vdiff);
+        inv_vdiff_s = (vdiff > 0) ? (1.0f / vdiff) : 0;
+    } else {
+        vmin_arr = sq.trained.data();
+        vdiff_arr = sq.trained.data() + d;
+    }
+
+    const int half = M2 / 2;
+    for (idx_t i = 0; i < n; i++) {
+        const float* xi = x + i * d;
+        uint8_t* dst = nibbles + i * half;
+        memset(dst, 0, half);
+        for (int m = 0; m + 1 < d; m += 2) {
+            float f0, f1;
+            if (is_uniform) {
+                f0 = (xi[m] - vmin_s) * inv_vdiff_s;
+                f1 = (xi[m + 1] - vmin_s) * inv_vdiff_s;
+            } else {
+                f0 = (vdiff_arr[m] > 0) ? (xi[m] - vmin_arr[m]) / vdiff_arr[m]
+                                        : 0;
+                f1 = (vdiff_arr[m + 1] > 0)
+                        ? (xi[m + 1] - vmin_arr[m + 1]) / vdiff_arr[m + 1]
+                        : 0;
+            }
+            uint8_t lo = (uint8_t)std::min(15, std::max(0, (int)(f0 * 15.0f)));
+            uint8_t hi = (uint8_t)std::min(15, std::max(0, (int)(f1 * 15.0f)));
+            dst[m / 2] = lo | (hi << 4);
+        }
+        if (d & 1) {
+            float f;
+            if (is_uniform) {
+                f = (xi[d - 1] - vmin_s) * inv_vdiff_s;
+            } else {
+                f = (vdiff_arr[d - 1] > 0)
+                        ? (xi[d - 1] - vmin_arr[d - 1]) / vdiff_arr[d - 1]
+                        : 0;
+            }
+            dst[(d - 1) / 2] =
+                    (uint8_t)std::min(15, std::max(0, (int)(f * 15.0f)));
+        }
+    }
+}
+
+} // namespace sq_fastscan
+} // namespace faiss

faiss/index_factory.cpp

@@ -179,8 +179,8 @@ std::map<std::string, ScalarQuantizer::QuantizerType> sq_types = {
 };
 const std::string sq_pattern =
         "(SQ0|SQ4|SQ8|SQ6|SQfp16|SQbf16|SQ8_direct_signed|SQ8_direct|SQtqmse1|SQtqmse2|SQtqmse3|SQtqmse4|SQtqmse8|SQtq2|SQtq3|SQtq4|SQtq5)";
-// Native 4-bit types with "fs" suffix for IndexSQFastScan, optional _bbs
-const std::string sq_fs_pattern = "(SQ4)fs(_[0-9]+)?";
+const std::string sq_fs_pattern =
+        "(SQ4|SQ6|SQ8|SQ8_direct_signed|SQ8_direct)fs(_[0-9]+)?";
 // All SQ types with "fs" suffix for IndexIVFSQFastScan, optional _bbs
 const std::string ivf_sq_fs_pattern =
         "(SQ0|SQ4|SQ8|SQ6|SQfp16|SQbf16|SQ8_direct_signed|SQ8_direct|SQtqmse1|SQtqmse2|SQtqmse3|SQtqmse4|SQtqmse8|SQtq2|SQtq3|SQtq4|SQtq5)fs(_[0-9]+)?";

tests/test_sq_fastscan.py

@@ -48,12 +48,10 @@ class TestSQFastScanConstruction(unittest.TestCase):
         self.assertEqual(index.d, 0)
         self.assertEqual(index.ntotal, 0)
 
-    def test_rejects_non_4bit(self):
+    def test_accepts_reranked_types_rejects_fallback(self):
         d = 64
-        with self.assertRaises(RuntimeError):
-            faiss.IndexSQFastScan(d, SQ.QT_8bit)
-        with self.assertRaises(RuntimeError):
-            faiss.IndexSQFastScan(d, SQ.QT_6bit)
+        for qtype in (SQ.QT_8bit, SQ.QT_6bit):
+            faiss.IndexSQFastScan(d, qtype)
         with self.assertRaises(RuntimeError):
             faiss.IndexSQFastScan(d, SQ.QT_fp16)
         with self.assertRaises(RuntimeError):
@@ -177,11 +175,17 @@ class TestSQFastScanConversionConstructor(unittest.TestCase):
                 r1_fs = (I_fs[:, 0] >= 0).mean()
                 self.assertEqual(r1_sq, r1_fs)
 
-    def test_conversion_rejects_non_4bit(self):
+    def test_conversion_accepts_8bit_and_rejects_fallback(self):
         sq = faiss.IndexScalarQuantizer(self.d, SQ.QT_8bit)
         sq.train(self.ds.get_train())
+        sq.add(self.ds.get_database())
+        converted = faiss.IndexSQFastScan(sq)
+        self.assertEqual(converted.ntotal, sq.ntotal)
+
+        fp16 = faiss.IndexScalarQuantizer(self.d, SQ.QT_fp16)
+        fp16.train(self.ds.get_train())
         with self.assertRaises(RuntimeError):
-            faiss.IndexSQFastScan(sq)
+            faiss.IndexSQFastScan(fp16)
 
 
 class TestSQFastScanReset(unittest.TestCase):
@@ -586,3 +590,120 @@ class TestRaBitQFastScanBiasBound(unittest.TestCase):
         self.assertGreaterEqual(
             fs, ref - 0.1, f"fast scan R@1 {fs:.2f} trails the reference {ref:.2f}"
         )
+
+
+class TestSQFastScanSplitCodePaths(unittest.TestCase):
+    """Paths that rebuild a full code from the split halves.
+
+    QT_4bit hides every bug here, because its code is already one nibble per
+    dimension and sq.code_size equals the fast-scan code size.
+    """
+
+    def _pair(self, qtype, d=32, n=200):
+        rng = np.random.default_rng(0)
+        xb = rng.standard_normal((n, d), dtype=np.float32)
+        ref = faiss.IndexScalarQuantizer(d, qtype)
+        ref.train(xb)
+        ref.add(xb)
+        fs = faiss.IndexSQFastScan(d, qtype)
+        fs.train(xb)
+        fs.add(xb)
+        return xb, ref, fs
+
+    def test_reconstruct_matches_reference(self):
+        for qtype in (SQ.QT_4bit, SQ.QT_6bit, SQ.QT_8bit):
+            _, ref, fs = self._pair(qtype)
+            for i in (0, 7, 199):
+                np.testing.assert_array_equal(
+                    ref.reconstruct(i), fs.reconstruct(i)
+                )
+
+    def test_range_search_matches_reference(self):
+        for qtype in (SQ.QT_4bit, SQ.QT_6bit, SQ.QT_8bit):
+            xb, ref, fs = self._pair(qtype)
+            _, _, iref = ref.range_search(xb[:5], 20.0)
+            _, _, ifs = fs.range_search(xb[:5], 20.0)
+            self.assertEqual(sorted(iref.tolist()), sorted(ifs.tolist()))
+
+    def test_permute_entries_moves_both_halves(self):
+        for qtype in (SQ.QT_6bit, SQ.QT_8bit):
+            _, _, fs = self._pair(qtype)
+            n = fs.ntotal
+            before = fs.reconstruct(5).copy()
+            perm = np.arange(n - 1, -1, -1, dtype="int64")
+            fs.permute_entries(perm)
+            np.testing.assert_array_equal(before, fs.reconstruct(n - 6))
+
+    def test_direct_qtypes_match_reference(self):
+        rng = np.random.default_rng(0)
+        d, n = 32, 200
+        for qtype, lo in (
+            (SQ.QT_8bit_direct, 0),
+            (SQ.QT_8bit_direct_signed, -128),
+        ):
+            xb = rng.integers(lo, lo + 256, size=(n, d)).astype(np.float32)
+            ref = faiss.IndexScalarQuantizer(d, qtype)
+            ref.train(xb)
+            ref.add(xb)
+            fs = faiss.IndexSQFastScan(d, qtype)
+            fs.train(xb)
+            fs.add(xb)
+            _, iref = ref.search(xb[:10], 5)
+            _, ifs = fs.search(xb[:10], 5)
+            hits = sum(len(set(a) & set(b)) for a, b in zip(iref, ifs))
+            self.assertGreaterEqual(hits, 0.9 * iref.size)
+
+    def test_lo_codes_track_ntotal_across_add_paths(self):
+        rng = np.random.default_rng(0)
+        d, qtype = 8, SQ.QT_8bit
+        # over the blocking size of IndexFastScan::add, which re-enters add()
+        n = 70000
+        xb = rng.standard_normal((n, d), dtype=np.float32)
+        ref = faiss.IndexScalarQuantizer(d, qtype)
+        ref.train(xb)
+        ref.add(xb)
+
+        one = faiss.IndexSQFastScan(d, qtype)
+        one.train(xb)
+        one.add(xb)
+        self.assertEqual(one.ntotal, n)
+
+        split = faiss.IndexSQFastScan(d, qtype)
+        split.train(xb)
+        split.add(xb[:40000])
+        split.add(xb[40000:])
+        self.assertEqual(split.ntotal, n)
+
+        for i in (0, 39999, 65536, n - 1):
+            expect = ref.reconstruct(i)
+            np.testing.assert_array_equal(expect, one.reconstruct(i))
+            np.testing.assert_array_equal(expect, split.reconstruct(i))
+
+    def test_add_sa_codes_fills_the_rerank_half(self):
+        xb, ref, _ = self._pair(SQ.QT_8bit)
+        fs = faiss.IndexSQFastScan(ref.d, SQ.QT_8bit)
+        fs.train(xb)
+        fs.add_sa_codes(ref.sa_encode(xb))
+        self.assertEqual(fs.ntotal, ref.ntotal)
+        for i in (0, 7, 199):
+            np.testing.assert_array_equal(ref.reconstruct(i), fs.reconstruct(i))
+
+    def test_merge_from_moves_the_rerank_half(self):
+        rng = np.random.default_rng(0)
+        d, n = 16, 100
+        xb = rng.standard_normal((2 * n, d), dtype=np.float32)
+        parts = []
+        for lo, hi in ((0, n), (n, 2 * n)):
+            idx = faiss.IndexSQFastScan(d, SQ.QT_8bit)
+            idx.train(xb)
+            idx.add(xb[lo:hi])
+            parts.append(idx)
+        parts[0].merge_from(parts[1])
+        self.assertEqual(parts[0].ntotal, 2 * n)
+
+        whole = faiss.IndexSQFastScan(d, SQ.QT_8bit)
+        whole.train(xb)
+        whole.add(xb)
+        np.testing.assert_array_equal(
+            whole.reconstruct_n(0, 2 * n), parts[0].reconstruct_n(0, 2 * n)
+        )