From 0b2b71b126e8b68b94b89de3f668a4d8fdae2577 Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Sun, 25 Oct 2026 00:28:17 +0800 Subject: [PATCH] refactor(core): share quantization math and reuse residuals --- .../rabitqlib/core/detail/quantized_level.hpp | 23 + include/rabitqlib/core/detail/rabitq_impl.hpp | 393 +++++++++--------- src/core/api.cpp | 46 +- src/simd/quantization_avx2.cpp | 8 +- src/simd/quantization_avx512.cpp | 8 +- src/simd/rescale_search.hpp | 82 +--- .../rabitqlib/quantization/factor_test.cpp | 163 ++++++++ .../full_code_preparation_test.cpp | 81 ++++ 8 files changed, 511 insertions(+), 293 deletions(-) create mode 100644 include/rabitqlib/core/detail/quantized_level.hpp create mode 100644 tests/unit/rabitqlib/quantization/factor_test.cpp create mode 100644 tests/unit/rabitqlib/quantization/full_code_preparation_test.cpp diff --git a/include/rabitqlib/core/detail/quantized_level.hpp b/include/rabitqlib/core/detail/quantized_level.hpp new file mode 100644 index 00000000..8ed4615a --- /dev/null +++ b/include/rabitqlib/core/detail/quantized_level.hpp @@ -0,0 +1,23 @@ +#pragma once + +#include + +namespace rabitqlib::quant::detail { + +// Preconditions: magnitude >= 0 and finite; scale >= 0 and finite; max_code >= 0. +// Keep this helper local to each translation unit: callers use different ISA flags. +// Search evaluation and final encoding must use the same division-based thresholds. +static inline int quantized_level_at_scale(double magnitude, double scale, int max_code) { + if (magnitude == 0) + return 0; + int code = static_cast(std::min(scale * magnitude, static_cast(max_code))); + // Multiplication can round to either side of an event; correct against the + // exact thresholds rather than adding an epsilon that could cross other events. + if (code < max_code && (code + 1.0) / magnitude <= scale) + ++code; + else if (code > 0 && static_cast(code) / magnitude > scale) + --code; + return code; +} + +} // namespace rabitqlib::quant::detail diff --git a/include/rabitqlib/core/detail/rabitq_impl.hpp b/include/rabitqlib/core/detail/rabitq_impl.hpp index 31742690..79ca6c8e 100644 --- a/include/rabitqlib/core/detail/rabitq_impl.hpp +++ b/include/rabitqlib/core/detail/rabitq_impl.hpp @@ -13,6 +13,7 @@ #include "rabitqlib/core/data_layout.hpp" #include "rabitqlib/core/detail/eigen.hpp" +#include "rabitqlib/core/detail/quantized_level.hpp" #include "rabitqlib/core/detail/scratch.hpp" #include "rabitqlib/core/detail/simd/quantization_dispatch.hpp" #include "rabitqlib/core/metric.hpp" @@ -29,6 +30,66 @@ namespace rabitqlib::quant::rabitq_impl { constexpr float kConstEpsilon = 1.9; +// Arithmetic only: callers own residual/centered storage and choose rescale arithmetic. +template +struct CodeFactorTerms { + T add; + T error; + T residual_norm; + T residual_code_ip; +}; + +template +inline CodeFactorTerms code_factor_terms( + const T* residual, + const T* centroid, + const T* centered, + size_t dim, + T l2_sqr, + MetricType metric_type, + const char* invalid_metric_message +) { + // A zero residual needs no centered code or inner-product divisions. + if (l2_sqr == 0) { + if (metric_type == METRIC_L2) { + return {0, 0, 0, 0}; + } + if (metric_type == METRIC_IP) { + return {1, 0, 0, 0}; + } + throw std::invalid_argument(invalid_metric_message); + } + const T l2_norm = std::sqrt(l2_sqr); + const T ip_resi_xucb = dot_product(residual, centered, dim); + const T ip_cent_xucb = dot_product(centroid, centered, dim); + assert(ip_resi_xucb > 0); + + // Cauchy-Schwarz makes this nonnegative; collinear vectors can round below zero. + const T normalized_error = + (((l2_sqr * l2norm_sqr(centered, dim)) / (ip_resi_xucb * ip_resi_xucb)) - 1) / + (dim - 1); + const T tmp_error = + l2_norm * kConstEpsilon * std::sqrt(std::max(normalized_error, T{0})); + + // Metric-specific derivation: docs/docs/rabitq/estimator.md. + if (metric_type == METRIC_L2) { + return { + l2_sqr + (2 * l2_sqr * ip_cent_xucb / ip_resi_xucb), + 2 * tmp_error, + l2_norm, + ip_resi_xucb}; + } + if (metric_type == METRIC_IP) { + return { + 1 - dot_product(residual, centroid, dim) + + (l2_sqr * ip_cent_xucb / ip_resi_xucb), + 1 * tmp_error, + l2_norm, + ip_resi_xucb}; + } + throw std::invalid_argument(invalid_metric_message); +} + namespace one_bit { /** @@ -94,52 +155,21 @@ inline void one_bit_code_with_factor_into( xu_cb = x_u.template cast() + T{-0.5}; } - // distance to centroid - T l2_sqr = l2norm_sqr(residual, dim); - T l2_norm = std::sqrt(l2_sqr); - - if (l2_sqr == 0) { - if (metric_type == METRIC_L2) { - f_add = 0; - } else if (metric_type == METRIC_IP) { - f_add = 1; - } else { - throw std::invalid_argument("Unsupported metric type in quantization"); - } - f_rescale = 0; - f_error = 0; - return; - } - T ip_resi_xucb = dot_product(residual, centered, dim); - T ip_cent_xucb = dot_product(centroid, centered, dim); - - // A nonzero residual and its sign code have a strictly positive inner product. - assert(ip_resi_xucb > 0); - - // tan^2(angle(residual, centered)) / (dim - 1), using unnormalized vectors. - // Cauchy-Schwarz makes this nonnegative; collinear vectors can round below zero. - const T normalized_error = - (((l2_sqr * l2norm_sqr(centered, dim)) / (ip_resi_xucb * ip_resi_xucb)) - 1) / - (dim - 1); - T tmp_error = l2_norm * kConstEpsilon * std::sqrt(std::max(normalized_error, T{0})); - - // Factor derivation: docs/docs/rabitq/estimator.md, metric-specific estimators. - // For f_rescale and 2nd item of f_add, we use the dot product of raw residual (rather - // than the normalized one) as the denominator, thus we need to multiply another l2norm. - // For (ip_cent_xucb / ip_resi_xucb), the norm of xucb does not matter since it is also - // in numerator. - 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 { - throw std::invalid_argument("Unsupported metric type in quantization"); - } + const T l2_sqr = l2norm_sqr(residual, dim); + const auto factors = code_factor_terms( + residual, + centroid, + centered, + dim, + l2_sqr, + metric_type, + "Unsupported metric type in quantization" + ); + f_add = factors.add; + f_rescale = l2_sqr == 0 ? T{0} + : metric_type == METRIC_L2 ? -2 * l2_sqr / factors.residual_code_ip + : -l2_sqr / factors.residual_code_ip; + f_error = factors.error; } template @@ -320,21 +350,7 @@ constexpr std::array kTightStart = { 0.81, }; -// Preconditions: magnitude >= 0 and finite; t >= 0 and finite; max_code >= 0. -// Compare thresholds exactly as the queue does. Multiplication alone can round -// to the wrong side of an event; adding a fixed epsilon can cross other events. -inline int quantized_level_at_scale(double magnitude, double t, int max_code) { - if (magnitude == 0) - return 0; - int code = static_cast(std::min(t * magnitude, static_cast(max_code))); - // Clamping only bounds the code; correct rounding in either direction to match - // the queue's division-based thresholds. - if (code < max_code && (code + 1.0) / magnitude <= t) - ++code; - else if (code > 0 && static_cast(code) / magnitude > t) - --code; - return code; -} +using ::rabitqlib::quant::detail::quantized_level_at_scale; template inline double event_sweep( @@ -556,8 +572,9 @@ T faster_quantize_ex(const T* o_abs, TP* code, size_t dim, size_t ex_bits, doubl } template -inline T ex_bits_code_with_scratch( +inline T ex_bits_code_with_norm( const T* residual, + T residual_norm, size_t dim, size_t ex_bits, TP* ex_code, @@ -566,7 +583,6 @@ inline T ex_bits_code_with_scratch( ) { ConstRowMajorArrayMap res_arr(residual, 1, dim); T ipnorm_inv = 1; - const T residual_norm = std::sqrt(l2norm_sqr(residual, dim)); if (residual_norm == 0) { std::fill(ex_code, ex_code + dim, static_cast(0)); } else { @@ -599,14 +615,22 @@ template inline T ex_bits_code( const T* residual, size_t dim, size_t ex_bits, TP* ex_code, double t_const = -1 ) { - return ex_bits_code_with_scratch( - residual, dim, ex_bits, ex_code, t_const, static_cast*>(nullptr) + return ex_bits_code_with_norm( + residual, + std::sqrt(l2norm_sqr(residual, dim)), + dim, + ex_bits, + ex_code, + t_const, + static_cast*>(nullptr) ); } -template -inline void ex_bits_code_with_factor_with_scratch( - const T* data, +// l2_sqr belongs to residual. FullCode returns the combined code used for factors. +template +inline void code_with_factor_from_residual( + const T* residual, + T l2_sqr, const T* centroid, size_t dim, size_t ex_bits, @@ -618,39 +642,38 @@ inline void ex_bits_code_with_factor_with_scratch( double t_const, Scratch* scratch ) { - ConstRowMajorArrayMap data_arr(data, 1, dim); - ConstRowMajorArrayMap cent_arr(centroid, 1, dim); - RowMajorArray owned_residual; - auto residual_arr = scratch_array( - scratch == nullptr ? nullptr : scratch->residual, dim, owned_residual - ); - residual_arr = data_arr - cent_arr; - - const T l2_sqr = l2norm_sqr(residual_arr.data(), dim); if (l2_sqr == 0) { std::fill(ex_code, ex_code + dim, static_cast(0)); - if (metric_type == METRIC_L2) { - f_add_ex = 0; - } else if (metric_type == METRIC_IP) { - f_add_ex = 1; - } else { - throw std::invalid_argument("Unsupported metric for ex_bits_code()"); - } + const auto factors = code_factor_terms( + residual, + centroid, + nullptr, + dim, + l2_sqr, + metric_type, + "Unsupported metric for ex_bits_code()" + ); + f_add_ex = factors.add; f_rescale_ex = 0; - f_error_ex = 0; + f_error_ex = factors.error; + if constexpr (FullCode) { + // A tiny nonzero residual can have a squared norm of zero. + for (size_t i = 0; i < dim; ++i) { + ex_code[i] = static_cast(residual[i] > 0) << ex_bits; + } + } return; } - T ipnorm_inv = ex_bits_code_with_scratch( - residual_arr.data(), dim, ex_bits, ex_code, t_const, scratch + T ipnorm_inv = ex_bits_code_with_norm( + residual, std::sqrt(l2_sqr), dim, ex_bits, ex_code, t_const, scratch ); RowMajorArray owned_total; auto total_code = scratch_array(scratch == nullptr ? nullptr : scratch->total, dim, owned_total); 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) - << ex_bits; + total_code(0, static_cast(i)) += static_cast(residual[i] > 0) << ex_bits; } // Full-code factors use the metric-specific derivation in @@ -662,32 +685,23 @@ inline void ex_bits_code_with_factor_with_scratch( ); xu_cb = total_code.template cast() + cb; - T l2_norm = std::sqrt(l2_sqr); - - T ip_resi_xucb = dot_product(residual_arr.data(), xu_cb.data(), dim); - T ip_cent_xucb = dot_product(centroid, xu_cb.data(), dim); - - // A nonzero residual and its quantized code have a strictly positive inner product. - assert(ip_resi_xucb > 0); - - // Cauchy-Schwarz makes this nonnegative; collinear vectors can round below zero. - 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_ex = l2_sqr + (2 * l2_sqr * ip_cent_xucb / ip_resi_xucb); - f_rescale_ex = ipnorm_inv * -2 * l2_norm; - f_error_ex = 2 * tmp_error; - } else if (metric_type == METRIC_IP) { - f_add_ex = 1 - dot_product(residual_arr.data(), centroid, dim) + - (l2_sqr * ip_cent_xucb / ip_resi_xucb); - f_rescale_ex = ipnorm_inv * -l2_norm; - f_error_ex = 1 * tmp_error; - } else { - throw std::invalid_argument("Unsupported metric for ex_bits_code()"); + const auto factors = code_factor_terms( + residual, + centroid, + xu_cb.data(), + dim, + l2_sqr, + metric_type, + "Unsupported metric for ex_bits_code()" + ); + f_add_ex = factors.add; + // Retain the normalized-code rescale expression and its rounding order. + f_rescale_ex = metric_type == METRIC_L2 ? ipnorm_inv * -2 * factors.residual_norm + : ipnorm_inv * -factors.residual_norm; + f_error_ex = factors.error; + if constexpr (FullCode) { + RowMajorArrayMap(ex_code, 1, static_cast(dim)) = + total_code.template cast(); } } @@ -704,8 +718,12 @@ inline void ex_bits_code_with_factor( MetricType metric_type = METRIC_L2, double t_const = -1 ) { - ex_bits_code_with_factor_with_scratch( - data, + RowMajorArray residual = + ConstRowMajorArrayMap(data, 1, static_cast(dim)) - + ConstRowMajorArrayMap(centroid, 1, static_cast(dim)); + code_with_factor_from_residual( + residual.data(), + l2norm_sqr(residual.data(), dim), centroid, dim, ex_bits, @@ -832,65 +850,56 @@ static inline void rabitq_full_impl_with_scratch( double t_const, Scratch* scratch ) { + const size_t ex_bits = total_bits - 1; + if (ex_bits > 0) { + RowMajorArray owned_residual; + auto residual = scratch_array( + scratch == nullptr ? nullptr : scratch->residual, dim, owned_residual + ); + residual = ConstRowMajorArrayMap(data, 1, static_cast(dim)) - + ConstRowMajorArrayMap(centroid, 1, static_cast(dim)); + ex_bits::code_with_factor_from_residual( + residual.data(), + l2norm_sqr(residual.data(), dim), + centroid, + dim, + ex_bits, + total_code, + f_add, + f_rescale, + f_error, + metric_type, + t_const, + scratch + ); + return; + } + std::vector owned_binary; int* binary_code = nullptr; if (scratch == nullptr) { owned_binary.resize(dim); binary_code = owned_binary.data(); + one_bit::one_bit_code_with_factor( + data, centroid, dim, binary_code, f_add, f_rescale, f_error, metric_type + ); } else { binary_code = scratch->binary; - } - size_t ex_bits = total_bits - 1; - - if (ex_bits > 0) { - if (scratch == nullptr) { - rabitq_impl::one_bit::one_bit_code(data, centroid, dim, binary_code); - } else { - ConstRowMajorArrayMap input(data, 1, static_cast(dim)); - ConstRowMajorArrayMap center(centroid, 1, static_cast(dim)); - RowMajorArrayMap residual(scratch->residual, 1, static_cast(dim)); - residual = input - center; - RowMajorArrayMap signs(binary_code, 1, static_cast(dim)); - signs = (residual > 0).template cast(); - } - - ex_bits::ex_bits_code_with_factor_with_scratch( + one_bit::one_bit_code_with_factor_into( data, centroid, dim, - ex_bits, - total_code, + binary_code, f_add, f_rescale, f_error, metric_type, - t_const, - scratch + scratch->residual, + scratch->centered ); - } else { - if (scratch == nullptr) { - one_bit::one_bit_code_with_factor( - data, centroid, dim, binary_code, f_add, f_rescale, f_error, metric_type - ); - } else { - one_bit::one_bit_code_with_factor_into( - data, - centroid, - dim, - binary_code, - f_add, - f_rescale, - f_error, - metric_type, - scratch->residual, - scratch->centered - ); - } } - 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; + total_code[i] = static_cast(binary_code[i]); } } template @@ -951,12 +960,8 @@ inline void validate_bit_size(size_t base_bits, size_t ex_bits) { * the code's integer range alone does not guarantee this. A zero residual * produces zero rescale/error factors and a metric-specific additive term. * - * 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. + * Centers the supplied code before applying the shared factor arithmetic. + * With cb = -0.5 and a sign-bit code, it matches the one-bit factors. * * @param residual data - centroid * @param code integer code, values in [0, 2^bits) @@ -977,49 +982,21 @@ inline void code_factors( 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 { - throw std::invalid_argument("Unsupported metric type in code_factors()"); - } - 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); - - // A nonzero residual and its quantized code have a strictly positive inner product. - assert(ip_resi_xucb > 0); - - // 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); - 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 { - throw std::invalid_argument("Unsupported metric type in code_factors()"); - } + const T l2_sqr = l2norm_sqr(residual, dim); + const auto factors = code_factor_terms( + residual, + centroid, + xu_cb.data(), + dim, + l2_sqr, + metric_type, + "Unsupported metric type in code_factors()" + ); + f_add = factors.add; + f_rescale = l2_sqr == 0 ? T{0} + : metric_type == METRIC_L2 ? -2 * l2_sqr / factors.residual_code_ip + : -l2_sqr / factors.residual_code_ip; + f_error = factors.error; } template diff --git a/src/core/api.cpp b/src/core/api.cpp index ee52c837..2eda58f2 100644 --- a/src/core/api.cpp +++ b/src/core/api.cpp @@ -279,23 +279,41 @@ Factors encode( RowMajorArrayMap residual(s.scratch.residual, 1, static_cast(p.dim)); residual = ConstRowMajorArrayMap(input, 1, static_cast(p.dim)) - ConstRowMajorArrayMap(centroid, 1, static_cast(p.dim)); - if (!std::isfinite(l2norm_sqr(s.scratch.residual, p.dim))) { + const float residual_squared_norm = l2norm_sqr(s.scratch.residual, p.dim); + if (!std::isfinite(residual_squared_norm)) { throw std::overflow_error("Residual squared norm is not representable in float32"); } Factors factors; - rabitq_impl::total_bits::rabitq_full_impl_with_scratch( - input, - centroid, - p.dim, - p.total_bits, - s.codes, - factors.add, - factors.rescale, - factors.error, - p.metric, - p.t_const, - &s.scratch - ); + if (p.total_bits > 1) { + rabitq_impl::ex_bits::code_with_factor_from_residual( + s.scratch.residual, + residual_squared_norm, + centroid, + p.dim, + p.total_bits - 1, + s.codes, + factors.add, + factors.rescale, + factors.error, + p.metric, + p.t_const, + &s.scratch + ); + } else { + rabitq_impl::total_bits::rabitq_full_impl_with_scratch( + input, + centroid, + p.dim, + p.total_bits, + s.codes, + factors.add, + factors.rescale, + factors.error, + p.metric, + p.t_const, + &s.scratch + ); + } if (!std::isfinite(factors.add) || !std::isfinite(factors.rescale) || !std::isfinite(factors.error)) { throw std::overflow_error("Quantization factors are not representable in float32"); diff --git a/src/simd/quantization_avx2.cpp b/src/simd/quantization_avx2.cpp index 88da2919..c85c925f 100644 --- a/src/simd/quantization_avx2.cpp +++ b/src/simd/quantization_avx2.cpp @@ -64,7 +64,7 @@ static RescaleSearchState evaluate_scale_state_avx2( _mm256_store_pd(corrected, c); for (unsigned mask = boundary; mask; mask &= mask - 1) { unsigned k = bitops::countr_zero32(mask); - corrected[k] = quantized_code_at_scale(magnitudes[i + k], scale, max_code); + corrected[k] = quantized_level_at_scale(magnitudes[i + k], scale, max_code); } c = _mm256_load_pd(corrected); } @@ -123,18 +123,18 @@ static RescaleSearchState evaluate_scale_state_avx2( s.first = std::max( s.first, static_cast( - quantized_code_at_scale(magnitudes[fi[fk]], scale, max_code) + quantized_level_at_scale(magnitudes[fi[fk]], scale, max_code) ) / magnitudes[fi[fk]] ); if (ni[nk] != std::numeric_limits::max()) s.next = std::min( s.next, - (quantized_code_at_scale(magnitudes[ni[nk]], scale, max_code) + 1.0) / + (quantized_level_at_scale(magnitudes[ni[nk]], scale, max_code) + 1.0) / magnitudes[ni[nk]] ); // Handle the remaining zero to three coordinates with the same definitions. for (; i < dim; ++i) { - int c = quantized_code_at_scale(magnitudes[i], scale, max_code); + int c = quantized_level_at_scale(magnitudes[i], scale, max_code); double q = c + .5; s.numerator += magnitudes[i] * q; s.squared_norm += q * q; diff --git a/src/simd/quantization_avx512.cpp b/src/simd/quantization_avx512.cpp index c0a06d72..bab967b9 100644 --- a/src/simd/quantization_avx512.cpp +++ b/src/simd/quantization_avx512.cpp @@ -59,7 +59,7 @@ static RescaleSearchState evaluate_scale_state_avx512( _mm512_store_pd(corrected, c); for (unsigned mask = boundary; mask; mask &= mask - 1) { unsigned k = bitops::countr_zero32(mask); - corrected[k] = quantized_code_at_scale(magnitudes[i + k], scale, max_code); + corrected[k] = quantized_level_at_scale(magnitudes[i + k], scale, max_code); } c = _mm512_load_pd(corrected); } @@ -115,18 +115,18 @@ static RescaleSearchState evaluate_scale_state_avx512( s.first = std::max( s.first, static_cast( - quantized_code_at_scale(magnitudes[fi[fk]], scale, max_code) + quantized_level_at_scale(magnitudes[fi[fk]], scale, max_code) ) / magnitudes[fi[fk]] ); if (ni[nk] != std::numeric_limits::max()) s.next = std::min( s.next, - (quantized_code_at_scale(magnitudes[ni[nk]], scale, max_code) + 1.0) / + (quantized_level_at_scale(magnitudes[ni[nk]], scale, max_code) + 1.0) / magnitudes[ni[nk]] ); // Handle the remaining zero to seven coordinates with the same definitions. for (; i < dim; ++i) { - int c = quantized_code_at_scale(magnitudes[i], scale, max_code); + int c = quantized_level_at_scale(magnitudes[i], scale, max_code); double q = c + .5; s.numerator += magnitudes[i] * q; s.squared_norm += q * q; diff --git a/src/simd/rescale_search.hpp b/src/simd/rescale_search.hpp index ad402840..b9b2a999 100644 --- a/src/simd/rescale_search.hpp +++ b/src/simd/rescale_search.hpp @@ -7,8 +7,12 @@ #include #include +#include "rabitqlib/core/detail/quantized_level.hpp" + namespace rabitqlib::simd::detail { +using ::rabitqlib::quant::detail::quantized_level_at_scale; + struct RescaleScratch { std::vector magnitudes; std::vector reciprocals; @@ -33,55 +37,28 @@ struct RescaleSearchState { // Keep executable helpers local to each translation unit. This header is // compiled with different ISA flags; merging external inline/template symbols // could make the AVX2 backend call an AVX512-compiled helper. -static inline int quantized_code_at_scale(double magnitude, double scale, int max_code) { - // Keep the division-based boundary correction identical to the event sweep. - if (magnitude == 0) - return 0; - int code = static_cast(std::min(scale * magnitude, static_cast(max_code))); - if (code < max_code && (code + 1.0) / magnitude <= scale) - ++code; - else if (code > 0 && static_cast(code) / magnitude > scale) - --code; - return code; -} -// Return the next representable double greater than value, e.g. -// next_double_up(1.0) = 1.0000000000000002. This is one floating-point step, not +1. -// All bound operands here are finite. After a rounded arithmetic operation, -// stepping outward gives a conservative bound even if the result was exact. -// Use upward steps for upper bounds and downward steps for lower bounds so -// rounding cannot make an interval look worse and cause an incorrect prune. -// These helpers do not change the processor's floating-point rounding mode. +// Finite IEEE-754 binary64 neighbors for outward-rounded bounds. +// Adjusting the representation leaves the floating-point rounding mode unchanged. static inline double next_double_up(double value) { - // The bit ordering below relies on the IEEE-754 binary64 representation. static_assert(std::numeric_limits::is_iec559 && sizeof(double) == 8); - // Both +0 and -0 step to the smallest positive subnormal, not min(), - // which is the smallest positive *normal* double and would skip values. + // Both signed zeros step to the smallest positive subnormal. if (value == 0) return std::numeric_limits::denorm_min(); uint64_t bits = 0; - // Copy the representation without numeric conversion or pointer aliasing. std::memcpy(&bits, &value, sizeof(bits)); - // For positive doubles, increasing the bits increases the value. For - // negative doubles, decreasing the bits moves toward zero (a larger value). - // Converting -1 to uint64_t makes this unsigned addition subtract one. + // Negative values have reversed bit ordering. bits += value > 0 ? uint64_t{1} : static_cast(-1); std::memcpy(&value, &bits, sizeof(value)); return value; } -// Return the next representable double less than value, e.g. -// next_double_down(1.0) = 0.9999999999999999. This supplies the downward step used -// in lower bounds; its spacing depends on value, unlike a fixed epsilon. -// It uses the same finite-input and binary64 assumptions as next_double_up above. static inline double next_double_down(double value) { // Both signed zeros step to the smallest-magnitude negative subnormal. if (value == 0) return -std::numeric_limits::denorm_min(); uint64_t bits = 0; std::memcpy(&bits, &value, sizeof(bits)); - // Reverse next_double_up: positive values move toward zero by subtracting one - // from the bits; negative values become more negative by adding one. bits += value > 0 ? static_cast(-1) : uint64_t{1}; std::memcpy(&value, &bits, sizeof(value)); return value; @@ -200,25 +177,12 @@ static inline double interval_squared_cosine_upper_bound( ); } -// SIMD branch-and-bound alternative to the heap sweep in rabitq_impl.hpp. -// Evaluating a scale scans independent coordinates, which maps well to SIMD. -// Recursively split the unexplored scale range, discarding constant-code spans -// and intervals whose score upper bound cannot beat an evaluated candidate. -// Unlike binary search, either or both halves may need further exploration. -// -// Search outline: -// 1. Evaluate the first and last allowed scales and save the best candidate. -// 2. Bound every unevaluated code between two endpoint plateaus. -// 3. Skip the gap if it has no other codes or cannot improve the best score. -// 4. Otherwise evaluate a midpoint, split around its entire plateau, and recur. -// No fixed step size is used: interval bounds determine how large a range can -// be skipped. The supplied [t_start, t_end) limits are the same as the heap's. -// -// Return false to use the exact event search when the interval cannot be -// resolved within bounded work. The evaluator must emit legal states and -// division-based first/next thresholds; only its numerator sum may round. -// False discards the tentative SIMD result: the caller runs the heap sweep, -// rather than accepting an approximate winner after a budget or precision limit. +// SIMD branch-and-bound over the heap sweep's [t_start, t_end) range. +// Evaluate endpoint plateaus, then split gaps whose bounds can beat the best score. +// Evaluators return legal states and division-based plateau boundaries; +// only numerator accumulation may round. +// Return false on ambiguous scores or work limits: the caller discards the +// tentative result and runs the exact heap sweep, preserving its tie order. template static bool try_find_best_scale_by_interval_search( const float* input, @@ -275,8 +239,7 @@ static bool try_find_best_scale_by_interval_search( return false; ++evaluations; result.state = evaluator(magnitudes, reciprocals, dim, max_code, t_start, t_end, t); - // A valid evaluator must return a plateau containing the requested t. - // Every t in that plateau produces the same codes and score. + // Reject evaluator states whose plateau does not contain t. if (result.state.first > t || result.state.next <= t) return false; // If the computed numerator is N_hat and its relative error is at most @@ -305,8 +268,6 @@ static bool try_find_best_scale_by_interval_search( result.score_upper >= best.score_lower && result.state.squared_norm != best.state.squared_norm ) { - // The event sweep resolves numerically indistinguishable winners - // with its existing arithmetic and tie order. return false; } return true; @@ -334,15 +295,12 @@ static bool try_find_best_scale_by_interval_search( if (left.state.next > right.state.first || depth == kMaxDepth) return false; - // The endpoint plateaus are already evaluated. Split only the gap - // between them, not the constant-code spans surrounding their scales. + // Split only [left.next, right.first); stay below the right plateau + // despite rounding. double midpoint = left.state.next + (right.state.first - left.state.next) * 0.5; midpoint = std::max( left.state.next, std::min(midpoint, next_double_down(right.state.first)) ); - // Keep the trial inside [left.next, right.first); rounding must not - // send us back to the already evaluated right plateau. If no new state - // can be established, let the heap resolve the remaining events. if (midpoint >= right.state.first) return false; BoundedRescaleState middle; @@ -351,12 +309,10 @@ static bool try_find_best_scale_by_interval_search( middle.state.squared_norm >= right.state.squared_norm) return false; - // The middle code covers [middle.first, middle.next). These two bounds - // exclude that whole plateau, so recursion never searches inside it. + // Exclude the middle plateau from both recursive gaps. const double left_upper = interval_squared_cosine_upper_bound(left, middle); const double right_upper = interval_squared_cosine_upper_bound(middle, right); - // Explore the more promising half first; a better candidate found there - // raises best_lower and can let us prune the other half immediately. + // Explore the stronger bound first to improve subsequent pruning. if (left_upper > right_upper) { return self(self, left, middle, left_upper, depth + 1) && self(self, middle, right, right_upper, depth + 1); diff --git a/tests/unit/rabitqlib/quantization/factor_test.cpp b/tests/unit/rabitqlib/quantization/factor_test.cpp new file mode 100644 index 00000000..a297e494 --- /dev/null +++ b/tests/unit/rabitqlib/quantization/factor_test.cpp @@ -0,0 +1,163 @@ +#include + +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/core/detail/rabitq_impl.hpp" +#include "rabitqlib/core/encoder.hpp" +#include "rabitqlib/core/metric.hpp" + +namespace rabitqlib::quant { +namespace { + +template +void check_collinear_factors() { + constexpr size_t kDim = 64; + constexpr T kCenter = T{0.25}; + std::vector data(kDim), centroid(kDim, kCenter), residual(kDim); + std::vector codes(kDim); + for (const auto metric : {METRIC_L2, METRIC_IP}) { + for (const T value : {T{0}, T{0.5}, T{-0.5}}) { + std::fill(data.begin(), data.end(), kCenter + value); + std::fill(residual.begin(), residual.end(), value); + for (const size_t bits : {1U, 2U, 4U, 9U}) { + for (const double scale : {-1.0, 7.0}) { + RabitqConfig config; + config.t_const = scale; + T add, rescale, error; + quantize_full_single( + data.data(), + centroid.data(), + kDim, + bits, + codes.data(), + add, + rescale, + error, + metric, + config + ); + ASSERT_TRUE(std::all_of(codes.begin(), codes.end(), [&](int code) { + return code == codes.front(); + })); + const T centered_code = static_cast(codes.front()) - + static_cast((1U << bits) - 1) / 2; + // Constant residual/code vectors are collinear: the error is zero. + // L2 add = d*r^2 + 2*d*r*c; IP's centroid terms cancel to 1. + const T expected_add = + metric == METRIC_L2 + ? kDim * value * value + 2 * kDim * value * kCenter + : T{1}; + const T expected_rescale = value == 0 ? T{0} + : metric == METRIC_L2 + ? -2 * value / centered_code + : -value / centered_code; + const auto check = [&] { + EXPECT_EQ(add, expected_add); + EXPECT_NEAR( + rescale, + expected_rescale, + 8 * std::numeric_limits::epsilon() * + std::max(T{1}, std::abs(expected_rescale)) + ); + EXPECT_EQ(error, T{0}); + }; + check(); + rabitq_impl::xy_bits::code_factors( + residual.data(), + centroid.data(), + kDim, + codes.data(), + -static_cast((1U << bits) - 1) / 2, + add, + rescale, + error, + metric + ); + check(); + } + } + } + } +} + +TEST(RabitqFactorTest, CollinearMetricFactorsFloat) { check_collinear_factors(); } + +TEST(RabitqFactorTest, CollinearMetricFactorsDouble) { check_collinear_factors(); } + +TEST(RabitqFactorTest, InvalidMetricMessagesArePreserved) { + constexpr size_t kDim = 64; + const auto invalid_metric = static_cast(2); + std::vector data(kDim), centroid(kDim, 0), residual(kDim), centered(kDim); + std::vector codes(kDim); + float add, rescale, error; + const auto check = [](auto&& operation, const char* message) { + try { + operation(); + FAIL() << "Expected invalid metric rejection"; + } catch (const std::invalid_argument& exception) { + EXPECT_STREQ(exception.what(), message); + } + }; + for (const float value : {0.0F, 0.5F}) { + std::fill(data.begin(), data.end(), value); + std::fill(residual.begin(), residual.end(), value); + check( + [&] { + rabitq_impl::one_bit::one_bit_code_with_factor_into( + data.data(), + centroid.data(), + kDim, + codes.data(), + add, + rescale, + error, + invalid_metric, + residual.data(), + centered.data() + ); + }, + "Unsupported metric type in quantization" + ); + check( + [&] { + rabitq_impl::ex_bits::ex_bits_code_with_factor( + data.data(), + centroid.data(), + kDim, + 3, + codes.data(), + add, + rescale, + error, + invalid_metric + ); + }, + "Unsupported metric for ex_bits_code()" + ); + std::fill(codes.begin(), codes.end(), 1); + check( + [&] { + rabitq_impl::xy_bits::code_factors( + residual.data(), + centroid.data(), + kDim, + codes.data(), + -0.5F, + add, + rescale, + error, + invalid_metric + ); + }, + "Unsupported metric type in code_factors()" + ); + } +} + +} // namespace +} // namespace rabitqlib::quant diff --git a/tests/unit/rabitqlib/quantization/full_code_preparation_test.cpp b/tests/unit/rabitqlib/quantization/full_code_preparation_test.cpp new file mode 100644 index 00000000..bf682567 --- /dev/null +++ b/tests/unit/rabitqlib/quantization/full_code_preparation_test.cpp @@ -0,0 +1,81 @@ +#include + +#include +#include +#include + +#include "rabitqlib/core/api.hpp" +#include "rabitqlib/core/encoder.hpp" +#include "rabitqlib/core/metric.hpp" + +namespace rabitqlib::quant { +namespace { + +TEST(FullCodePreparation, SquaredNormUnderflowPreservesSigns) { + namespace api = rabitqlib::quant::v1; + for (const size_t dim : {32U, 64U, 96U}) { + std::vector data(dim), centroid(dim, 0); + for (size_t i = 0; i < dim; ++i) { + data[i] = i % 3 == 0 ? 0.0F : i % 3 == 1 ? 1e-30F : -1e-30F; + } + // Values remain nonzero floats, but every squared term underflows to zero. + ASSERT_GT(data[1], 0.0F); + ASSERT_EQ(data[1] * data[1], 0.0F); + for (const auto metric : {METRIC_L2, METRIC_IP}) { + for (const size_t bits : {1U, 2U, 4U, 9U}) { + for (const double scale : {-1.0, 7.0}) { + std::vector code(dim); + RabitqConfig config; + config.t_const = scale; + float add, rescale, error; + quantize_full_single( + data.data(), + centroid.data(), + dim, + bits, + code.data(), + add, + rescale, + error, + metric, + config + ); + EXPECT_EQ(add, metric == METRIC_L2 ? 0.0F : 1.0F); + EXPECT_EQ(rescale, 0.0F); + EXPECT_EQ(error, 0.0F); + for (size_t i = 0; i < dim; ++i) { + EXPECT_EQ(code[i], data[i] > 0 ? 1U << (bits - 1) : 0U); + } + + const api::Parameters parameters{dim, bits, metric, scale}; + const auto layout = api::CodeLayout::UnpackedU16; + const auto requirements = api::requirements(parameters, layout); + std::vector bytes(requirements.code_bytes); + std::vector workspace( + (requirements.workspace_bytes + sizeof(std::max_align_t) - 1) / + sizeof(std::max_align_t) + ); + const auto factors = api::encode( + parameters, + data.data(), + centroid.data(), + {layout, {bytes.data(), bytes.size()}, {}}, + {workspace.data(), workspace.size() * sizeof(std::max_align_t)} + ); + EXPECT_EQ(factors.add, add); + EXPECT_EQ(factors.rescale, rescale); + EXPECT_EQ(factors.error, error); + for (size_t i = 0; i < dim; ++i) { + const uint16_t value = static_cast( + bytes[2 * i] | (uint16_t{bytes[2 * i + 1]} << 8U) + ); + EXPECT_EQ(value, code[i]); + } + } + } + } + } +} + +} // namespace +} // namespace rabitqlib::quant