faiss 0.6.1 → 0.6.2

data/vendor/faiss/faiss/impl/scalar_quantizer/sq-dispatch.h

@@ -7,11 +7,15 @@
 
 #pragma once
 
+#include <faiss/impl/RaBitQUtils.h>
 #include <faiss/impl/scalar_quantizer/codecs.h>
 #include <faiss/impl/scalar_quantizer/distance_computers.h>
 #include <faiss/impl/scalar_quantizer/quantizers.h>
 #include <faiss/impl/scalar_quantizer/scanners.h>
 #include <faiss/impl/scalar_quantizer/similarities.h>
+#include <faiss/utils/distances.h>
+#include <faiss/utils/rabitq_simd.h>
+#include <limits>
 
 #ifndef THE_LEVEL_TO_DISPATCH
 #error "THE_LEVEL_TO_DISPATCH should be set on input to this header"
@@ -24,10 +28,324 @@ namespace scalar_quantizer {
 // Define SL as alias for THE_LEVEL_TO_DISPATCH for use in this file
 constexpr SIMDLevel SL = THE_LEVEL_TO_DISPATCH;
 
+/*******************************************************************
+ * TurboQuant SIMD kernel: masked_sum
+ * Compute sum of arr[j] where bit j of the bitmask is set.
+ * NONE specialization is inline; AVX2/AVX512/NEON specializations
+ * live in sq-avx2.cpp / sq-avx512.cpp / sq-neon.cpp.
+ *******************************************************************/
+
+template <SIMDLevel SL0>
+float turboq_masked_sum(const float* arr, const uint8_t* bits, size_t d);
+
+template <>
+inline float turboq_masked_sum<SIMDLevel::NONE>(
+        const float* arr,
+        const uint8_t* bits,
+        size_t d) {
+    float result = 0;
+    for (size_t byte_idx = 0; byte_idx < (d + 7) / 8; byte_idx++) {
+        uint8_t b = bits[byte_idx];
+        size_t base = byte_idx * 8;
+        size_t end = std::min(base + 8, d);
+        for (size_t j = base; j < end; j++) {
+            if (b & (1 << (j - base))) {
+                result += arr[j];
+            }
+        }
+    }
+    return result;
+}
+
+/*******************************************************************
+ * Full TurboQuant DC — lives here because it needs both
+ * quantizers.h (QuantizerTurboQuantFull, SQTurboQFactors) and
+ * similarities.h (Similarity::metric_type). distance_computers.h
+ * can't include quantizers.h due to header ordering.
+ *******************************************************************/
+template <int NBits, class Similarity, SIMDLevel SL2>
+struct DCTurboQuantFull : ScalarQuantizer::TurboQuantRefine::DistanceComputer {
+    using Sim = Similarity;
+    QuantizerTurboQuantFull<NBits, SIMDLevel::NONE> quant;
+    std::vector<float> query;
+    std::vector<float> query_proj;
+    float q_norm_sq = 0;
+    float qjl_coeff = 0;
+    float total_qproj_sum = 0;
+
+    // Pre-screening state
+    const float* threshold_ptr = nullptr;
+    bool prescreen_l2 = false;
+    float qjl_error_coeff = 0;
+    mutable size_t n_total = 0;
+    mutable size_t n_skipped = 0;
+
+    // Integer popcount state
+    uint8_t qb = 0;
+    bool int_qjl = false;
+    std::vector<uint8_t> rearranged_q;
+    float mse_base = 0;
+    float mse_int_scale = 0;
+    float mse_popcnt_scale = 0;
+
+    // Integer QJL popcount state
+    std::vector<uint8_t> rearranged_qproj;
+    float qjl_int_scale = 0;
+    float qjl_popcnt_scale = 0;
+
+    // Scaled centroids for 1-bit MSE fast path (NBits==2)
+    float scaled_c0 = 0;
+    float scaled_c1 = 0;
+    float delta_centroid = 0;
+    float total_q_sum = 0;
+
+    // Multi-bit MSE decomposed coefficients (NBits==3, kMSEBits==2)
+    float mse_multi_base = 0;
+    float mse_coeff_s0 = 0;
+    float mse_coeff_s1 = 0;
+    float mse_coeff_s01 = 0;
+    mutable std::vector<uint8_t> scratch_and;
+
+    DCTurboQuantFull(size_t d, const std::vector<float>& trained)
+            : quant(d, trained) {
+        qjl_coeff = std::sqrt(M_PI / 2.0f) / static_cast<float>(d);
+    }
+
+    void configure(uint8_t qb_in, bool int_qjl_in) override {
+        qb = qb_in;
+        int_qjl = int_qjl_in;
+    }
+
+    void set_prescreen_threshold(const float* ptr, bool l2) override {
+        threshold_ptr = ptr;
+        prescreen_l2 = l2;
+    }
+
+    void clear_prescreen_threshold() override {
+        threshold_ptr = nullptr;
+    }
+
+    void set_query(const float* x) final {
+        q = x;
+        size_t d = quant.d;
+        query.assign(x, x + d);
+        q_norm_sq = fvec_norm_L2sqr(x, d);
+
+        // Project query
+        query_proj.resize(d);
+        quant.project_forward(x, query_proj.data());
+        float inv_sqrt_pd =
+                1.0f / std::sqrt(static_cast<float>(quant.padded_d));
+        for (size_t j = 0; j < d; j++) {
+            query_proj[j] *= inv_sqrt_pd;
+        }
+
+        total_qproj_sum = 0;
+        for (size_t j = 0; j < d; j++) {
+            total_qproj_sum += query_proj[j];
+        }
+
+        // Pre-screening: worst-case L1 bound on QJL error
+        float qproj_l1 = 0;
+        for (size_t j = 0; j < d; j++) {
+            qproj_l1 += std::abs(query_proj[j]);
+        }
+        qjl_error_coeff = qjl_coeff * qproj_l1;
+
+        // Pre-compute for 1-bit MSE fast path
+        if constexpr (NBits == 2) {
+            float inv_sqrt_d = 1.0f / std::sqrt(static_cast<float>(d));
+            scaled_c0 = quant.centroids[0] * inv_sqrt_d;
+            scaled_c1 = quant.centroids[1] * inv_sqrt_d;
+            delta_centroid = scaled_c1 - scaled_c0;
+            total_q_sum = 0;
+            for (size_t j = 0; j < d; j++) {
+                total_q_sum += query[j];
+            }
+
+            // Integer popcount setup
+            if (qb > 0) {
+                size_t byte_size = (d + 7) / 8;
+                float q_min = *std::min_element(query.begin(), query.end());
+                float q_max = *std::max_element(query.begin(), query.end());
+                float q_range = q_max - q_min;
+                if (q_range < 1e-30f) {
+                    q_range = 1e-30f;
+                }
+                float max_val = static_cast<float>((1 << qb) - 1);
+                float scale = max_val / q_range;
+                float delta_q = q_range / max_val;
+
+                rearranged_q.assign(byte_size * qb, 0);
+                for (size_t j = 0; j < d; j++) {
+                    int qval = static_cast<int>(
+                            std::round((query[j] - q_min) * scale));
+                    qval = std::max(
+                            0, std::min(static_cast<int>(max_val), qval));
+                    for (int b = 0; b < qb; b++) {
+                        if (qval & (1 << b)) {
+                            rearranged_q[b * byte_size + j / 8] |=
+                                    (1 << (j % 8));
+                        }
+                    }
+                }
+                mse_base = scaled_c0 * total_q_sum;
+                mse_int_scale = delta_centroid * delta_q;
+                mse_popcnt_scale = delta_centroid * q_min;
+            }
+        }
+
+        // Pre-compute for 2-bit MSE decomposed path (NBits==3)
+        if constexpr (NBits == 3) {
+            float inv_sqrt_d = 1.0f / std::sqrt(static_cast<float>(d));
+            const float* c = quant.centroids;
+            total_q_sum = 0;
+            for (size_t j = 0; j < d; j++) {
+                total_q_sum += query[j];
+            }
+            mse_multi_base = c[0] * inv_sqrt_d * total_q_sum;
+            mse_coeff_s0 = (c[1] - c[0]) * inv_sqrt_d;
+            mse_coeff_s1 = (c[2] - c[0]) * inv_sqrt_d;
+            mse_coeff_s01 = (c[3] - c[2] - c[1] + c[0]) * inv_sqrt_d;
+            scratch_and.resize((d + 7) / 8);
+        }
+
+        // Integer QJL: quantize projected query into bit-planes
+        if (qb > 0 && int_qjl) {
+            size_t byte_size = (d + 7) / 8;
+            float qp_min =
+                    *std::min_element(query_proj.begin(), query_proj.end());
+            float qp_max =
+                    *std::max_element(query_proj.begin(), query_proj.end());
+            float qp_range = qp_max - qp_min;
+            if (qp_range < 1e-30f) {
+                qp_range = 1e-30f;
+            }
+            float max_val = static_cast<float>((1 << qb) - 1);
+            float qp_scale = max_val / qp_range;
+            float delta_qp = qp_range / max_val;
+
+            rearranged_qproj.assign(byte_size * qb, 0);
+            for (size_t j = 0; j < d; j++) {
+                int qval = static_cast<int>(
+                        std::round((query_proj[j] - qp_min) * qp_scale));
+                qval = std::max(0, std::min(static_cast<int>(max_val), qval));
+                for (int b = 0; b < qb; b++) {
+                    if (qval & (1 << b)) {
+                        rearranged_qproj[b * byte_size + j / 8] |=
+                                (1 << (j % 8));
+                    }
+                }
+            }
+            qjl_popcnt_scale = qp_min;
+            qjl_int_scale = delta_qp;
+        }
+
+        n_total = 0;
+        n_skipped = 0;
+    }
+
+    float query_to_code(const uint8_t* code) const final {
+        size_t d = quant.d;
+        float inv_sqrt_d = 1.0f / std::sqrt(static_cast<float>(d));
+        const auto* factors = reinterpret_cast<const SQTurboQFactors*>(
+                code + quant.mse_total_bytes + quant.qjl_plane_bytes);
+        float norm = factors->norm;
+        float gamma = factors->gamma;
+
+        // Stage 1: MSE dot product
+        float mse_dot = 0;
+        if constexpr (NBits == 2) {
+            if (qb > 0) {
+                // Integer popcount path for 1-bit MSE
+                size_t byte_size = (d + 7) / 8;
+                uint64_t and_result = rabitq::bitwise_and_dot_product<SL2>(
+                        rearranged_q.data(), code, byte_size, qb);
+                uint64_t pop = rabitq::popcount<SL2>(code, byte_size);
+                mse_dot = mse_base +
+                        mse_int_scale * static_cast<float>(and_result) +
+                        mse_popcnt_scale * static_cast<float>(pop);
+            } else {
+                // Float path: masked accumulation
+                float pos_sum = turboq_masked_sum<SL2>(query.data(), code, d);
+                mse_dot = scaled_c0 * total_q_sum + delta_centroid * pos_sum;
+            }
+        } else if constexpr (NBits == 3) {
+            // 2-bit MSE: decompose into 3 masked sums over bit-planes.
+            size_t pb = quant.mse_plane_bytes;
+            float s0 = turboq_masked_sum<SL2>(query.data(), code, d);
+            float s1 = turboq_masked_sum<SL2>(query.data(), code + pb, d);
+            for (size_t i = 0; i < pb; i++) {
+                scratch_and[i] = code[i] & code[pb + i];
+            }
+            float s01 =
+                    turboq_masked_sum<SL2>(query.data(), scratch_and.data(), d);
+            mse_dot = mse_multi_base + mse_coeff_s0 * s0 + mse_coeff_s1 * s1 +
+                    mse_coeff_s01 * s01;
+        } else {
+            // kMSEBits > 2: per-dimension fallback
+            for (size_t j = 0; j < d; j++) {
+                uint8_t idx = quant.load_mse_index(code, j);
+                mse_dot += query[j] * quant.centroids[idx] * inv_sqrt_d;
+            }
+        }
+
+        // Pre-screening
+        if (threshold_ptr != nullptr) {
+            n_total++;
+            float bound = qjl_error_coeff * gamma * norm;
+            float mse_ip = norm * mse_dot;
+
+            if constexpr (Similarity::metric_type == METRIC_INNER_PRODUCT) {
+                if (mse_ip + bound <= *threshold_ptr) {
+                    n_skipped++;
+                    return -std::numeric_limits<float>::infinity();
+                }
+            } else {
+                float best_possible =
+                        q_norm_sq + norm * norm - 2.0f * (mse_ip + bound);
+                if (best_possible >= *threshold_ptr) {
+                    n_skipped++;
+                    return std::numeric_limits<float>::infinity();
+                }
+            }
+        }
+
+        // Stage 2: QJL dot product
+        const uint8_t* qjl_code = code + quant.mse_total_bytes;
+        float qjl_dot;
+        if (qb > 0 && int_qjl) {
+            size_t byte_size = (d + 7) / 8;
+            uint64_t and_result = rabitq::bitwise_and_dot_product<SL2>(
+                    rearranged_qproj.data(), qjl_code, byte_size, qb);
+            uint64_t pop = rabitq::popcount<SL2>(qjl_code, byte_size);
+            float pos_sum = qjl_popcnt_scale * static_cast<float>(pop) +
+                    qjl_int_scale * static_cast<float>(and_result);
+            qjl_dot = qjl_coeff * gamma * (2.0f * pos_sum - total_qproj_sum);
+        } else {
+            float pos_sum =
+                    turboq_masked_sum<SL2>(query_proj.data(), qjl_code, d);
+            qjl_dot = qjl_coeff * gamma * (2.0f * pos_sum - total_qproj_sum);
+        }
+
+        float estimated_ip = norm * (mse_dot + qjl_dot);
+
+        if constexpr (Similarity::metric_type == METRIC_INNER_PRODUCT) {
+            return estimated_ip;
+        } else {
+            return q_norm_sq + norm * norm - 2.0f * estimated_ip;
+        }
+    }
+
+    float symmetric_dis(idx_t, idx_t) override {
+        FAISS_THROW_MSG("Not implemented");
+    }
+};
+
 // Returns true if dimension d is compatible with the given SIMD level
 template <SIMDLevel SL2>
 constexpr bool is_dimension_compatible(size_t d) {
-    if constexpr (SL2 == SIMDLevel::AVX512) {
+    if constexpr (SL2 == SIMDLevel::AVX512 || SL2 == SIMDLevel::AVX512_SPR) {
         return d % 16 == 0;
     } else if constexpr (SL2 == SIMDLevel::AVX2 || SL2 == SIMDLevel::ARM_NEON) {
         return d % 8 == 0;
@@ -98,6 +416,14 @@ ScalarQuantizer::SQuantizer* sq_select_quantizer<THE_LEVEL_TO_DISPATCH>(
             return new QuantizerTurboQuantMSE<4, SL>(d, trained);
         case ScalarQuantizer::QT_8bit_tqmse:
             return new QuantizerTurboQuantMSE<8, SL>(d, trained);
+        case ScalarQuantizer::QT_2bit_tq:
+            return new QuantizerTurboQuantFull<2, SL>(d, trained);
+        case ScalarQuantizer::QT_3bit_tq:
+            return new QuantizerTurboQuantFull<3, SL>(d, trained);
+        case ScalarQuantizer::QT_4bit_tq:
+            return new QuantizerTurboQuantFull<4, SL>(d, trained);
+        case ScalarQuantizer::QT_5bit_tq:
+            return new QuantizerTurboQuantFull<5, SL>(d, trained);
         default:
             FAISS_THROW_MSG("unknown qtype");
     }
@@ -171,7 +497,8 @@ SQDistanceComputer* select_distance_computer_body(
             return new DCTemplate<QuantizerBF16<SL2>, Sim, SL2>(d, trained);
 
         case ScalarQuantizer::QT_8bit_direct:
-            if constexpr (SL2 == SIMDLevel::AVX512) {
+            if constexpr (
+                    SL2 == SIMDLevel::AVX512 || SL2 == SIMDLevel::AVX512_SPR) {
                 if (d % 32 == 0) {
                     return new DistanceComputerByte<Sim, SL2>(
                             static_cast<int>(d), trained);
@@ -186,6 +513,12 @@ SQDistanceComputer* select_distance_computer_body(
                     d, trained);
 
         case ScalarQuantizer::QT_8bit_direct_signed:
+            if constexpr (SL2 == SIMDLevel::AVX512_SPR) {
+                if (d % 64 == 0) {
+                    return new DistanceComputerByteSigned<Sim, SL2>(
+                            static_cast<int>(d), trained);
+                }
+            }
             return new DCTemplate<Quantizer8bitDirectSigned<SL2>, Sim, SL2>(
                     d, trained);
         case ScalarQuantizer::QT_0bit:
@@ -206,6 +539,16 @@ SQDistanceComputer* select_distance_computer_body(
         case ScalarQuantizer::QT_8bit_tqmse:
             return new DCTemplate<QuantizerTurboQuantMSE<8, SL2>, Sim, SL2>(
                     d, trained);
+        case ScalarQuantizer::QT_2bit_tq:
+            // FRICTION: bypasses DCTemplate entirely — custom DC
+            // that doesn't fit the Quantizer+Similarity decomposition
+            return new DCTurboQuantFull<2, Sim, SL2>(d, trained);
+        case ScalarQuantizer::QT_3bit_tq:
+            return new DCTurboQuantFull<3, Sim, SL2>(d, trained);
+        case ScalarQuantizer::QT_4bit_tq:
+            return new DCTurboQuantFull<4, Sim, SL2>(d, trained);
+        case ScalarQuantizer::QT_5bit_tq:
+            return new DCTurboQuantFull<5, Sim, SL2>(d, trained);
         default:
             FAISS_THROW_MSG("unknown qtype");
     }
@@ -320,7 +663,9 @@ InvertedListScanner* sq_select_InvertedListScanner<THE_LEVEL_TO_DISPATCH>(
                 return scan.template
                 operator()<DCTemplate<QuantizerBF16<SL2>, Similarity, SL2>>();
             case ScalarQuantizer::QT_8bit_direct:
-                if constexpr (SL2 == SIMDLevel::AVX512) {
+                if constexpr (
+                        SL2 == SIMDLevel::AVX512 ||
+                        SL2 == SIMDLevel::AVX512_SPR) {
                     if (d % 32 == 0) {
                         return scan.template
                         operator()<DistanceComputerByte<Similarity, SL2>>();
@@ -336,6 +681,12 @@ InvertedListScanner* sq_select_InvertedListScanner<THE_LEVEL_TO_DISPATCH>(
                         Similarity,
                         SL2>>();
             case ScalarQuantizer::QT_8bit_direct_signed:
+                if constexpr (SL2 == SIMDLevel::AVX512_SPR) {
+                    if (d % 64 == 0) {
+                        return scan.template operator()<
+                                DistanceComputerByteSigned<Similarity, SL2>>();
+                    }
+                }
                 return scan.template operator()<DCTemplate<
                         Quantizer8bitDirectSigned<SL2>,
                         Similarity,
@@ -368,6 +719,18 @@ InvertedListScanner* sq_select_InvertedListScanner<THE_LEVEL_TO_DISPATCH>(
                         QuantizerTurboQuantMSE<8, SL2>,
                         Similarity,
                         SL2>>();
+            case ScalarQuantizer::QT_2bit_tq:
+                return scan.template
+                operator()<DCTurboQuantFull<2, Similarity, SL2>>();
+            case ScalarQuantizer::QT_3bit_tq:
+                return scan.template
+                operator()<DCTurboQuantFull<3, Similarity, SL2>>();
+            case ScalarQuantizer::QT_4bit_tq:
+                return scan.template
+                operator()<DCTurboQuantFull<4, Similarity, SL2>>();
+            case ScalarQuantizer::QT_5bit_tq:
+                return scan.template
+                operator()<DCTurboQuantFull<5, Similarity, SL2>>();
             default:
                 FAISS_THROW_MSG("unknown qtype");
         }

data/vendor/faiss/faiss/impl/scalar_quantizer/sq-neon.cpp

@@ -9,6 +9,7 @@
 
 #include <faiss/impl/simdlib/simdlib_neon.h>
 
+#include <algorithm>
 #include <cstring>
 
 #include <faiss/impl/scalar_quantizer/codecs.h>
@@ -180,6 +181,12 @@ struct QuantizerTemplate<
                                 xi.data.val[1],
                                 this->vdiff)});
     }
+
+    /// Raw codec decode without denormalization (for pre-decode opt)
+    FAISS_ALWAYS_INLINE simd8float32
+    decode_8_raw(const uint8_t* code, int i) const {
+        return Codec::decode_8_components(code, i);
+    }
 };
 
 template <class Codec>
@@ -219,23 +226,34 @@ struct QuantizerTemplate<
  * TurboQuant MSE quantizer
  **********************************************************/
 
-#define DEFINE_TQMSE_NEON_SPECIALIZATION(NBITS, UNPACK_FN)                  \
-    template <>                                                             \
-    struct QuantizerTurboQuantMSE<NBITS, SIMDLevel::ARM_NEON>               \
-            : QuantizerTurboQuantMSE<NBITS, SIMDLevel::NONE> {              \
-        using Base = QuantizerTurboQuantMSE<NBITS, SIMDLevel::NONE>;        \
-                                                                            \
-        QuantizerTurboQuantMSE(size_t d, const std::vector<float>& trained) \
-                : Base(d, trained) {                                        \
-            assert(d % 8 == 0);                                             \
-        }                                                                   \
-                                                                            \
-        FAISS_ALWAYS_INLINE simd8float32                                    \
-        reconstruct_8_components(const uint8_t* code, int i) const {        \
-            uint8_t indices[8];                                             \
-            UNPACK_FN(code, i, indices);                                    \
-            return gather_8_components(this->centroids, indices);           \
-        }                                                                   \
+// NEON TurboQuantMSE: decode via gather, encode stays scalar.
+// NEON doesn't have movemask so 1-bit encode is also scalar.
+#define DEFINE_TQMSE_NEON_SPECIALIZATION(NBITS, UNPACK_FN)                   \
+    template <>                                                              \
+    struct QuantizerTurboQuantMSE<NBITS, SIMDLevel::ARM_NEON>                \
+            : QuantizerTurboQuantMSE<NBITS, SIMDLevel::NONE> {               \
+        using Base = QuantizerTurboQuantMSE<NBITS, SIMDLevel::NONE>;         \
+                                                                             \
+        QuantizerTurboQuantMSE(size_t d, const std::vector<float>& trained)  \
+                : Base(d, trained) {                                         \
+            assert(d % 8 == 0);                                              \
+        }                                                                    \
+                                                                             \
+        FAISS_ALWAYS_INLINE simd8float32                                     \
+        reconstruct_8_components(const uint8_t* code, int i) const {         \
+            uint8_t indices[8];                                              \
+            UNPACK_FN(code, i, indices);                                     \
+            return gather_8_components(this->centroids, indices);            \
+        }                                                                    \
+                                                                             \
+        void decode_vector(const uint8_t* code, float* x) const final {      \
+            for (size_t i = 0; i < this->d; i += 8) {                        \
+                simd8float32 xi =                                            \
+                        reconstruct_8_components(code, static_cast<int>(i)); \
+                vst1q_f32(x + i, xi.data.val[0]);                            \
+                vst1q_f32(x + i + 4, xi.data.val[1]);                        \
+            }                                                                \
+        }                                                                    \
     }
 
 DEFINE_TQMSE_NEON_SPECIALIZATION(1, unpack_8x1bit_to_u8);
@@ -261,6 +279,15 @@ struct QuantizerTurboQuantMSE<8, SIMDLevel::ARM_NEON>
         std::memcpy(indices, code + static_cast<size_t>(i), sizeof(indices));
         return gather_8_components(this->centroids, indices);
     }
+
+    void decode_vector(const uint8_t* code, float* x) const final {
+        for (size_t i = 0; i < this->d; i += 8) {
+            simd8float32 xi =
+                    reconstruct_8_components(code, static_cast<int>(i));
+            vst1q_f32(x + i, xi.data.val[0]);
+            vst1q_f32(x + i + 4, xi.data.val[1]);
+        }
+    }
 };
 
 /**********************************************************
@@ -397,6 +424,22 @@ struct SimilarityL2<SIMDLevel::ARM_NEON> {
     FAISS_ALWAYS_INLINE float result_8() {
         return horizontal_add(accu8);
     }
+
+    static void adjust_query_for_raw_decode(
+            const float* x,
+            float* q_adj,
+            size_t d,
+            float vmin,
+            float vdiff,
+            float& scale_factor,
+            float& bias) {
+        float inv_vdiff = (vdiff != 0) ? 1.0f / vdiff : 0.0f;
+        for (size_t i = 0; i < d; i++) {
+            q_adj[i] = (x[i] - vmin) * inv_vdiff;
+        }
+        scale_factor = vdiff * vdiff;
+        bias = 0;
+    }
 };
 
 template <>
@@ -431,6 +474,23 @@ struct SimilarityIP<SIMDLevel::ARM_NEON> {
     FAISS_ALWAYS_INLINE float result_8() {
         return horizontal_add(accu8);
     }
+
+    static void adjust_query_for_raw_decode(
+            const float* x,
+            float* q_adj,
+            size_t d,
+            float vmin,
+            float vdiff,
+            float& scale_factor,
+            float& bias) {
+        float sum_q = 0;
+        for (size_t i = 0; i < d; i++) {
+            q_adj[i] = x[i];
+            sum_q += x[i];
+        }
+        scale_factor = vdiff;
+        bias = vmin * sum_q;
+    }
 };
 
 /**********************************************************
@@ -444,8 +504,23 @@ struct DCTemplate<Quantizer, Similarity, SIMDLevel::ARM_NEON>
 
     Quantizer quant;
 
+    // Pre-adjusted query buffer for uniform quantizers
+    std::vector<float> q_adj;
+    float scale_factor = 0;
+    float bias = 0;
+
+    static constexpr bool has_decode_raw() {
+        return requires(const Quantizer& q, const uint8_t* c, int i) {
+            { q.decode_8_raw(c, i) };
+        };
+    }
+
     DCTemplate(size_t d, const std::vector<float>& trained)
-            : quant(d, trained) {}
+            : quant(d, trained) {
+        if constexpr (has_decode_raw()) {
+            q_adj.resize(d);
+        }
+    }
 
     float compute_distance(const float* x, const uint8_t* code) const {
         Similarity sim(x);
@@ -471,6 +546,26 @@ struct DCTemplate<Quantizer, Similarity, SIMDLevel::ARM_NEON>
 
     void set_query(const float* x) final {
         q = x;
+        if constexpr (has_decode_raw()) {
+            Sim::adjust_query_for_raw_decode(
+                    x,
+                    q_adj.data(),
+                    quant.d,
+                    quant.vmin,
+                    quant.vdiff,
+                    scale_factor,
+                    bias);
+        }
+    }
+
+    float query_to_code_predecoded(const uint8_t* code) const {
+        Similarity sim(q_adj.data());
+        sim.begin_8();
+        for (size_t i = 0; i < quant.d; i += 8) {
+            simd8float32 xi = quant.decode_8_raw(code, i);
+            sim.add_8_components(xi);
+        }
+        return bias + scale_factor * sim.result_8();
     }
 
     float symmetric_dis(idx_t i, idx_t j) override {
@@ -479,7 +574,11 @@ struct DCTemplate<Quantizer, Similarity, SIMDLevel::ARM_NEON>
     }
 
     float query_to_code(const uint8_t* code) const final {
-        return compute_distance(q, code);
+        if constexpr (has_decode_raw()) {
+            return query_to_code_predecoded(code);
+        } else {
+            return compute_distance(q, code);
+        }
     }
 
     void query_to_codes_batch_4(
@@ -564,6 +663,32 @@ struct DistanceComputerByte<Similarity, SIMDLevel::ARM_NEON>
     }
 };
 
+/**********************************************************
+ * TurboQuant masked_sum NEON specialization (scalar fallback)
+ **********************************************************/
+
+template <SIMDLevel SL0>
+float turboq_masked_sum(const float* arr, const uint8_t* bits, size_t d);
+
+template <>
+float turboq_masked_sum<SIMDLevel::ARM_NEON>(
+        const float* arr,
+        const uint8_t* bits,
+        size_t d) {
+    float result = 0;
+    for (size_t byte_idx = 0; byte_idx < (d + 7) / 8; byte_idx++) {
+        uint8_t b = bits[byte_idx];
+        size_t base = byte_idx * 8;
+        size_t end = std::min(base + 8, d);
+        for (size_t j = base; j < end; j++) {
+            if (b & (1 << (j - base))) {
+                result += arr[j];
+            }
+        }
+    }
+    return result;
+}
+
 } // namespace scalar_quantizer
 } // namespace faiss