diff --git a/CMakeLists.txt b/CMakeLists.txt index c856696..0c84ac3 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -50,6 +50,7 @@ set(RABITQ_COMMON_SOURCES ) set(RABITQ_AVX2_SOURCES + src/simd/quantization_avx2.cpp src/simd/pack_excode_avx2.cpp src/simd/space_excode_avx2.cpp src/simd/space_avx2.cpp @@ -59,6 +60,7 @@ set(RABITQ_AVX2_SOURCES ) set(RABITQ_AVX512_SOURCES + src/simd/quantization_avx512.cpp src/simd/pack_excode_avx512.cpp src/simd/space_excode_avx512.cpp src/simd/space_avx512.cpp diff --git a/include/rabitqlib/quantization/rabitq_impl.hpp b/include/rabitqlib/quantization/rabitq_impl.hpp index ec8f398..ef860ae 100644 --- a/include/rabitqlib/quantization/rabitq_impl.hpp +++ b/include/rabitqlib/quantization/rabitq_impl.hpp @@ -9,8 +9,8 @@ #include #include #include -#include #include +#include #include #include @@ -18,6 +18,7 @@ #include "rabitqlib/fastscan/fastscan.hpp" #include "rabitqlib/quantization/data_layout.hpp" #include "rabitqlib/quantization/pack_excode.hpp" +#include "rabitqlib/simd/quantization_dispatch.hpp" #include "rabitqlib/utils/space.hpp" namespace rabitqlib::quant::rabitq_impl { @@ -321,46 +322,90 @@ inline double best_rescale_factor(const T* o_abs, size_t dim, size_t ex_bits) { const double t_end = static_cast(max_code + 10) / max_o; const double t_start = t_end * kTightStart[ex_bits]; + // Both paths maximize cosine similarity between the input and quantized + // directions over the configured scale interval. Here a_i = o_abs[i] is a + // rotated, normalized residual magnitude, and q_i = code_i + 0.5. + // With ||a|| = 1, cosine(a, q) = dot(a, q / ||q||) = N / sqrt(S), where + // N = sum(a_i * q_i) and S = sum(q_i * q_i). + // SIMD evaluates whole vectors at selected scales and prunes intervals; + // its search and bounds are in src/simd/rescale_search.hpp. + if constexpr (std::is_same_v) { + const double t = simd::best_rescale_factor(o_abs, dim, max_code, t_start, t_end); + if (t >= 0) + return t; + } + + // Scalar event sweep, also used if SIMD is unavailable or cannot certify a + // winner within its work limit. Each heap entry holds a coordinate's next + // code-changing scale (code_i + 1) / a_i. Visiting only these events skips + // ranges with constant codes and updates N and S in O(1) per coordinate. + // Heap traversal is sequential and branch-dependent, so it offers little + // SIMD parallelism; the SIMD path instead prunes ranges of candidate codes. using Event = std::pair; - std::priority_queue, std::greater> next_t; + std::vector next_t; + next_t.reserve(dim); std::vector cur_o_bar(dim); double sqr_denominator = static_cast(dim) * 0.25; double numerator = 0; - auto enqueue_next = [&](size_t i) { - const double magnitude = static_cast(o_abs[i]); - // Never increment a saturated coordinate; skip zero coordinates. - if (magnitude > 0 && cur_o_bar[i] < max_code) { - const double next = static_cast(cur_o_bar[i] + 1) / magnitude; - if (next < t_end) - next_t.emplace(next, i); - } - }; - + // init quantization codes for each coordinate by t_start for (size_t i = 0; i < dim; ++i) { const double magnitude = static_cast(o_abs[i]); const int cur = quantized_level_at_scale(magnitude, t_start, max_code); cur_o_bar[i] = cur; sqr_denominator += (cur * cur) + cur; numerator += (cur + 0.5) * magnitude; - enqueue_next(i); + // Never increment a saturated coordinate; skip zero coordinates. + if (magnitude > 0 && cur < max_code) { + const double next = static_cast(cur + 1) / magnitude; + if (next < t_end) + next_t.emplace_back(next, i); + } } + std::make_heap(next_t.begin(), next_t.end(), std::greater{}); // The initial state may already be the best state inside the interval. double max_ip = numerator / std::sqrt(sqr_denominator); double best_t = t_start; while (!next_t.empty()) { - const double cur_t = next_t.top().first; + const double cur_t = next_t.front().first; // All coordinates crossing the same threshold change together. do { - const size_t i = next_t.top().second; - next_t.pop(); + const size_t i = next_t.front().second; ++cur_o_bar[i]; + // With the new code k, q changes from k-0.5 to k+0.5: delta S = 2k. sqr_denominator += 2.0 * cur_o_bar[i]; - numerator += static_cast(o_abs[i]); - enqueue_next(i); - } while (!next_t.empty() && next_t.top().first == cur_t); + const double magnitude = static_cast(o_abs[i]); + numerator += magnitude; + + Event next{t_end, i}; + if (cur_o_bar[i] < max_code) + next.first = static_cast(cur_o_bar[i] + 1) / magnitude; + if (next.first >= t_end) { + // No further event for this coordinate: fill the root's hole + // with the last heap entry, then restore the heap below. + next = next_t.back(); + next_t.pop_back(); + } + if (!next_t.empty()) { + // Replace the consumed root event with `next`. Its threshold + // cannot precede the old minimum, so one sift-down suffices; + // priority_queue pop() followed by push() would repair twice. + size_t parent = 0; + size_t child = 1; + while (child < next_t.size()) { + if (child + 1 < next_t.size() && next_t[child + 1] < next_t[child]) + ++child; + if (!(next_t[child] < next)) + break; + next_t[parent] = next_t[child]; + parent = child; + child = 2 * parent + 1; + } + next_t[parent] = next; + } + } while (!next_t.empty() && next_t.front().first == cur_t); const double cur_ip = numerator / std::sqrt(sqr_denominator); if (cur_ip > max_ip) { diff --git a/include/rabitqlib/simd/quantization_dispatch.hpp b/include/rabitqlib/simd/quantization_dispatch.hpp new file mode 100644 index 0000000..bdf3b21 --- /dev/null +++ b/include/rabitqlib/simd/quantization_dispatch.hpp @@ -0,0 +1,20 @@ +#pragma once + +#include + +namespace rabitqlib::simd { + +// Search normalized, nonnegative float magnitudes with max_code in [0, 255]. +// A negative result requests the event sweep when the bounded search cannot +// certify a winner within its work limit. +double best_rescale_factor( + const float* magnitudes, size_t dim, int max_code, double start, double end +); +double best_rescale_factor_avx2( + const float* magnitudes, size_t dim, int max_code, double start, double end +); +double best_rescale_factor_avx512( + const float* magnitudes, size_t dim, int max_code, double start, double end +); + +} // namespace rabitqlib::simd diff --git a/src/simd/dispatch.cpp b/src/simd/dispatch.cpp index 1a3eed6..56f379f 100644 --- a/src/simd/dispatch.cpp +++ b/src/simd/dispatch.cpp @@ -7,13 +7,31 @@ #include "rabitqlib/simd/fastscan_dispatch.hpp" #include "rabitqlib/simd/pack_excode_dispatch.hpp" +#include "rabitqlib/simd/quantization_dispatch.hpp" #include "rabitqlib/simd/rotator_dispatch.hpp" #include "rabitqlib/simd/space_dispatch.hpp" #include "rabitqlib/simd/warmup_dispatch.hpp" #include "rabitqlib/utils/cpu_features.hpp" +#include "rescale_search.hpp" namespace rabitqlib::simd { +namespace detail { + +RescaleScratch& get_thread_local_rescale_scratch(size_t dim) { + // Search evaluation is synchronous and does not re-enter the quantizer. + // Retain the largest buffers on each worker; smaller vectors overwrite only + // their active prefix without shrinking or zero-initializing the storage. + thread_local RescaleScratch scratch; + if (scratch.magnitudes.size() < dim) + scratch.magnitudes.resize(dim); + if (scratch.reciprocals.size() < dim) + scratch.reciprocals.resize(dim); + return scratch; +} + +} // namespace detail + [[noreturn]] static void missing_feature(const char* feature_name) { throw std::runtime_error( std::string(feature_name) + " requires AVX2/FMA or AVX512 support" @@ -29,6 +47,27 @@ static float ip_fxu0( return 0.0F; } +static double request_scalar_rescale_search(const float*, size_t, int, double, double) { + return -1; +} + +using BestRescaleFactorFn = double (*)(const float*, size_t, int, double, double); +const BestRescaleFactorFn kBestRescaleFactorFn = [] { + if (cpu::has_avx512_core()) { + return best_rescale_factor_avx512; + } else if (cpu::has_avx2()) { + return best_rescale_factor_avx2; + } else { + return request_scalar_rescale_search; + } +}(); + +double best_rescale_factor( + const float* magnitudes, size_t dim, int max_code, double start, double end +) { + return kBestRescaleFactorFn(magnitudes, dim, max_code, start, end); +} + static float missing_excode_ip(const float*, const uint8_t*, size_t) { missing_feature("excode ip functions"); } diff --git a/src/simd/quantization_avx2.cpp b/src/simd/quantization_avx2.cpp new file mode 100644 index 0000000..70e2f2a --- /dev/null +++ b/src/simd/quantization_avx2.cpp @@ -0,0 +1,167 @@ +#include + +#include +#include +#include + +#include "rabitqlib/simd/quantization_dispatch.hpp" +#include "rescale_search.hpp" +namespace rabitqlib::simd::detail { + +// Evaluate one candidate scale, processing four coordinates per iteration. +// For a_i = magnitudes[i], c_i = min(floor(scale*a_i), max_code), q_i = c_i+0.5, +// return N = sum(a_i*q_i), S = sum(q_i*q_i), and the constant-code interval +// [first, next). The shared search uses N^2/S to compare and bound squared cosine. +// No code vector is stored here; the caller emits codes once the search finishes. +// +// Magnitudes are float32 values promoted exactly to double. For levels <= 255, +// unequal comparable level/magnitude ratios differ by at least 2^-33 relatively. +// Reciprocal multiplication has much less error, so its extreme indices remain +// valid. Recompute those two endpoints by division to retain exact event scales. +static RescaleSearchState evaluate_scale_state_avx2( + const double* magnitudes, + const double* reciprocals, + size_t dim, + int max_code, + double t_start, + double t_end, + double scale +) { + __m256d z = _mm256_setzero_pd(), one = _mm256_set1_pd(1), half = _mm256_set1_pd(.5), + vt = _mm256_set1_pd(scale), vcap = _mm256_set1_pd(max_code); + __m256d guard = _mm256_set1_pd( + 2 * std::numeric_limits::epsilon() * (max_code + 1.0) + ), + vn = z, vd = z, vfirst = z, + vnext = _mm256_set1_pd(std::numeric_limits::infinity()); + __m256d positions = _mm256_setr_pd(0, 1, 2, 3), first_idx = _mm256_set1_pd(-1), + next_idx = first_idx; + size_t i = 0; + for (; i + 4 <= dim; i += 4) { + __m256d v = _mm256_loadu_pd(magnitudes + i), vi = _mm256_loadu_pd(reciprocals + i), + product = _mm256_mul_pd(v, vt), + c = _mm256_min_pd(_mm256_floor_pd(product), vcap); + // Floor is sufficient away from integer boundaries. The guard covers + // multiplication rounding near those boundaries; saturation removes + // the upper boundary. Zero magnitudes always keep code zero. + __m256d interior = _mm256_and_pd( + _mm256_cmp_pd(product, _mm256_add_pd(c, guard), _CMP_GT_OQ), + _mm256_or_pd( + _mm256_cmp_pd(c, vcap, _CMP_EQ_OQ), + _mm256_cmp_pd( + product, _mm256_sub_pd(_mm256_add_pd(c, one), guard), _CMP_LT_OQ + ) + ) + ); + unsigned boundary = + (~static_cast(_mm256_movemask_pd(interior))) & + static_cast(_mm256_movemask_pd(_mm256_cmp_pd(v, z, _CMP_GT_OQ))) & 15; + // Repair only flagged lanes using the heap's division-based thresholds. + // This ensures an event at k/a_i belongs to code k in both algorithms. + if (boundary) { + alignas(32) double corrected[4]; + _mm256_store_pd(corrected, c); + for (unsigned mask = boundary; mask; mask &= mask - 1) { + unsigned k = __builtin_ctz(mask); + corrected[k] = quantized_code_at_scale(magnitudes[i + k], scale, max_code); + } + c = _mm256_load_pd(corrected); + } + // vn/vd hold lane-local contributions to N/S. Accumulate in double; + // the shared search encloses rounding in N, while S remains exact. + __m256d q = _mm256_add_pd(c, half); + vn = _mm256_add_pd(vn, _mm256_mul_pd(v, q)); + vd = _mm256_add_pd(vd, _mm256_mul_pd(q, q)); + // A coordinate enters its current code at c_i/a_i (only if c_i>0) + // and leaves at (c_i+1)/a_i (only if a_i>0 and not saturated). + // Inactive coordinates contribute neutral values: zero to max(first) + // and infinity to min(next). Reciprocals avoid per-coordinate division. + __m256d af = _mm256_cmp_pd(c, z, _CMP_GT_OQ), + an = _mm256_and_pd( + _mm256_cmp_pd(c, vcap, _CMP_LT_OQ), _mm256_cmp_pd(v, z, _CMP_GT_OQ) + ); + __m256d f = _mm256_blendv_pd(z, _mm256_mul_pd(c, vi), af), + n = _mm256_blendv_pd( + _mm256_set1_pd(std::numeric_limits::infinity()), + _mm256_mul_pd(_mm256_add_pd(c, one), vi), + an + ); + __m256d fm = _mm256_and_pd(af, _mm256_cmp_pd(f, vfirst, _CMP_GT_OQ)), + nm = _mm256_and_pd(an, _mm256_cmp_pd(n, vnext, _CMP_LT_OQ)); + // Retain original coordinate indices along with each lane's extrema, + // so the final boundary can be recomputed by division below. + first_idx = _mm256_blendv_pd(first_idx, positions, fm); + next_idx = _mm256_blendv_pd(next_idx, positions, nm); + vfirst = _mm256_max_pd(vfirst, f); + vnext = _mm256_min_pd(vnext, n); + positions = _mm256_add_pd(positions, _mm256_set1_pd(4)); + } + // Reduce the four partial sums and select the global boundary coordinates. + // Index -1 means this lane never had an active boundary candidate. + alignas(32) double ns[4], ds[4], fs[4], ts[4]; + _mm256_store_pd(ns, vn); + _mm256_store_pd(ds, vd); + _mm256_store_pd(fs, vfirst); + _mm256_store_pd(ts, vnext); + RescaleSearchState s{ + ns[0] + ns[1] + ns[2] + ns[3], ds[0] + ds[1] + ds[2] + ds[3], t_start, t_end}; + alignas(32) double first_positions[4], next_positions[4]; + _mm256_store_pd(first_positions, first_idx); + _mm256_store_pd(next_positions, next_idx); + size_t fi[4], ni[4]; + for (size_t k = 0; k < 4; ++k) { + fi[k] = first_positions[k] < 0 ? std::numeric_limits::max() + : static_cast(first_positions[k]); + ni[k] = next_positions[k] < 0 ? std::numeric_limits::max() + : static_cast(next_positions[k]); + } + size_t fk = std::max_element(fs, fs + 4) - fs, nk = std::min_element(ts, ts + 4) - ts; + // Reciprocal products selected the coordinates; division now recovers the + // exact heap event scales, clipped to the configured search interval. + if (fi[fk] != std::numeric_limits::max()) + s.first = std::max( + s.first, + static_cast(quantized_code_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) / + 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); + double q = c + .5; + s.numerator += magnitudes[i] * q; + s.squared_norm += q * q; + if (c > 0) + s.first = std::max(s.first, static_cast(c) / magnitudes[i]); + if (c < max_code && magnitudes[i] > 0) + s.next = std::min(s.next, (c + 1.0) / magnitudes[i]); + } + return s; +} + +} // namespace rabitqlib::simd::detail + +namespace rabitqlib::simd { +double best_rescale_factor_avx2( + const float* magnitudes, size_t dim, int max_code, double start, double end +) { + // A negative result asks the caller to rerun the scalar heap search. + double best_t = 0; + return detail::try_find_best_scale_by_interval_search( + magnitudes, + dim, + max_code, + start, + end, + detail::evaluate_scale_state_avx2, + best_t + ) + ? best_t + : -1; +} +} // namespace rabitqlib::simd diff --git a/src/simd/quantization_avx512.cpp b/src/simd/quantization_avx512.cpp new file mode 100644 index 0000000..8be32f0 --- /dev/null +++ b/src/simd/quantization_avx512.cpp @@ -0,0 +1,158 @@ +#include + +#include +#include +#include + +#include "rabitqlib/simd/quantization_dispatch.hpp" +#include "rescale_search.hpp" +namespace rabitqlib::simd::detail { + +// Evaluate one candidate scale, processing eight coordinates per iteration. +// For a_i = magnitudes[i], c_i = min(floor(scale*a_i), max_code), q_i = c_i+0.5, +// return N = sum(a_i*q_i), S = sum(q_i*q_i), and the constant-code interval +// [first, next). The shared search uses N^2/S to compare and bound squared cosine. +// No code vector is stored here; the caller emits codes once the search finishes. +// +// Magnitudes are float32 values promoted exactly to double. For levels <= 255, +// unequal comparable level/magnitude ratios differ by at least 2^-33 relatively. +// Reciprocal multiplication has much less error, so its extreme indices remain +// valid. Recompute those two endpoints by division to retain exact event scales. +static RescaleSearchState evaluate_scale_state_avx512( + const double* magnitudes, + const double* reciprocals, + size_t dim, + int max_code, + double t_start, + double t_end, + double scale +) { + __m512d z = _mm512_setzero_pd(), one = _mm512_set1_pd(1), half = _mm512_set1_pd(.5), + vt = _mm512_set1_pd(scale), vcap = _mm512_set1_pd(max_code); + __m512d guard = _mm512_set1_pd( + 2 * std::numeric_limits::epsilon() * (max_code + 1.0) + ), + vn = z, vd = z, vfirst = z, + vnext = _mm512_set1_pd(std::numeric_limits::infinity()); + __m512i positions = _mm512_setr_epi64(0, 1, 2, 3, 4, 5, 6, 7), + first_idx = _mm512_set1_epi64(-1), next_idx = first_idx; + size_t i = 0; + for (; i + 8 <= dim; i += 8) { + __m512d v = _mm512_loadu_pd(magnitudes + i), vi = _mm512_loadu_pd(reciprocals + i), + product = _mm512_mul_pd(v, vt), + c = _mm512_min_pd(_mm512_floor_pd(product), vcap); + // Floor is sufficient away from integer boundaries. The guard covers + // multiplication rounding near those boundaries; saturation removes + // the upper boundary. Zero magnitudes always keep code zero. + __mmask8 interior = + _mm512_cmp_pd_mask(product, _mm512_add_pd(c, guard), _CMP_GT_OQ) & + (_mm512_cmp_pd_mask(c, vcap, _CMP_EQ_OQ) | + _mm512_cmp_pd_mask( + product, _mm512_sub_pd(_mm512_add_pd(c, one), guard), _CMP_LT_OQ + )); + __mmask8 boundary = ~interior & _mm512_cmp_pd_mask(v, z, _CMP_GT_OQ); + // Repair only flagged lanes using the heap's division-based thresholds. + // This ensures an event at k/a_i belongs to code k in both algorithms. + if (boundary) { + alignas(64) double corrected[8]; + _mm512_store_pd(corrected, c); + for (unsigned mask = boundary; mask; mask &= mask - 1) { + unsigned k = __builtin_ctz(mask); + corrected[k] = quantized_code_at_scale(magnitudes[i + k], scale, max_code); + } + c = _mm512_load_pd(corrected); + } + // vn/vd hold lane-local contributions to N/S. Accumulate in double; + // the shared search encloses rounding in N, while S remains exact. + __m512d q = _mm512_add_pd(c, half); + vn = _mm512_add_pd(vn, _mm512_mul_pd(v, q)); + vd = _mm512_add_pd(vd, _mm512_mul_pd(q, q)); + // A coordinate enters its current code at c_i/a_i (only if c_i>0) + // and leaves at (c_i+1)/a_i (only if a_i>0 and not saturated). + // Inactive coordinates contribute neutral values: zero to max(first) + // and infinity to min(next). Reciprocals avoid per-coordinate division. + __mmask8 activefirst = _mm512_cmp_pd_mask(c, z, _CMP_GT_OQ), + activenext = _mm512_cmp_pd_mask(c, vcap, _CMP_LT_OQ) & + _mm512_cmp_pd_mask(v, z, _CMP_GT_OQ); + __m512d f = _mm512_maskz_mul_pd(activefirst, c, vi), + n = _mm512_mask_mul_pd( + _mm512_set1_pd(std::numeric_limits::infinity()), + activenext, + _mm512_add_pd(c, one), + vi + ); + __mmask8 fm = activefirst & _mm512_cmp_pd_mask(f, vfirst, _CMP_GT_OQ), + nm = activenext & _mm512_cmp_pd_mask(n, vnext, _CMP_LT_OQ); + // Retain original coordinate indices along with each lane's extrema, + // so the final boundary can be recomputed by division below. + first_idx = _mm512_mask_mov_epi64(first_idx, fm, positions); + next_idx = _mm512_mask_mov_epi64(next_idx, nm, positions); + vfirst = _mm512_max_pd(vfirst, f); + vnext = _mm512_min_pd(vnext, n); + positions = _mm512_add_epi64(positions, _mm512_set1_epi64(8)); + } + // Reduce the eight partial sums and select the global boundary coordinates. + // Index -1 becomes size_t::max(), marking lanes without an active candidate. + RescaleSearchState s{ + _mm512_reduce_add_pd(vn), _mm512_reduce_add_pd(vd), t_start, t_end}; + alignas(64) size_t fi[8], ni[8]; + _mm512_store_si512(fi, first_idx); + _mm512_store_si512(ni, next_idx); + double f = _mm512_reduce_max_pd(vfirst), n = _mm512_reduce_min_pd(vnext); + // At least one lane equals each reduced extremum. Select its first set bit; + // tied boundary ratios describe the same event scale. + unsigned fk = __builtin_ctz(static_cast( + _mm512_cmp_pd_mask(vfirst, _mm512_set1_pd(f), _CMP_EQ_OQ) + )), + nk = __builtin_ctz(static_cast( + _mm512_cmp_pd_mask(vnext, _mm512_set1_pd(n), _CMP_EQ_OQ) + )); + // Reciprocal products selected the coordinates; division now recovers the + // exact heap event scales, clipped to the configured search interval. + if (fi[fk] != std::numeric_limits::max()) + s.first = std::max( + s.first, + static_cast(quantized_code_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) / + 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); + double q = c + .5; + s.numerator += magnitudes[i] * q; + s.squared_norm += q * q; + if (c > 0) + s.first = std::max(s.first, static_cast(c) / magnitudes[i]); + if (c < max_code && magnitudes[i] > 0) + s.next = std::min(s.next, (c + 1.0) / magnitudes[i]); + } + return s; +} + +} // namespace rabitqlib::simd::detail + +namespace rabitqlib::simd { +double best_rescale_factor_avx512( + const float* magnitudes, size_t dim, int max_code, double start, double end +) { + // A negative result asks the caller to rerun the scalar heap search. + double best_t = 0; + return detail::try_find_best_scale_by_interval_search( + magnitudes, + dim, + max_code, + start, + end, + detail::evaluate_scale_state_avx512, + best_t + ) + ? best_t + : -1; +} +} // namespace rabitqlib::simd diff --git a/src/simd/rescale_search.hpp b/src/simd/rescale_search.hpp new file mode 100644 index 0000000..22dbe19 --- /dev/null +++ b/src/simd/rescale_search.hpp @@ -0,0 +1,378 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +namespace rabitqlib::simd::detail { + +struct RescaleScratch { + std::vector magnitudes; + std::vector reciprocals; +}; + +// Keep TLS initialization/destruction in generic code shared by both backends. +RescaleScratch& get_thread_local_rescale_scratch(size_t dim); + +struct RescaleSearchState { + double numerator; // N = sum(a_i * q_i), where q_i = code_i + 0.5. + double squared_norm; // S = sum(q_i * q_i); the squared score is N^2 / S. + // This code vector is unchanged throughout [first, next), clipped to the + // configured search range. These boundaries let us skip constant-code spans. + // For the current codes c_i and maximum code M: + // first = max(t_start, max_{c_i > 0} c_i / a_i), + // next = min(t_end, min_{a_i > 0, c_i < M} (c_i + 1) / a_i). + // Zero coordinates never change; saturated coordinates have no next event. + double first; + double next; +}; + +// 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. +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. + 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. + 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; +} + +struct BoundedRescaleState { + RescaleSearchState state{}; + // Enclose the true dot product despite SIMD summation rounding. S is exact + // under the dimension guard, so these also give bounds on the true N^2 / S. + double numerator_lower = 0; + double numerator_upper = 0; + double score_lower = 0; + double score_upper = 0; +}; + +// Upper bound on squared cosine similarity for every code between two evaluated +// states. Here a_i = input[i] is a rotated, normalized residual magnitude, and +// q_i = code_i + 0.5. With ||a|| = 1, cosine(a, q)^2 = N^2 / S, where +// N = sum(a_i * q_i) and S = sum(q_i * q_i). All magnitudes are nonnegative, +// so maximizing squared cosine also maximizes cosine, the heap sweep's objective. +// This bounds the whole interval; comparing endpoint scores alone cannot +// exclude a better interior code because the score need not be monotonic. +static inline double interval_squared_cosine_upper_bound( + const BoundedRescaleState& left, const BoundedRescaleState& right +) { + const double endpoints = std::max(left.score_upper, right.score_upper); + // Adjacent plateaus have no intermediate codes. If the numerator upper + // bounds are reversed, every interior numerator is already <= the left + // upper bound, while its squared norm is >= the left norm; endpoints suffice. + if (left.state.next == right.state.first || + left.numerator_upper >= right.numerator_upper) + return endpoints; + + // Derive a bound without enumerating the events inside this interval: + // + // 1. An event at tau = k/a_i changes code_i from k-1 to k. It adds a_i + // to N and (k+0.5)^2 - (k-0.5)^2 = 2k to S. Thus delta S = 2*tau*delta N. + // 2. All intervening event scales lie between l = left.next and + // u = right.first. Therefore 2l*delta N <= delta S <= 2u*delta N. + // 3. For any intermediate numerator x, summing events forward from the + // left endpoint gives S(x) >= S0 + 2l*(x-N0). Summing the remaining + // events to the right endpoint gives S(x) >= S1 - 2u*(N1-x). + // + // N0/N1 are the true endpoint numerators; S0/S1 are their exact squared + // norms. Replacing N0/N1 by their upper bounds U0/U1 only lowers these two + // lines, so the denominator bound stays conservative: + // S(x) >= max(S0 + 2l*(x-U0), S1 - 2u*(U1-x)). + // Hence x^2 divided by this maximum is an upper bound on the score. + // Values x <= U0 are covered by left.score_upper; only [U0,U1] remains. + // lower_t/upper_t and all subsequent arithmetic are rounded outward. + const double lower_t = next_double_down(left.state.next); + const double upper_t = next_double_up(right.state.first); + const double delta_n_lower = + next_double_down(right.numerator_upper - left.numerator_upper); + const double delta_n_upper = + next_double_up(right.numerator_upper - left.numerator_upper); + const double delta_d_lower = + next_double_down(right.state.squared_norm - left.state.squared_norm); + const double delta_d_upper = + next_double_up(right.state.squared_norm - left.state.squared_norm); + + // Each active denominator line has the form A+B*x. Where it is positive, + // f(x) = x^2/(A+B*x) is convex: f''(x) = 2*A^2/(A+B*x)^3 >= 0. + // Its maximum on a section is at a section endpoint. We therefore need + // only the two endpoint scores and the intersection of the two lines. + // Solving for the intersection's offset p from U0 gives + // p = (2u*(U1-U0) - (S1-S0)) / (2*(u-l)). + // delta_d_* above enclose S1-S0; intersection_n/d_* below enclose the + // numerator and denominator of this expression, not the cosine score. + const double intersection_n_lower = + next_double_down(next_double_down(2 * upper_t * delta_n_lower) - delta_d_upper); + const double intersection_n_upper = + next_double_up(next_double_up(2 * upper_t * delta_n_upper) - delta_d_lower); + const double intersection_d_lower = 2 * next_double_down(upper_t - lower_t); + const double intersection_d_upper = 2 * next_double_up(upper_t - lower_t); + if (intersection_d_lower <= 0) { + // The event-scale gap is too small to divide by reliably. Use the + // looser monotonic bound N <= U1 and S >= S0 instead of pruning unsafely. + return next_double_up( + next_double_up(right.numerator_upper * right.numerator_upper) / + left.state.squared_norm + ); + } + // Clamp the enclosed intersection to [0, U1-U0]. An intersection outside + // that range introduces no interior maximum; the endpoints already cover it. + const double offset_lower = std::max( + 0.0, + std::min( + delta_n_lower, + next_double_down(std::max(0.0, intersection_n_lower) / intersection_d_upper) + ) + ); + const double offset_upper = std::max( + 0.0, + std::min( + delta_n_upper, + next_double_up(std::max(0.0, intersection_n_upper) / intersection_d_lower) + ) + ); + // Both denominator lines increase with x. Use the larger possible offset + // for the numerator and the smaller one for the denominator, producing an + // upper score bound even when the intersection cannot be represented exactly. + // S0 is also a valid lower bound on every intermediate squared norm. + const double numerator = next_double_up(left.numerator_upper + offset_upper); + const double left_denominator = next_double_down( + left.state.squared_norm + next_double_down(2 * lower_t * offset_lower) + ); + const double right_denominator = next_double_down( + right.state.squared_norm - + next_double_up(2 * upper_t * next_double_up(delta_n_upper - offset_lower)) + ); + const double denominator = + std::max({left.state.squared_norm, left_denominator, right_denominator}); + return std::max( + endpoints, next_double_up(next_double_up(numerator * numerator) / denominator) + ); +} + +// 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. +template +static bool try_find_best_scale_by_interval_search( + const float* input, + size_t dim, + int max_code, + double t_start, + double t_end, + Evaluator evaluator, + double& best_t +) { + if (dim == 0 || max_code == 0) { + best_t = t_start; + return true; + } + // S is a sum of quarter integers. This guard makes every addition exact. + const uint64_t largest_odd = static_cast(2 * max_code + 1); + if (dim > (uint64_t{1} << 53) / (largest_odd * largest_odd)) + return false; + + // Promote float inputs exactly once and cache 1/a_i for repeated candidate + // evaluations. The thread-local buffers are reused across vectors; each + // call overwrites its active prefix. Zero reciprocals are masked by evaluators. + auto& scratch = get_thread_local_rescale_scratch(dim); + auto& magnitudes = scratch.magnitudes; + auto& reciprocals = scratch.reciprocals; + for (size_t i = 0; i < dim; ++i) { + magnitudes[i] = static_cast(input[i]); + reciprocals[i] = input[i] > 0 ? 1.0 / magnitudes[i] : 0; + } + + // Float magnitudes times half-integer codes are exact in double. Positive + // summation has relative error gamma_(dim+4), bounded by this delta under + // the dimension guard above, including the SIMD horizontal reduction. + const double delta = + static_cast(dim + 8) * std::numeric_limits::epsilon(); + const double lower_divisor = next_double_up(1 + delta); + const double upper_divisor = next_double_down(1 - delta); + // These cap search work, not quantization accuracy: hitting either cap + // requests the heap fallback instead of returning the best-so-far scale. + constexpr size_t kEvaluationBudget = 512; + constexpr size_t kMaxDepth = 64; + size_t evaluations = 0; + double best_lower = 0; // Largest certified score lower bound seen so far. + BoundedRescaleState best; + bool have_best = false; + + const auto try_evaluate_scale_with_score_bounds = [&](double t, + BoundedRescaleState& result) { + if (evaluations == kEvaluationBudget) + return false; + ++evaluations; + result.state = evaluator( + magnitudes.data(), reciprocals.data(), 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. + 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 + // delta, then N_hat/(1+delta) <= N <= N_hat/(1-delta). Round outward + // again when squaring/dividing so pruning never relies on a rounded-up + // candidate score or a rounded-down interval bound. + result.numerator_lower = next_double_down(result.state.numerator / lower_divisor); + result.numerator_upper = next_double_up(result.state.numerator / upper_divisor); + result.score_lower = next_double_down( + next_double_down(result.numerator_lower * result.numerator_lower) / + result.state.squared_norm + ); + result.score_upper = next_double_up( + next_double_up(result.numerator_upper * result.numerator_upper) / + result.state.squared_norm + ); + best_lower = std::max(best_lower, result.score_lower); + + // Replace the winner only when the score enclosures are disjoint. + // Equal S means the same code plateau, since codes are monotonic in t. + // Overlapping scores from different plateaus require the heap's tie order. + if (!have_best || result.score_lower > best.score_upper) { + best = result; + have_best = true; + } else if (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; + }; + + BoundedRescaleState first; + BoundedRescaleState last; + // t_end is excluded by the heap search too. Evaluate the preceding double + // so an event exactly at t_end is not accidentally admitted here. + if (!try_evaluate_scale_with_score_bounds(t_start, first) || + !try_evaluate_scale_with_score_bounds(next_double_down(t_end), last)) + return false; + + const auto try_search_interval = [&](auto&& self, + const BoundedRescaleState& left, + const BoundedRescaleState& right, + double upper_bound, + size_t depth) -> bool { + // Codes increase coordinate-wise with scale, so equal squared norms + // mean equal codes. Adjacent plateaus have no intermediate states; + // otherwise prune only when the entire interval is provably worse. + if (left.state.squared_norm == right.state.squared_norm || + left.state.next == right.state.first || upper_bound < best_lower) + return true; + 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. + 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; + if (!try_evaluate_scale_with_score_bounds(midpoint, middle) || + middle.state.squared_norm <= left.state.squared_norm || + 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. + 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. + if (left_upper > right_upper) { + return self(self, left, middle, left_upper, depth + 1) && + self(self, middle, right, right_upper, depth + 1); + } + return self(self, middle, right, right_upper, depth + 1) && + self(self, left, middle, left_upper, depth + 1); + }; + if (!try_search_interval( + try_search_interval, + first, + last, + interval_squared_cosine_upper_bound(first, last), + 0 + )) + return false; + // Return the first scale of the winning plateau, not the sampled midpoint, + // to preserve the event sweep's canonical scale convention. + best_t = best.state.first; + return true; +} + +} // namespace rabitqlib::simd::detail diff --git a/tests/unit/rabitqlib/quantization/rabitq_test.cpp b/tests/unit/rabitqlib/quantization/rabitq_test.cpp index 6bd694a..d5cd5be 100644 --- a/tests/unit/rabitqlib/quantization/rabitq_test.cpp +++ b/tests/unit/rabitqlib/quantization/rabitq_test.cpp @@ -2,15 +2,141 @@ #include +#include #include #include #include #include #include +#include +#include namespace rabitqlib::quant { namespace { +int level_from_thresholds(double magnitude, double t, int max_code) { + int result = 0; + if (magnitude > 0) { + for (int level = 1; level <= max_code; ++level) { + if (static_cast(level) / magnitude <= t) + result = level; + } + } + return result; +} + +template +double sorted_event_rescale_factor(const std::vector& magnitudes, size_t bits) { + if (magnitudes.empty()) + return 0; + const double max_o = *std::max_element(magnitudes.begin(), magnitudes.end()); + if (max_o == 0) + return 0; + + const int max_code = (1 << bits) - 1; + const double end = static_cast(max_code + 10) / max_o; + const double start = end * rabitq_impl::ex_bits::kTightStart[bits]; + std::vector levels(magnitudes.size()); + std::vector> events; + double denominator = static_cast(magnitudes.size()) * 0.25; + double numerator = 0; + for (size_t i = 0; i < magnitudes.size(); ++i) { + const double magnitude = magnitudes[i]; + const int level = level_from_thresholds(magnitude, start, max_code); + levels[i] = level; + denominator += (level * level) + level; + numerator += (level + 0.5) * magnitude; + if (magnitude > 0) { + for (int next = level + 1; next <= max_code; ++next) { + const double threshold = static_cast(next) / magnitude; + if (threshold < end) + events.emplace_back(threshold, i); + } + } + } + // Sorting all legal events independently checks heap ordering, including ties. + std::sort(events.begin(), events.end()); + double best_t = start; + double best_ip = numerator / std::sqrt(denominator); + for (size_t event = 0; event < events.size();) { + const double threshold = events[event].first; + do { + const size_t i = events[event++].second; + denominator += 2.0 * ++levels[i]; + numerator += static_cast(magnitudes[i]); + } while (event < events.size() && events[event].first == threshold); + const double ip = numerator / std::sqrt(denominator); + if (ip > best_ip) { + best_ip = ip; + best_t = threshold; + } + } + return best_t; +} + +template +void check_rescale_search_against_sorted_events() { + for (size_t dim : {0U, 1U, 2U, 3U, 7U, 31U, 64U, 65U}) { + for (int pattern = 0; pattern < 4; ++pattern) { + std::vector magnitudes(dim); + uint32_t state = 42; + double norm_sq = 0; + for (size_t i = 0; i < dim; ++i) { + double value = 0; + if (pattern == 1) + value = i == dim / 2 ? 1 : 0; + else if (pattern == 2) + value = static_cast(i % 4); + else if (pattern == 3) { + state = state * 1664525U + 1013904223U; + value = i % 5 == 0 ? 0 : 1.0 + (state >> 16U); + } + magnitudes[i] = static_cast(value); + norm_sq += value * value; + } + if (norm_sq > 0) { + const double norm = std::sqrt(norm_sq); + for (T& magnitude : magnitudes) + magnitude = static_cast(static_cast(magnitude) / norm); + } + for (size_t bits = 0; bits <= 8; ++bits) { + SCOPED_TRACE( + ::testing::Message() << "sizeof(T)=" << sizeof(T) << " dim=" << dim + << " pattern=" << pattern << " bits=" << bits + ); + const double expected_t = sorted_event_rescale_factor(magnitudes, bits); + EXPECT_EQ( + rabitq_impl::ex_bits::best_rescale_factor(magnitudes.data(), dim, bits), + expected_t + ); + std::vector expected_code(dim); + double ipnorm = 0; + for (size_t i = 0; i < dim; ++i) { + const double magnitude = magnitudes[i]; + const int level = + level_from_thresholds(magnitude, expected_t, (1 << bits) - 1); + expected_code[i] = static_cast(level); + ipnorm += (level + 0.5) * magnitude; + } + T expected_factor = ipnorm == 0 ? T{1} : static_cast(1.0 / ipnorm); + if (!std::isnormal(expected_factor)) + expected_factor = T{1}; + std::vector code(dim, 0xFF); + const T factor = rabitq_impl::ex_bits::quantize_ex( + magnitudes.data(), code.data(), dim, bits + ); + EXPECT_EQ(code, expected_code); + EXPECT_EQ(factor, expected_factor); + } + } + } +} + +TEST(RabitqRescaleSearchTest, MatchesSortedLegalEventsExactly) { + check_rescale_search_against_sorted_events(); + check_rescale_search_against_sorted_events(); +} + TEST(RabitqQuantizedLevelTest, MatchesThresholdsAroundRoundingBoundaries) { const std::array magnitudes = { 0.1, diff --git a/tests/unit/rabitqlib/quantization/rescale_search_test.cpp b/tests/unit/rabitqlib/quantization/rescale_search_test.cpp new file mode 100644 index 0000000..f6d53d8 --- /dev/null +++ b/tests/unit/rabitqlib/quantization/rescale_search_test.cpp @@ -0,0 +1,258 @@ +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/quantization/rabitq_impl.hpp" +#include "rabitqlib/simd/quantization_dispatch.hpp" +#include "rabitqlib/utils/cpu_features.hpp" + +namespace rabitqlib::quant { +namespace { + +using RescaleSearch = double (*)(const float*, size_t, int, double, double); + +std::vector levels_at_threshold( + const std::vector& magnitudes, double scale, int max_code +) { + std::vector levels(magnitudes.size()); + for (size_t i = 0; i < magnitudes.size(); ++i) { + if (magnitudes[i] == 0) + continue; + for (int level = 1; level <= max_code; ++level) { + if (static_cast(level) / magnitudes[i] <= scale) + levels[i] = level; + } + } + return levels; +} + +long double squared_cosine( + const std::vector& magnitudes, const std::vector& levels +) { + long double numerator = 0; + long double denominator = 0; + long double input_norm = 0; + for (size_t i = 0; i < magnitudes.size(); ++i) { + const long double magnitude = magnitudes[i]; + const long double value = levels[i] + 0.5L; + numerator += value * magnitude; + denominator += value * value; + input_norm += magnitude * magnitude; + } + return input_norm == 0 ? 0 : numerator * numerator / (denominator * input_norm); +} + +long double exhaustive_best_squared_cosine( + const std::vector& magnitudes, int max_code, double start, double end +) { + std::vector levels = levels_at_threshold(magnitudes, start, max_code); + std::vector> events; + for (size_t i = 0; i < magnitudes.size(); ++i) { + if (magnitudes[i] == 0) + continue; + for (int level = levels[i] + 1; level <= max_code; ++level) { + const double threshold = static_cast(level) / magnitudes[i]; + if (threshold < end) + events.emplace_back(threshold, i); + } + } + std::sort(events.begin(), events.end()); + long double best = squared_cosine(magnitudes, levels); + for (size_t event = 0; event < events.size();) { + const double threshold = events[event].first; + do { + ++levels[events[event++].second]; + } while (event < events.size() && events[event].first == threshold); + // Recompute in long double rather than sharing the search's reductions + // or accumulating floating-point error across thousands of events. + best = std::max(best, squared_cosine(magnitudes, levels)); + } + return best; +} + +void normalize_magnitudes(std::vector& magnitudes) { + long double norm_squared = 0; + for (float magnitude : magnitudes) + norm_squared += static_cast(magnitude) * magnitude; + if (norm_squared == 0) + return; + const long double norm = std::sqrt(norm_squared); + for (float& magnitude : magnitudes) + magnitude = static_cast(magnitude / norm); +} + +std::pair search_interval( + const std::vector& magnitudes, size_t bits +) { + const double maximum = *std::max_element(magnitudes.begin(), magnitudes.end()); + if (maximum == 0) + return {0, 1}; + const double end = static_cast((1 << bits) + 9) / maximum; + return {end * rabitq_impl::ex_bits::kTightStart[bits], end}; +} + +void expect_optimal_canonical_scale( + const std::vector& magnitudes, int max_code, double start, double end, double t +) { + ASSERT_TRUE(std::isfinite(t)); + ASSERT_GE(t, start); + ASSERT_LT(t, end); + const auto levels = levels_at_threshold(magnitudes, t, max_code); + double first = start; + for (size_t i = 0; i < magnitudes.size(); ++i) { + if (levels[i] > 0) + first = std::max(first, static_cast(levels[i]) / magnitudes[i]); + } + EXPECT_EQ(t, first); + const long double actual = squared_cosine(magnitudes, levels); + const long double optimum = + exhaustive_best_squared_cosine(magnitudes, max_code, start, end); + // Float magnitudes times half-integer levels fit exactly in double. The + // tolerance covers reduction/comparison rounding, without permitting a + // meaningful loss in the optimized cosine objective. + constexpr long double kTolerance = 8 * std::numeric_limits::epsilon(); + EXPECT_GE(actual + kTolerance, optimum); +} + +class RabitqRescaleBackendTest : public ::testing::TestWithParam { + protected: + void SetUp() override { + if (GetParam() ? !cpu::has_avx512_core() : !cpu::has_avx2()) + GTEST_SKIP() << "Requested SIMD backend is unavailable"; + } + + RescaleSearch search() const { + return GetParam() ? simd::best_rescale_factor_avx512 + : simd::best_rescale_factor_avx2; + } +}; + +TEST_P(RabitqRescaleBackendTest, MatchesExhaustiveObjectiveAndCanonicalThreshold) { + for (size_t dim : + {1U, 2U, 3U, 4U, 5U, 7U, 8U, 9U, 15U, 16U, 17U, 31U, 33U, 65U, 129U}) { + for (int pattern = 0; pattern < 7; ++pattern) { + std::vector magnitudes(dim); + std::mt19937 random(8143); + std::normal_distribution gaussian; + for (size_t i = 0; i < dim; ++i) { + if (pattern == 1) + magnitudes[i] = i == dim / 2 ? 1.0F : 0; + else if (pattern == 2) + magnitudes[i] = static_cast(i % 4); + else if (pattern == 3) + magnitudes[i] = 1.0F + static_cast(i % 7) * + std::numeric_limits::epsilon(); + else if (pattern == 4) + magnitudes[i] = i % 5 == 0 ? 0 : std::abs(gaussian(random)); + else if (pattern == 5) + magnitudes[i] = i + 1 == dim + ? std::numeric_limits::denorm_min() + : std::ldexp(1.0F, -static_cast(i % 150)); + else if (pattern == 6) + magnitudes[i] = std::abs(gaussian(random)); + } + normalize_magnitudes(magnitudes); + for (size_t bits = 1; bits <= 8; ++bits) { + SCOPED_TRACE( + ::testing::Message() + << "dim=" << dim << " pattern=" << pattern << " bits=" << bits + ); + const auto [start, end] = search_interval(magnitudes, bits); + const int max_code = (1 << bits) - 1; + const double t = search()(magnitudes.data(), dim, max_code, start, end); + ASSERT_TRUE(std::isfinite(t)); + if (t >= 0) + expect_optimal_canonical_scale(magnitudes, max_code, start, end, t); + } + } + } +} + +TEST_P(RabitqRescaleBackendTest, CompletesGaussianSearchWithoutFallback) { + std::vector magnitudes(128); + std::mt19937 random(42); + std::normal_distribution gaussian; + for (float& magnitude : magnitudes) + magnitude = std::abs(gaussian(random)); + normalize_magnitudes(magnitudes); + const auto [start, end] = search_interval(magnitudes, 7); + const double t = search()(magnitudes.data(), magnitudes.size(), 127, start, end); + ASSERT_GE(t, 0); + expect_optimal_canonical_scale(magnitudes, 127, start, end, t); +} + +TEST_P(RabitqRescaleBackendTest, ReusesScratchAcrossDimensionsAndThreads) { + std::vector> inputs; + std::vector> intervals; + std::vector expected; + const auto backend = search(); + std::mt19937 random(71); + std::normal_distribution gaussian; + for (size_t dim : {129U, 7U, 257U, 1U, 65U, 513U, 3U, 128U, 0U, 17U}) { + std::vector magnitudes(dim); + for (float& magnitude : magnitudes) + magnitude = dim == 17 ? 0 : std::abs(gaussian(random)); + normalize_magnitudes(magnitudes); + const auto interval = + dim == 0 ? std::pair{0, 1} : search_interval(magnitudes, 7); + expected.push_back( + backend(magnitudes.data(), dim, 127, interval.first, interval.second) + ); + intervals.push_back(interval); + inputs.push_back(std::move(magnitudes)); + } + + size_t mismatches = 0; +#pragma omp parallel for num_threads(8) schedule(dynamic, 1) reduction(+ : mismatches) + for (size_t call = 0; call < 160; ++call) { + const size_t index = call % inputs.size(); + const auto& magnitudes = inputs[index]; + const auto [start, end] = intervals[index]; + const double actual = + backend(magnitudes.data(), magnitudes.size(), 127, start, end); + mismatches += actual != expected[index]; + } + EXPECT_EQ(mismatches, 0U); +} + +TEST_P(RabitqRescaleBackendTest, FallsBackWhenNearlyCollinearWinnersAreAmbiguous) { + std::vector magnitudes(129); + for (size_t i = 0; i < magnitudes.size(); ++i) + magnitudes[i] = + 1.0F + static_cast(i) * std::numeric_limits::epsilon(); + normalize_magnitudes(magnitudes); + const auto [start, end] = search_interval(magnitudes, 8); + const double t = search()(magnitudes.data(), magnitudes.size(), 255, start, end); + ASSERT_TRUE(std::isfinite(t)); + ASSERT_LT(t, 0); + + std::vector code(magnitudes.size()); + const float factor = rabitq_impl::ex_bits::quantize_ex( + magnitudes.data(), code.data(), magnitudes.size(), 8 + ); + EXPECT_TRUE(std::isnormal(factor)); + EXPECT_GT(factor, 0); + const std::vector levels(code.begin(), code.end()); + const long double actual = squared_cosine(magnitudes, levels); + const long double optimum = exhaustive_best_squared_cosine(magnitudes, 255, start, end); + EXPECT_GE(actual + 8 * std::numeric_limits::epsilon(), optimum); +} + +INSTANTIATE_TEST_SUITE_P( + ExplicitSimd, + RabitqRescaleBackendTest, + ::testing::Bool(), + [](const ::testing::TestParamInfo& info) { + return info.param ? "Avx512" : "Avx2"; + } +); + +} // namespace +} // namespace rabitqlib::quant