From ee3a97e053076f751b136af4c95e744618243373 Mon Sep 17 00:00:00 2001 From: Tamim Ehsan Date: Mon, 31 Aug 2026 17:51:48 +0800 Subject: [PATCH 1/4] feat: add base+ex quantization support --- .../rabitqlib/quantization/data_layout.hpp | 47 ++++ include/rabitqlib/quantization/rabitq.hpp | 24 ++ .../rabitqlib/quantization/rabitq_impl.hpp | 219 ++++++++++++++++-- 3 files changed, 276 insertions(+), 14 deletions(-) diff --git a/include/rabitqlib/quantization/data_layout.hpp b/include/rabitqlib/quantization/data_layout.hpp index ddb3b6a..83e0ffa 100644 --- a/include/rabitqlib/quantization/data_layout.hpp +++ b/include/rabitqlib/quantization/data_layout.hpp @@ -143,6 +143,53 @@ struct ConstExDataMap { const T& f_recale_ex_; }; +template +struct BaseDataMap { + public: + explicit BaseDataMap(char* data, size_t padded_dim, size_t base_bits) + : base_code_(reinterpret_cast(data)) + , f_add_(*reinterpret_cast(data + (padded_dim * base_bits / 8))) + , f_rescale_(*(reinterpret_cast(data + (padded_dim * base_bits / 8)) + 1)) + , f_error_(*(reinterpret_cast(data + (padded_dim * base_bits / 8)) + 2)) {} + + static size_t data_bytes(size_t padded_dim, size_t base_bits) { + return (padded_dim * base_bits / 8) + (sizeof(T) * 3); + } + + [[nodiscard]] uint8_t* base_code() { return base_code_; } + [[nodiscard]] T& f_add() { return f_add_; } + [[nodiscard]] T& f_rescale() { return f_rescale_; } + [[nodiscard]] T& f_error() { return f_error_; } + + private: + uint8_t* base_code_; + T& f_add_; + T& f_rescale_; + T& f_error_; +}; + +template +struct ConstBaseDataMap { + public: + explicit ConstBaseDataMap(const char* data, size_t padded_dim, size_t base_bits) + : base_code_(reinterpret_cast(data)) + , f_add_(*reinterpret_cast(data + (padded_dim * base_bits / 8))) + , f_rescale_(*(reinterpret_cast(data + (padded_dim * base_bits / 8)) + 1)) + , f_error_(*(reinterpret_cast(data + (padded_dim * base_bits / 8)) + 2)) { + } + + [[nodiscard]] const uint8_t* base_code() const { return base_code_; } + [[nodiscard]] const T& f_add() const { return f_add_; } + [[nodiscard]] const T& f_rescale() const { return f_rescale_; } + [[nodiscard]] const T& f_error() const { return f_error_; } + + private: + const uint8_t* base_code_; + const T& f_add_; + const T& f_rescale_; + const T& f_error_; +}; + template struct BinDataMap { public: diff --git a/include/rabitqlib/quantization/rabitq.hpp b/include/rabitqlib/quantization/rabitq.hpp index bbd2581..229b184 100644 --- a/include/rabitqlib/quantization/rabitq.hpp +++ b/include/rabitqlib/quantization/rabitq.hpp @@ -265,6 +265,30 @@ inline void quantize_split_single( } } +inline void quantize_xy_single( + const float* data, + const float* centroid, + size_t padded_dim, + size_t base_bits, + size_t ex_bits, + char* base_data, + char* ex_data, + MetricType metric_type = METRIC_L2, + RabitqConfig config = RabitqConfig() +) { + rabitq_impl::xy_bits::split_code_with_factor( + data, + centroid, + padded_dim, + base_bits, + ex_bits, + base_data, + ex_data, + metric_type, + config.t_const + ); +} + template inline void quantize_full_single( const T* data, diff --git a/include/rabitqlib/quantization/rabitq_impl.hpp b/include/rabitqlib/quantization/rabitq_impl.hpp index 4d11461..98d2921 100644 --- a/include/rabitqlib/quantization/rabitq_impl.hpp +++ b/include/rabitqlib/quantization/rabitq_impl.hpp @@ -10,6 +10,7 @@ #include "rabitqlib/defines.hpp" #include "rabitqlib/fastscan/fastscan.hpp" +#include "rabitqlib/quantization/data_layout.hpp" #include "rabitqlib/quantization/pack_excode.hpp" #include "rabitqlib/utils/space.hpp" @@ -558,6 +559,25 @@ inline void ex_bits_compact_code( } // namespace ex_bits namespace total_bits { + +template +inline void combined_code( + const T* residual, size_t dim, size_t total_bits, TP* total_code, double t_const = -1 +) { + const size_t ex_bits = total_bits - 1; + + ConstRowMajorArrayMap res_arr(residual, 1, static_cast(dim)); + RowMajorArrayMap code_arr(total_code, 1, static_cast(dim)); + + if (ex_bits > 0) { + ex_bits::ex_bits_code(residual, dim, ex_bits, total_code, t_const); + } else { + code_arr.setZero(); + } + + code_arr += (res_arr > 0).template cast() * static_cast(1 << ex_bits); +} + template static inline void rabitq_scalar_impl( const T* data, @@ -570,11 +590,11 @@ static inline void rabitq_scalar_impl( double t_const = -1, ScalarQuantizerType scalar_quantizer_type = ScalarQuantizerType::RECONSTRUCTION ) { - std::vector binary_code(dim); size_t ex_bits = total_bits - 1; - RowMajorArray residual_arr = - rabitq_impl::one_bit::one_bit_code(data, centroid, dim, binary_code.data()); + ConstRowMajorArrayMap data_arr(data, 1, static_cast(dim)); + ConstRowMajorArrayMap cent_arr(centroid, 1, static_cast(dim)); + RowMajorArray residual_arr = data_arr - cent_arr; if (l2norm_sqr(residual_arr.data(), dim) == 0) { std::fill(total_code, total_code + dim, static_cast(0)); @@ -583,17 +603,7 @@ static inline void rabitq_scalar_impl( return; } - if (ex_bits > 0) { - ex_bits::ex_bits_code( - residual_arr.data(), dim, ex_bits, total_code, t_const - ); - } - - // merge 2 one_bit code and ex_bits code - for (size_t i = 0; i < dim; ++i) { - const TP sign_code = static_cast(binary_code[i]) << ex_bits; - total_code[i] = ex_bits > 0 ? total_code[i] + sign_code : sign_code; - } + combined_code(residual_arr.data(), dim, total_bits, total_code, t_const); float cb = -(static_cast(1 << ex_bits) - 0.5F); RowMajorArrayMap total_code_arr(total_code, 1, dim); @@ -658,4 +668,185 @@ static inline void rabitq_full_impl( } } } // namespace total_bits + +namespace xy_bits { + +constexpr size_t kMaxCombinedBits = 9; + +static_assert( + ex_bits::kTightStart.size() >= kMaxCombinedBits, + "kTightStart must cover magnitude widths up to kMaxCombinedBits - 1" +); + +inline void validate_bit_size(size_t base_bits, size_t ex_bits) { + if (base_bits < 1 || base_bits > 8) { + std::cerr << "base_bits must be in [1, 8]\n" << std::flush; + exit(1); + } + if (ex_bits > 8) { + std::cerr << "ex_bits must be in [0, 8]\n" << std::flush; + exit(1); + } + if (base_bits + ex_bits > kMaxCombinedBits) { + std::cerr << "base_bits + ex_bits must be in [1, " << kMaxCombinedBits << "]\n" + << std::flush; + exit(1); + } +} + +/** + * @brief Derive the 3 estimation factors for an arbitrary integer code. + * + * one_bit_code_with_factor() and ex_bits_code_with_factor() both compute the + * same three quantities from xu_cb = code + cb, differing only in how the + * code and cb are produced. This is that shared derivation, so a base-only + * code and a combined code can both be turned into factors without + * duplicating the algebra. With cb = -0.5 and a sign-bit code it reproduces + * one_bit_code_with_factor()'s outputs bit-for-bit. + * + * @param residual data - centroid + * @param code integer code, values in [0, 2^bits) + * @param cb offset-binary constant for that code width, -(2^bits - 1)/2 + */ +template +inline void code_factors( + const T* residual, + const T* centroid, + size_t dim, + const int* code, + float cb, + T& f_add, + T& f_rescale, + T& f_error, + MetricType metric_type = METRIC_L2 +) { + ConstRowMajorArrayMap code_arr(code, 1, static_cast(dim)); + RowMajorArray xu_cb = code_arr.template cast() + cb; + + T l2_sqr = l2norm_sqr(residual, dim); + + if (l2_sqr == 0) { + if (metric_type == METRIC_L2) { + f_add = 0; + } else if (metric_type == METRIC_IP) { + f_add = 1; + } else { + std::cerr << "Unsupported metric type in code_factors()\n" << std::flush; + exit(1); + } + f_rescale = 0; + f_error = 0; + return; + } + + T l2_norm = std::sqrt(l2_sqr); + + T ip_resi_xucb = dot_product(residual, xu_cb.data(), dim); + T ip_cent_xucb = dot_product(centroid, xu_cb.data(), dim); + + // corner case + if (ip_resi_xucb == 0) { + ip_resi_xucb = std::numeric_limits::infinity(); + } + + T tmp_error = + l2_norm * kConstEpsilon * + std::sqrt( + (((l2_sqr * l2norm_sqr(xu_cb.data(), dim)) / (ip_resi_xucb * ip_resi_xucb)) - + 1) / + (dim - 1) + ); + + if (metric_type == METRIC_L2) { + f_add = l2_sqr + (2 * l2_sqr * ip_cent_xucb / ip_resi_xucb); + f_rescale = -2 * l2_sqr / ip_resi_xucb; + f_error = 2 * tmp_error; + } else if (metric_type == METRIC_IP) { + f_add = 1 - dot_product(residual, centroid, dim) + + (l2_sqr * ip_cent_xucb / ip_resi_xucb); + f_rescale = -l2_sqr / ip_resi_xucb; + f_error = 1 * tmp_error; + } else { + std::cerr << "Unsupported metric type in code_factors()\n" << std::flush; + exit(1); + } +} + +template +inline void split_code_with_factor( + const T* data, + const T* centroid, + size_t dim, + size_t base_bits, + size_t ex_bits, + char* base_data, + char* ex_data, + MetricType metric_type = METRIC_L2, + double t_const = -1 +) { + validate_bit_size(base_bits, ex_bits); + + ConstRowMajorArrayMap data_arr(data, 1, static_cast(dim)); + ConstRowMajorArrayMap cent_arr(centroid, 1, static_cast(dim)); + RowMajorArray residual_arr = data_arr - cent_arr; + const T* residual = residual_arr.data(); + + const size_t total_bits = base_bits + ex_bits; + std::vector total_code(dim); + total_bits::combined_code( + residual, dim, total_bits, total_code.data(), t_const + ); + + const auto ex_mask = (ex_bits > 0) ? static_cast((1U << ex_bits) - 1) : 0; + std::vector base_code_int(dim); + std::vector ex_raw(ex_bits > 0 ? dim : 0, 0); + for (size_t i = 0; i < dim; ++i) { + base_code_int[i] = total_code[i] >> ex_bits; + if (ex_bits > 0) { + ex_raw[i] = static_cast(total_code[i] & ex_mask); + } + } + + std::vector base_raw(dim); + for (size_t i = 0; i < dim; ++i) { + base_raw[i] = static_cast(base_code_int[i]); + } + + // Base layer first: its factors describe the filter code on its own. + BaseDataMap base_map(base_data, dim, base_bits); + code_factors( + residual, + centroid, + dim, + base_code_int.data(), + -(static_cast((1U << base_bits) - 1) / 2.F), + base_map.f_add(), + base_map.f_rescale(), + base_map.f_error(), + metric_type + ); + ex_bits::packing_rabitqplus_code(base_raw.data(), base_map.base_code(), dim, base_bits); + + if (ex_bits == 0) { + return; + } + + // Refine layer: the same derivation over the combined code. + ExDataMap ex_map(ex_data, dim, ex_bits); + T f_error_ex = 0; + code_factors( + residual, + centroid, + dim, + total_code.data(), + -(static_cast((1U << total_bits) - 1) / 2.F), + ex_map.f_add_ex(), + ex_map.f_rescale_ex(), + f_error_ex, + metric_type + ); + + ex_bits::packing_rabitqplus_code(ex_raw.data(), ex_map.ex_code(), dim, ex_bits); +} +} // namespace xy_bits } // namespace rabitqlib::quant::rabitq_impl From 90fa7d732b1757f54f4e3e6308dc184d662f21a3 Mon Sep 17 00:00:00 2001 From: Tamim Ehsan Date: Mon, 31 Aug 2026 17:52:24 +0800 Subject: [PATCH 2/4] feat: add distance estimators for base and ex quantization --- include/rabitqlib/index/estimator.hpp | 57 +++++++++++++++++++++++++++ include/rabitqlib/index/query.hpp | 14 ++++--- 2 files changed, 66 insertions(+), 5 deletions(-) diff --git a/include/rabitqlib/index/estimator.hpp b/include/rabitqlib/index/estimator.hpp index 6d93dff..f75e87a 100644 --- a/include/rabitqlib/index/estimator.hpp +++ b/include/rabitqlib/index/estimator.hpp @@ -307,4 +307,61 @@ inline void split_single_estdist_direct( low_dist = est_dist - (cur_bin.f_error() * g_error); }; +inline void xy_single_base_dist( + const char* base_data, + float (*base_ip_func)(const float*, const uint8_t*, size_t), + const SplitSingleQuery& q_obj, + size_t padded_dim, + size_t base_bits, + float& ip_base, + float& est_dist, + float& low_dist, + float g_add = 0, + float g_error = 0 +) { + ConstBaseDataMap cur_base(base_data, padded_dim, base_bits); + + ip_base = base_ip_func(q_obj.rotated_query(), cur_base.base_code(), padded_dim); + + est_dist = + cur_base.f_add() + g_add + (cur_base.f_rescale() * (ip_base + q_obj.kbase_sumq())); + + low_dist = est_dist - (cur_base.f_error() * g_error); +} + +inline void xy_single_full_dist( + const char* base_data, + const char* ex_data, + float (*base_ip_func)(const float*, const uint8_t*, size_t), + float (*ex_ip_func)(const float*, const uint8_t*, size_t), + const SplitSingleQuery& q_obj, + size_t padded_dim, + size_t base_bits, + size_t ex_bits, + float& est_dist, + float& low_dist, + float& ip_base, + float g_add = 0, + float g_error = 0 +) { + ConstBaseDataMap cur_base(base_data, padded_dim, base_bits); + ip_base = base_ip_func(q_obj.rotated_query(), cur_base.base_code(), padded_dim); + + if (ex_bits == 0) { + est_dist = cur_base.f_add() + g_add + + (cur_base.f_rescale() * (ip_base + q_obj.kbase_sumq())); + low_dist = est_dist - (cur_base.f_error() * g_error); + return; + } + + ConstExDataMap cur_ex(ex_data, padded_dim, ex_bits); + + float ip = (static_cast(1 << ex_bits) * ip_base) + + ex_ip_func(q_obj.rotated_query(), cur_ex.ex_code(), padded_dim); + + est_dist = cur_ex.f_add_ex() + g_add + (cur_ex.f_rescale_ex() * (ip + q_obj.kbxsumq())); + + low_dist = est_dist - (cur_base.f_error() * g_error / static_cast(1 << ex_bits)); +} + } // namespace rabitqlib diff --git a/include/rabitqlib/index/query.hpp b/include/rabitqlib/index/query.hpp index 6e49bf2..754dcd8 100644 --- a/include/rabitqlib/index/query.hpp +++ b/include/rabitqlib/index/query.hpp @@ -116,6 +116,7 @@ class SplitSingleQuery { std::vector QueryBin_; T G_add_; T G_k1xSumq_; + T G_kbaseSumq_; T G_kbxSumq_; T G_error_; T delta_; @@ -129,15 +130,18 @@ class SplitSingleQuery { size_t padded_dim, size_t ex_bits, quant::RabitqConfig config, - size_t metric_type = METRIC_L2 + size_t metric_type = METRIC_L2, + size_t base_bits = 1 ) : rotated_query_(rotated_query), QueryBin_(padded_dim * kNumBits / 64, 0) { float c_1 = -static_cast((1 << 1) - 1) / 2.F; - float c_b = -static_cast((1 << (ex_bits + 1)) - 1) / 2.F; + float c_base = -static_cast((1U << base_bits) - 1) / 2.F; + float c_b = -static_cast((1U << (base_bits + ex_bits)) - 1) / 2.F; T sumq = std::accumulate(rotated_query, rotated_query + padded_dim, static_cast(0)); G_k1xSumq_ = sumq * c_1; + G_kbaseSumq_ = sumq * c_base; G_kbxSumq_ = sumq * c_b; metric_type_ = (metric_type == METRIC_IP) ? METRIC_IP : METRIC_L2; @@ -153,9 +157,6 @@ class SplitSingleQuery { rabitqlib::new_transpose_bin_512( quant_query.data(), QueryBin_.data(), padded_dim, kNumBits ); - - // new_transpose_bin_512 already stores the query in the bit-plane/chunk - // layout consumed by warmup_ip_x0_q_512. } [[nodiscard]] size_t num_bits() const { return kNumBits; } @@ -170,6 +171,9 @@ class SplitSingleQuery { [[nodiscard]] T k1xsumq() const { return G_k1xSumq_; } + // Equals k1xsumq() when base_bits == 1. + [[nodiscard]] T kbase_sumq() const { return G_kbaseSumq_; } + [[nodiscard]] T kbxsumq() const { return G_kbxSumq_; } [[nodiscard]] T g_add() const { return G_add_; } From abd9ddb61c9267b547e7486a3f526336fd1177e2 Mon Sep 17 00:00:00 2001 From: Tamim Ehsan Date: Mon, 31 Aug 2026 17:52:35 +0800 Subject: [PATCH 3/4] feat: add unit tests for XY quantization functionality --- .../quantization/xy_quantization_test.cpp | 571 ++++++++++++++++++ 1 file changed, 571 insertions(+) create mode 100644 tests/unit/rabitqlib/quantization/xy_quantization_test.cpp diff --git a/tests/unit/rabitqlib/quantization/xy_quantization_test.cpp b/tests/unit/rabitqlib/quantization/xy_quantization_test.cpp new file mode 100644 index 0000000..a39458c --- /dev/null +++ b/tests/unit/rabitqlib/quantization/xy_quantization_test.cpp @@ -0,0 +1,571 @@ +#include + +#include +#include +#include + +#include "rabitqlib/defines.hpp" +#include "rabitqlib/index/estimator.hpp" +#include "rabitqlib/index/query.hpp" +#include "rabitqlib/quantization/data_layout.hpp" +#include "rabitqlib/quantization/rabitq.hpp" +#include "rabitqlib/quantization/rabitq_impl.hpp" +#include "rabitqlib/utils/space.hpp" +#include "test_helpers.hpp" + +using namespace rabitqlib; +using namespace rabitq_test; + +namespace { + +std::vector RandomVec(size_t dim, std::mt19937& gen) { + std::normal_distribution dist(0.0F, 1.0F); + std::vector v(dim); + for (auto& x : v) { + x = dist(gen); + } + return v; +} + +// The function writes packed codes into the data blocks and no unpacker exists. +// Recover the raw per-dimension codes through the library's own reader instead +// of reimplementing seven SIMD bit layouts: a one-hot query makes the ex-code +// inner product return exactly code[j]. This asserts the codes as the search +// path actually interprets them, which is the property that matters. +std::vector UnpackViaIpKernel(const uint8_t* packed, size_t dim, size_t bits) { + std::vector out(dim, 0); + if (bits == 0) { + return out; + } + auto ip = select_excode_ipfunc(bits); + std::vector probe(dim, 0.0F); + for (size_t j = 0; j < dim; ++j) { + probe[j] = 1.0F; + out[j] = static_cast(std::lround(ip(probe.data(), packed, dim))); + probe[j] = 0.0F; + } + return out; +} + +// Convenience wrapper: split_code_with_factor with both factor triplets. +struct XySplitResult { + std::vector base_code; + std::vector extra_code; + float f_add_base = 0; + float f_rescale_base = 0; + float f_error_base = 0; + float f_add_full = 0; + float f_rescale_full = 0; +}; + +XySplitResult SplitCode( + const std::vector& data, + const std::vector& centroid, + size_t dim, + size_t base_bits, + size_t extra_bits +) { + XySplitResult res; + res.base_code.resize(dim); + res.extra_code.resize(dim); + + std::vector base_block(BaseDataMap::data_bytes(dim, base_bits), 0); + std::vector extra_block( + extra_bits > 0 ? ExDataMap::data_bytes(dim, extra_bits) : 1, 0 + ); + + // Through the public entry point: rabitq_impl now takes raw buffers, and + // rabitq.hpp is what maps them onto BaseDataMap/ExDataMap. + quant::quantize_xy_single( + data.data(), + centroid.data(), + dim, + base_bits, + extra_bits, + base_block.data(), + extra_bits > 0 ? extra_block.data() : nullptr, + METRIC_L2 + ); + + BaseDataMap base_map(base_block.data(), dim, base_bits); + res.f_add_base = base_map.f_add(); + res.f_rescale_base = base_map.f_rescale(); + res.f_error_base = base_map.f_error(); + res.base_code = UnpackViaIpKernel(base_map.base_code(), dim, base_bits); + + if (extra_bits > 0) { + ExDataMap extra_map(extra_block.data(), dim, extra_bits); + res.f_add_full = extra_map.f_add_ex(); + res.f_rescale_full = extra_map.f_rescale_ex(); + res.extra_code = UnpackViaIpKernel(extra_map.ex_code(), dim, extra_bits); + } else { + res.f_add_full = res.f_add_base; + res.f_rescale_full = res.f_rescale_base; + } + return res; +} + +} // namespace + +// Backward-compat anchor: at base_bits == 1, split_code_with_factor must +// match today's one_bit_code_with_factor + ex_bits_code_with_factor -- codes +// *and* both factor triplets. The base layer's factors are what BinDataMap +// holds in the 1+y path, the full layer's are what ExDataMap holds. +TEST(XyQuantization, BaseBitsOneMatchesExistingOneBitPlusExBits) { + constexpr size_t kDim = 64; + constexpr size_t kExtraBits = 3; + + std::mt19937 gen(42); + std::vector data = RandomVec(kDim, gen); + std::vector centroid(kDim, 0.0F); + + // Reference: today's separate one_bit_code_with_factor (sign bit + the + // 1-bit factors) + ex_bits_code_with_factor (magnitude + full factors). + std::vector ref_binary_code(kDim); + float ref_f_add_bin = 0; + float ref_f_rescale_bin = 0; + float ref_f_error_bin = 0; + quant::rabitq_impl::one_bit::one_bit_code_with_factor( + data.data(), + centroid.data(), + kDim, + ref_binary_code.data(), + ref_f_add_bin, + ref_f_rescale_bin, + ref_f_error_bin, + METRIC_L2 + ); + + std::vector ref_ex_code(kDim); + float ref_f_add = 0; + float ref_f_rescale = 0; + float ref_f_error = 0; + quant::rabitq_impl::ex_bits::ex_bits_code_with_factor( + data.data(), + centroid.data(), + kDim, + kExtraBits, + ref_ex_code.data(), + ref_f_add, + ref_f_rescale, + ref_f_error, + METRIC_L2 + ); + + // New path: base_bits=1, extra_bits=kExtraBits. + XySplitResult res = SplitCode(data, centroid, kDim, /*base_bits=*/1, kExtraBits); + + for (size_t i = 0; i < kDim; ++i) { + EXPECT_EQ(res.base_code[i], static_cast(ref_binary_code[i])) + << "dim " << i; + EXPECT_EQ(res.extra_code[i], ref_ex_code[i]) << "dim " << i; + } + + EXPECT_FLOAT_NEARLY_EQUAL(res.f_add_base, ref_f_add_bin, 1e-4F); + EXPECT_FLOAT_NEARLY_EQUAL(res.f_rescale_base, ref_f_rescale_bin, 1e-4F); + EXPECT_FLOAT_NEARLY_EQUAL(res.f_error_base, ref_f_error_bin, 1e-4F); + + EXPECT_FLOAT_NEARLY_EQUAL(res.f_add_full, ref_f_add, 1e-4F); + EXPECT_FLOAT_NEARLY_EQUAL(res.f_rescale_full, ref_f_rescale, 1e-4F); + // The combined code's f_error is not exposed by the function, so it is not + // asserted here; ref_f_error stays as documentation of the reference value. +} + +// ip(base)*2^extra_bits + ip(extra) must equal ip(total_code) -- the +// linearity the boosting step relies on. +// The sign bit must be assigned the same way here as in the 1-bit path, +// including for an exactly-zero residual component. one_bit_code() uses +// `residual > 0`, so a zero lands on the 0 side; combined_code() has to agree +// or the same vector gets different sign bits from the two paths. +// +// Random gaussian data never produces an exact zero, so the other tests here +// cannot catch a divergence -- this one constructs the case directly, by +// putting the centroid exactly on the data value in half the dimensions. +TEST(XyQuantization, SignBitMatchesOneBitPathOnZeroResiduals) { + constexpr size_t kDim = 64; + constexpr size_t kExtraBits = 3; + + std::vector centroid(kDim, 1.0F); + std::vector data(kDim); + size_t zero_dims = 0; + for (size_t j = 0; j < kDim; ++j) { + if (j % 2 == 0) { + data[j] = 1.0F; // exact zero residual + ++zero_dims; + } else { + data[j] = 1.0F + ((static_cast(j % 3) - 1.0F) * 0.7F); + } + } + ASSERT_GT(zero_dims, 0U); + + std::vector ref_bin(kDim, -1); + quant::rabitq_impl::one_bit::one_bit_code( + data.data(), centroid.data(), kDim, ref_bin.data() + ); + + std::vector residual(kDim); + for (size_t j = 0; j < kDim; ++j) { + residual[j] = data[j] - centroid[j]; + } + + std::vector total_code(kDim, -1); + quant::rabitq_impl::total_bits::combined_code( + residual.data(), kDim, 1 + kExtraBits, total_code.data() + ); + + for (size_t j = 0; j < kDim; ++j) { + EXPECT_EQ(total_code[j] >> kExtraBits, ref_bin[j]) << "dim " << j; + } +} + +// A point sitting exactly on its centroid is normal (an IVF cluster of one, +// or a duplicate of the centroid). code_factors must return the same finite +// zeros one_bit_code_with_factor and ex_bits_code_with_factor return, not a +// NaN f_error -- a NaN propagates into low_dist and makes every pruning +// comparison false, so such a node is never pruned. +// +// Goes through quantize_xy_single rather than code_factors directly: the +// combined_code-level test above cannot reach this, which is how it was +// missed. +TEST(XyQuantization, ZeroResidualGivesFiniteFactors) { + constexpr size_t kDim = 64; + std::vector centroid(kDim, 0.7F); + std::vector data = centroid; // residual is exactly zero + + for (size_t base_bits : {1U, 2U, 4U, 8U}) { + for (size_t ex_bits : {0U, 1U, 3U}) { + if (base_bits + ex_bits > quant::rabitq_impl::xy_bits::kMaxCombinedBits) { + continue; + } + std::vector base_block( + BaseDataMap::data_bytes(kDim, base_bits), 0 + ); + std::vector ex_block( + ex_bits > 0 ? ExDataMap::data_bytes(kDim, ex_bits) : 1, 0 + ); + quant::quantize_xy_single( + data.data(), + centroid.data(), + kDim, + base_bits, + ex_bits, + base_block.data(), + ex_bits > 0 ? ex_block.data() : nullptr, + METRIC_L2 + ); + + ConstBaseDataMap base_map(base_block.data(), kDim, base_bits); + EXPECT_TRUE(std::isfinite(base_map.f_add())) << base_bits << "+" << ex_bits; + EXPECT_TRUE(std::isfinite(base_map.f_rescale())) << base_bits << "+" << ex_bits; + EXPECT_TRUE(std::isfinite(base_map.f_error())) << base_bits << "+" << ex_bits; + EXPECT_FLOAT_EQ(base_map.f_error(), 0.0F) << base_bits << "+" << ex_bits; + + if (ex_bits > 0) { + ConstExDataMap ex_map(ex_block.data(), kDim, ex_bits); + EXPECT_TRUE(std::isfinite(ex_map.f_add_ex())) + << base_bits << "+" << ex_bits; + EXPECT_TRUE(std::isfinite(ex_map.f_rescale_ex())) + << base_bits << "+" << ex_bits; + } + } + } +} + +TEST(XyQuantization, SplitInnerProductRecombinesExactly) { + constexpr size_t kDim = 64; + constexpr size_t kBaseBits = 3; + constexpr size_t kExtraBits = 4; + + std::mt19937 gen(7); + std::vector data = RandomVec(kDim, gen); + std::vector centroid(kDim, 0.0F); + std::vector query = RandomVec(kDim, gen); + + XySplitResult res = SplitCode(data, centroid, kDim, kBaseBits, kExtraBits); + + double split_ip = 0; + double direct_ip = 0; + for (size_t i = 0; i < kDim; ++i) { + uint32_t total = (static_cast(res.base_code[i]) << kExtraBits) | + static_cast(res.extra_code[i]); + direct_ip += static_cast(query[i]) * static_cast(total); + } + for (size_t i = 0; i < kDim; ++i) { + split_ip += static_cast(query[i]) * static_cast(res.base_code[i]) * + static_cast(1U << kExtraBits); + split_ip += static_cast(query[i]) * static_cast(res.extra_code[i]); + } + + EXPECT_NEAR(split_ip, direct_ip, 1e-6 * std::abs(direct_ip) + 1e-6); +} + +// The base code must be stored exactly once: total storage is +// (base_bits + extra_bits) bits per dimension plus 5 floats (3 base factors + +// 2 full factors), never (2*base_bits + extra_bits) as an independently +// quantized second layer would cost. +TEST(XyQuantization, StorageCostsBasePlusExtraBitsOnly) { + constexpr size_t kDim = 128; + constexpr size_t kBaseBits = 3; + constexpr size_t kExtraBits = 4; + + size_t base_bytes = BaseDataMap::data_bytes(kDim, kBaseBits); + size_t extra_bytes = ExDataMap::data_bytes(kDim, kExtraBits); + + EXPECT_EQ(base_bytes + extra_bytes, (kDim * (kBaseBits + kExtraBits) / 8) + (5 * 4)); + EXPECT_EQ(ExDataMap::data_bytes(kDim, 0), 0U); +} + +// End-to-end backward-compat check: at base_bits=1, quantize_xy_single + +// xy_single_base_dist + split_distance_boosting must agree with today's production +// formula ((1< data = RandomVec(kDim, gen); + std::vector centroid = RandomVec(kDim, gen); + std::vector query = RandomVec(kDim, gen); + + constexpr float kGAdd = 3.5F; + constexpr float kGError = 0.75F; + + // --- Reference: today's production formula, built from the existing, + // unchanged one_bit_code_with_factor + ex_bits_code_with_factor. --- + std::vector ref_binary_code(kDim); + float ref_f_add_bin = 0; + float ref_f_rescale_bin = 0; + float ref_f_error_bin = 0; + quant::rabitq_impl::one_bit::one_bit_code_with_factor( + data.data(), + centroid.data(), + kDim, + ref_binary_code.data(), + ref_f_add_bin, + ref_f_rescale_bin, + ref_f_error_bin, + METRIC_L2 + ); + + std::vector ref_ex_code(kDim); + float ref_f_add_ex = 0; + float ref_f_rescale_ex = 0; + float ref_f_error_ex = 0; + quant::rabitq_impl::ex_bits::ex_bits_code_with_factor( + data.data(), + centroid.data(), + kDim, + kExtraBits, + ref_ex_code.data(), + ref_f_add_ex, + ref_f_rescale_ex, + ref_f_error_ex, + METRIC_L2 + ); + + double ip_x0_qr = 0; + double ex_ip = 0; + double sumq = 0; + for (size_t i = 0; i < kDim; ++i) { + ip_x0_qr += static_cast(query[i]) * static_cast(ref_binary_code[i]); + ex_ip += static_cast(query[i]) * static_cast(ref_ex_code[i]); + sumq += query[i]; + } + double ref_c_b = -(static_cast(1 << (kExtraBits + 1)) - 1) / 2.0; + double ref_kbxsumq = sumq * ref_c_b; + + // 1-bit stage, as split_single_estdist_direct computes it. + float ref_base_est = static_cast( + ref_f_add_bin + kGAdd + (ref_f_rescale_bin * (ip_x0_qr - (0.5 * sumq))) + ); + + // Boosted stage, as split_distance_boosting computes it. + float ref_full_est = static_cast( + ref_f_add_ex + kGAdd + + (ref_f_rescale_ex * + (static_cast(1 << kExtraBits) * ip_x0_qr + ex_ip + ref_kbxsumq)) + ); + + // --- New path --- + std::vector base_data(BaseDataMap::data_bytes(kDim, 1)); + std::vector extra_data(ExDataMap::data_bytes(kDim, kExtraBits)); + quant::quantize_xy_single( + data.data(), + centroid.data(), + kDim, + /*base_bits=*/1, + kExtraBits, + base_data.data(), + extra_data.data() + ); + + SplitSingleQuery q_obj( + query.data(), kDim, kExtraBits, quant::RabitqConfig(), METRIC_L2, /*base_bits=*/1 + ); + + auto base_ip_func = select_excode_ipfunc(1); + auto extra_ip_func = select_excode_ipfunc(kExtraBits); + + float ip_base = 0; + float base_est = 0; + float base_low = 0; + xy_single_base_dist( + base_data.data(), + base_ip_func, + q_obj, + kDim, + /*base_bits=*/1, + ip_base, + base_est, + base_low, + kGAdd, + kGError + ); + + EXPECT_FLOAT_NEARLY_EQUAL(ip_base, static_cast(ip_x0_qr), 1e-3F); + EXPECT_FLOAT_NEARLY_EQUAL(base_est, ref_base_est, 1e-2F); + EXPECT_FLOAT_NEARLY_EQUAL(base_low, base_est - (ref_f_error_bin * kGError), 1e-4F); + + // split_distance_boosting reads g_add off the query rather than taking + // it as a parameter; there is no x+y-specific boosting function. + q_obj.set_g_add(std::sqrt(kGAdd)); + float boosted_est = split_distance_boosting( + extra_data.data(), extra_ip_func, q_obj, kDim, kExtraBits, ip_base + ); + + EXPECT_FLOAT_NEARLY_EQUAL(boosted_est, ref_full_est, 1e-2F); + + // The one-shot path recomputes ip_base itself and must agree exactly with + // the boosted path, plus supply the lower bound boosting omits. + float full_est = 0; + float full_low = 0; + float full_ip_base = 0; + xy_single_full_dist( + base_data.data(), + extra_data.data(), + base_ip_func, + extra_ip_func, + q_obj, + kDim, + /*base_bits=*/1, + kExtraBits, + full_est, + full_low, + full_ip_base, + kGAdd, + kGError + ); + + EXPECT_FLOAT_EQ(full_ip_base, ip_base); + EXPECT_FLOAT_EQ(full_est, boosted_est); + EXPECT_FLOAT_NEARLY_EQUAL( + full_low, + full_est - (ref_f_error_bin * kGError / static_cast(1 << kExtraBits)), + 1e-4F + ); +} + +// The boosted estimate must equal the estimate you would get by reading the +// whole (base_bits+extra_bits)-bit code directly -- i.e. reusing the filter +// stage's base inner product loses nothing. This is the property the old +// two-layer implementation could not have, since its two layers held +// different codes. +TEST(XyQuantization, BoostedEstimateMatchesUnsplitCombinedCode) { + constexpr size_t kDim = 128; + constexpr size_t kBaseBits = 3; + constexpr size_t kExtraBits = 4; + + std::mt19937 gen(2024); + std::vector data = RandomVec(kDim, gen); + std::vector centroid = RandomVec(kDim, gen); + std::vector query = RandomVec(kDim, gen); + + // g_add / g_error as the L2 search path actually supplies them + // (||q - centroid||^2 and ||q - centroid||), so the filter's lower bound + // below is the real one rather than an arbitrary scaling of f_error. + float q_to_cent_sqr = 0; + for (size_t i = 0; i < kDim; ++i) { + float d = query[i] - centroid[i]; + q_to_cent_sqr += d * d; + } + const float kGError = std::sqrt(q_to_cent_sqr); + const float kGAdd = q_to_cent_sqr; + + XySplitResult res = SplitCode(data, centroid, kDim, kBaseBits, kExtraBits); + + double total_ip = 0; + double sumq = 0; + for (size_t i = 0; i < kDim; ++i) { + uint32_t total = (static_cast(res.base_code[i]) << kExtraBits) | + static_cast(res.extra_code[i]); + total_ip += static_cast(query[i]) * static_cast(total); + sumq += query[i]; + } + double c_b = -(static_cast(1 << (kBaseBits + kExtraBits)) - 1) / 2.0; + float ref_est = static_cast( + res.f_add_full + kGAdd + (res.f_rescale_full * (total_ip + (sumq * c_b))) + ); + + std::vector base_data(BaseDataMap::data_bytes(kDim, kBaseBits)); + std::vector extra_data(ExDataMap::data_bytes(kDim, kExtraBits)); + quant::quantize_xy_single( + data.data(), + centroid.data(), + kDim, + kBaseBits, + kExtraBits, + base_data.data(), + extra_data.data() + ); + + SplitSingleQuery q_obj( + query.data(), kDim, kExtraBits, quant::RabitqConfig(), METRIC_L2, kBaseBits + ); + + float ip_base = 0; + float base_est = 0; + float base_low = 0; + xy_single_base_dist( + base_data.data(), + select_excode_ipfunc(kBaseBits), + q_obj, + kDim, + kBaseBits, + ip_base, + base_est, + base_low, + kGAdd, + kGError + ); + + q_obj.set_g_add(std::sqrt(kGAdd)); + float full_est = split_distance_boosting( + extra_data.data(), + select_excode_ipfunc(kExtraBits), + q_obj, + kDim, + kExtraBits, + ip_base + ); + + EXPECT_FLOAT_NEARLY_EQUAL(full_est, ref_est, 1e-2F); + // The cheap filter's lower bound must not exclude the refined estimate, + // otherwise the base layer would prune candidates the refine layer would + // have kept. + EXPECT_LE(base_low, full_est); +} + +TEST(XyQuantization, CombinedBitsBeyondCapAborts) { + constexpr size_t kDim = 64; + std::mt19937 gen(1); + std::vector data = RandomVec(kDim, gen); + std::vector centroid(kDim, 0.0F); + + auto call_with_bad_bits = [&]() { + SplitCode(data, centroid, kDim, /*base_bits=*/8, /*extra_bits=*/8); + }; + EXPECT_DEATH(call_with_bad_bits(), ""); +} From 5db366788b911dba0224f8a6d5cea44126c48ab9 Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Tue, 1 Sep 2026 22:11:44 +0800 Subject: [PATCH 4/4] fix: harden two-level quantization edge cases --- include/rabitqlib/index/query.hpp | 8 ++ .../rabitqlib/quantization/rabitq_impl.hpp | 23 ++-- .../quantization/xy_quantization_test.cpp | 115 +++++++++++++++++- 3 files changed, 132 insertions(+), 14 deletions(-) diff --git a/include/rabitqlib/index/query.hpp b/include/rabitqlib/index/query.hpp index 754dcd8..ad1e4fd 100644 --- a/include/rabitqlib/index/query.hpp +++ b/include/rabitqlib/index/query.hpp @@ -3,6 +3,7 @@ #include #include #include +#include #include #include "rabitqlib/defines.hpp" @@ -134,6 +135,13 @@ class SplitSingleQuery { size_t base_bits = 1 ) : rotated_query_(rotated_query), QueryBin_(padded_dim * kNumBits / 64, 0) { + if (base_bits < 1 || base_bits > 8 || ex_bits > 8 || base_bits + ex_bits > 9) { + throw std::invalid_argument( + "SplitSingleQuery requires base_bits in [1, 8], ex_bits in [0, 8], " + "and base_bits + ex_bits <= 9" + ); + } + float c_1 = -static_cast((1 << 1) - 1) / 2.F; float c_base = -static_cast((1U << base_bits) - 1) / 2.F; float c_b = -static_cast((1U << (base_bits + ex_bits)) - 1) / 2.F; diff --git a/include/rabitqlib/quantization/rabitq_impl.hpp b/include/rabitqlib/quantization/rabitq_impl.hpp index 98d2921..a59a698 100644 --- a/include/rabitqlib/quantization/rabitq_impl.hpp +++ b/include/rabitqlib/quantization/rabitq_impl.hpp @@ -2,6 +2,7 @@ #include +#include #include #include #include @@ -486,7 +487,7 @@ inline void ex_bits_code_with_factor( RowMajorArray total_code = RowMajorArrayMap(ex_code, 1, dim).template cast(); for (size_t i = 0; i < dim; ++i) { - total_code(0, static_cast(i)) += static_cast(residual_arr.data()[i] >= 0) + total_code(0, static_cast(i)) += static_cast(residual_arr.data()[i] > 0) << ex_bits; } @@ -744,18 +745,16 @@ inline void code_factors( T ip_resi_xucb = dot_product(residual, xu_cb.data(), dim); T ip_cent_xucb = dot_product(centroid, xu_cb.data(), dim); - // corner case - if (ip_resi_xucb == 0) { - ip_resi_xucb = std::numeric_limits::infinity(); - } + // A nonzero residual and its quantized code have a strictly positive inner product. + assert(ip_resi_xucb > 0); - T tmp_error = - l2_norm * kConstEpsilon * - std::sqrt( - (((l2_sqr * l2norm_sqr(xu_cb.data(), dim)) / (ip_resi_xucb * ip_resi_xucb)) - - 1) / - (dim - 1) - ); + // Cauchy-Schwarz guarantees this value is nonnegative. Clamp small negative + // round-off errors for collinear residual/code vectors before taking the square root. + const T normalized_error = + (((l2_sqr * l2norm_sqr(xu_cb.data(), dim)) / (ip_resi_xucb * ip_resi_xucb)) - 1 + ) / + (dim - 1); + T tmp_error = l2_norm * kConstEpsilon * std::sqrt(std::max(normalized_error, T{0})); if (metric_type == METRIC_L2) { f_add = l2_sqr + (2 * l2_sqr * ip_cent_xucb / ip_resi_xucb); diff --git a/tests/unit/rabitqlib/quantization/xy_quantization_test.cpp b/tests/unit/rabitqlib/quantization/xy_quantization_test.cpp index a39458c..b2d3d1a 100644 --- a/tests/unit/rabitqlib/quantization/xy_quantization_test.cpp +++ b/tests/unit/rabitqlib/quantization/xy_quantization_test.cpp @@ -63,7 +63,8 @@ XySplitResult SplitCode( const std::vector& centroid, size_t dim, size_t base_bits, - size_t extra_bits + size_t extra_bits, + MetricType metric = METRIC_L2 ) { XySplitResult res; res.base_code.resize(dim); @@ -84,7 +85,7 @@ XySplitResult SplitCode( extra_bits, base_block.data(), extra_bits > 0 ? extra_block.data() : nullptr, - METRIC_L2 + metric ); BaseDataMap base_map(base_block.data(), dim, base_bits); @@ -218,6 +219,66 @@ TEST(XyQuantization, SignBitMatchesOneBitPathOnZeroResiduals) { } } +TEST(XyQuantization, BaseBitsOneMatchesExistingFactorsOnZeroResidualComponents) { + constexpr size_t kDim = 64; + constexpr size_t kExtraBits = 3; + + std::vector centroid(kDim, 1.0F); + std::vector data(kDim); + for (size_t dim = 0; dim < kDim; ++dim) { + data[dim] = dim % 2 == 0 + ? centroid[dim] + : centroid[dim] + + ((static_cast(static_cast(dim % 3) - 1)) * 0.7F); + } + + for (MetricType metric : {METRIC_L2, METRIC_IP}) { + std::vector reference_base_code(kDim); + float reference_base_add = 0; + float reference_base_rescale = 0; + float reference_base_error = 0; + quant::rabitq_impl::one_bit::one_bit_code_with_factor( + data.data(), + centroid.data(), + kDim, + reference_base_code.data(), + reference_base_add, + reference_base_rescale, + reference_base_error, + metric + ); + + std::vector reference_extra_code(kDim); + float reference_full_add = 0; + float reference_full_rescale = 0; + float reference_full_error = 0; + quant::rabitq_impl::ex_bits::ex_bits_code_with_factor( + data.data(), + centroid.data(), + kDim, + kExtraBits, + reference_extra_code.data(), + reference_full_add, + reference_full_rescale, + reference_full_error, + metric + ); + + XySplitResult result = + SplitCode(data, centroid, kDim, /*base_bits=*/1, kExtraBits, metric); + + for (size_t dim = 0; dim < kDim; ++dim) { + EXPECT_EQ(result.base_code[dim], reference_base_code[dim]); + EXPECT_EQ(result.extra_code[dim], reference_extra_code[dim]); + } + EXPECT_FLOAT_NEARLY_EQUAL(result.f_add_base, reference_base_add, 1e-4F); + EXPECT_FLOAT_NEARLY_EQUAL(result.f_rescale_base, reference_base_rescale, 1e-4F); + EXPECT_FLOAT_NEARLY_EQUAL(result.f_error_base, reference_base_error, 1e-4F); + EXPECT_FLOAT_NEARLY_EQUAL(result.f_add_full, reference_full_add, 1e-4F); + EXPECT_FLOAT_NEARLY_EQUAL(result.f_rescale_full, reference_full_rescale, 1e-4F); + } +} + // A point sitting exactly on its centroid is normal (an IVF cluster of one, // or a duplicate of the centroid). code_factors must return the same finite // zeros one_bit_code_with_factor and ex_bits_code_with_factor return, not a @@ -271,6 +332,40 @@ TEST(XyQuantization, ZeroResidualGivesFiniteFactors) { } } +TEST(XyQuantization, CollinearResidualGivesFiniteErrorFactor) { + constexpr size_t kDim = 64; + constexpr size_t kBaseBits = 2; + + std::mt19937 gen(314); + std::uniform_real_distribution scale_dist(0.001F, 10000.0F); + std::uniform_int_distribution code_dist(0, 3); + + const float scale = scale_dist(gen); + std::vector centroid(kDim, 0.0F); + std::vector data(kDim); + for (float& value : data) { + value = (static_cast(code_dist(gen)) - 1.5F) * scale; + } + + std::vector base_data(BaseDataMap::data_bytes(kDim, kBaseBits), 0); + quant::quantize_xy_single( + data.data(), + centroid.data(), + kDim, + kBaseBits, + /*ex_bits=*/0, + base_data.data(), + /*ex_data=*/nullptr, + METRIC_L2 + ); + + ConstBaseDataMap base_map(base_data.data(), kDim, kBaseBits); + EXPECT_TRUE(std::isfinite(base_map.f_add())); + EXPECT_TRUE(std::isfinite(base_map.f_rescale())); + EXPECT_TRUE(std::isfinite(base_map.f_error())); + EXPECT_FLOAT_EQ(base_map.f_error(), 0.0F); +} + TEST(XyQuantization, SplitInnerProductRecombinesExactly) { constexpr size_t kDim = 64; constexpr size_t kBaseBits = 3; @@ -569,3 +664,19 @@ TEST(XyQuantization, CombinedBitsBeyondCapAborts) { }; EXPECT_DEATH(call_with_bad_bits(), ""); } + +TEST(XyQuantization, SplitSingleQueryRejectsInvalidBitWidths) { + constexpr size_t kDim = 64; + std::vector query(kDim, 1.0F); + + for (const auto& [base_bits, ex_bits] : + {std::pair{0, 0}, {9, 0}, {1, 9}, {8, 2}, {32, 0}}) { + EXPECT_THROW( + SplitSingleQuery( + query.data(), kDim, ex_bits, quant::RabitqConfig(), METRIC_L2, base_bits + ), + std::invalid_argument + ) << "base_bits=" + << base_bits << ", ex_bits=" << ex_bits; + } +}