From ccf6344e3eb4ca7a91b7242f616568cf236611e9 Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Sat, 31 Oct 2026 06:15:34 +0800 Subject: [PATCH 1/4] fix: allow overlapping input in AVX FHT rotation Rotators may be called in place, and the generic and NEON backends copy the input with memmove. The shared AVX2/AVX-512 template used memcpy, which is undefined for overlapping ranges. Use memmove and test overlapping calls on every available backend. --- src/simd/rotator_kernels.hpp | 2 +- tests/unit/rabitqlib/utils/rotator_test.cpp | 62 +++++++++++++++++++++ 2 files changed, 63 insertions(+), 1 deletion(-) diff --git a/src/simd/rotator_kernels.hpp b/src/simd/rotator_kernels.hpp index 4befb336..70f9b7af 100644 --- a/src/simd/rotator_kernels.hpp +++ b/src/simd/rotator_kernels.hpp @@ -69,7 +69,7 @@ void fht_rotate_impl( } #undef RABITQ_FHT_CASE - std::memcpy(rotated_data, data, sizeof(float) * dim); + std::memmove(rotated_data, data, sizeof(float) * dim); std::fill(rotated_data + dim, rotated_data + padded_dim, 0.0F); if (trunc_dim == padded_dim) { diff --git a/tests/unit/rabitqlib/utils/rotator_test.cpp b/tests/unit/rabitqlib/utils/rotator_test.cpp index 2d0f3262..5eb2ef5c 100644 --- a/tests/unit/rabitqlib/utils/rotator_test.cpp +++ b/tests/unit/rabitqlib/utils/rotator_test.cpp @@ -242,6 +242,68 @@ TEST(FhtDispatchTest, BackendsMatchScalarButterfliesAndPadding) { } } +// Rotators allow in-place calls, so every backend must accept an output that +// overlaps its input. Run under ASan to catch overlapping memcpy. +TEST(FhtDispatchTest, BackendsAcceptOverlappingInputAndOutput) { + std::vector backends{ + simd::fht_rotate_generic, simd::fht_rotate}; +#if defined(__x86_64__) || defined(_M_X64) + if (cpu::has_avx2()) { + backends.push_back(simd::fht_rotate_avx2); + } + if (cpu::has_avx512_core()) { + backends.push_back(simd::fht_rotate_avx512); + } +#endif +#if defined(__aarch64__) || defined(_M_ARM64) + if (cpu::has_neon()) { + backends.push_back(simd::fht_rotate_neon); + } +#endif + for (size_t dim : {64U, 65U, 127U, 128U, 193U, 960U, 4097U}) { + SCOPED_TRACE(dim); + const size_t padded = (dim + 31) / 32 * 32; + size_t trunc = 64; + while (trunc * 2 <= dim) { + trunc *= 2; + } + const float fac = 1.0F / std::sqrt(static_cast(trunc)); + std::vector flip(padded / 2); + for (size_t i = 0; i < flip.size(); ++i) { + flip[i] = static_cast(i * 73 + i / 7); + } + std::vector input(dim); + for (size_t i = 0; i < dim; ++i) { + input[i] = std::sin(static_cast(i) * 0.13F); + } + for (auto backend : backends) { + std::vector expected(padded); + backend(input.data(), expected.data(), dim, padded, trunc, fac, flip.data()); + for (size_t offset : {0U, 1U}) { + SCOPED_TRACE(offset); + std::vector actual(padded + 2, 12345); + std::copy(input.begin(), input.end(), actual.begin() + 1); + backend( + actual.data() + 1, + actual.data() + offset, + dim, + padded, + trunc, + fac, + flip.data() + ); + EXPECT_EQ( + std::memcmp( + actual.data() + offset, expected.data(), padded * sizeof(float) + ), + 0 + ); + EXPECT_EQ(actual.back(), 12345); + } + } + } +} + TEST(FhtDispatchTest, PreservesZeroNormAndInnerProduct) { std::vector backends{ simd::fht_rotate_generic, simd::fht_rotate}; From e3e7539aa6f42f03a6de40d715ba6d8515b3790c Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Sat, 31 Oct 2026 06:15:34 +0800 Subject: [PATCH 2/4] fix: reject search and save on uninitialized HNSW A default-constructed HierarchicalNSW has no rotator, so search and save dereferenced a null pointer. Reuse the initialization check from resize and throw std::logic_error before touching the rotator or opening the output file. --- src/index/hnsw.cpp | 6 ++++++ tests/unit/rabitqlib/index/hnsw_test.cpp | 23 +++++++++++++++++++++++ 2 files changed, 29 insertions(+) diff --git a/src/index/hnsw.cpp b/src/index/hnsw.cpp index aef1adcc..67736a6d 100644 --- a/src/index/hnsw.cpp +++ b/src/index/hnsw.cpp @@ -152,6 +152,9 @@ HierarchicalNSW::HierarchicalNSW( } void HierarchicalNSW::save(const char* filename) const { + if (level0_record_bytes_ == 0) { + throw std::logic_error("HNSW index must be initialized before save"); + } std::ofstream output(rabitqlib::io_impl::filesystem_path(filename), std::ios::binary); if (!output.is_open()) { throw std::runtime_error("HNSW: cannot open index file for writing"); @@ -957,6 +960,9 @@ PID HierarchicalNSW::mutually_connect_new_element( std::vector>> HierarchicalNSW::search( const float* queries, size_t num_queries, size_t k, size_t ef_search, size_t num_threads ) { + if (level0_record_bytes_ == 0) { + throw std::logic_error("HNSW index must be initialized before search"); + } if (k == 0) { throw std::invalid_argument("HNSW k must be greater than zero"); } diff --git a/tests/unit/rabitqlib/index/hnsw_test.cpp b/tests/unit/rabitqlib/index/hnsw_test.cpp index c42ae916..c979a844 100644 --- a/tests/unit/rabitqlib/index/hnsw_test.cpp +++ b/tests/unit/rabitqlib/index/hnsw_test.cpp @@ -555,6 +555,29 @@ TEST(HnswResizeTest, RejectsUninitializedIndex) { } } +TEST(HnswUninitializedTest, RejectsSearchAndSave) { + HierarchicalNSW index; + const std::vector query(64, 1.0F); + for (const size_t num_queries : {1U, 0U}) { + SCOPED_TRACE(num_queries); + try { + index.search(query.data(), num_queries, 1, 10, 1); + FAIL() << "Searching an uninitialized index must fail"; + } catch (const std::logic_error& error) { + EXPECT_STREQ(error.what(), "HNSW index must be initialized before search"); + } + } + + const rabitq_test::TempFile path_file("rabitq_hnsw_uninitialized_save.index"); + try { + index.save(path_file.path().c_str()); + FAIL() << "Saving an uninitialized index must fail"; + } catch (const std::logic_error& error) { + EXPECT_STREQ(error.what(), "HNSW index must be initialized before save"); + } + EXPECT_FALSE(std::filesystem::exists(std::filesystem::u8path(path_file.path()))); +} + TEST(HnswResizeTest, AllowsResizeBeforeConstruction) { constexpr size_t kBuilt = 64; constexpr size_t kCapacity = 128; From 77a0e0cfa5d69a258bdf9dc029e4575921fd95aa Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Sat, 31 Oct 2026 06:15:34 +0800 Subject: [PATCH 3/4] fix: align zero-residual xy codes with other encoders For a residual with zero norm, ex_bits_code_with_norm zeroed the magnitude codes and then inverted them, so quantize_xy_single wrote 2^ex_bits - 1 while the full, split and v1 encoders write 0. Return before the inversion. All factors are zero on this path, so estimates are unchanged. --- include/rabitqlib/core/detail/rabitq_impl.hpp | 2 + .../quantization/xy_quantization_test.cpp | 49 +++++++++++++++++++ 2 files changed, 51 insertions(+) diff --git a/include/rabitqlib/core/detail/rabitq_impl.hpp b/include/rabitqlib/core/detail/rabitq_impl.hpp index ef49bc64..16de73e5 100644 --- a/include/rabitqlib/core/detail/rabitq_impl.hpp +++ b/include/rabitqlib/core/detail/rabitq_impl.hpp @@ -584,7 +584,9 @@ inline T ex_bits_code_with_norm( ConstRowMajorArrayMap res_arr(residual, 1, dim); T ipnorm_inv = 1; if (residual_norm == 0) { + // Match code_with_factor_from_residual(): no magnitude bits, so skip the revert. std::fill(ex_code, ex_code + dim, static_cast(0)); + return ipnorm_inv; } else { RowMajorArray owned_absolute; auto abs_res = scratch_array( diff --git a/tests/unit/rabitqlib/quantization/xy_quantization_test.cpp b/tests/unit/rabitqlib/quantization/xy_quantization_test.cpp index 77545f07..a3fa8cfd 100644 --- a/tests/unit/rabitqlib/quantization/xy_quantization_test.cpp +++ b/tests/unit/rabitqlib/quantization/xy_quantization_test.cpp @@ -219,6 +219,55 @@ TEST(XyQuantization, SignBitMatchesOneBitPathOnZeroResiduals) { } } +// A residual whose squared norm is zero takes the zero-residual path of the other +// encoders; the split code must reassemble to the same full code. The second case +// has nonzero components whose float squares underflow to zero. +TEST(XyQuantization, ZeroNormResidualMatchesFullCode) { + constexpr size_t kDim = 64; + constexpr size_t kExtraBits = 3; + + std::vector data(kDim, 1.0F); + std::vector tiny(kDim, 0.0F); + for (size_t j = 0; j < kDim; j += 2) { + tiny[j] = 1e-30F; + } + ASSERT_EQ(l2norm_sqr(tiny.data(), kDim), 0.0F); + const std::vector, std::vector>> cases{ + {data, data}, {tiny, std::vector(kDim, 0.0F)}}; + + for (const auto& [input, centroid] : cases) { + for (size_t base_bits : {1U, 2U}) { + for (MetricType metric : {METRIC_L2, METRIC_IP}) { + SCOPED_TRACE(base_bits); + SCOPED_TRACE(metric); + XySplitResult res = + SplitCode(input, centroid, kDim, base_bits, kExtraBits, metric); + + std::vector full_code(kDim); + float full_add = 0, full_rescale = 0, full_error = 0; + quant::quantize_full_single( + input.data(), + centroid.data(), + kDim, + base_bits + kExtraBits, + full_code.data(), + full_add, + full_rescale, + full_error, + metric + ); + for (size_t j = 0; j < kDim; ++j) { + const uint32_t total = + (static_cast(res.base_code[j]) << kExtraBits) | + static_cast(res.extra_code[j]); + EXPECT_EQ(total, full_code[j]) << "dim " << j; + } + EXPECT_EQ(res.f_rescale_full, 0.0F); + } + } + } +} + TEST(XyQuantization, BaseBitsOneMatchesExistingFactorsOnZeroResidualComponents) { constexpr size_t kDim = 64; constexpr size_t kExtraBits = 3; From d3f0cd2eb93e1c0ed36edb4c014b281b320a8dc0 Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Sat, 31 Oct 2026 06:15:34 +0800 Subject: [PATCH 4/4] fix: accept the maximum point count in PiPNN seeding QuantizedGraph and IVF accept up to kSearchBufferMaxPointCount points, but the PiPNN seed rejected a count equal to it. Use the same bound. --- src/index/pipnn.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/index/pipnn.cpp b/src/index/pipnn.cpp index 0d96608f..9080f34d 100644 --- a/src/index/pipnn.cpp +++ b/src/index/pipnn.cpp @@ -454,7 +454,7 @@ InitialGraph build_initial_graph( ) { validate_metric_type(metric_type); if (data == nullptr || dim == 0 || degree == 0 || degree >= count || - count >= buffer::kSearchBufferMaxPointCount || degree % 32 != 0) { + count > buffer::kSearchBufferMaxPointCount || degree % 32 != 0) { throw std::invalid_argument( "PiPNN seed requires data, positive dim, and a degree multiple of 32 below " "count"