From 0cb67e54efe59514ed0f90b7245949f7f3360915 Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Wed, 7 Oct 2026 13:45:41 +0800 Subject: [PATCH 1/3] Support 32-dimension padding and simplify optimized kernels --- AGENTS.md | 14 +- docs/docs/index/hnsw.md | 5 + docs/docs/index/ivf.md | 18 +- docs/docs/index/qg.md | 4 + docs/docs/rabitq/rotator.md | 18 +- include/rabitqlib/index/hnsw/hnsw.hpp | 14 +- include/rabitqlib/index/hnsw/hnsw_quant.hpp | 4 +- include/rabitqlib/index/ivf/ivf.hpp | 50 ++- include/rabitqlib/index/query.hpp | 3 +- .../rabitqlib/quantization/pack_excode.hpp | 3 +- .../rabitqlib/quantization/rabitq_impl.hpp | 13 +- include/rabitqlib/utils/bitops.hpp | 24 +- include/rabitqlib/utils/rotator.hpp | 6 +- include/rabitqlib/utils/space.hpp | 65 ++- include/rabitqlib/utils/warmup_space.hpp | 36 +- src/index/hnsw_search_avx2_kernels.hpp | 219 +--------- src/index/hnsw_search_avx512_kernels.hpp | 109 +---- src/index/qg.cpp | 13 +- src/simd/fht_kernels.hpp | 29 +- src/simd/hnsw_neon_kernels.hpp | 24 +- src/simd/pack_excode_avx2.cpp | 44 +- src/simd/pack_excode_avx512.cpp | 45 ++- src/simd/pack_excode_kernels.hpp | 216 +++++++--- src/simd/packed_tail_avx2.hpp | 145 +++++++ src/simd/portable_generic.cpp | 84 +++- src/simd/query_neon.cpp | 9 +- src/simd/rotator_avx2.cpp | 2 +- src/simd/rotator_avx512.cpp | 46 ++- src/simd/rotator_kernels.hpp | 96 ++--- src/simd/space_avx2.cpp | 171 ++++---- src/simd/space_avx512.cpp | 95 +++-- src/simd/space_excode_avx2.cpp | 26 +- src/simd/space_excode_avx512.cpp | 109 +++-- src/simd/space_packed_neon.cpp | 18 +- src/simd/warmup_avx2.cpp | 159 +------- src/simd/warmup_avx512.cpp | 63 +-- src/simd/warmup_kernels.hpp | 208 ++++++++++ tests/python/fixtures/README.md | 17 +- .../fixtures/hnsw_legacy_padding65_4bit.index | Bin 0 -> 1105 bytes .../fixtures/ivf_legacy_padding65_1bit.index | Bin 0 -> 1526 bytes .../fixtures/ivf_raw_v1_padding65.index | Bin 0 -> 2318 bytes tests/python/test_hnsw.py | 25 ++ tests/python/test_ivf.py | 49 ++- tests/python/test_portable_indexes.py | 34 +- tests/python/test_symqg.py | 6 +- tests/unit/rabitqlib/index/hnsw_test.cpp | 63 +++ tests/unit/rabitqlib/index/ivf_test.cpp | 6 +- tests/unit/rabitqlib/index/qg_test.cpp | 57 +++ .../rabitqlib/quantization/rabitq_test.cpp | 42 ++ tests/unit/rabitqlib/utils/excode_ip_test.cpp | 8 +- tests/unit/rabitqlib/utils/mask_ip_test.cpp | 17 +- tests/unit/rabitqlib/utils/neon_hnsw_test.cpp | 10 +- .../rabitqlib/utils/portable_kernel_test.cpp | 108 ++++- .../utils/query_preparation_test.cpp | 11 +- tests/unit/rabitqlib/utils/rotator_test.cpp | 375 ++++++++++++------ tests/unit/rabitqlib/utils/space_test.cpp | 125 +++++- 56 files changed, 2100 insertions(+), 1060 deletions(-) create mode 100644 src/simd/packed_tail_avx2.hpp create mode 100644 src/simd/warmup_kernels.hpp create mode 100644 tests/python/fixtures/hnsw_legacy_padding65_4bit.index create mode 100644 tests/python/fixtures/ivf_legacy_padding65_1bit.index create mode 100644 tests/python/fixtures/ivf_raw_v1_padding65.index diff --git a/AGENTS.md b/AGENTS.md index b8caacb0..fbe3e5a1 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -147,7 +147,8 @@ Recommended: pack/unpack, and estimation coverage. Test both `METRIC_L2` and `METRIC_IP` where supported. - IVF/HNSW quantized total bits are one sign bit plus `ex_bits`, with totals 1 through 9. IVF also accepts `bits == 32`: one-bit filtering plus owned original float32 vectors - in place of extra-bit codes; its raw-mode persistence has a magic/version header. + in place of extra-bit codes. New IVF files have a magic/version header and an explicit + padded dimension; legacy quantized files and raw v1 files infer 64-dimension padding. IVF search defaults to HACC for 4–9 bits and standard FastScan for 1–3 bits or raw storage. SymphonyQG supports raw storage (`quantization_bits == 0`) and quantized storage at 4 or 8 bits. @@ -170,13 +171,20 @@ do not use it where callers expect in-place mutation or pointer identity. ### Rotation and padded dimensions -Indexes quantize and search in the rotated, padded domain. FHT/Kac rotation pads to a multiple of -64; callers provide vectors in the original `dim`, while internal code and quantization buffers +Indexes quantize and search in the rotated, padded domain. FHT/Kac, RaBitQKMeans, QGKMeans, +SymphonyQG, IVF (all bit widths), and HNSW pad to multiples of 32. Packed extra-bit +and sign codes retain the historical full 64-coordinate blocks and compact final +32-coordinate tails. Binary query planes round their scratch storage up to +whole uint64_t words; 32-coordinate padding keeps HNSW record strides PID-aligned. Callers +provide vectors in the original `dim`, while internal code and quantization buffers generally use `padded_dim`. A buffer allocated for `dim` must never receive `padded_dim` values. Centroids, data, and queries compared by the same estimator must be in the same domain. Rotator state is part of persisted index state. Loading an index must restore the exact rotation used at construction; generating a new random rotator produces silently incorrect distances. +Honor stored padded dimensions instead of recomputing them from current defaults. Old IVF formats +omit this field and must retain their historical 64-dimension padding when loaded. +Unreleased pad16 files whose stored dimension is not divisible by 32 are rejected. ### Distance conventions diff --git a/docs/docs/index/hnsw.md b/docs/docs/index/hnsw.md index b5351931..62bae2c8 100644 --- a/docs/docs/index/hnsw.md +++ b/docs/docs/index/hnsw.md @@ -242,3 +242,8 @@ A file that has removals cannot be read by earlier releases. Releases 0.3.7 through 0.5.1 check every cluster ID on `load` and reject the file as invalid instead of returning removed points. Releases before 0.3.7 do not check cluster IDs, so do not open such a file with them. + +New HNSW indexes pad dimensions to multiples of 32 for all supported bit widths. +Loading preserves the saved padded dimension and rotation, so older indexes +using multiples of 64 remain readable. Older library versions cannot load new +files whose padded dimension is not a multiple of 64. diff --git a/docs/docs/index/ivf.md b/docs/docs/index/ivf.md index 48cf6b14..123a1397 100644 --- a/docs/docs/index/ivf.md +++ b/docs/docs/index/ivf.md @@ -87,14 +87,24 @@ Construction reuses a rotation buffer of at most 32 vectors per worker. This temporary storage uses at most `workers * 32 * padded_dim * sizeof(float)` bytes, independently of cluster sizes. +New indexes pad the rotated dimension to a multiple of 32 for every supported +bit width, including raw mode (32). This applies to both rotator types. +Extra-bit layouts retain full 64-coordinate blocks and store any final +32-coordinate tail compactly. + After construction, you can directly save the index file to disk: ```c++ ivf.save(output_index_file); ``` -Raw-mode files use a magic/version header and include both the original vectors -and the rotation state. Loading restores the mode automatically, so querying -does not require an external dataset. Existing quantized files retain their -format and remain compatible; versions before 0.3.3 cannot load raw-mode files. +New files use the `RABQIDX1` magic, version 1, storage flags, and an explicit +padded dimension before the index metadata. Loading restores the storage mode, +rotation state, and padded dimension; raw files also contain the original vectors, +so querying does not require an external dataset. + +Legacy unversioned quantized files and raw v1 files remain readable and retain +their original 64-dimension padding. Saving a loaded legacy index writes the new +header without changing its padded dimension or encoded data. Older library +versions cannot read the new format. ### Data Layout The main data layout for our IVF is organized as follows: diff --git a/docs/docs/index/qg.md b/docs/docs/index/qg.md index abbe3428..e2a73761 100644 --- a/docs/docs/index/qg.md +++ b/docs/docs/index/qg.md @@ -7,6 +7,10 @@ FastScan estimates neighbor distances; visited vertices are scored using stored raw vectors or quantized codes. Set `quantization_bits=0` for raw vectors (default), or `4`/`8` for quantized SymphonyQG. +All three modes pad the rotated domain to a multiple of 32. Loading an older +index preserves its saved padded dimension, including historical multiples of 64. +Older library versions that enforce 64-dimension padding cannot load new files +whose padded dimension is not a multiple of 64. Quantized vectors share a rotated global centroid. Refinement temporarily reconstructs source vectors and estimates distances to stored target codes; reverse edges are rescored because these estimates are directional. Raw refinement diff --git a/docs/docs/rabitq/rotator.md b/docs/docs/rabitq/rotator.md index f6159a55..7d91c6e3 100644 --- a/docs/docs/rabitq/rotator.md +++ b/docs/docs/rabitq/rotator.md @@ -6,10 +6,20 @@ RaBitQLib provides two types of random rotation. All implementations sample and By default, the library uses the `FFHT + Kac’s Walk` method. The default rotator accepts input dimensions from 64 through 65,536 inclusive, -including 16,384. It pads to a multiple of 64 and uses AVX or NEON kernels +including 16,384. It pads to a multiple of 32 and uses AVX or NEON kernels for power-of-two blocks through `2^16`. Allocate the output using `rotator->size()`, which can exceed the input dimension. +RaBitQKMeans, QGKMeans, SymphonyQG, IVF, and HNSW also use multiples of 32 +for all their supported bit widths. +Index loading preserves the padded domain used when the index was built. +Reducing the padded dimension changes the sampled rotation and quantization even +with the same seed. Newly built indexes can therefore have different recall and +latency; smaller padding does not guarantee faster search or unchanged accuracy. +When restoring a standalone rotator saved with earlier 64-dimension padding, +pass that original padded dimension to `choose_rotator` before loading its state. +The saved sign bytes do not contain dimension metadata. + The implementation can be found in `rotator.hpp`. ```css @@ -30,12 +40,12 @@ The implementation can be found in `rotator.hpp`. #include "rabitqlib/utils/rotator.hpp" int main() { - const size_t dim = 769; // deliberately not a multiple of 64 + const size_t dim = 769; // deliberately not a multiple of 32 std::unique_ptr> rotator( rabitqlib::choose_rotator(dim, rabitqlib::RotatorType::FhtKacRotator) ); std::vector x(dim, 1.0F); - std::vector x_prime(rotator->size()); // 832 elements + std::vector x_prime(rotator->size()); // 800 elements rotator->rotate(x.data(), x_prime.data()); } ``` @@ -46,7 +56,7 @@ to quantization and estimation routines that operate on rotated vectors. For a random orthogonal transformation, select `RotatorType::MatrixRotator`. It keeps the input dimension by default; pass a padded dimension explicitly if -subsequent packed-code operations require a multiple of 64. Its storage and +subsequent packed-code operations require a multiple of 32. Its storage and computation cost are quadratic in the dimension. ## FFHT + Kac’s Walk diff --git a/include/rabitqlib/index/hnsw/hnsw.hpp b/include/rabitqlib/index/hnsw/hnsw.hpp index 8e8ed6f3..d25781be 100644 --- a/include/rabitqlib/index/hnsw/hnsw.hpp +++ b/include/rabitqlib/index/hnsw/hnsw.hpp @@ -464,11 +464,11 @@ inline HierarchicalNSW::HierarchicalNSW( max_elements_ = max_elements; dim_ = dim; rotator_.reset(choose_rotator( - dim, RotatorType::FhtKacRotator, round_up_to_multiple(dim_, 64) + dim, RotatorType::FhtKacRotator, round_up_to_multiple(dim_, 32) )); padded_dim_ = rotator_->size(); /* check size */ - assert(padded_dim_ % 64 == 0); + assert(padded_dim_ % 32 == 0); assert(padded_dim_ >= dim_); ex_bits_ = total_bits - 1; @@ -476,7 +476,7 @@ inline HierarchicalNSW::HierarchicalNSW( throw std::invalid_argument("HNSW quantization bits must be in [1, 9]"); }; - assert(padded_dim_ % 64 == 0); + assert(padded_dim_ % 32 == 0); ip_func_ = select_excode_ipfunc(ex_bits_); @@ -500,9 +500,7 @@ inline HierarchicalNSW::HierarchicalNSW( sizeof(PID); // (# of edges + edges) + (cluster_id) + (external label) offsetExData_ = offsetBinData_ + size_bin_data_; // (# of edges + edges) + (cluster_id) // + (external label) + (BinData) - size_data_per_element_ = - offsetExData_ + size_ex_data_; // (# of edges + edges) + (cluster_id) + (external - // label) + (BinData) + (ExData) + size_data_per_element_ = offsetExData_ + size_ex_data_; data_level0_memory_ = memory::huge_page_allocate(max_elements_ * size_data_per_element_); if (data_level0_memory_ == nullptr) { @@ -652,8 +650,8 @@ inline void HierarchicalNSW::load(const char* filename) { loaded.num_cluster_ > buffer::kSearchBufferMaxPointCount || (element_count != 0 && loaded.num_cluster_ == 0) || loaded.dim_ < 64 || loaded.dim_ > rotator_impl::FhtKacRotator::kMaxDim || - loaded.padded_dim_ != round_up_to_multiple(loaded.dim_, 64) || - loaded.ex_bits_ > 8 || + loaded.padded_dim_ < loaded.dim_ || loaded.padded_dim_ % 32 != 0 || + loaded.padded_dim_ > rotator_impl::FhtKacRotator::kMaxDim || loaded.ex_bits_ > 8 || (loaded.metric_type_ != METRIC_L2 && loaded.metric_type_ != METRIC_IP) || loaded.M_ < 2 || loaded.M_ > 10000 || loaded.maxM_ != loaded.M_ || loaded.maxM0_ != 2 * loaded.M_ || loaded.ef_construction_ < loaded.M_ || diff --git a/include/rabitqlib/index/hnsw/hnsw_quant.hpp b/include/rabitqlib/index/hnsw/hnsw_quant.hpp index 4558ff29..a138240f 100644 --- a/include/rabitqlib/index/hnsw/hnsw_quant.hpp +++ b/include/rabitqlib/index/hnsw/hnsw_quant.hpp @@ -60,11 +60,11 @@ inline void HierarchicalNSW::reconstruct_rotated(PID id, float* out) const { const uint8_t* sign_code = bin.bin_code(); for (size_t word = 0; word * 64 < padded_dim_; ++word) { uint64_t bits = 0; - std::memcpy(&bits, sign_code + (word * sizeof(uint64_t)), sizeof(bits)); const size_t base = word * 64; const size_t span = std::min(64, padded_dim_ - base); + std::memcpy(&bits, sign_code + (word * sizeof(uint64_t)), span / 8); for (size_t j = 0; j < span; ++j) { - combined[base + j] = static_cast((bits >> (63 - j)) & 1ULL); + combined[base + j] = static_cast((bits >> (span - 1 - j)) & 1ULL); } } diff --git a/include/rabitqlib/index/ivf/ivf.hpp b/include/rabitqlib/index/ivf/ivf.hpp index 5168c20c..97e7672a 100644 --- a/include/rabitqlib/index/ivf/ivf.hpp +++ b/include/rabitqlib/index/ivf/ivf.hpp @@ -67,6 +67,8 @@ class IVF { static constexpr uint64_t kRawFormatMagic = 0x3157415251464252ULL; static constexpr uint32_t kRawFormatVersion = 1; + static constexpr uint64_t kFormatMagic = 0x3158444951424152ULL; // "RABQIDX1" + static constexpr uint32_t kFormatVersion = 1; void quantize_cluster(Cluster&, const std::vector&, const float*, const float*, float*, float*, const quant::RabitqConfig&); @@ -297,7 +299,7 @@ inline IVF::IVF( if (dim == 0 || dim > (std::numeric_limits::max() / 32) - 64) { throw std::invalid_argument("IVF dimension is invalid or too large"); } - padded_dim_ = round_up_to_multiple(dim_, 64); + padded_dim_ = round_up_to_multiple(dim_, 32); const size_t max_size = std::numeric_limits::max(); if (exceeds_storage(n) || cluster_num > max_size / sizeof(float) / padded_dim_ || (type == RotatorType::MatrixRotator && dim > max_size / sizeof(float) / padded_dim_ @@ -306,7 +308,7 @@ inline IVF::IVF( } rotator_.reset(choose_rotator(dim, type, padded_dim_)); /* check size */ - assert(padded_dim_ % 64 == 0); + assert(padded_dim_ % 32 == 0); assert(padded_dim_ >= dim_); } @@ -780,14 +782,14 @@ inline void IVF::save(const char* filename) const { std::ofstream output(rabitqlib::io_impl::filesystem_path(filename), std::ios::binary); output.exceptions(std::ios::failbit | std::ios::badbit); - if (raw_reranking_) { - output.write( - reinterpret_cast(&kRawFormatMagic), sizeof(kRawFormatMagic) - ); - output.write( - reinterpret_cast(&kRawFormatVersion), sizeof(kRawFormatVersion) - ); - } + const uint32_t flags = raw_reranking_ ? 1 : 0; + const uint64_t stored_padded_dim = padded_dim_; + output.write(reinterpret_cast(&kFormatMagic), sizeof(kFormatMagic)); + output.write(reinterpret_cast(&kFormatVersion), sizeof(kFormatVersion)); + output.write(reinterpret_cast(&flags), sizeof(flags)); + output.write( + reinterpret_cast(&stored_padded_dim), sizeof(stored_padded_dim) + ); /* Save meta data */ output.write(reinterpret_cast(&num_), sizeof(size_t)); @@ -842,8 +844,23 @@ inline void IVF::load(const char* filename) { input.seekg(0); uint64_t magic = 0; input.read(reinterpret_cast(&magic), sizeof(magic)); - loaded.raw_reranking_ = magic == kRawFormatMagic; - if (loaded.raw_reranking_) { + const bool explicit_padding = magic == kFormatMagic; + if (explicit_padding) { + uint32_t version = 0, flags = 0; + uint64_t stored_padded_dim = 0; + input.read(reinterpret_cast(&version), sizeof(version)); + input.read(reinterpret_cast(&flags), sizeof(flags)); + input.read(reinterpret_cast(&stored_padded_dim), sizeof(stored_padded_dim)); + if (version != kFormatVersion) { + throw std::runtime_error("Unsupported IVF index version"); + } + if (flags > 1 || stored_padded_dim > std::numeric_limits::max()) { + throw std::runtime_error("Invalid IVF index metadata"); + } + loaded.raw_reranking_ = flags == 1; + loaded.padded_dim_ = static_cast(stored_padded_dim); + } else if (magic == kRawFormatMagic) { + loaded.raw_reranking_ = true; uint32_t version = 0; input.read(reinterpret_cast(&version), sizeof(version)); if (version != kRawFormatVersion) { @@ -868,7 +885,14 @@ inline void IVF::load(const char* filename) { loaded.dim_ > (std::numeric_limits::max() / 32) - 64) { throw std::runtime_error("Invalid IVF index metadata"); } - loaded.padded_dim_ = round_up_to_multiple(loaded.dim_, 64); + if (!explicit_padding) { + // Both historical formats omit padded_dim and always used 64-coordinate blocks. + loaded.padded_dim_ = round_up_to_multiple(loaded.dim_, 64); + } + if (loaded.padded_dim_ < loaded.dim_ || loaded.padded_dim_ % 32 != 0 || + loaded.padded_dim_ > (std::numeric_limits::max() / 32) - 64) { + throw std::runtime_error("Invalid padded dimension in IVF index file"); + } // Bound every payload by the actual file before allocating or multiplying sizes. size_t remaining = file_bytes - static_cast(input.tellg()); const auto consume = [&remaining](size_t count, size_t bytes) { diff --git a/include/rabitqlib/index/query.hpp b/include/rabitqlib/index/query.hpp index ec4f1460..51ddf47b 100644 --- a/include/rabitqlib/index/query.hpp +++ b/include/rabitqlib/index/query.hpp @@ -161,7 +161,8 @@ class SplitSingleQuery { size_t metric_type = METRIC_L2, size_t base_bits = 1 ) - : rotated_query_(rotated_query), QueryBin_(padded_dim * kNumBits / 64, 0) { + : rotated_query_(rotated_query) + , QueryBin_(div_round_up(padded_dim, 64) * kNumBits, 0) { if (base_bits < 1 || base_bits > 8 || ex_bits > 8 || base_bits + ex_bits > 9) { throw std::invalid_argument( "SplitSingleQuery requires base_bits in [1, 8], ex_bits in [0, 8], " diff --git a/include/rabitqlib/quantization/pack_excode.hpp b/include/rabitqlib/quantization/pack_excode.hpp index 75b39719..15b8ae02 100644 --- a/include/rabitqlib/quantization/pack_excode.hpp +++ b/include/rabitqlib/quantization/pack_excode.hpp @@ -58,8 +58,7 @@ inline void packing_8bit_excode(const uint8_t* o_raw, uint8_t* o_compact, size_t * * @param o_raw unpacked code, code for each dim is represented by uint8 * @param o_compact compact format of code - * @param dim dimension of code, NOTICE: different num of bits requried different - * dimension padding, dim should obey the corresponding requirement + * @param dim dimension of code, a multiple of 32 for every supported bit width * @param ex_bits number of bits used for code */ inline void packing_rabitqplus_code( diff --git a/include/rabitqlib/quantization/rabitq_impl.hpp b/include/rabitqlib/quantization/rabitq_impl.hpp index d6a2b81d..2425f9c9 100644 --- a/include/rabitqlib/quantization/rabitq_impl.hpp +++ b/include/rabitqlib/quantization/rabitq_impl.hpp @@ -229,7 +229,7 @@ inline void one_bit_compact_code( ); } -// Requires a positive padded_dim divisible by sizeof(TC) * 8. +// Requires a positive padded_dim divisible by 8. template < typename T, typename TC, @@ -248,19 +248,17 @@ inline void one_bit_compact_codes( FError f_error, MetricType metric_type = METRIC_L2 ) { - constexpr size_t kTypeBits = sizeof(TC) * 8; - #pragma omp parallel for if (Parallel) for (std::ptrdiff_t index = 0; index < static_cast(num); ++index) { const size_t i = static_cast(index); T add; T rescale; T error; - one_bit_compact_code( + one_bit_compact_code_to_bytes( data + (padded_dim * i), centroid, padded_dim, - compact_code + (padded_dim / kTypeBits * i), + reinterpret_cast(compact_code) + (padded_dim / 8 * i), add, rescale, error, @@ -272,8 +270,9 @@ inline void one_bit_compact_codes( } } -// Encoding requires a positive padded_dim divisible by 8; index pipelines use 64. -// packed_code needs ceil(num / 32) * 32 * (padded_dim / 8) bytes, including tail padding. +// Encoding requires a positive padded_dim divisible by 8; index pipelines use multiples +// of 32. packed_code needs ceil(num / 32) * 32 * (padded_dim / 8) bytes, including tail +// padding. template < typename T, bool Parallel = false, diff --git a/include/rabitqlib/utils/bitops.hpp b/include/rabitqlib/utils/bitops.hpp index 610df2e2..0702a666 100644 --- a/include/rabitqlib/utils/bitops.hpp +++ b/include/rabitqlib/utils/bitops.hpp @@ -1,6 +1,8 @@ #pragma once +#include #include +#include #if defined(_MSC_VER) #include @@ -8,6 +10,23 @@ namespace rabitqlib::bitops { +// Load 1..8 bytes of a compact word, zero-extending a final partial word. +inline uint64_t load_word(const void* data, size_t bytes) { + if (bytes == 8) { + uint64_t word; + std::memcpy(&word, data, sizeof(word)); + return word; + } + if (bytes == 4) { + uint32_t word; + std::memcpy(&word, data, sizeof(word)); + return word; + } + uint64_t word = 0; + std::memcpy(&word, data, bytes); + return word; +} + inline unsigned popcount32(uint32_t value) { #if defined(_MSC_VER) value -= (value >> 1) & 0x55555555U; @@ -19,7 +38,10 @@ inline unsigned popcount32(uint32_t value) { } inline unsigned popcount64(uint64_t value) { -#if defined(_MSC_VER) +#if defined(_MSC_VER) || \ + ((defined(__x86_64__) || defined(__i386__)) && !defined(__POPCNT__)) + // Inline the software fallback on generic x86 instead of calling libgcc + // once per word. ISA-specific builds still use their native instruction. value -= (value >> 1) & 0x5555555555555555ULL; value = (value & 0x3333333333333333ULL) + ((value >> 2) & 0x3333333333333333ULL); return static_cast( diff --git a/include/rabitqlib/utils/rotator.hpp b/include/rabitqlib/utils/rotator.hpp index 0e5f3215..47dbfd10 100644 --- a/include/rabitqlib/utils/rotator.hpp +++ b/include/rabitqlib/utils/rotator.hpp @@ -52,10 +52,10 @@ inline size_t padding_requirement(size_t dim, RotatorType type) { return dim; } if (type == RotatorType::FhtKacRotator) { - if (dim > std::numeric_limits::max() - 63) { + if (dim > std::numeric_limits::max() - 31) { throw std::invalid_argument("Rotator dimension is too large to pad"); } - return round_up_to_multiple(dim, 64); + return round_up_to_multiple(dim, 32); } throw std::invalid_argument("Invalid rotator type in padding_requirement()"); } @@ -147,7 +147,7 @@ class FhtKacRotator : public Rotator { size_t dim, size_t padded_dim, uint32_t seed = std::random_device{}() ) : Rotator(dim, padded_dim) { - if (dim < 64 || padded_dim < dim || padded_dim % 64 != 0 || padded_dim > kMaxDim) { + if (dim < 64 || padded_dim < dim || padded_dim % 32 != 0 || padded_dim > kMaxDim) { throw std::invalid_argument("Unsupported dimension for FhtKacRotator"); } flip_.resize((padded_dim / kByteLen) * 4); diff --git a/include/rabitqlib/utils/space.hpp b/include/rabitqlib/utils/space.hpp index 9ddfa432..3c2633a8 100644 --- a/include/rabitqlib/utils/space.hpp +++ b/include/rabitqlib/utils/space.hpp @@ -146,7 +146,9 @@ inline T normalize_vec( } // Pack 0/1 data into words without requiring the destination byte storage to -// contain aligned, lifetime-started T objects. +// contain aligned, lifetime-started T objects. Length is a multiple of 8; the +// final partial word stores only its used bytes, with the first coordinate at +// its highest used bit. Full words retain their original layout. template inline void pack_binary_to_bytes( const int* __restrict__ binary_code, uint8_t* __restrict__ compact_code, size_t length @@ -154,12 +156,17 @@ inline void pack_binary_to_bytes( constexpr size_t kTypeBits = sizeof(T) * 8; for (size_t i = 0; i < length; i += kTypeBits) { + const size_t width = std::min(kTypeBits, length - i); T cur = 0; - for (size_t j = 0; j < kTypeBits; ++j) { - cur |= (static_cast(binary_code[i + j]) << (kTypeBits - 1 - j)); + for (size_t j = 0; j < width; j += 8) { + uint8_t byte = 0; + for (size_t bit = 0; bit < 8; ++bit) { + byte |= static_cast(binary_code[i + j + bit] << (7 - bit)); + } + cur = static_cast((cur << 8) | byte); } - std::memcpy(compact_code, &cur, sizeof(cur)); - compact_code += sizeof(cur); + std::memcpy(compact_code, &cur, width / 8); + compact_code += width / 8; } } @@ -332,6 +339,8 @@ static inline uint64_t reverse_bits_u64(uint64_t n) { return n; } +// padded_dim is a multiple of 32. Transposed queries need ceil(padded_dim / 64) * b_query +// uint64_t words, including a zero-extended final partial word in each bit plane. void new_transpose_bin(const uint16_t* q, uint64_t* tq, size_t padded_dim, size_t b_query); void new_transpose_bin_512( @@ -354,14 +363,14 @@ inline float ip_x0_q( size_t b_query ) { auto num_blk = padded_dim / 64; - const auto* it_data = data; + const auto* it_data = reinterpret_cast(data); const auto* it_query = query; size_t ip = 0; size_t ppc = 0; for (size_t i = 0; i < num_blk; ++i) { - uint64_t x = *static_cast(it_data); + const uint64_t x = rabitqlib::bitops::load_word(it_data, sizeof(uint64_t)); ppc += bitops::popcount64(x); for (size_t j = 0; j < b_query; ++j) { @@ -369,7 +378,16 @@ inline float ip_x0_q( ip += (bitops::popcount64(x & y) << j); it_query++; } - it_data++; + it_data += sizeof(uint64_t); + } + + if (padded_dim % 64 != 0) { + const uint64_t x = rabitqlib::bitops::load_word(it_data, (padded_dim % 64) / 8); + ppc += bitops::popcount64(x); + for (size_t j = 0; j < b_query; ++j) { + ip += (bitops::popcount64(x & *it_query) << j); + ++it_query; + } } return (delta * static_cast(ip)) + (vl * static_cast(ppc)); @@ -377,11 +395,20 @@ inline float ip_x0_q( static inline uint32_t ip_bin_bin(const uint64_t* q, const uint64_t* d, size_t padded_dim) { uint64_t ret = 0; - size_t iter = padded_dim / 64; - for (size_t i = 0; i < iter; ++i) { - ret += bitops::popcount64((*d) & (*q)); - q++; - d++; + const size_t words = padded_dim / 64; + const auto* query = reinterpret_cast(q); + const auto* data = reinterpret_cast(d); + for (size_t i = 0; i < words; ++i) { + ret += bitops::popcount64( + bitops::load_word(data + i * 8, 8) & bitops::load_word(query + i * 8, 8) + ); + } + if (padded_dim % 64 != 0) { + const size_t bytes = (padded_dim % 64) / 8; + ret += bitops::popcount64( + bitops::load_word(data + words * 8, bytes) & + bitops::load_word(query + words * 8, bytes) + ); } return ret; } @@ -390,7 +417,7 @@ inline uint32_t ip_byte_bin( const uint64_t* q, const uint64_t* d, size_t padded_dim, size_t b_query ) { uint32_t ret = 0; - size_t offset = (padded_dim / 64); + size_t offset = (padded_dim + 63) / 64; for (size_t i = 0; i < b_query; i++) { ret += (ip_bin_bin(q, d, padded_dim) << i); q += offset; @@ -400,9 +427,13 @@ inline uint32_t ip_byte_bin( inline size_t popcount(const uint64_t* __restrict__ d, size_t length) { size_t ret = 0; - for (size_t i = 0; i < length / 64; ++i) { - ret += bitops::popcount64((*d)); - ++d; + const size_t words = length / 64; + const auto* data = reinterpret_cast(d); + for (size_t i = 0; i < words; ++i) { + ret += bitops::popcount64(bitops::load_word(data + i * 8, 8)); + } + if (length % 64 != 0) { + ret += bitops::popcount64(bitops::load_word(data + words * 8, (length % 64) / 8)); } return ret; } diff --git a/include/rabitqlib/utils/warmup_space.hpp b/include/rabitqlib/utils/warmup_space.hpp index 9bbe73be..749ac28e 100644 --- a/include/rabitqlib/utils/warmup_space.hpp +++ b/include/rabitqlib/utils/warmup_space.hpp @@ -35,14 +35,14 @@ inline float warmup_ip_x0_q( [[maybe_unused]] size_t _b_query = 0 // not used ) { auto num_blk = padded_dim / 64; - const auto* it_data = data; + const auto* it_data = reinterpret_cast(data); const auto* it_query = query; size_t ip = 0; size_t ppc = 0; for (size_t i = 0; i < num_blk; ++i) { - uint64_t x = *static_cast(it_data); + const uint64_t x = rabitqlib::bitops::load_word(it_data, sizeof(uint64_t)); ppc += bitops::popcount64(x); for (size_t j = 0; j < b_query; ++j) { @@ -50,7 +50,16 @@ inline float warmup_ip_x0_q( ip += (bitops::popcount64(x & y) << j); it_query++; } - it_data++; + it_data += sizeof(uint64_t); + } + + if (padded_dim % 64 != 0) { + const uint64_t x = rabitqlib::bitops::load_word(it_data, (padded_dim % 64) / 8); + ppc += bitops::popcount64(x); + for (size_t j = 0; j < b_query; ++j) { + ip += (bitops::popcount64(x & *it_query) << j); + ++it_query; + } } return (delta * static_cast(ip)) + (vl * static_cast(ppc)); @@ -67,24 +76,5 @@ inline float warmup_ip_x0_q( size_t _padded_dim = 0, // not used size_t _b_query = 0 // not used ) { - auto num_blk = padded_dim / 64; - const auto* it_data = data; - const auto* it_query = query; - - size_t ip = 0; - size_t ppc = 0; - - for (size_t i = 0; i < num_blk; ++i) { - uint64_t x = *static_cast(it_data); - ppc += rabitqlib::bitops::popcount64(x); - - for (size_t j = 0; j < b_query; ++j) { - uint64_t y = *static_cast(it_query); - ip += (rabitqlib::bitops::popcount64(x & y) << j); - it_query++; - } - it_data++; - } - - return (delta * static_cast(ip)) + (vl * static_cast(ppc)); + return rabitqlib::warmup_ip_x0_q(data, query, delta, vl, padded_dim); } diff --git a/src/index/hnsw_search_avx2_kernels.hpp b/src/index/hnsw_search_avx2_kernels.hpp index a1dfb37a..4b8f6924 100644 --- a/src/index/hnsw_search_avx2_kernels.hpp +++ b/src/index/hnsw_search_avx2_kernels.hpp @@ -1,112 +1,20 @@ #pragma once -#include - #include #include -#include +#include "../simd/packed_tail_avx2.hpp" +#include "../simd/warmup_kernels.hpp" #include "rabitqlib/index/query.hpp" namespace rabitqlib::hnsw::detail { -// Helper: AVX2 64-bit Popcount; Mula's method -static inline __m256i hnsw_popcount_avx2(__m256i v) { - // Lookup table for population count of 0-15 - const __m256i lookup = _mm256_setr_epi8( - 0, - 1, - 1, - 2, - 1, - 2, - 2, - 3, - 1, - 2, - 2, - 3, - 2, - 3, - 3, - 4, - 0, - 1, - 1, - 2, - 1, - 2, - 2, - 3, - 1, - 2, - 2, - 3, - 2, - 3, - 3, - 4 - ); - const __m256i low_mask = _mm256_set1_epi8(0x0f); - - // Count low nibbles - __m256i lo = _mm256_and_si256(v, low_mask); - __m256i cnt_lo = _mm256_shuffle_epi8(lookup, lo); - - // Count high nibbles - __m256i hi = _mm256_and_si256(_mm256_srli_epi16(v, 4), low_mask); - __m256i cnt_hi = _mm256_shuffle_epi8(lookup, hi); - - // Add counts (bytes) - __m256i cnt_bytes = _mm256_add_epi8(cnt_lo, cnt_hi); - - // Sum bytes horizontally into 64-bit integers (SAD against 0) - return _mm256_sad_epu8(cnt_bytes, _mm256_setzero_si256()); -} - static inline float hnsw_mask_ip_x0_q_avx2( const float* query, const uint64_t* data, size_t padded_dim ) { - const size_t num_blk = padded_dim / 64; - const auto* it_data = reinterpret_cast(data); - const float* it_query = query; - const __m256i shifts0 = _mm256_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7); - const __m256i shifts1 = _mm256_setr_epi32(8, 9, 10, 11, 12, 13, 14, 15); - const __m256i shifts2 = _mm256_setr_epi32(16, 17, 18, 19, 20, 21, 22, 23); - const __m256i shifts3 = _mm256_setr_epi32(24, 25, 26, 27, 28, 29, 30, 31); - __m256 sum0 = _mm256_setzero_ps(); - __m256 sum1 = _mm256_setzero_ps(); - __m256 sum2 = _mm256_setzero_ps(); - __m256 sum3 = _mm256_setzero_ps(); - - for (size_t i = 0; i < num_blk; ++i) { - // Stored coordinates run from bit 63 to bit 0: high word first, - // with each selected bit shifted into a maskload lane's sign bit. - for (size_t half = 0; half < 2; ++half) { - int32_t word; - std::memcpy(&word, it_data + (1 - half) * sizeof(word), sizeof(word)); - const __m256i bits = _mm256_set1_epi32(word); - sum0 = _mm256_add_ps( - sum0, _mm256_maskload_ps(it_query, _mm256_sllv_epi32(bits, shifts0)) - ); - sum1 = _mm256_add_ps( - sum1, _mm256_maskload_ps(it_query + 8, _mm256_sllv_epi32(bits, shifts1)) - ); - sum2 = _mm256_add_ps( - sum2, _mm256_maskload_ps(it_query + 16, _mm256_sllv_epi32(bits, shifts2)) - ); - sum3 = _mm256_add_ps( - sum3, _mm256_maskload_ps(it_query + 24, _mm256_sllv_epi32(bits, shifts3)) - ); - it_query += 32; - } - it_data += sizeof(uint64_t); - } - - const __m256 sum = _mm256_add_ps(_mm256_add_ps(sum0, sum1), _mm256_add_ps(sum2, sum3)); - __m128 lanes = _mm_add_ps(_mm256_castps256_ps128(sum), _mm256_extractf128_ps(sum, 1)); - lanes = _mm_add_ps(lanes, _mm_movehl_ps(lanes, lanes)); - return _mm_cvtss_f32(_mm_add_ss(lanes, _mm_movehdup_ps(lanes))); + return simd::detail::mask_ip_avx2( + query, reinterpret_cast(data), padded_dim + ); } static inline float hnsw_warmup_ip_x0_q_512_avx2( @@ -119,111 +27,22 @@ static inline float hnsw_warmup_ip_x0_q_512_avx2( ) { size_t ip_scalar = 0; size_t ppc_scalar = 0; - - __m256i acc_ip = _mm256_setzero_si256(); - __m256i acc_ppc = _mm256_setzero_si256(); - - size_t i = 0; - // Step by 512 bits at a time (64 bytes = 16 elements of 32-bit integers) - size_t dim_end_512 = (padded_dim / 512) * 512; - - __m256i acc_bits[SplitSingleQuery::kNumBits]; - for (size_t j = 0; j < b_query; ++j) { - acc_bits[j] = _mm256_setzero_si256(); + if (padded_dim < 512) { + simd::detail::accumulate_warmup_tail( + reinterpret_cast(data), + query, + padded_dim, + b_query, + ip_scalar, + ppc_scalar + ); + return (delta * static_cast(ip_scalar)) + + (vl * static_cast(ppc_scalar)); } - for (; i < dim_end_512; i += 512) { - // Load 64 bytes of data using paired 32-byte loads - __m256i data_vec_lo = _mm256_loadu_si256(reinterpret_cast(data)); - __m256i data_vec_hi = - _mm256_loadu_si256(reinterpret_cast(data + 4)); - data += 8; // Advance 8 x 64-bit ints (64 bytes) - - acc_ppc = _mm256_add_epi64(acc_ppc, hnsw_popcount_avx2(data_vec_lo)); - acc_ppc = _mm256_add_epi64(acc_ppc, hnsw_popcount_avx2(data_vec_hi)); - - for (size_t j = 0; j < b_query; ++j) { - // Load 64 bytes of transposed query matching the 512-bit block layout - __m256i query_vec_lo = - _mm256_loadu_si256(reinterpret_cast(query)); - __m256i query_vec_hi = - _mm256_loadu_si256(reinterpret_cast(query + 4)); - query += 8; // Advance 8 x 64-bit ints (64 bytes) - - __m256i pop_lo = - hnsw_popcount_avx2(_mm256_and_si256(data_vec_lo, query_vec_lo)); - __m256i pop_hi = - hnsw_popcount_avx2(_mm256_and_si256(data_vec_hi, query_vec_hi)); - - acc_bits[j] = _mm256_add_epi64(acc_bits[j], pop_lo); - acc_bits[j] = _mm256_add_epi64(acc_bits[j], pop_hi); - } - } - - // Remainder block: handles leftovers less than 512 bits wide (e.g., last 448 bits) - size_t remaining_dim = padded_dim - i; - if (remaining_dim > 0) { - size_t num_chunks_64 = remaining_dim / 64; - size_t num_chunks_32 = remaining_dim / 32; - - size_t chunks_lo = (num_chunks_32 > 8) ? 8 : num_chunks_32; - size_t chunks_hi = (num_chunks_32 > 8) ? (num_chunks_32 - 8) : 0; - - // 1. Create a baseline sequence register - __m256i sequence = _mm256_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7); - - // 2. Generate masks in-register using Greater-Than comparisons - // If chunks_lo is 3, limit will be [3,3,3,3,3,3,3,3]. - // 3 > seq results in [-1, -1, -1, 0, 0, 0, 0, 0], which is the exact mask needed. - __m256i limit_lo = _mm256_set1_epi32(static_cast(chunks_lo)); - __m256i mask_lo = _mm256_cmpgt_epi32(limit_lo, sequence); - - __m256i limit_hi = _mm256_set1_epi32(static_cast(chunks_hi)); - __m256i mask_hi = _mm256_cmpgt_epi32(limit_hi, sequence); - - // 3. Vectorized execution continues with zero memory latency - __m256i data_vec_lo = - _mm256_maskload_epi32(reinterpret_cast(data), mask_lo); - __m256i data_vec_hi = - _mm256_maskload_epi32(reinterpret_cast(data + 4), mask_hi); - - acc_ppc = _mm256_add_epi64(acc_ppc, hnsw_popcount_avx2(data_vec_lo)); - acc_ppc = _mm256_add_epi64(acc_ppc, hnsw_popcount_avx2(data_vec_hi)); - - for (size_t j = 0; j < b_query; ++j) { - __m256i query_vec_lo = - _mm256_maskload_epi32(reinterpret_cast(query), mask_lo); - __m256i query_vec_hi = - _mm256_maskload_epi32(reinterpret_cast(query + 4), mask_hi); - query += num_chunks_64; - - __m256i pop_lo = - hnsw_popcount_avx2(_mm256_and_si256(data_vec_lo, query_vec_lo)); - __m256i pop_hi = - hnsw_popcount_avx2(_mm256_and_si256(data_vec_hi, query_vec_hi)); - - acc_bits[j] = _mm256_add_epi64(acc_bits[j], pop_lo); - acc_bits[j] = _mm256_add_epi64(acc_bits[j], pop_hi); - } - } - - for (size_t j = 0; j < b_query; ++j) { - __m128i shift = _mm_cvtsi32_si128(static_cast(j)); - acc_ip = _mm256_add_epi64(acc_ip, _mm256_sll_epi64(acc_bits[j], shift)); - } - - // Standard reduction for a single __m256i - auto mm256_reduce_add_epi64 = [](__m256i v) { - __m128i low = _mm256_castsi256_si128(v); - __m128i high = _mm256_extracti128_si256(v, 1); - __m128i sum = _mm_add_epi64(low, high); - return _mm_extract_epi64(sum, 0) + _mm_extract_epi64(sum, 1); - }; - - ip_scalar += mm256_reduce_add_epi64(acc_ip); - ppc_scalar += mm256_reduce_add_epi64(acc_ppc); - - return (delta * static_cast(ip_scalar)) + (vl * static_cast(ppc_scalar)); + return simd::detail::warmup_blocks_avx2::kNumBits>( + reinterpret_cast(data), query, delta, vl, padded_dim, b_query + ); } } // namespace rabitqlib::hnsw::detail diff --git a/src/index/hnsw_search_avx512_kernels.hpp b/src/index/hnsw_search_avx512_kernels.hpp index f2773575..eecef63a 100644 --- a/src/index/hnsw_search_avx512_kernels.hpp +++ b/src/index/hnsw_search_avx512_kernels.hpp @@ -1,54 +1,20 @@ #pragma once -#include - #include #include +#include "../simd/warmup_kernels.hpp" #include "rabitqlib/index/query.hpp" -#include "rabitqlib/utils/space.hpp" +#include "rabitqlib/simd/space_dispatch.hpp" namespace rabitqlib::hnsw::detail { static inline float hnsw_mask_ip_x0_q_avx512( const float* query, const uint64_t* data, size_t padded_dim ) { - const size_t num_blk = padded_dim / 64; - const uint8_t* it_data = reinterpret_cast(data); - const float* it_query = query; - - // __m512 sum0 = _mm512_setzero_ps(); - // __m512 sum1 = _mm512_setzero_ps(); - // __m512 sum2 = _mm512_setzero_ps(); - // __m512 sum3 = _mm512_setzero_ps(); - - __m512 sum = _mm512_setzero_ps(); - for (size_t i = 0; i < num_blk; ++i) { - uint64_t bits = rabitqlib::reverse_bits_u64(rabitqlib::load_unaligned_u64(it_data)); - - auto mask0 = static_cast<__mmask16>(bits); - auto mask1 = static_cast<__mmask16>(bits >> 16); - auto mask2 = static_cast<__mmask16>(bits >> 32); - auto mask3 = static_cast<__mmask16>(bits >> 48); - - __m512 masked0 = _mm512_maskz_loadu_ps(mask0, it_query); - __m512 masked1 = _mm512_maskz_loadu_ps(mask1, it_query + 16); - __m512 masked2 = _mm512_maskz_loadu_ps(mask2, it_query + 32); - __m512 masked3 = _mm512_maskz_loadu_ps(mask3, it_query + 48); - - sum = _mm512_add_ps(sum, masked0); - sum = _mm512_add_ps(sum, masked1); - sum = _mm512_add_ps(sum, masked2); - sum = _mm512_add_ps(sum, masked3); - - // _mm_prefetch(reinterpret_cast(it_query + 128), _MM_HINT_T1); - - it_data += sizeof(uint64_t); - it_query += 64; - } - - // __m512 sum = _mm512_add_ps(_mm512_add_ps(sum0, sum1), _mm512_add_ps(sum2, sum3)); - return _mm512_reduce_add_ps(sum); + return simd::mask_ip_x0_q_avx2( + query, reinterpret_cast(data), padded_dim + ); } static inline float hnsw_warmup_ip_x0_q_512_avx512( @@ -62,58 +28,23 @@ static inline float hnsw_warmup_ip_x0_q_512_avx512( size_t ip_scalar = 0; size_t ppc_scalar = 0; - __m512i acc_ip = _mm512_setzero_si512(); - __m512i acc_ppc = _mm512_setzero_si512(); - - size_t i = 0; - size_t dim_end_512 = (padded_dim / 512) * 512; - - __m512i acc_bits[SplitSingleQuery::kNumBits]; - for (size_t j = 0; j < b_query; ++j) { - acc_bits[j] = _mm512_setzero_si512(); + // The fixed-width HNSW path favors vector popcount beyond two words. + if (padded_dim <= 128) { + simd::detail::accumulate_warmup_tail( + reinterpret_cast(data), + query, + padded_dim, + b_query, + ip_scalar, + ppc_scalar + ); + return (delta * static_cast(ip_scalar)) + + (vl * static_cast(ppc_scalar)); } - for (; i < dim_end_512; i += 512) { - __m512i data_vec = _mm512_loadu_si512(data); - data += 8; - - acc_ppc = _mm512_add_epi64(acc_ppc, _mm512_popcnt_epi64(data_vec)); - - for (size_t j = 0; j < b_query; ++j) { - __m512i query_vec = _mm512_loadu_si512(query); - query += 8; - - __m512i pop = _mm512_popcnt_epi64(_mm512_and_si512(data_vec, query_vec)); - acc_bits[j] = _mm512_add_epi64(acc_bits[j], pop); - } - } - - size_t remaining_dim = padded_dim - i; - if (remaining_dim > 0) { - size_t num_chunks = remaining_dim / 64; - auto valid_mask = static_cast<__mmask8>((1u << num_chunks) - 1u); - - __m512i data_vec = _mm512_maskz_loadu_epi64(valid_mask, data); - acc_ppc = _mm512_add_epi64(acc_ppc, _mm512_popcnt_epi64(data_vec)); - - for (size_t j = 0; j < b_query; ++j) { - __m512i query_vec = _mm512_maskz_loadu_epi64(valid_mask, query); - query += num_chunks; - - __m512i pop = _mm512_popcnt_epi64(_mm512_and_si512(data_vec, query_vec)); - acc_bits[j] = _mm512_add_epi64(acc_bits[j], pop); - } - } - - for (size_t j = 0; j < b_query; ++j) { - __m128i shift = _mm_cvtsi32_si128(static_cast(j)); - acc_ip = _mm512_add_epi64(acc_ip, _mm512_sll_epi64(acc_bits[j], shift)); - } - - ip_scalar += static_cast(_mm512_reduce_add_epi64(acc_ip)); - ppc_scalar += static_cast(_mm512_reduce_add_epi64(acc_ppc)); - - return (delta * static_cast(ip_scalar)) + (vl * static_cast(ppc_scalar)); + return simd::detail::warmup_blocks_avx512::kNumBits>( + reinterpret_cast(data), query, delta, vl, padded_dim, b_query + ); } } // namespace rabitqlib::hnsw::detail diff --git a/src/index/qg.cpp b/src/index/qg.cpp index fd15130c..67a98655 100644 --- a/src/index/qg.cpp +++ b/src/index/qg.cpp @@ -70,7 +70,8 @@ size_t QuantizedGraph::checked_multiply(size_t lhs, size_t rhs) { } size_t QuantizedGraph::padded_dimension(size_t dim) { - return (checked_add(dim, 63) / 64) * 64; + // Use the same rotated domain as the other indexes. + return (checked_add(dim, 31) / 32) * 32; } size_t QuantizedGraph::num_vertices() const { return this->num_points_; } @@ -122,6 +123,7 @@ QuantizedGraph::QuantizedGraph( , quantization_bits_(quantization_bits) , seed_(seed) { validate_configuration(); + padded_dim_ = padded_dimension(dim_); initialize(); } @@ -336,10 +338,12 @@ void QuantizedGraph::load(const char* filename) { loaded.raw_dist_func_ = (loaded.metric_type_ == METRIC_IP) ? dot_product_dis : euclidean_sqr; loaded.validate_configuration(); - loaded.padded_dim_ = padded_dimension(loaded.dim_); - if (stored_padded_dim != loaded.padded_dim_) { + if (stored_padded_dim < loaded.dim_ || stored_padded_dim % 32 != 0) { throw std::runtime_error("Invalid padded dimension in quantized graph file"); } + // Older files use 64-coordinate padding. Their saved rotation and codes must + // retain that dimension, including when initialize() allocates the row storage. + loaded.padded_dim_ = stored_padded_dim; loaded.initialize_layout(); const size_t centroid_bytes = loaded.quantization_bits_ == 0 @@ -652,10 +656,9 @@ void QuantizedGraph::initialize_layout() { } void QuantizedGraph::initialize() { - padded_dim_ = padded_dimension(dim_); rotator_.reset(choose_rotator(dim_, rotator_type_, padded_dim_, seed_)); - assert(padded_dim_ % 64 == 0); + assert(padded_dim_ % 32 == 0); assert(padded_dim_ >= dim_); initialize_layout(); diff --git a/src/simd/fht_kernels.hpp b/src/simd/fht_kernels.hpp index f717456d..3abcd9fc 100644 --- a/src/simd/fht_kernels.hpp +++ b/src/simd/fht_kernels.hpp @@ -39,8 +39,8 @@ inline void fht_butterfly(__m256& a, __m256& b) { // Fuse three stages so eight registers are loaded and stored only once. // Stride is eight for the 64-point leaf, or N/8 for a larger transform. -template -inline void fht_eight_registers(float* data, size_t stride) { +template +inline void fht_eight_registers(float* data, size_t stride, float factor = 1.0F) { __m256 a = _mm256_loadu_ps(data); __m256 b = _mm256_loadu_ps(data + stride); __m256 c = _mm256_loadu_ps(data + 2 * stride); @@ -89,6 +89,17 @@ inline void fht_eight_registers(float* data, size_t stride) { fht_butterfly(b, f); fht_butterfly(c, g); fht_butterfly(d, h); + if constexpr (normalize) { + const __m256 multiplier = _mm256_set1_ps(factor); + a = _mm256_mul_ps(a, multiplier); + b = _mm256_mul_ps(b, multiplier); + c = _mm256_mul_ps(c, multiplier); + d = _mm256_mul_ps(d, multiplier); + e = _mm256_mul_ps(e, multiplier); + f = _mm256_mul_ps(f, multiplier); + g = _mm256_mul_ps(g, multiplier); + h = _mm256_mul_ps(h, multiplier); + } _mm256_storeu_ps(data, a); _mm256_storeu_ps(data + stride, b); _mm256_storeu_ps(data + 2 * stride, c); @@ -99,12 +110,13 @@ inline void fht_eight_registers(float* data, size_t stride) { _mm256_storeu_ps(data + 7 * stride, h); } -template -void fht_intrinsics(float* data) { +// Only the outermost merge applies the rotation scale, after every butterfly. +template +void fht_intrinsics(float* data, float factor = 1.0F) { static_assert(log_dim >= 6 && log_dim <= 16); constexpr size_t kDim = size_t{1} << log_dim; if constexpr (log_dim == 6) { - fht_eight_registers(data, 8); + fht_eight_registers(data, 8, factor); } else if constexpr (log_dim < 9) { fht_intrinsics(data); fht_intrinsics(data + kDim / 2); @@ -112,6 +124,11 @@ void fht_intrinsics(float* data) { __m256 a = _mm256_loadu_ps(data + i); __m256 b = _mm256_loadu_ps(data + kDim / 2 + i); fht_butterfly(a, b); + if constexpr (normalize) { + const __m256 multiplier = _mm256_set1_ps(factor); + a = _mm256_mul_ps(a, multiplier); + b = _mm256_mul_ps(b, multiplier); + } _mm256_storeu_ps(data + i, a); _mm256_storeu_ps(data + kDim / 2 + i, b); } @@ -123,7 +140,7 @@ void fht_intrinsics(float* data) { fht_intrinsics(data + block); } for (size_t i = 0; i < kStride; i += 8) { - fht_eight_registers(data + i, kStride); + fht_eight_registers(data + i, kStride, factor); } } } diff --git a/src/simd/hnsw_neon_kernels.hpp b/src/simd/hnsw_neon_kernels.hpp index f7cde1d3..4fd800b2 100644 --- a/src/simd/hnsw_neon_kernels.hpp +++ b/src/simd/hnsw_neon_kernels.hpp @@ -8,12 +8,14 @@ #include #include +#include "rabitqlib/simd/space_dispatch.hpp" + namespace rabitqlib::simd::detail { inline float mask_ip_neon(const float* query, const uint8_t* data, size_t dim) { float32x4_t sum0 = vdupq_n_f32(0), sum1 = sum0, sum2 = sum0, sum3 = sum0; const int32x4_t shifts = {-31, -30, -29, -28}; - for (size_t block = 0; block < dim; block += 64) { + for (size_t block = 0; block < dim - dim % 64; block += 64) { uint64_t word; std::memcpy(&word, data + block / 8, sizeof(word)); // Sign-code coordinates run from bit 63 to bit 0 in each stored word. @@ -40,7 +42,12 @@ inline float mask_ip_neon(const float* query, const uint8_t* data, size_t dim) { word <<= 32; } } - return vaddvq_f32(vaddq_f32(vaddq_f32(sum0, sum1), vaddq_f32(sum2, sum3))); + const float result = + vaddvq_f32(vaddq_f32(vaddq_f32(sum0, sum1), vaddq_f32(sum2, sum3))); + return dim % 64 == 0 + ? result + : result + + mask_ip_x0_q_generic(query, data + (dim - dim % 64) / 8, dim % 64); } inline float warmup_ip_neon( @@ -56,10 +63,17 @@ inline float warmup_ip_neon( } uint64_t ip = 0, count = 0; for (size_t block = 0; block < dim; block += 512) { - const size_t chunks = std::min(size_t{512}, dim - block) / 64; + const size_t width = std::min(size_t{512}, dim - block); + const size_t chunks = (width + 63) / 64; + const uint8_t* block_data = data + block / 8; + alignas(16) uint8_t tail[64]{}; + if (width % 64 != 0) { + std::memcpy(tail, block_data, width / 8); + block_data = tail; + } size_t chunk = 0; for (; chunk + 2 <= chunks; chunk += 2) { - const uint8x16_t code = vld1q_u8(data + block / 8 + chunk * 8); + const uint8x16_t code = vld1q_u8(block_data + chunk * 8); count += vaddlvq_u8(vcntq_u8(code)); for (size_t b = 0; b < bits; ++b) { const uint8x16_t plane = @@ -68,7 +82,7 @@ inline float warmup_ip_neon( } } if (chunk < chunks) { - const uint8x8_t code = vld1_u8(data + block / 8 + chunk * 8); + const uint8x8_t code = vld1_u8(block_data + chunk * 8); count += vaddlv_u8(vcnt_u8(code)); for (size_t b = 0; b < bits; ++b) { const uint8x8_t plane = diff --git a/src/simd/pack_excode_avx2.cpp b/src/simd/pack_excode_avx2.cpp index 6bb88bb3..7cf15cf4 100644 --- a/src/simd/pack_excode_avx2.cpp +++ b/src/simd/pack_excode_avx2.cpp @@ -15,7 +15,7 @@ void packing_3bit_excode_avx2(const uint8_t* o_raw, uint8_t* o_compact, size_t d } void packing_4bit_excode_avx2(const uint8_t* o_raw, uint8_t* o_compact, size_t dim) { - detail::packing_4bit_excode_intrinsics(o_raw, o_compact, dim); + detail::packing_4bit_excode_intrinsics<64>(o_raw, o_compact, dim); } void packing_5bit_excode_avx2(const uint8_t* o_raw, uint8_t* o_compact, size_t dim) { @@ -23,7 +23,47 @@ void packing_5bit_excode_avx2(const uint8_t* o_raw, uint8_t* o_compact, size_t d } void packing_6bit_excode_avx2(const uint8_t* o_raw, uint8_t* o_compact, size_t dim) { - detail::packing_6bit_excode_intrinsics(o_raw, o_compact, dim); + const __m256i mask = _mm256_set1_epi8(0x3F); + const __m256i shifts = _mm256_setr_epi32(6, 6, 6, 6, 4, 4, 4, 4); + for (size_t i = 0; i < dim - dim % 64; i += 64) { + // Preserve the full-block layout: the last 16 codes contribute two + // bits to each of the preceding three groups. Load each input once. + const __m256i first = _mm256_loadu_si256(reinterpret_cast(o_raw)); + const __m256i second = + _mm256_loadu_si256(reinterpret_cast(o_raw + 32)); + const __m128i third = _mm256_castsi256_si128(second); + const __m128i last = _mm256_extracti128_si256(second, 1); + const __m256i high = + _mm256_sllv_epi32(_mm256_permute4x64_epi64(second, 0xEE), shifts); + // Quantized codes are in [0, 63], so the low six bits need no mask. + const __m256i packed = _mm256_or_si256(first, _mm256_andnot_si256(mask, high)); + const __m128i mask128 = _mm256_castsi256_si128(mask); + const __m128i end = + _mm_or_si128(third, _mm_andnot_si128(mask128, _mm_slli_epi16(last, 2))); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(o_compact), packed); + _mm_storeu_si128(reinterpret_cast<__m128i*>(o_compact + 32), end); + o_raw += 64; + o_compact += 48; + } + if (dim % 64 != 0) { + const __m256i values = _mm256_loadu_si256(reinterpret_cast(o_raw)); + const __m256i pairs = _mm256_maddubs_epi16(_mm256_set1_epi16(0x4001), values); + const __m256i groups = _mm256_madd_epi16(pairs, _mm256_set1_epi32(0x10000001)); + // Each lane holds 12 packed bytes. Remove the gap between lanes before + // writing the complete 24-byte tail. + const __m128i shuffle = + _mm_setr_epi8(0, 1, 2, 4, 5, 6, 8, 9, 10, 12, 13, 14, -128, -128, -128, -128); + const __m256i lanes = + _mm256_shuffle_epi8(groups, _mm256_broadcastsi128_si256(shuffle)); + const __m256i packed = + _mm256_permutevar8x32_epi32(lanes, _mm256_setr_epi32(0, 1, 2, 4, 5, 6, 3, 7)); + _mm_storeu_si128( + reinterpret_cast<__m128i*>(o_compact), _mm256_castsi256_si128(packed) + ); + _mm_storel_epi64( + reinterpret_cast<__m128i*>(o_compact + 16), _mm256_extracti128_si256(packed, 1) + ); + } } void packing_7bit_excode_avx2(const uint8_t* o_raw, uint8_t* o_compact, size_t dim) { diff --git a/src/simd/pack_excode_avx512.cpp b/src/simd/pack_excode_avx512.cpp index a3b6cb9a..77e3e6cd 100644 --- a/src/simd/pack_excode_avx512.cpp +++ b/src/simd/pack_excode_avx512.cpp @@ -6,28 +6,61 @@ namespace rabitqlib::simd { +namespace { +void pack_2bit_tail(const uint8_t* raw, uint8_t* compact, size_t dim) { + if (dim == 0) { + return; + } + // Each four input bytes become one output byte; both operations use the + // same lane mask. The tail contains exactly 32 coordinates. + constexpr __mmask16 mask = 0xFF; + const __m512i values = _mm512_maskz_loadu_epi32(mask, raw); + const __m512i pairs = _mm512_maddubs_epi16(_mm512_set1_epi16(0x0401), values); + const __m512i groups = _mm512_madd_epi16(pairs, _mm512_set1_epi32(0x00100001)); + _mm512_mask_cvtepi32_storeu_epi8(compact, mask, groups); +} + +void pack_6bit_tail(const uint8_t* raw, uint8_t* compact, size_t dim) { + if (dim == 0) { + return; + } + constexpr __mmask16 input_mask = 0xFF; + const __m512i values = _mm512_maskz_loadu_epi32(input_mask, raw); + const __m512i pairs = _mm512_maddubs_epi16(_mm512_set1_epi16(0x4001), values); + const __m512i groups = _mm512_madd_epi16(pairs, _mm512_set1_epi32(0x10000001)); + // Four codes occupy three bytes. Compact each 128-bit lane into three + // words, then compress those words into exactly 24 output bytes. + const __m128i shuffle = + _mm_setr_epi8(0, 1, 2, 4, 5, 6, 8, 9, 10, 12, 13, 14, -128, -128, -128, -128); + const __m512i lanes = _mm512_shuffle_epi8(groups, _mm512_broadcast_i32x4(shuffle)); + constexpr __mmask16 selected = 0x77; + _mm512_mask_compressstoreu_epi32(compact, selected, lanes); +} +} // namespace + void packing_2bit_excode_avx512(const uint8_t* o_raw, uint8_t* o_compact, size_t dim) { - detail::packing_2bit_excode_intrinsics(o_raw, o_compact, dim); + detail::packing_2bit_excode_intrinsics(o_raw, o_compact, dim); } +// Reuse the AVX2 odd-width packers: wider code generation adds overhead here. void packing_3bit_excode_avx512(const uint8_t* o_raw, uint8_t* o_compact, size_t dim) { - detail::packing_3bit_excode_intrinsics(o_raw, o_compact, dim); + packing_3bit_excode_avx2(o_raw, o_compact, dim); } void packing_4bit_excode_avx512(const uint8_t* o_raw, uint8_t* o_compact, size_t dim) { - detail::packing_4bit_excode_intrinsics(o_raw, o_compact, dim); + detail::packing_4bit_excode_intrinsics<128>(o_raw, o_compact, dim); } void packing_5bit_excode_avx512(const uint8_t* o_raw, uint8_t* o_compact, size_t dim) { - detail::packing_5bit_excode_intrinsics(o_raw, o_compact, dim); + packing_5bit_excode_avx2(o_raw, o_compact, dim); } void packing_6bit_excode_avx512(const uint8_t* o_raw, uint8_t* o_compact, size_t dim) { - detail::packing_6bit_excode_intrinsics(o_raw, o_compact, dim); + detail::packing_6bit_excode_intrinsics(o_raw, o_compact, dim); } void packing_7bit_excode_avx512(const uint8_t* o_raw, uint8_t* o_compact, size_t dim) { - detail::packing_7bit_excode_intrinsics(o_raw, o_compact, dim); + packing_7bit_excode_avx2(o_raw, o_compact, dim); } } // namespace rabitqlib::simd diff --git a/src/simd/pack_excode_kernels.hpp b/src/simd/pack_excode_kernels.hpp index b00ac5c9..2a250f23 100644 --- a/src/simd/pack_excode_kernels.hpp +++ b/src/simd/pack_excode_kernels.hpp @@ -2,17 +2,98 @@ #include +#include #include #include #include namespace rabitqlib::simd::detail { +template +inline void pack_excode_tail_intrinsics(const uint8_t* raw, uint8_t* compact, size_t dim) { + static_assert(Bits == 2 || Bits == 3 || Bits == 5 || Bits == 6 || Bits == 7); + if (dim == 0) + return; + // A partial block is always 32 coordinates. + const __m128i pair_weights = + _mm_set1_epi16(static_cast(1U | (1U << (Bits + 8)))); + const __m128i group_weights = + _mm_set1_epi32(static_cast(1U | (1U << (2 * Bits + 16)))); + alignas(16) static constexpr auto kPositions = [] { + std::array positions{}; + for (size_t i = 0; i < 16; ++i) { + positions[i] = static_cast( + Bits % 2 == 0 ? (i < 2 * Bits ? (i / (Bits / 2)) * 4 + i % (Bits / 2) : 128) + : (i < Bits ? i : (i < 2 * Bits ? 8 + i - Bits : 128)) + ); + } + return positions; + }(); + const __m128i shuffle = + _mm_load_si128(reinterpret_cast(kPositions.data())); + for (size_t i = 0; i < 32; i += 16) { + const __m128i values = _mm_loadu_si128(reinterpret_cast(raw + i)); + // Codes are at most 127: signed byte operands are nonnegative, and + // even a pair of 7-bit codes fits without saturating the 16-bit sum. + const __m128i pairs = _mm_maddubs_epi16(pair_weights, values); + const __m128i groups = _mm_madd_epi16(pairs, group_weights); + __m128i packed = groups; + // Even widths already form whole-byte groups of four codes; the shuffle + // compacts those bytes directly. Odd widths need bitwise concatenation. + if constexpr (Bits % 2 != 0) { + packed = _mm_or_si128( + _mm_and_si128(groups, _mm_set1_epi64x(0xFFFFFFFFLL)), + _mm_slli_epi64(_mm_srli_epi64(groups, 32), 4 * Bits) + ); + } + packed = _mm_shuffle_epi8(packed, shuffle); + // Write the packed bytes in exact 8/4/2-byte pieces. + if constexpr (Bits >= 4) { + _mm_storel_epi64(reinterpret_cast<__m128i*>(compact), packed); + } + if constexpr (Bits == 2 || Bits == 3) { + _mm_storeu_si32(compact, packed); + } else if constexpr (Bits == 6 || Bits == 7) { + _mm_storeu_si32(compact + 8, _mm_srli_si128(packed, 8)); + } + if constexpr (Bits % 2 != 0) { + const auto last = static_cast(_mm_extract_epi16(packed, Bits - 1)); + std::memcpy(compact + 2 * Bits - 2, &last, sizeof(last)); + } + compact += 2 * Bits; + } +} + +// Pack the selected bit from eight rows of eight bytes, preserving column order. +template +inline uint64_t pack_excode_high_bit(const uint8_t* raw) { + static_assert(Bit == 2 || Bit == 4 || Bit == 6); + const __m256i mask = _mm256_set1_epi64x(0x0101010101010101LL); + const __m256i low = _mm256_and_si256( + _mm256_srli_epi64(_mm256_loadu_si256(reinterpret_cast(raw)), Bit), + mask + ); + const __m256i high = _mm256_and_si256( + _mm256_srli_epi64( + _mm256_loadu_si256(reinterpret_cast(raw + 32)), Bit + ), + mask + ); + const __m256i pairs = _mm256_or_si256(low, _mm256_slli_epi64(high, 4)); + const __m256i placed = _mm256_sllv_epi64(pairs, _mm256_setr_epi64x(0, 1, 2, 3)); + const __m128i folded = + _mm_or_si128(_mm256_castsi256_si128(placed), _mm256_extracti128_si256(placed, 1)); + return static_cast( + _mm_cvtsi128_si64(_mm_or_si128(folded, _mm_srli_si128(folded, 8))) + ); +} + +template > inline void packing_2bit_excode_intrinsics( const uint8_t* o_raw, uint8_t* o_compact, size_t dim ) { - // ! require dim % 64 == 0 - for (size_t j = 0; j < dim; j += 64) { + // Full blocks preserve the existing byte layout. + for (size_t j = 0; j < dim - dim % 64; j += 64) { // pack 64 2-bit codes into 128 bits (16 bytes) // the lower 2 bits of each byte represent vec00 to vec04... __m128i vec_00_to_15 = _mm_loadu_si128(reinterpret_cast(o_raw)); @@ -33,14 +114,15 @@ inline void packing_2bit_excode_intrinsics( o_raw += 64; o_compact += 16; } + pack_tail(o_raw, o_compact, dim % 64); } inline void packing_3bit_excode_intrinsics( const uint8_t* o_raw, uint8_t* o_compact, size_t dim ) { - // ! require dim % 64 == 0 + // Full blocks preserve the existing byte layout. const __m128i mask = _mm_set1_epi8(0b11); - for (size_t d = 0; d < dim; d += 64) { + for (size_t d = 0; d < dim - dim % 64; d += 64) { // split 3-bit codes into 2 bits and 1 bit // for 2-bit part, compact it like 2-bit code // for 1-bit part, compact 64 1-bit code into a int64 @@ -51,6 +133,7 @@ inline void packing_3bit_excode_intrinsics( _mm_loadu_si128(reinterpret_cast(o_raw + 32)); __m128i vec_48_to_63 = _mm_loadu_si128(reinterpret_cast(o_raw + 48)); + const uint64_t top_bit = pack_excode_high_bit<2>(o_raw); vec_00_to_15 = _mm_and_si128(vec_00_to_15, mask); vec_16_to_31 = _mm_slli_epi16(_mm_and_si128(vec_16_to_31, mask), 2); @@ -67,49 +150,88 @@ inline void packing_3bit_excode_intrinsics( // from lower to upper, each bit in each byte represents vec00 to vec07, // ..., vec56 to vec63 - uint64_t top_bit = 0; - constexpr uint64_t kMask1 = 0x0101010101010101ULL; - for (size_t i = 0; i < 64; i += 8) { - uint64_t cur_codes; - std::memcpy(&cur_codes, o_raw + i, sizeof(cur_codes)); - top_bit |= ((cur_codes >> 2) & kMask1) << (i / 8); - } std::memcpy(o_compact, &top_bit, sizeof(uint64_t)); o_raw += 64; o_compact += 8; } + pack_excode_tail_intrinsics<3>(o_raw, o_compact, dim % 64); +} + +// Each 16-coordinate group stores the first eight codes in the low nibbles +// and the next eight in the high nibbles. Load whole groups before packing. +template +inline void pack_4bit_block(const uint8_t* raw, uint8_t* compact) { + static_assert(Dim == 16 || Dim == 32 || Dim == 64 || Dim == 128); + if constexpr (Dim == 128) { + const __m512i lo = _mm512_loadu_si512(raw); + const __m512i hi = _mm512_loadu_si512(raw + 64); + const __m512i even = + _mm512_permutex2var_epi64(lo, _mm512_setr_epi64(0, 2, 4, 6, 8, 10, 12, 14), hi); + const __m512i odd = + _mm512_permutex2var_epi64(lo, _mm512_setr_epi64(1, 3, 5, 7, 9, 11, 13, 15), hi); + _mm512_storeu_si512(compact, _mm512_or_si512(even, _mm512_slli_epi64(odd, 4))); + } else if constexpr (Dim == 64) { + const __m256i lo = _mm256_loadu_si256(reinterpret_cast(raw)); + const __m256i hi = _mm256_loadu_si256(reinterpret_cast(raw + 32)); + const __m256i even = _mm256_permute4x64_epi64(_mm256_unpacklo_epi64(lo, hi), 0xD8); + const __m256i odd = _mm256_permute4x64_epi64(_mm256_unpackhi_epi64(lo, hi), 0xD8); + _mm256_storeu_si256( + reinterpret_cast<__m256i*>(compact), + _mm256_or_si256(even, _mm256_slli_epi64(odd, 4)) + ); + } else if constexpr (Dim == 32) { + const __m128i lo = _mm_loadu_si128(reinterpret_cast(raw)); + const __m128i hi = _mm_loadu_si128(reinterpret_cast(raw + 16)); + const __m128i even = _mm_unpacklo_epi64(lo, hi); + const __m128i odd = _mm_unpackhi_epi64(lo, hi); + _mm_storeu_si128( + reinterpret_cast<__m128i*>(compact), _mm_or_si128(even, _mm_slli_epi64(odd, 4)) + ); + } else { + uint64_t lo, hi; + std::memcpy(&lo, raw, sizeof(lo)); + std::memcpy(&hi, raw + 8, sizeof(hi)); + const uint64_t packed = lo | (hi << 4); + std::memcpy(compact, &packed, sizeof(packed)); + } } +template inline void packing_4bit_excode_intrinsics( - const uint8_t* o_raw, uint8_t* o_compact, size_t dim + const uint8_t* raw, uint8_t* compact, size_t dim ) { - // Although this part only requires SSE, computing inner products for this organization - // requires AVX512F, similar for remaining functions - // ! require dim % 16 == 0 - for (size_t j = 0; j < dim; j += 16) { - // pack 16 4-bit codes into uint64 - // the lower 4 bits represent vec00 to vec07 - uint64_t code0; - std::memcpy(&code0, o_raw, sizeof(code0)); - // the upper 4 bits represent vec08 to vec15 - uint64_t code1; - std::memcpy(&code1, o_raw + 8, sizeof(code1)); - - uint64_t compact = (code1 << 4) | code0; - std::memcpy(o_compact, &compact, sizeof(compact)); - - o_raw += 16; - o_compact += 8; + static_assert(BlockDim == 64 || BlockDim == 128); + for (; dim >= BlockDim; dim -= BlockDim) { + pack_4bit_block(raw, compact); + raw += BlockDim; + compact += BlockDim / 2; + } + if constexpr (BlockDim == 128) { + if (dim >= 64) { + pack_4bit_block<64>(raw, compact); + raw += 64; + compact += 32; + dim -= 64; + } + } + if (dim >= 32) { + pack_4bit_block<32>(raw, compact); + raw += 32; + compact += 16; + dim -= 32; + } + if (dim != 0) { + pack_4bit_block<16>(raw, compact); } } inline void packing_5bit_excode_intrinsics( const uint8_t* o_raw, uint8_t* o_compact, size_t dim ) { - // ! require dim % 64 == 0 + // Full blocks preserve the existing byte layout. const __m128i mask = _mm_set1_epi8(0b1111); - for (size_t j = 0; j < dim; j += 64) { + for (size_t j = 0; j < dim - dim % 64; j += 64) { __m128i vec_00_to_15 = _mm_loadu_si128(reinterpret_cast(o_raw)); __m128i vec_16_to_31 = _mm_loadu_si128(reinterpret_cast(o_raw + 16)); @@ -117,6 +239,7 @@ inline void packing_5bit_excode_intrinsics( _mm_loadu_si128(reinterpret_cast(o_raw + 32)); __m128i vec_48_to_63 = _mm_loadu_si128(reinterpret_cast(o_raw + 48)); + const uint64_t top_bit = pack_excode_high_bit<4>(o_raw); vec_00_to_15 = _mm_and_si128(vec_00_to_15, mask); vec_16_to_31 = _mm_slli_epi16(_mm_and_si128(vec_16_to_31, mask), 4); @@ -133,28 +256,23 @@ inline void packing_5bit_excode_intrinsics( // from lower to upper, each bit in each byte represents vec00 to vec07, // ..., vec56 to vec63 - uint64_t top_bit = 0; - constexpr uint64_t kMask1 = 0x0101010101010101ULL; - for (size_t i = 0; i < 64; i += 8) { - uint64_t cur_codes; - std::memcpy(&cur_codes, o_raw + i, sizeof(cur_codes)); - top_bit |= ((cur_codes >> 4) & kMask1) << (i / 8); - } std::memcpy(o_compact, &top_bit, sizeof(uint64_t)); o_raw += 64; o_compact += 8; } + pack_excode_tail_intrinsics<5>(o_raw, o_compact, dim % 64); } +template > inline void packing_6bit_excode_intrinsics( const uint8_t* o_raw, uint8_t* o_compact, size_t dim ) { // for vec00 to vec47, split code into 6 // for vec48 to vec63, split code into 2 + 2 + 2 - const __m128i mask2 = _mm_set1_epi8(static_cast(0b11000000)); + // Complementary masks let AVX512 fuse each bit selection into one instruction. const __m128i mask6 = _mm_set1_epi8(0b00111111); - for (size_t d = 0; d < dim; d += 64) { + for (size_t d = 0; d < dim - dim % 64; d += 64) { __m128i vec_00_to_15 = _mm_loadu_si128(reinterpret_cast(o_raw)); __m128i vec_16_to_31 = _mm_loadu_si128(reinterpret_cast(o_raw + 16)); @@ -165,24 +283,25 @@ inline void packing_6bit_excode_intrinsics( __m128i compact = _mm_or_si128( _mm_and_si128(vec_00_to_15, mask6), - _mm_and_si128(_mm_slli_epi16(vec_48_to_63, 6), mask2) + _mm_andnot_si128(mask6, _mm_slli_epi16(vec_48_to_63, 6)) ); _mm_storeu_si128(reinterpret_cast<__m128i*>(o_compact), compact); compact = _mm_or_si128( _mm_and_si128(vec_16_to_31, mask6), - _mm_and_si128(_mm_slli_epi16(vec_48_to_63, 4), mask2) + _mm_andnot_si128(mask6, _mm_slli_epi16(vec_48_to_63, 4)) ); _mm_storeu_si128(reinterpret_cast<__m128i*>(o_compact + 16), compact); compact = _mm_or_si128( _mm_and_si128(vec_32_to_47, mask6), - _mm_and_si128(_mm_slli_epi16(vec_48_to_63, 2), mask2) + _mm_andnot_si128(mask6, _mm_slli_epi16(vec_48_to_63, 2)) ); _mm_storeu_si128(reinterpret_cast<__m128i*>(o_compact + 32), compact); o_compact += 48; o_raw += 64; } + pack_tail(o_raw, o_compact, dim % 64); } inline void packing_7bit_excode_intrinsics( @@ -192,7 +311,7 @@ inline void packing_7bit_excode_intrinsics( // for vec48 to vec63, split code into 2 + 2 + 2 + 1 const __m128i mask2 = _mm_set1_epi8(static_cast(0b11000000)); const __m128i mask6 = _mm_set1_epi8(0b00111111); - for (size_t d = 0; d < dim; d += 64) { + for (size_t d = 0; d < dim - dim % 64; d += 64) { __m128i vec_00_to_15 = _mm_loadu_si128(reinterpret_cast(o_raw)); __m128i vec_16_to_31 = _mm_loadu_si128(reinterpret_cast(o_raw + 16)); @@ -200,6 +319,7 @@ inline void packing_7bit_excode_intrinsics( _mm_loadu_si128(reinterpret_cast(o_raw + 32)); __m128i vec_48_to_63 = _mm_loadu_si128(reinterpret_cast(o_raw + 48)); + const uint64_t top_bit = pack_excode_high_bit<6>(o_raw); __m128i compact = _mm_or_si128( _mm_and_si128(vec_00_to_15, mask6), @@ -220,18 +340,12 @@ inline void packing_7bit_excode_intrinsics( _mm_storeu_si128(reinterpret_cast<__m128i*>(o_compact + 32), compact); o_compact += 48; - uint64_t top_bit = 0; - constexpr uint64_t top_mask = 0x0101010101010101ULL; - for (size_t i = 0; i < 64; i += 8) { - uint64_t cur_codes; - std::memcpy(&cur_codes, o_raw + i, sizeof(cur_codes)); - top_bit |= ((cur_codes >> 6) & top_mask) << (i / 8); - } std::memcpy(o_compact, &top_bit, sizeof(uint64_t)); o_compact += 8; o_raw += 64; } + pack_excode_tail_intrinsics<7>(o_raw, o_compact, dim % 64); } } // namespace rabitqlib::simd::detail diff --git a/src/simd/packed_tail_avx2.hpp b/src/simd/packed_tail_avx2.hpp new file mode 100644 index 00000000..7994e12f --- /dev/null +++ b/src/simd/packed_tail_avx2.hpp @@ -0,0 +1,145 @@ +#pragma once + +#include + +#include +#include +#include + +namespace rabitqlib::simd::detail { +// Read exactly 2 * Bits bytes for 16 codes. Staging short stores through a +// 16-byte stack buffer blocks store forwarding on the following vector load. +template +inline __m128i load_excode_tail(const uint8_t* compact) { + static_assert(Bits == 2 || Bits == 3 || Bits == 5 || Bits == 6 || Bits == 7); + if constexpr (Bits == 2) { + return _mm_loadu_si32(compact); + } else if constexpr (Bits == 3) { + uint16_t high; + std::memcpy(&high, compact + 4, sizeof(high)); + return _mm_insert_epi16(_mm_loadu_si32(compact), high, 2); + } else if constexpr (Bits == 5) { + uint16_t high; + std::memcpy(&high, compact + 8, sizeof(high)); + return _mm_insert_epi16(_mm_loadu_si64(compact), high, 4); + } else if constexpr (Bits == 6) { + return _mm_unpacklo_epi64(_mm_loadu_si64(compact), _mm_loadu_si32(compact + 8)); + } else { + uint16_t high; + std::memcpy(&high, compact + 12, sizeof(high)); + const __m128i upper = _mm_insert_epi16(_mm_loadu_si32(compact + 8), high, 2); + return _mm_unpacklo_epi64(_mm_loadu_si64(compact), upper); + } +} + +template +inline void accumulate_excode_tail_avx2( + const float* query, const uint8_t* compact, size_t dim, __m256& sum, __m256& sum_hi +) { + if (dim == 0) + return; + // The fixed 32-coordinate tail contains two 16-code groups. + // Select the two bytes spanning each code, then shift/mask eight codes at once. + alignas(16) uint8_t positions[32]; + for (size_t i = 0; i < 16; ++i) { + positions[2 * i] = static_cast(i * Bits / 8); + positions[2 * i + 1] = static_cast(i * Bits / 8 + 1); + } + const __m128i low = _mm_load_si128(reinterpret_cast(positions)); + const __m128i high = _mm_load_si128(reinterpret_cast(positions + 16)); + constexpr size_t kShiftPeriod = Bits == 2 ? 32 : 8; + const __m256i shifts = _mm256_setr_epi32( + 0, + Bits % kShiftPeriod, + (2 * Bits) % kShiftPeriod, + (3 * Bits) % kShiftPeriod, + (4 * Bits) % kShiftPeriod, + (5 * Bits) % kShiftPeriod, + (6 * Bits) % kShiftPeriod, + (7 * Bits) % kShiftPeriod + ); + const __m256i mask = _mm256_set1_epi32((1U << Bits) - 1); + for (size_t i = 0; i < 32; i += 16) { + const __m128i packed = load_excode_tail(compact); + __m256i words_lo, words_hi; + if constexpr (Bits == 2) { + // All 16 codes fit in one word; broadcast it instead of gathering bytes. + words_lo = _mm256_broadcastd_epi32(packed); + words_hi = _mm256_srli_epi32(words_lo, 16); + } else { + words_lo = _mm256_cvtepu16_epi32(_mm_shuffle_epi8(packed, low)); + words_hi = _mm256_cvtepu16_epi32(_mm_shuffle_epi8(packed, high)); + } + const auto decode = [&](const __m256i words) { + return _mm256_cvtepi32_ps( + _mm256_and_si256(_mm256_srlv_epi32(words, shifts), mask) + ); + }; + sum = _mm256_fmadd_ps(decode(words_lo), _mm256_loadu_ps(query + i), sum); + sum_hi = _mm256_fmadd_ps(decode(words_hi), _mm256_loadu_ps(query + i + 8), sum_hi); + compact += 2 * Bits; + } +} + +static inline float mask_ip_avx2( + const float* query, const uint8_t* data, size_t padded_dim +) { + const size_t num_blk = padded_dim / 64; + const auto* it_data = data; + const float* it_query = query; + const __m256i shifts0 = _mm256_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7); + const __m256i shifts1 = _mm256_setr_epi32(8, 9, 10, 11, 12, 13, 14, 15); + const __m256i shifts2 = _mm256_setr_epi32(16, 17, 18, 19, 20, 21, 22, 23); + const __m256i shifts3 = _mm256_setr_epi32(24, 25, 26, 27, 28, 29, 30, 31); + __m256 sum0 = _mm256_setzero_ps(); + __m256 sum1 = _mm256_setzero_ps(); + __m256 sum2 = _mm256_setzero_ps(); + __m256 sum3 = _mm256_setzero_ps(); + + for (size_t i = 0; i < num_blk; ++i) { + // Stored coordinates run from bit 63 to bit 0: high word first, + // with each selected bit shifted into a maskload lane's sign bit. + for (size_t half = 0; half < 2; ++half) { + int32_t word; + std::memcpy(&word, it_data + (1 - half) * sizeof(word), sizeof(word)); + const __m256i bits = _mm256_set1_epi32(word); + sum0 = _mm256_add_ps( + sum0, _mm256_maskload_ps(it_query, _mm256_sllv_epi32(bits, shifts0)) + ); + sum1 = _mm256_add_ps( + sum1, _mm256_maskload_ps(it_query + 8, _mm256_sllv_epi32(bits, shifts1)) + ); + sum2 = _mm256_add_ps( + sum2, _mm256_maskload_ps(it_query + 16, _mm256_sllv_epi32(bits, shifts2)) + ); + sum3 = _mm256_add_ps( + sum3, _mm256_maskload_ps(it_query + 24, _mm256_sllv_epi32(bits, shifts3)) + ); + it_query += 32; + } + it_data += sizeof(uint64_t); + } + + if (padded_dim % 64 != 0) { + const auto accumulate = [](__m256 sum, + const float* values, + __m256i bits, + __m256i shifts) { + const __m256 mask = + _mm256_castsi256_ps(_mm256_srai_epi32(_mm256_sllv_epi32(bits, shifts), 31)); + return _mm256_add_ps(sum, _mm256_and_ps(_mm256_loadu_ps(values), mask)); + }; + int32_t signed_bits; + std::memcpy(&signed_bits, it_data, sizeof(signed_bits)); + const __m256i bits = _mm256_set1_epi32(signed_bits); + sum0 = accumulate(sum0, it_query, bits, shifts0); + sum1 = accumulate(sum1, it_query + 8, bits, shifts1); + sum2 = accumulate(sum2, it_query + 16, bits, shifts2); + sum3 = accumulate(sum3, it_query + 24, bits, shifts3); + } + const __m256 sum = _mm256_add_ps(_mm256_add_ps(sum0, sum1), _mm256_add_ps(sum2, sum3)); + __m128 lanes = _mm_add_ps(_mm256_castps256_ps128(sum), _mm256_extractf128_ps(sum, 1)); + lanes = _mm_add_ps(lanes, _mm_movehl_ps(lanes, lanes)); + return _mm_cvtss_f32(_mm_add_ss(lanes, _mm_movehdup_ps(lanes))); +} +} // namespace rabitqlib::simd::detail diff --git a/src/simd/portable_generic.cpp b/src/simd/portable_generic.cpp index d4804165..0f8acf57 100644 --- a/src/simd/portable_generic.cpp +++ b/src/simd/portable_generic.cpp @@ -13,7 +13,7 @@ namespace rabitqlib::simd { namespace { -// Coordinate-to-bit mapping of the existing x86 packed format. Odd widths +// Coordinate-to-bit mapping of the historical full-block layout. Odd widths // store the high bit in eight bytes, with coordinate i at byte i%8, bit i/8. template size_t packed_bit(size_t i, size_t bit) { @@ -42,13 +42,38 @@ size_t packed_bit(size_t i, size_t bit) { } } +// Read one code from the compact little-endian tail, without overreading its end. +template +unsigned tail_code(const uint8_t* compact, size_t i) { + const size_t offset = i * Bits; + unsigned code = compact[offset / 8] >> (offset % 8); + if (offset % 8 + Bits > 8) { + code |= unsigned{compact[offset / 8 + 1]} << (8 - offset % 8); + } + return code & ((1U << Bits) - 1); +} + template void pack(const uint8_t* raw, uint8_t* compact, size_t dim) { - std::fill_n(compact, dim * Bits / 8, uint8_t{0}); - for (size_t i = 0; i < dim; ++i) { - for (size_t b = 0; b < Bits; ++b) { + const size_t full = (Bits == 1 || Bits == 4 || Bits == 8) ? dim : dim - dim % 64; + std::fill_n(compact, full * Bits / 8, uint8_t{0}); + for (size_t i = 0; i < full; ++i) { + for (size_t b = 0; b < Bits; b += 2) { const size_t pos = packed_bit(i, b); - compact[pos / 8] |= ((raw[i] >> b) & 1U) << (pos % 8); + const unsigned mask = b + 1 < Bits ? 3U : 1U; + compact[pos / 8] |= ((raw[i] >> b) & mask) << (pos % 8); + } + } + // Every legal tail occupies whole bytes; at most one byte is emitted per code. + unsigned pending = 0, count = 0; + size_t output = full * Bits / 8; + for (size_t i = full; i < dim; ++i) { + pending |= (unsigned{raw[i]} & ((1U << Bits) - 1)) << count; + count += Bits; + if (count >= 8) { + compact[output++] = static_cast(pending); + pending >>= 8; + count -= 8; } } } @@ -56,20 +81,29 @@ void pack(const uint8_t* raw, uint8_t* compact, size_t dim) { // Inverse of pack, over the same mapping. Not a search path, so scalar only. template void unpack(const uint8_t* compact, uint8_t* raw, size_t dim) { - for (size_t i = 0; i < dim; ++i) { + const size_t full = (Bits == 1 || Bits == 4 || Bits == 8) ? dim : dim - dim % 64; + for (size_t i = 0; i < full; ++i) { unsigned code = 0; - for (size_t b = 0; b < Bits; ++b) { + // Every pair occupies adjacent bits, even in the split 6/7-bit layout. + // Odd widths leave the separate high bit for the final iteration. + for (size_t b = 0; b < Bits; b += 2) { const size_t pos = packed_bit(i, b); - code |= ((compact[pos / 8] >> (pos % 8)) & 1U) << b; + unsigned pair = compact[pos / 8] >> (pos % 8); + const unsigned mask = b + 1 < Bits ? 3U : 1U; + code |= (pair & mask) << b; } raw[i] = static_cast(code); } + for (size_t i = full; i < dim; ++i) { + raw[i] = static_cast(tail_code(compact + full * Bits / 8, i - full)); + } } template float excode_ip(const float* query, const uint8_t* compact, size_t dim) { double sum = 0; - for (size_t i = 0; i < dim; ++i) { + const size_t full = (Bits == 1 || Bits == 4 || Bits == 8) ? dim : dim - dim % 64; + for (size_t i = 0; i < full; ++i) { unsigned code = 0; for (size_t b = 0; b < Bits; ++b) { const size_t pos = packed_bit(i, b); @@ -77,6 +111,10 @@ float excode_ip(const float* query, const uint8_t* compact, size_t dim) { } sum += static_cast(query[i]) * code; } + for (size_t i = full; i < dim; ++i) { + sum += static_cast(query[i]) * + tail_code(compact + full * Bits / 8, i - full); + } return static_cast(sum); } @@ -181,10 +219,11 @@ void new_transpose_bin_generic( const uint16_t* query, uint64_t* transposed, size_t dim, size_t bits ) { for (size_t block = 0; block < dim; block += 64) { + const size_t width = std::min(size_t{64}, dim - block); for (size_t b = 0; b < bits; ++b) { uint64_t word = 0; - for (size_t i = 0; i < 64; ++i) { - word |= uint64_t{(query[block + i] >> b) & 1U} << (63 - i); + for (size_t i = 0; i < width; ++i) { + word |= uint64_t{(query[block + i] >> b) & 1U} << (width - 1 - i); } *transposed++ = word; } @@ -195,12 +234,14 @@ void new_transpose_bin_512_generic( const uint8_t* query, uint64_t* transposed, size_t dim, size_t bits ) { for (size_t block = 0; block < dim; block += 512) { - const size_t chunks = std::min(size_t{512}, dim - block) / 64; + const size_t chunks = (std::min(size_t{512}, dim - block) + 63) / 64; for (size_t b = 0; b < bits; ++b) { for (size_t chunk = 0; chunk < chunks; ++chunk) { + const size_t width = std::min(size_t{64}, dim - block - chunk * 64); uint64_t word = 0; - for (size_t i = 0; i < 64; ++i) { - word |= uint64_t{(query[block + chunk * 64 + i] >> b) & 1U} << (63 - i); + for (size_t i = 0; i < width; ++i) { + word |= uint64_t{(query[block + chunk * 64 + i] >> b) & 1U} + << (width - 1 - i); } *transposed++ = word; } @@ -211,10 +252,10 @@ void new_transpose_bin_512_generic( float mask_ip_x0_q_generic(const float* query, const uint8_t* data, size_t dim) { double sum = 0; for (size_t block = 0; block < dim; block += 64) { - uint64_t word; - std::memcpy(&word, data + block / 8, sizeof(word)); - for (size_t i = 0; i < 64; ++i) { - if ((word >> (63 - i)) & 1U) + const size_t width = std::min(size_t{64}, dim - block); + const uint64_t word = bitops::load_word(data + block / 8, width / 8); + for (size_t i = 0; i < width; ++i) { + if ((word >> (width - 1 - i)) & 1U) sum += query[block + i]; } } @@ -236,10 +277,11 @@ float warmup_ip_x0_q_512_generic( throw std::invalid_argument("warmup_ip_x0_q_512 requires at most 8 query bits"); uint64_t ip = 0, count = 0; for (size_t block = 0; block < dim; block += 512) { - const size_t chunks = std::min(size_t{512}, dim - block) / 64; + const size_t chunks = (std::min(size_t{512}, dim - block) + 63) / 64; for (size_t chunk = 0; chunk < chunks; ++chunk) { - uint64_t word; - std::memcpy(&word, data + block / 8 + chunk * 8, sizeof(word)); + const size_t width = std::min(size_t{64}, dim - block - chunk * 64); + const uint64_t word = + bitops::load_word(data + block / 8 + chunk * 8, width / 8); count += bitops::popcount64(word); for (size_t b = 0; b < bits; ++b) { ip += uint64_t{bitops::popcount64(word & query[b * chunks + chunk])} << b; diff --git a/src/simd/query_neon.cpp b/src/simd/query_neon.cpp index 7d244a1a..664692b1 100644 --- a/src/simd/query_neon.cpp +++ b/src/simd/query_neon.cpp @@ -69,7 +69,7 @@ void scalar_quantize_uint16_neon( void new_transpose_bin_neon( const uint16_t* query, uint64_t* transposed, size_t dim, size_t bits ) { - for (size_t block = 0; block < dim; block += 64) { + for (size_t block = 0; block < dim - dim % 64; block += 64) { uint16x8_t values[8]; for (size_t c = 0; c < 8; ++c) { values[c] = vld1q_u16(query + block + c * 8); @@ -83,12 +83,19 @@ void new_transpose_bin_neon( *transposed++ = word; } } + if (dim % 64 != 0) { + new_transpose_bin_generic(query + dim - dim % 64, transposed, dim % 64, bits); + } } void new_transpose_bin_512_neon( const uint8_t* query, uint64_t* transposed, size_t dim, size_t bits ) { for (size_t block = 0; block < dim; block += 512) { + if ((dim - block) < 512 && dim % 64 != 0) { + new_transpose_bin_512_generic(query + block, transposed, dim - block, bits); + break; + } const size_t chunks = std::min(size_t{512}, dim - block) / 64; for (size_t c = 0; c < chunks; ++c) { const uint8_t* row = query + block + c * 64; diff --git a/src/simd/rotator_avx2.cpp b/src/simd/rotator_avx2.cpp index 5b9be3d4..0f30eb94 100644 --- a/src/simd/rotator_avx2.cpp +++ b/src/simd/rotator_avx2.cpp @@ -40,7 +40,7 @@ void flip_sign_avx2(const uint8_t* flip, float* data, size_t dim) { } void kacs_walk_avx2(float* data, size_t len) { - // ! len % 16 == 0; + // ! len % 32 == 0; for (size_t i = 0; i < len / 2; i += 8) { __m256 x = _mm256_loadu_ps(&data[i]); __m256 y = _mm256_loadu_ps(&data[i + (len / 2)]); diff --git a/src/simd/rotator_avx512.cpp b/src/simd/rotator_avx512.cpp index 2b34d5a2..c04b0173 100644 --- a/src/simd/rotator_avx512.cpp +++ b/src/simd/rotator_avx512.cpp @@ -9,6 +9,30 @@ namespace rabitqlib::simd { +namespace { + +// Apply the next round's signs before storing the Kac outputs. +void kacs_walk_flip(const uint8_t* flip, float* data, size_t len) { + const size_t half = len / 2; + const __m512 sign = + _mm512_castsi512_ps(_mm512_set1_epi32(std::numeric_limits::min())); + for (size_t i = 0; i < half; i += 16) { + const __m512 x = _mm512_loadu_ps(data + i); + const __m512 y = _mm512_loadu_ps(data + half + i); + const __m512 sum = _mm512_add_ps(x, y); + const __m512 difference = _mm512_sub_ps(x, y); + uint16_t left_bits, right_bits; + std::memcpy(&left_bits, flip + i / 8, sizeof(left_bits)); + std::memcpy(&right_bits, flip + (half + i) / 8, sizeof(right_bits)); + _mm512_storeu_ps(data + i, _mm512_mask_xor_ps(sum, left_bits, sum, sign)); + _mm512_storeu_ps( + data + half + i, _mm512_mask_xor_ps(difference, right_bits, difference, sign) + ); + } +} + +} // namespace + void flip_sign_avx512(const uint8_t* flip, float* data, size_t dim) { constexpr size_t kFloatsPerChunk = 64; // Process 64 floats per iteration @@ -17,7 +41,8 @@ void flip_sign_avx512(const uint8_t* flip, float* data, size_t dim) { "floats_per_chunk must be divisible by AVX512 register width" ); - for (size_t i = 0; i < dim; i += kFloatsPerChunk) { + size_t i = 0; + for (; i + kFloatsPerChunk <= dim; i += kFloatsPerChunk) { // Load 64 bits (8 bytes) from the bit sequence uint64_t mask_bits; std::memcpy(&mask_bits, &flip[i / 8], sizeof(mask_bits)); @@ -52,10 +77,25 @@ void flip_sign_avx512(const uint8_t* flip, float* data, size_t dim) { vec3 = _mm512_mask_xor_ps(vec3, mask3, vec3, sign_flip); _mm512_storeu_ps(&data[i + 48], vec3); } + if (i < dim) { + uint32_t bits; + std::memcpy(&bits, flip + i / 8, sizeof(bits)); + const __m512 sign = + _mm512_castsi512_ps(_mm512_set1_epi32(std::numeric_limits::min())); + const __m512 lo = _mm512_loadu_ps(data + i); + const __m512 hi = _mm512_loadu_ps(data + i + 16); + _mm512_storeu_ps( + data + i, _mm512_mask_xor_ps(lo, static_cast<__mmask16>(bits), lo, sign) + ); + _mm512_storeu_ps( + data + i + 16, + _mm512_mask_xor_ps(hi, static_cast<__mmask16>(bits >> 16), hi, sign) + ); + } } void kacs_walk_avx512(float* data, size_t len) { - // ! len % 32 == 0; + // Each half contains whole 16-float registers. for (size_t i = 0; i < len / 2; i += 16) { __m512 x = _mm512_loadu_ps(&data[i]); __m512 y = _mm512_loadu_ps(&data[i + (len / 2)]); @@ -77,7 +117,7 @@ void fht_rotate_avx512( float fac, const uint8_t* flip ) { - fht_rotate_impl( + fht_rotate_impl( data, rotated_vec, dim, padded_dim, trunc_dim, fac, flip ); } diff --git a/src/simd/rotator_kernels.hpp b/src/simd/rotator_kernels.hpp index 1602b5e7..1128d1a4 100644 --- a/src/simd/rotator_kernels.hpp +++ b/src/simd/rotator_kernels.hpp @@ -9,13 +9,35 @@ namespace rabitqlib::simd { namespace { -inline void rescale(float* data, size_t dim, float factor) { - for (size_t i = 0; i < dim; ++i) { - data[i] *= factor; +// Preserve the final Kac arithmetic and scale, storing each result only once. +inline void kacs_walk_rescale(float* data, size_t dim, float factor) { + const size_t half = dim / 2; +#if defined(__AVX512F__) + const __m512 multiplier512 = _mm512_set1_ps(factor); + for (size_t i = 0; i < half; i += 16) { + const __m512 left = _mm512_loadu_ps(data + i); + const __m512 right = _mm512_loadu_ps(data + half + i); + _mm512_storeu_ps( + data + i, _mm512_mul_ps(_mm512_add_ps(left, right), multiplier512) + ); + _mm512_storeu_ps( + data + half + i, _mm512_mul_ps(_mm512_sub_ps(left, right), multiplier512) + ); } +#else + const __m256 multiplier = _mm256_set1_ps(factor); + for (size_t i = 0; i < half; i += 8) { + const __m256 left = _mm256_loadu_ps(data + i); + const __m256 right = _mm256_loadu_ps(data + half + i); + _mm256_storeu_ps(data + i, _mm256_mul_ps(_mm256_add_ps(left, right), multiplier)); + _mm256_storeu_ps( + data + half + i, _mm256_mul_ps(_mm256_sub_ps(left, right), multiplier) + ); + } +#endif } // Both backends and all compilers use the same AVX butterfly ordering. -template +template void fht_rotate_impl( const float* data, float* rotated_vec, @@ -25,10 +47,10 @@ void fht_rotate_impl( float fac, const uint8_t* flip ) { - void (*fht)(float*) = nullptr; -#define RABITQ_FHT_CASE(log_dim) \ - case size_t{1} << (log_dim): \ - fht = fht_intrinsics; \ + void (*fht)(float*, float) = nullptr; +#define RABITQ_FHT_CASE(log_dim) \ + case size_t{1} << (log_dim): \ + fht = fht_intrinsics; \ break; switch (trunc_dim) { RABITQ_FHT_CASE(6) @@ -51,50 +73,32 @@ void fht_rotate_impl( std::fill(rotated_vec + dim, rotated_vec + padded_dim, 0.0F); if (trunc_dim == padded_dim) { - flip_sign(flip, rotated_vec, padded_dim); - fht(rotated_vec); - rescale(rotated_vec, trunc_dim, fac); - - flip_sign(flip + (padded_dim / 8), rotated_vec, padded_dim); - fht(rotated_vec); - rescale(rotated_vec, trunc_dim, fac); - - flip_sign(flip + (2 * padded_dim / 8), rotated_vec, padded_dim); - fht(rotated_vec); - rescale(rotated_vec, trunc_dim, fac); - - flip_sign(flip + (3 * padded_dim / 8), rotated_vec, padded_dim); - fht(rotated_vec); - rescale(rotated_vec, trunc_dim, fac); - + for (size_t round = 0; round < 4; ++round) { + flip_sign(flip + round * (padded_dim / 8), rotated_vec, padded_dim); + fht(rotated_vec, fac); + } return; } + const auto walk_and_flip = [&](const uint8_t* signs) { + if constexpr (kacs_flip != nullptr) { + kacs_flip(signs, rotated_vec, padded_dim); + } else { + kacs_walk(rotated_vec, padded_dim); + flip_sign(signs, rotated_vec, padded_dim); + } + }; size_t start = padded_dim - trunc_dim; flip_sign(flip, rotated_vec, padded_dim); - fht(rotated_vec); - rescale(rotated_vec, trunc_dim, fac); - kacs_walk(rotated_vec, padded_dim); - - flip_sign(flip + (padded_dim / 8), rotated_vec, padded_dim); - fht(rotated_vec + start); - rescale(rotated_vec + start, trunc_dim, fac); - kacs_walk(rotated_vec, padded_dim); - - flip_sign(flip + (2 * padded_dim / 8), rotated_vec, padded_dim); - fht(rotated_vec); - rescale(rotated_vec, trunc_dim, fac); - kacs_walk(rotated_vec, padded_dim); - - flip_sign(flip + (3 * padded_dim / 8), rotated_vec, padded_dim); - fht(rotated_vec + start); - rescale(rotated_vec + start, trunc_dim, fac); - kacs_walk(rotated_vec, padded_dim); - - // This can be removed if we don't care about the absolute value of - // similarities. - rescale(rotated_vec, padded_dim, 0.25F); + fht(rotated_vec, fac); + walk_and_flip(flip + (padded_dim / 8)); + fht(rotated_vec + start, fac); + walk_and_flip(flip + (2 * padded_dim / 8)); + fht(rotated_vec, fac); + walk_and_flip(flip + (3 * padded_dim / 8)); + fht(rotated_vec + start, fac); + kacs_walk_rescale(rotated_vec, padded_dim, 0.25F); } } // namespace } // namespace rabitqlib::simd diff --git a/src/simd/space_avx2.cpp b/src/simd/space_avx2.cpp index 334094f0..8f624b8d 100644 --- a/src/simd/space_avx2.cpp +++ b/src/simd/space_avx2.cpp @@ -5,6 +5,7 @@ #include #include +#include "packed_tail_avx2.hpp" #include "rabitqlib/simd/space_dispatch.hpp" #include "rabitqlib/utils/space.hpp" #include "space_float_kernels.hpp" @@ -91,92 +92,123 @@ void scalar_quantize_uint16_avx2( void new_transpose_bin_avx2( const uint16_t* q, uint64_t* tq, size_t padded_dim, size_t b_query ) { - for (size_t i = 0; i < padded_dim; i += 64) { - __m256i vec_00_to_15 = _mm256_loadu_si256((__m256i const*)(q)); - __m256i vec_16_to_31 = _mm256_loadu_si256((__m256i const*)(q + 16)); - __m256i vec_32_to_47 = _mm256_loadu_si256((__m256i const*)(q + 32)); - __m256i vec_48_to_63 = _mm256_loadu_si256((__m256i const*)(q + 48)); + // Reverse coordinates once; movemask then produces the stored bit order. + const __m256i order = _mm256_broadcastsi128_si256( + _mm_setr_epi8(14, 15, 12, 13, 10, 11, 8, 9, 6, 7, 4, 5, 2, 3, 0, 1) + ); + for (size_t i = 0; i < padded_dim - padded_dim % 64; i += 64) { + const __m256i a = _mm256_loadu_si256(reinterpret_cast(q)); + const __m256i b = _mm256_loadu_si256(reinterpret_cast(q + 16)); + const __m256i c = _mm256_loadu_si256(reinterpret_cast(q + 32)); + const __m256i d = _mm256_loadu_si256(reinterpret_cast(q + 48)); + // Arrange eight-coordinate groups so packs produces the final bit order. + __m256i hi_odd = _mm256_shuffle_epi8(_mm256_permute2x128_si256(b, a, 0x20), order); + __m256i hi_even = _mm256_shuffle_epi8(_mm256_permute2x128_si256(b, a, 0x31), order); + __m256i lo_odd = _mm256_shuffle_epi8(_mm256_permute2x128_si256(d, c, 0x20), order); + __m256i lo_even = _mm256_shuffle_epi8(_mm256_permute2x128_si256(d, c, 0x31), order); // the first (16 - b_query) bits are empty const int shift = static_cast(16 - b_query); - vec_00_to_15 = _mm256_slli_epi32(vec_00_to_15, shift); - vec_16_to_31 = _mm256_slli_epi32(vec_16_to_31, shift); - vec_32_to_47 = _mm256_slli_epi32(vec_32_to_47, shift); - vec_48_to_63 = _mm256_slli_epi32(vec_48_to_63, shift); + hi_odd = _mm256_slli_epi32(hi_odd, shift); + hi_even = _mm256_slli_epi32(hi_even, shift); + lo_odd = _mm256_slli_epi32(lo_odd, shift); + lo_even = _mm256_slli_epi32(lo_even, shift); for (size_t j = 0; j < b_query; ++j) { // pack two 16-bit vectors to 8-bit interleaved vectors - __m256i p0 = _mm256_packs_epi16(vec_00_to_15, vec_16_to_31); - __m256i p1 = _mm256_packs_epi16(vec_32_to_47, vec_48_to_63); + __m256i p0 = _mm256_packs_epi16(lo_even, lo_odd); + __m256i p1 = _mm256_packs_epi16(hi_even, hi_odd); - uint32_t m0 = _mm256_movemask_epi8(p0); - uint32_t m1 = _mm256_movemask_epi8(p1); + const uint32_t m0 = _mm256_movemask_epi8(p0); + const uint32_t m1 = _mm256_movemask_epi8(p1); - // Fix AVX2 Lane Ordering of the interleaved mask - auto fix_avx2_mask = [](uint32_t m) { - return (m & 0xFF0000FF) | ((m & 0x00FF0000) >> 8) | ((m & 0x0000FF00) << 8); - }; - - m0 = fix_avx2_mask(m0); - m1 = fix_avx2_mask(m1); - - m0 = reverse_bits(m0); - m1 = reverse_bits(m1); - - uint64_t v = (static_cast(m0) << 32) | m1; + const uint64_t v = uint64_t{m0} | (uint64_t{m1} << 32); tq[b_query - j - 1] = v; - vec_00_to_15 = _mm256_slli_epi16(vec_00_to_15, 1); - vec_16_to_31 = _mm256_slli_epi16(vec_16_to_31, 1); - vec_32_to_47 = _mm256_slli_epi16(vec_32_to_47, 1); - vec_48_to_63 = _mm256_slli_epi16(vec_48_to_63, 1); + hi_odd = _mm256_slli_epi16(hi_odd, 1); + hi_even = _mm256_slli_epi16(hi_even, 1); + lo_odd = _mm256_slli_epi16(lo_odd, 1); + lo_even = _mm256_slli_epi16(lo_even, 1); } tq += b_query; q += 64; } + const size_t tail_dim = padded_dim % 64; + // The partial word contains exactly 32 coordinates. + if (tail_dim != 0) { + const __m256i a = _mm256_loadu_si256(reinterpret_cast(q)); + const __m256i b = _mm256_loadu_si256(reinterpret_cast(q + 16)); + __m256i odd = _mm256_shuffle_epi8(_mm256_permute2x128_si256(b, a, 0x20), order); + __m256i even = _mm256_shuffle_epi8(_mm256_permute2x128_si256(b, a, 0x31), order); + odd = _mm256_slli_epi16(odd, static_cast(16 - b_query)); + even = _mm256_slli_epi16(even, static_cast(16 - b_query)); + for (size_t bit = 0; bit < b_query; ++bit) { + const __m256i bytes = _mm256_packs_epi16(even, odd); + tq[b_query - bit - 1] = static_cast(_mm256_movemask_epi8(bytes)); + odd = _mm256_slli_epi16(odd, 1); + even = _mm256_slli_epi16(even, 1); + } + } } void new_transpose_bin_512_avx2( const uint8_t* q, uint64_t* tq, size_t padded_dim, size_t b_query ) { + const __m128i reverse_bytes = + _mm_setr_epi8(15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0); + const auto load_reversed = [&](const uint8_t* values) { + const __m256i loaded = _mm256_loadu_si256(reinterpret_cast(values)); + return _mm256_permute4x64_epi64( + _mm256_shuffle_epi8(loaded, _mm256_broadcastsi128_si256(reverse_bytes)), 0x4E + ); + }; + // Shift each plane to byte sign bits. Across at most eight planes, carries + // from the neighboring byte never reach those sign bits. for (size_t i = 0; i < padded_dim;) { size_t block_size = 512; if (i + 512 > padded_dim) { block_size = padded_dim - i; } // Each chunk represents 64 bytes (512 bits) of dimensions - size_t num_chunks = block_size / 64; + const size_t full_chunks = block_size / 64; + const size_t tail_dim = block_size % 64; + const size_t num_chunks = full_chunks + (tail_dim != 0); - for (size_t k = 0; k < num_chunks; ++k) { + for (size_t k = 0; k < full_chunks; ++k) { // Load 64 bytes using two sequential 32-byte AVX2 registers const uint8_t* current_q_lo = q + i + k * 64; const uint8_t* current_q_hi = q + i + k * 64 + 32; - __m256i vec_lo = - _mm256_loadu_si256(reinterpret_cast(current_q_lo)); - __m256i vec_hi = - _mm256_loadu_si256(reinterpret_cast(current_q_hi)); + __m256i vec_lo = _mm256_slli_epi16( + load_reversed(current_q_lo), static_cast(8 - b_query) + ); + __m256i vec_hi = _mm256_slli_epi16( + load_reversed(current_q_hi), static_cast(8 - b_query) + ); for (size_t j = 0; j < b_query; ++j) { - int bit_idx = static_cast(b_query - 1 - j); - __m256i mask_vec = _mm256_set1_epi8(static_cast(1 << bit_idx)); - - // Process lower 32 bytes - __m256i res_lo = _mm256_and_si256(vec_lo, mask_vec); - __m256i eq_lo = _mm256_cmpeq_epi8(res_lo, _mm256_setzero_si256()); - uint32_t m_lo = ~static_cast(_mm256_movemask_epi8(eq_lo)); - - // Process upper 32 bytes - __m256i res_hi = _mm256_and_si256(vec_hi, mask_vec); - __m256i eq_hi = _mm256_cmpeq_epi8(res_hi, _mm256_setzero_si256()); - uint32_t m_hi = ~static_cast(_mm256_movemask_epi8(eq_hi)); + const auto m_lo = static_cast(_mm256_movemask_epi8(vec_lo)); + const auto m_hi = static_cast(_mm256_movemask_epi8(vec_hi)); // Combine both 32-bit masks into a single 64-bit mask - uint64_t m = (static_cast(m_hi) << 32) | m_lo; + uint64_t m = (static_cast(m_lo) << 32) | m_hi; // Write into the 64-bit structured macro-layout - tq[(b_query - j - 1) * num_chunks + k] = reverse_bits_u64(m); + tq[(b_query - j - 1) * num_chunks + k] = m; + vec_lo = _mm256_slli_epi16(vec_lo, 1); + vec_hi = _mm256_slli_epi16(vec_hi, 1); + } + } + + if (tail_dim != 0) { + __m256i values = _mm256_slli_epi16( + load_reversed(q + i + full_chunks * 64), static_cast(8 - b_query) + ); + for (size_t bit = 0; bit < b_query; ++bit) { + tq[(b_query - bit - 1) * num_chunks + full_chunks] = + static_cast(_mm256_movemask_epi8(values)); + values = _mm256_slli_epi16(values, 1); } } @@ -186,46 +218,7 @@ void new_transpose_bin_512_avx2( } float mask_ip_x0_q_avx2(const float* query, const uint8_t* data, size_t padded_dim) { - const size_t num_blk = padded_dim / 64; - const auto* it_data = data; - const float* it_query = query; - const __m256i shifts0 = _mm256_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7); - const __m256i shifts1 = _mm256_setr_epi32(8, 9, 10, 11, 12, 13, 14, 15); - const __m256i shifts2 = _mm256_setr_epi32(16, 17, 18, 19, 20, 21, 22, 23); - const __m256i shifts3 = _mm256_setr_epi32(24, 25, 26, 27, 28, 29, 30, 31); - __m256 sum0 = _mm256_setzero_ps(); - __m256 sum1 = _mm256_setzero_ps(); - __m256 sum2 = _mm256_setzero_ps(); - __m256 sum3 = _mm256_setzero_ps(); - - for (size_t i = 0; i < num_blk; ++i) { - // Stored coordinates run from bit 63 to bit 0: high word first, - // with each selected bit shifted into a maskload lane's sign bit. - for (size_t half = 0; half < 2; ++half) { - int32_t word; - std::memcpy(&word, it_data + (1 - half) * sizeof(word), sizeof(word)); - const __m256i bits = _mm256_set1_epi32(word); - sum0 = _mm256_add_ps( - sum0, _mm256_maskload_ps(it_query, _mm256_sllv_epi32(bits, shifts0)) - ); - sum1 = _mm256_add_ps( - sum1, _mm256_maskload_ps(it_query + 8, _mm256_sllv_epi32(bits, shifts1)) - ); - sum2 = _mm256_add_ps( - sum2, _mm256_maskload_ps(it_query + 16, _mm256_sllv_epi32(bits, shifts2)) - ); - sum3 = _mm256_add_ps( - sum3, _mm256_maskload_ps(it_query + 24, _mm256_sllv_epi32(bits, shifts3)) - ); - it_query += 32; - } - it_data += sizeof(uint64_t); - } - - const __m256 sum = _mm256_add_ps(_mm256_add_ps(sum0, sum1), _mm256_add_ps(sum2, sum3)); - __m128 lanes = _mm_add_ps(_mm256_castps256_ps128(sum), _mm256_extractf128_ps(sum, 1)); - lanes = _mm_add_ps(lanes, _mm_movehl_ps(lanes, lanes)); - return _mm_cvtss_f32(_mm_add_ss(lanes, _mm_movehdup_ps(lanes))); + return detail::mask_ip_avx2(query, data, padded_dim); } float mask_ip_x0_q_avx2(const float* query, const uint64_t* data, size_t padded_dim) { diff --git a/src/simd/space_avx512.cpp b/src/simd/space_avx512.cpp index b0f1c271..9fbf1dbc 100644 --- a/src/simd/space_avx512.cpp +++ b/src/simd/space_avx512.cpp @@ -84,10 +84,18 @@ void scalar_quantize_uint16_avx512( void new_transpose_bin_avx512( const uint16_t* q, uint64_t* tq, size_t padded_dim, size_t b_query ) { + // Reverse coordinates once; mask extraction then produces the stored bit order. + const auto reverse_words = [](__m512i values) { + const __m512i order = _mm512_broadcast_i32x4( + _mm_setr_epi8(14, 15, 12, 13, 10, 11, 8, 9, 6, 7, 4, 5, 2, 3, 0, 1) + ); + values = _mm512_shuffle_epi8(values, order); + return _mm512_shuffle_i64x2(values, values, 0x1B); + }; // 512 / 16 = 32 - for (size_t i = 0; i < padded_dim; i += 64) { - __m512i vec_00_to_31 = _mm512_loadu_si512(q); - __m512i vec_32_to_63 = _mm512_loadu_si512(q + 32); + for (size_t i = 0; i < padded_dim - padded_dim % 64; i += 64) { + __m512i vec_00_to_31 = reverse_words(_mm512_loadu_si512(q)); + __m512i vec_32_to_63 = reverse_words(_mm512_loadu_si512(q + 32)); // the first (16 - b_query) bits are empty vec_00_to_31 = _mm512_slli_epi32(vec_00_to_31, (16 - b_query)); @@ -96,10 +104,7 @@ void new_transpose_bin_avx512( for (size_t j = 0; j < b_query; ++j) { uint32_t v0 = _mm512_movepi16_mask(vec_00_to_31); // get most significant bit uint32_t v1 = _mm512_movepi16_mask(vec_32_to_63); // get most significant bit - // [TODO: remove all reverse_bits] - v0 = reverse_bits(v0); - v1 = reverse_bits(v1); - uint64_t v = (static_cast(v0) << 32) + v1; + const uint64_t v = uint64_t{v1} | (uint64_t{v0} << 32); tq[b_query - j - 1] = v; @@ -109,11 +114,26 @@ void new_transpose_bin_avx512( tq += b_query; q += 64; } + if (padded_dim % 64 != 0) { + __m512i values = reverse_words(_mm512_loadu_si512(q)); + values = _mm512_slli_epi16(values, static_cast(16 - b_query)); + for (size_t bit = 0; bit < b_query; ++bit) { + tq[b_query - bit - 1] = _mm512_movepi16_mask(values); + values = _mm512_slli_epi16(values, 1); + } + } } void new_transpose_bin_512_avx512( const uint8_t* q, uint64_t* tq, size_t padded_dim, size_t b_query ) { + const auto reverse_bytes = [](__m512i values) { + const __m512i order = _mm512_broadcast_i32x4( + _mm_setr_epi8(15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0) + ); + values = _mm512_shuffle_epi8(values, order); + return _mm512_shuffle_i64x2(values, values, 0x1B); + }; // Keep full 512-dim blocks as 8 chunks, but store the tail as compact // [b_query x num_chunks] so runtime can use maskz loads without query padding. for (size_t i = 0; i < padded_dim;) { @@ -121,11 +141,13 @@ void new_transpose_bin_512_avx512( if (i + 512 > padded_dim) { block_size = padded_dim - i; } - size_t num_chunks = block_size / 64; + const size_t full_chunks = block_size / 64; + const size_t tail_dim = block_size % 64; + const size_t num_chunks = full_chunks + (tail_dim != 0); - for (size_t k = 0; k < num_chunks; ++k) { + for (size_t k = 0; k < full_chunks; ++k) { const uint8_t* current_q = q + i + k * 64; - __m512i vec = _mm512_loadu_si512(current_q); + const __m512i vec = reverse_bytes(_mm512_loadu_si512(current_q)); for (size_t j = 0; j < b_query; ++j) { int bit_idx = static_cast(b_query - 1 - j); @@ -134,8 +156,18 @@ void new_transpose_bin_512_avx512( 1U << bit_idx ); // NOLINT(bugprone-narrowing-conversions) __mmask64 m = _mm512_test_epi8_mask(vec, _mm512_set1_epi8(bit_mask)); - tq[(b_query - j - 1) * num_chunks + k] = - reverse_bits_u64(static_cast(m)); + tq[(b_query - j - 1) * num_chunks + k] = static_cast(m); + } + } + + if (tail_dim != 0) { + constexpr __mmask64 valid = 0xFFFFFFFFULL; + const __m512i values = + reverse_bytes(_mm512_maskz_loadu_epi8(valid, q + i + full_chunks * 64)); + for (size_t bit = 0; bit < b_query; ++bit) { + const __m512i mask = _mm512_set1_epi8(static_cast(1U << bit)); + const uint64_t selected = _mm512_test_epi8_mask(values, mask); + tq[bit * num_chunks + full_chunks] = selected >> 32; } } @@ -145,42 +177,9 @@ void new_transpose_bin_512_avx512( } float mask_ip_x0_q_avx512(const float* query, const uint8_t* data, size_t padded_dim) { - const size_t num_blk = padded_dim / 64; - const uint8_t* it_data = data; - const float* it_query = query; - - // __m512 sum0 = _mm512_setzero_ps(); - // __m512 sum1 = _mm512_setzero_ps(); - // __m512 sum2 = _mm512_setzero_ps(); - // __m512 sum3 = _mm512_setzero_ps(); - - __m512 sum = _mm512_setzero_ps(); - for (size_t i = 0; i < num_blk; ++i) { - uint64_t bits = reverse_bits_u64(load_unaligned_u64(it_data)); - - auto mask0 = static_cast<__mmask16>(bits); - auto mask1 = static_cast<__mmask16>(bits >> 16); - auto mask2 = static_cast<__mmask16>(bits >> 32); - auto mask3 = static_cast<__mmask16>(bits >> 48); - - __m512 masked0 = _mm512_maskz_loadu_ps(mask0, it_query); - __m512 masked1 = _mm512_maskz_loadu_ps(mask1, it_query + 16); - __m512 masked2 = _mm512_maskz_loadu_ps(mask2, it_query + 32); - __m512 masked3 = _mm512_maskz_loadu_ps(mask3, it_query + 48); - - sum = _mm512_add_ps(sum, masked0); - sum = _mm512_add_ps(sum, masked1); - sum = _mm512_add_ps(sum, masked2); - sum = _mm512_add_ps(sum, masked3); - - // _mm_prefetch(reinterpret_cast(it_query + 128), _MM_HINT_T1); - - it_data += sizeof(uint64_t); - it_query += 64; - } - - // __m512 sum = _mm512_add_ps(_mm512_add_ps(sum0, sum1), _mm512_add_ps(sum2, sum3)); - return _mm512_reduce_add_ps(sum); + // Reuse the AVX2 implementation: independent accumulators avoid the long + // reduction dependency of the wider masked-load loop. + return mask_ip_x0_q_avx2(query, data, padded_dim); } float mask_ip_x0_q_avx512(const float* query, const uint64_t* data, size_t padded_dim) { diff --git a/src/simd/space_excode_avx2.cpp b/src/simd/space_excode_avx2.cpp index 3d478c68..0ac1ac67 100644 --- a/src/simd/space_excode_avx2.cpp +++ b/src/simd/space_excode_avx2.cpp @@ -3,6 +3,7 @@ #include #include +#include "packed_tail_avx2.hpp" #include "rabitqlib/simd/space_dispatch.hpp" namespace rabitqlib::simd::excode_ipimpl { @@ -94,7 +95,7 @@ float ip64_fxu2_avx2( float result = 0; const __m128i mask = _mm_set1_epi8(0b00000011); - for (size_t i = 0; i < dim; i += 64) { + for (size_t i = 0; i < dim - dim % 64; i += 64) { __m128i compact = _mm_loadu_si128(reinterpret_cast(compact_code)); __m128i vec_00_to_15 = _mm_and_si128(compact, mask); @@ -109,6 +110,9 @@ float ip64_fxu2_avx2( compact_code += 16; } + simd::detail::accumulate_excode_tail_avx2<2>( + query + dim - dim % 64, compact_code, dim % 64, sum0, sum1 + ); result = mm256_reduce_add_ps( _mm256_add_ps(_mm256_add_ps(sum0, sum1), _mm256_add_ps(sum2, sum3)) ); @@ -125,7 +129,7 @@ float ip64_fxu3_avx2( const __m128i mask = _mm_set1_epi8(0b11); const __m128i top_mask = _mm_set1_epi8(0b100); - for (size_t i = 0; i < dim; i += 64) { + for (size_t i = 0; i < dim - dim % 64; i += 64) { __m128i compact2 = _mm_loadu_si128(reinterpret_cast(compact_code)); compact_code += 16; @@ -156,6 +160,9 @@ float ip64_fxu3_avx2( contribute_ip(vec_48_to_63, &query[i + 48], sum2, sum3); } + simd::detail::accumulate_excode_tail_avx2<3>( + query + dim - dim % 64, compact_code, dim % 64, sum0, sum1 + ); result = mm256_reduce_add_ps( _mm256_add_ps(_mm256_add_ps(sum0, sum1), _mm256_add_ps(sum2, sum3)) ); @@ -196,7 +203,7 @@ float ip64_fxu5_avx2( const __m128i mask = _mm_set1_epi8(0b1111); const __m128i top_mask = _mm_set1_epi8(0b10000); - for (size_t i = 0; i < dim; i += 64) { + for (size_t i = 0; i < dim - dim % 64; i += 64) { __m128i compact4_1 = _mm_loadu_si128(reinterpret_cast(compact_code)); __m128i compact4_2 = @@ -230,6 +237,9 @@ float ip64_fxu5_avx2( contribute_ip(vec_32_to_47, &query[i + 32], sum0, sum1); contribute_ip(vec_48_to_63, &query[i + 48], sum2, sum3); } + simd::detail::accumulate_excode_tail_avx2<5>( + query + dim - dim % 64, compact_code, dim % 64, sum0, sum1 + ); result = mm256_reduce_add_ps( _mm256_add_ps(_mm256_add_ps(sum0, sum1), _mm256_add_ps(sum2, sum3)) ); @@ -246,7 +256,7 @@ float ip64_fxu6_avx2( const __m128i mask6 = _mm_set1_epi8(0b00111111); const __m128i mask2 = _mm_set1_epi8(static_cast(0b11000000)); - for (size_t i = 0; i < dim; i += 64) { + for (size_t i = 0; i < dim - dim % 64; i += 64) { __m128i cpt1 = _mm_loadu_si128(reinterpret_cast(compact_code)); __m128i cpt2 = _mm_loadu_si128(reinterpret_cast(compact_code + 16)); __m128i cpt3 = _mm_loadu_si128(reinterpret_cast(compact_code + 32)); @@ -269,6 +279,9 @@ float ip64_fxu6_avx2( contribute_ip(vec_32_to_47, &query[i + 32], sum0, sum1); contribute_ip(vec_48_to_63, &query[i + 48], sum2, sum3); } + simd::detail::accumulate_excode_tail_avx2<6>( + query + dim - dim % 64, compact_code, dim % 64, sum0, sum1 + ); result = mm256_reduce_add_ps( _mm256_add_ps(_mm256_add_ps(sum0, sum1), _mm256_add_ps(sum2, sum3)) ); @@ -286,7 +299,7 @@ float ip64_fxu7_avx2( const __m128i mask2 = _mm_set1_epi8(static_cast(0b11000000)); const __m128i top_mask = _mm_set1_epi8(0b1000000); - for (size_t i = 0; i < dim; i += 64) { + for (size_t i = 0; i < dim - dim % 64; i += 64) { __m128i cpt1 = _mm_loadu_si128(reinterpret_cast(compact_code)); __m128i cpt2 = _mm_loadu_si128(reinterpret_cast(compact_code + 16)); __m128i cpt3 = _mm_loadu_si128(reinterpret_cast(compact_code + 32)); @@ -326,6 +339,9 @@ float ip64_fxu7_avx2( contribute_ip(vec_48_to_63, &query[i + 48], sum2, sum3); } + simd::detail::accumulate_excode_tail_avx2<7>( + query + dim - dim % 64, compact_code, dim % 64, sum0, sum1 + ); result = mm256_reduce_add_ps( _mm256_add_ps(_mm256_add_ps(sum0, sum1), _mm256_add_ps(sum2, sum3)) ); diff --git a/src/simd/space_excode_avx512.cpp b/src/simd/space_excode_avx512.cpp index e1bc2107..31b3233b 100644 --- a/src/simd/space_excode_avx512.cpp +++ b/src/simd/space_excode_avx512.cpp @@ -1,13 +1,58 @@ #include +#include #include #include +#include "packed_tail_avx2.hpp" #include "rabitqlib/simd/space_dispatch.hpp" namespace rabitqlib::simd::excode_ipimpl { namespace { +template +inline __m512 excode_ip_tail_sum_avx512( + const float* query, const uint8_t* compact, __m512 sum +) { + // Static constants avoid short stack stores followed by a wider SIMD reload. + alignas(64) static constexpr auto kPositions = [] { + std::array positions{}; + // Form 32-bit words directly; 0x80 shuffle indices zero their upper bytes. + for (size_t i = 0; i < 16; ++i) { + positions[4 * i] = static_cast(i * Bits / 8); + positions[4 * i + 1] = static_cast(i * Bits / 8 + 1); + positions[4 * i + 2] = 0x80; + positions[4 * i + 3] = 0x80; + } + return positions; + }(); + alignas(64) static constexpr auto kOffsets = [] { + std::array offsets{}; + for (size_t i = 0; i < 16; ++i) { + offsets[i] = static_cast(i * Bits % (Bits == 2 ? 32 : 8)); + } + return offsets; + }(); + const __m512i shuffle = _mm512_load_si512(kPositions.data()); + const __m512i shifts = _mm512_load_si512(kOffsets.data()); + const __m512i mask = _mm512_set1_epi32((1U << Bits) - 1); + for (size_t i = 0; i < 32; i += 16) { + const __m128i packed = simd::detail::load_excode_tail(compact); + __m512i words; + if constexpr (Bits == 2) { + // All 16 codes fit in one word; broadcast it instead of gathering bytes. + words = _mm512_broadcastd_epi32(packed); + } else { + words = _mm512_shuffle_epi8(_mm512_broadcast_i32x4(packed), shuffle); + } + const __m512 values = + _mm512_cvtepi32_ps(_mm512_and_si512(_mm512_srlv_epi32(words, shifts), mask)); + sum = _mm512_fmadd_ps(values, _mm512_loadu_ps(query + i), sum); + compact += 2 * Bits; + } + return sum; +} + [[nodiscard]] inline uint64_t load_u64(const uint8_t* data) noexcept { uint64_t value = 0; std::memcpy(&value, data, sizeof(value)); @@ -57,10 +102,9 @@ float ip64_fxu2_avx512( ) { __m512 sum0 = _mm512_setzero_ps(), sum1 = sum0, sum2 = sum0, sum3 = sum0; - float result = 0; const __m128i mask = _mm_set1_epi8(0b00000011); - for (size_t i = 0; i < dim; i += 64) { + for (size_t i = 0; i < dim - dim % 64; i += 64) { __m128i compact = _mm_loadu_si128(reinterpret_cast(compact_code)); __m128i vec_00_to_15 = _mm_and_si128(compact, mask); @@ -89,11 +133,12 @@ float ip64_fxu2_avx512( compact_code += 16; } - result = _mm512_reduce_add_ps( - _mm512_add_ps(_mm512_add_ps(sum0, sum1), _mm512_add_ps(sum2, sum3)) - ); - - return result; + __m512 combined = _mm512_add_ps(_mm512_add_ps(sum0, sum1), _mm512_add_ps(sum2, sum3)); + if (dim % 64 != 0) { + combined = + excode_ip_tail_sum_avx512<2>(query + dim - dim % 64, compact_code, combined); + } + return _mm512_reduce_add_ps(combined); } float ip64_fxu3_avx512( @@ -101,11 +146,10 @@ float ip64_fxu3_avx512( ) { __m512 sum = _mm512_setzero_ps(); - float result = 0; const __m128i mask = _mm_set1_epi8(0b11); const __m128i top_mask = _mm_set1_epi8(0b100); - for (size_t i = 0; i < dim; i += 64) { + for (size_t i = 0; i < dim - dim % 64; i += 64) { __m128i compact2 = _mm_loadu_si128(reinterpret_cast(compact_code)); compact_code += 16; @@ -150,9 +194,12 @@ float ip64_fxu3_avx512( sum = _mm512_fmadd_ps(q, cf, sum); } - result = _mm512_reduce_add_ps(sum); - - return result; + __m512 combined = sum; + if (dim % 64 != 0) { + combined = + excode_ip_tail_sum_avx512<3>(query + dim - dim % 64, compact_code, combined); + } + return _mm512_reduce_add_ps(combined); } float ip16_fxu4_avx512( @@ -188,11 +235,10 @@ float ip64_fxu5_avx512( ) { __m512 sum = _mm512_setzero_ps(); - float result = 0.0F; const __m128i mask = _mm_set1_epi8(0b1111); const __m128i top_mask = _mm_set1_epi8(0b10000); - for (size_t i = 0; i < dim; i += 64) { + for (size_t i = 0; i < dim - dim % 64; i += 64) { __m128i compact4_1 = _mm_loadu_si128(reinterpret_cast(compact_code)); __m128i compact4_2 = @@ -240,9 +286,12 @@ float ip64_fxu5_avx512( cf = _mm512_cvtepi32_ps(_mm512_cvtepu8_epi32(vec_48_to_63)); sum = _mm512_fmadd_ps(q, cf, sum); } - result = _mm512_reduce_add_ps(sum); - - return result; + __m512 combined = sum; + if (dim % 64 != 0) { + combined = + excode_ip_tail_sum_avx512<5>(query + dim - dim % 64, compact_code, combined); + } + return _mm512_reduce_add_ps(combined); } float ip64_fxu6_avx512( @@ -250,11 +299,10 @@ float ip64_fxu6_avx512( ) { __m512 sum0 = _mm512_setzero_ps(), sum1 = sum0, sum2 = sum0, sum3 = sum0; - float result = 0.0F; const __m128i mask6 = _mm_set1_epi8(0b00111111); const __m128i mask2 = _mm_set1_epi8(static_cast(0b11000000)); - for (size_t i = 0; i < dim; i += 64) { + for (size_t i = 0; i < dim - dim % 64; i += 64) { __m128i cpt1 = _mm_loadu_si128(reinterpret_cast(compact_code)); __m128i cpt2 = _mm_loadu_si128(reinterpret_cast(compact_code + 16)); __m128i cpt3 = _mm_loadu_si128(reinterpret_cast(compact_code + 32)); @@ -291,11 +339,12 @@ float ip64_fxu6_avx512( cf = _mm512_cvtepi32_ps(_mm512_cvtepu8_epi32(vec_48_to_63)); sum3 = _mm512_fmadd_ps(q, cf, sum3); } - result = _mm512_reduce_add_ps( - _mm512_add_ps(_mm512_add_ps(sum0, sum1), _mm512_add_ps(sum2, sum3)) - ); - - return result; + __m512 combined = _mm512_add_ps(_mm512_add_ps(sum0, sum1), _mm512_add_ps(sum2, sum3)); + if (dim % 64 != 0) { + combined = + excode_ip_tail_sum_avx512<6>(query + dim - dim % 64, compact_code, combined); + } + return _mm512_reduce_add_ps(combined); } float ip64_fxu7_avx512( @@ -303,12 +352,11 @@ float ip64_fxu7_avx512( ) { __m512 sum = _mm512_setzero_ps(); - float result = 0.0F; const __m128i mask6 = _mm_set1_epi8(0b00111111); const __m128i mask2 = _mm_set1_epi8(static_cast(0b11000000)); const __m128i top_mask = _mm_set1_epi8(0b1000000); - for (size_t i = 0; i < dim; i += 64) { + for (size_t i = 0; i < dim - dim % 64; i += 64) { __m128i cpt1 = _mm_loadu_si128(reinterpret_cast(compact_code)); __m128i cpt2 = _mm_loadu_si128(reinterpret_cast(compact_code + 16)); __m128i cpt3 = _mm_loadu_si128(reinterpret_cast(compact_code + 32)); @@ -362,9 +410,12 @@ float ip64_fxu7_avx512( sum = _mm512_fmadd_ps(q, cf, sum); } - result = _mm512_reduce_add_ps(sum); - - return result; + __m512 combined = sum; + if (dim % 64 != 0) { + combined = + excode_ip_tail_sum_avx512<7>(query + dim - dim % 64, compact_code, combined); + } + return _mm512_reduce_add_ps(combined); } float ip16_fxu8_avx512( diff --git a/src/simd/space_packed_neon.cpp b/src/simd/space_packed_neon.cpp index 4142970f..a39da5fa 100644 --- a/src/simd/space_packed_neon.cpp +++ b/src/simd/space_packed_neon.cpp @@ -67,7 +67,7 @@ float packed_ip(const float* query, const uint8_t* compact, size_t dim) { compact += 2 * Bits; } } else { - for (size_t i = 0; i < dim; i += 64) { + for (size_t i = 0; i < dim - dim % 64; i += 64) { uint8x16_t c0, c1, c2, c3; if constexpr (Bits <= 3) { const uint8x16_t packed = vld1q_u8(compact), mask = vdupq_n_u8(3); @@ -107,6 +107,22 @@ float packed_ip(const float* query, const uint8_t* compact, size_t dim) { compact += 8 * Bits; } } + if constexpr (Bits != 1 && Bits != 4 && Bits != 8) { + if (dim % 64 != 0) { + // The final 32 codes form a contiguous little-endian bit stream. + double tail = 0; + for (size_t i = 0; i < 32; ++i) { + const size_t offset = i * Bits; + unsigned code = compact[offset / 8] >> (offset % 8); + if (offset % 8 + Bits > 8) { + code |= unsigned{compact[offset / 8 + 1]} << (8 - offset % 8); + } + tail += + static_cast(query[dim - 32 + i]) * (code & ((1U << Bits) - 1)); + } + return sum.result() + static_cast(tail); + } + } return sum.result(); } } // namespace diff --git a/src/simd/warmup_avx2.cpp b/src/simd/warmup_avx2.cpp index 031712dd..2153072d 100644 --- a/src/simd/warmup_avx2.cpp +++ b/src/simd/warmup_avx2.cpp @@ -1,68 +1,12 @@ -#include - #include #include -#include #include #include "rabitqlib/simd/warmup_dispatch.hpp" +#include "warmup_kernels.hpp" namespace rabitqlib::simd { -// Helper: AVX2 64-bit Popcount; Mula's method -static inline __m256i popcount_avx2(__m256i v) { - // Lookup table for population count of 0-15 - const __m256i lookup = _mm256_setr_epi8( - 0, - 1, - 1, - 2, - 1, - 2, - 2, - 3, - 1, - 2, - 2, - 3, - 2, - 3, - 3, - 4, - 0, - 1, - 1, - 2, - 1, - 2, - 2, - 3, - 1, - 2, - 2, - 3, - 2, - 3, - 3, - 4 - ); - const __m256i low_mask = _mm256_set1_epi8(0x0f); - - // Count low nibbles - __m256i lo = _mm256_and_si256(v, low_mask); - __m256i cnt_lo = _mm256_shuffle_epi8(lookup, lo); - - // Count high nibbles - __m256i hi = _mm256_and_si256(_mm256_srli_epi16(v, 4), low_mask); - __m256i cnt_hi = _mm256_shuffle_epi8(lookup, hi); - - // Add counts (bytes) - __m256i cnt_bytes = _mm256_add_epi8(cnt_lo, cnt_hi); - - // Sum bytes horizontally into 64-bit integers (SAD against 0) - return _mm256_sad_epu8(cnt_bytes, _mm256_setzero_si256()); -} - float warmup_ip_x0_q_512_avx2( const uint8_t* data, const uint64_t* query, @@ -71,101 +15,20 @@ float warmup_ip_x0_q_512_avx2( size_t padded_dim, size_t b_query ) { - size_t ip_scalar = 0; - size_t ppc_scalar = 0; - - __m256i acc_ip = _mm256_setzero_si256(); - __m256i acc_ppc = _mm256_setzero_si256(); - - size_t i = 0; - // Step by 512 bits at a time (64 bytes = 16 elements of 32-bit integers) - size_t dim_end_512 = (padded_dim / 512) * 512; - - constexpr size_t kMaxQueryBits = 8; - if (b_query > kMaxQueryBits) { + if (b_query > 8) { throw std::invalid_argument("warmup_ip_x0_q_512 requires at most 8 query bits"); } - __m256i acc_bits[kMaxQueryBits]; - for (size_t j = 0; j < b_query; ++j) { - acc_bits[j] = _mm256_setzero_si256(); - } - - for (; i < dim_end_512; i += 512) { - // Load 64 bytes of data using paired 32-byte loads - __m256i data_vec_lo = _mm256_loadu_si256(reinterpret_cast(data)); - __m256i data_vec_hi = - _mm256_loadu_si256(reinterpret_cast(data + 32)); - data += 64; - - acc_ppc = _mm256_add_epi64(acc_ppc, popcount_avx2(data_vec_lo)); - acc_ppc = _mm256_add_epi64(acc_ppc, popcount_avx2(data_vec_hi)); - - for (size_t j = 0; j < b_query; ++j) { - // Load 64 bytes of transposed query matching the 512-bit block layout - __m256i query_vec_lo = - _mm256_loadu_si256(reinterpret_cast(query)); - __m256i query_vec_hi = - _mm256_loadu_si256(reinterpret_cast(query + 4)); - query += 8; // Advance 8 x 64-bit ints (64 bytes) - - __m256i pop_lo = popcount_avx2(_mm256_and_si256(data_vec_lo, query_vec_lo)); - __m256i pop_hi = popcount_avx2(_mm256_and_si256(data_vec_hi, query_vec_hi)); - - acc_bits[j] = _mm256_add_epi64(acc_bits[j], pop_lo); - acc_bits[j] = _mm256_add_epi64(acc_bits[j], pop_hi); - } - } - - // Remainder block: handles leftovers less than 512 bits wide (e.g., last 448 bits) - size_t remaining_dim = padded_dim - i; - if (remaining_dim > 0) { - size_t num_chunks_64 = remaining_dim / 64; - size_t remaining_bytes = num_chunks_64 * sizeof(uint64_t); - - alignas(32) uint8_t data_tail[64]{}; - std::memcpy(data_tail, data, remaining_bytes); - __m256i data_vec_lo = - _mm256_load_si256(reinterpret_cast(data_tail)); - __m256i data_vec_hi = - _mm256_load_si256(reinterpret_cast(data_tail + 32)); - - acc_ppc = _mm256_add_epi64(acc_ppc, popcount_avx2(data_vec_lo)); - acc_ppc = _mm256_add_epi64(acc_ppc, popcount_avx2(data_vec_hi)); - - for (size_t j = 0; j < b_query; ++j) { - alignas(32) uint8_t query_tail[64]{}; - std::memcpy(query_tail, query, remaining_bytes); - __m256i query_vec_lo = - _mm256_load_si256(reinterpret_cast(query_tail)); - __m256i query_vec_hi = - _mm256_load_si256(reinterpret_cast(query_tail + 32)); - query += num_chunks_64; - - __m256i pop_lo = popcount_avx2(_mm256_and_si256(data_vec_lo, query_vec_lo)); - __m256i pop_hi = popcount_avx2(_mm256_and_si256(data_vec_hi, query_vec_hi)); - - acc_bits[j] = _mm256_add_epi64(acc_bits[j], pop_lo); - acc_bits[j] = _mm256_add_epi64(acc_bits[j], pop_hi); - } - } - - for (size_t j = 0; j < b_query; ++j) { - __m128i shift = _mm_cvtsi32_si128(static_cast(j)); - acc_ip = _mm256_add_epi64(acc_ip, _mm256_sll_epi64(acc_bits[j], shift)); + size_t ip_scalar = 0; + size_t ppc_scalar = 0; + if (padded_dim < 512) { + detail::accumulate_warmup_tail( + data, query, padded_dim, b_query, ip_scalar, ppc_scalar + ); + return (delta * static_cast(ip_scalar)) + + (vl * static_cast(ppc_scalar)); } - // Standard reduction for a single __m256i - auto mm256_reduce_add_epi64 = [](__m256i v) { - __m128i low = _mm256_castsi256_si128(v); - __m128i high = _mm256_extracti128_si256(v, 1); - __m128i sum = _mm_add_epi64(low, high); - return _mm_extract_epi64(sum, 0) + _mm_extract_epi64(sum, 1); - }; - - ip_scalar += mm256_reduce_add_epi64(acc_ip); - ppc_scalar += mm256_reduce_add_epi64(acc_ppc); - - return (delta * static_cast(ip_scalar)) + (vl * static_cast(ppc_scalar)); + return detail::warmup_blocks_avx2<8>(data, query, delta, vl, padded_dim, b_query); } float warmup_ip_x0_q_512_avx2( diff --git a/src/simd/warmup_avx512.cpp b/src/simd/warmup_avx512.cpp index ded79983..c4b1b9a6 100644 --- a/src/simd/warmup_avx512.cpp +++ b/src/simd/warmup_avx512.cpp @@ -1,10 +1,9 @@ -#include - #include #include #include #include "rabitqlib/simd/warmup_dispatch.hpp" +#include "warmup_kernels.hpp" namespace rabitqlib::simd { @@ -19,62 +18,22 @@ float warmup_ip_x0_q_512_avx512( size_t ip_scalar = 0; size_t ppc_scalar = 0; - __m512i acc_ip = _mm512_setzero_si512(); - __m512i acc_ppc = _mm512_setzero_si512(); - - size_t i = 0; - size_t dim_end_512 = (padded_dim / 512) * 512; - constexpr size_t kMaxQueryBits = 8; if (b_query > kMaxQueryBits) { throw std::invalid_argument("warmup_ip_x0_q_512 requires at most 8 query bits"); } - __m512i acc_bits[kMaxQueryBits]; - for (size_t j = 0; j < b_query; ++j) { - acc_bits[j] = _mm512_setzero_si512(); + // For at most four words, scalar POPCNT avoids vector setup and reduction. + if (padded_dim <= 256) { + detail::accumulate_warmup_tail( + data, query, padded_dim, b_query, ip_scalar, ppc_scalar + ); + return (delta * static_cast(ip_scalar)) + + (vl * static_cast(ppc_scalar)); } - for (; i < dim_end_512; i += 512) { - __m512i data_vec = _mm512_loadu_si512(data); - data += 64; - - acc_ppc = _mm512_add_epi64(acc_ppc, _mm512_popcnt_epi64(data_vec)); - - for (size_t j = 0; j < b_query; ++j) { - __m512i query_vec = _mm512_loadu_si512(query); - query += 8; - - __m512i pop = _mm512_popcnt_epi64(_mm512_and_si512(data_vec, query_vec)); - acc_bits[j] = _mm512_add_epi64(acc_bits[j], pop); - } - } - - size_t remaining_dim = padded_dim - i; - if (remaining_dim > 0) { - size_t num_chunks = remaining_dim / 64; - auto valid_mask = static_cast<__mmask8>((1u << num_chunks) - 1u); - - __m512i data_vec = _mm512_maskz_loadu_epi64(valid_mask, data); - acc_ppc = _mm512_add_epi64(acc_ppc, _mm512_popcnt_epi64(data_vec)); - - for (size_t j = 0; j < b_query; ++j) { - __m512i query_vec = _mm512_maskz_loadu_epi64(valid_mask, query); - query += num_chunks; - - __m512i pop = _mm512_popcnt_epi64(_mm512_and_si512(data_vec, query_vec)); - acc_bits[j] = _mm512_add_epi64(acc_bits[j], pop); - } - } - - for (size_t j = 0; j < b_query; ++j) { - __m128i shift = _mm_cvtsi32_si128(static_cast(j)); - acc_ip = _mm512_add_epi64(acc_ip, _mm512_sll_epi64(acc_bits[j], shift)); - } - - ip_scalar += static_cast(_mm512_reduce_add_epi64(acc_ip)); - ppc_scalar += static_cast(_mm512_reduce_add_epi64(acc_ppc)); - - return (delta * static_cast(ip_scalar)) + (vl * static_cast(ppc_scalar)); + return detail::warmup_blocks_avx512( + data, query, delta, vl, padded_dim, b_query + ); } float warmup_ip_x0_q_512_avx512( diff --git a/src/simd/warmup_kernels.hpp b/src/simd/warmup_kernels.hpp new file mode 100644 index 00000000..70bcfe34 --- /dev/null +++ b/src/simd/warmup_kernels.hpp @@ -0,0 +1,208 @@ +#pragma once + +#include + +#include +#include +#include + +#include "rabitqlib/utils/bitops.hpp" + +namespace rabitqlib::simd::detail { + +// Keep ISA-specific instantiations local to each backend translation unit. + +// A partial 512-coordinate block contains at most eight words per query plane. +// Scalar POPCNT avoids staging/zeroing two full vectors for each short plane. +static inline void accumulate_warmup_tail( + const uint8_t* data, + const uint64_t* query, + size_t dim, + size_t bits, + size_t& ip, + size_t& count +) { + const size_t full_words = dim / 64; + const size_t words = (dim + 63) / 64; + const auto accumulate = [&](uint64_t word, size_t index) { + count += bitops::popcount64(word); + for (size_t bit = 0; bit < bits; ++bit) { + ip += size_t{bitops::popcount64(word & query[bit * words + index])} << bit; + } + }; + for (size_t i = 0; i < full_words; ++i) { + uint64_t word; + std::memcpy(&word, data + i * sizeof(word), sizeof(word)); + accumulate(word, i); + } + if (dim % 64 != 0) { + accumulate( + bitops::load_word(data + full_words * sizeof(uint64_t), sizeof(uint32_t)), + full_words + ); + } +} + +static inline __m256i popcount_avx2(__m256i v) { + // Lookup table for population count of 0-15 + const __m256i lookup = _mm256_broadcastsi128_si256( + _mm_setr_epi8(0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4) + ); + const __m256i low_mask = _mm256_set1_epi8(0x0f); + + // Count low nibbles + __m256i lo = _mm256_and_si256(v, low_mask); + __m256i cnt_lo = _mm256_shuffle_epi8(lookup, lo); + + // Count high nibbles + __m256i hi = _mm256_and_si256(_mm256_srli_epi16(v, 4), low_mask); + __m256i cnt_hi = _mm256_shuffle_epi8(lookup, hi); + + // Add counts (bytes) + __m256i cnt_bytes = _mm256_add_epi8(cnt_lo, cnt_hi); + + // Sum bytes horizontally into 64-bit integers (SAD against 0) + return _mm256_sad_epu8(cnt_bytes, _mm256_setzero_si256()); +} + +template +static inline float warmup_blocks_avx2( + const uint8_t* data, + const uint64_t* query, + float delta, + float vl, + size_t padded_dim, + size_t b_query +) { + size_t ip_scalar = 0; + size_t ppc_scalar = 0; + + __m256i acc_ip = _mm256_setzero_si256(); + __m256i acc_ppc = _mm256_setzero_si256(); + + size_t i = 0; + // Step by 512 bits at a time (64 bytes = 16 elements of 32-bit integers) + size_t dim_end_512 = (padded_dim / 512) * 512; + + __m256i acc_bits[MaxQueryBits]; + for (size_t j = 0; j < b_query; ++j) { + acc_bits[j] = _mm256_setzero_si256(); + } + + for (; i < dim_end_512; i += 512) { + // Load 64 bytes of data using paired 32-byte loads + __m256i data_vec_lo = _mm256_loadu_si256(reinterpret_cast(data)); + __m256i data_vec_hi = + _mm256_loadu_si256(reinterpret_cast(data + 32)); + data += 64; + + acc_ppc = _mm256_add_epi64(acc_ppc, popcount_avx2(data_vec_lo)); + acc_ppc = _mm256_add_epi64(acc_ppc, popcount_avx2(data_vec_hi)); + + for (size_t j = 0; j < b_query; ++j) { + // Load 64 bytes of transposed query matching the 512-bit block layout + __m256i query_vec_lo = + _mm256_loadu_si256(reinterpret_cast(query)); + __m256i query_vec_hi = + _mm256_loadu_si256(reinterpret_cast(query + 4)); + query += 8; // Advance 8 x 64-bit ints (64 bytes) + + __m256i pop_lo = popcount_avx2(_mm256_and_si256(data_vec_lo, query_vec_lo)); + __m256i pop_hi = popcount_avx2(_mm256_and_si256(data_vec_hi, query_vec_hi)); + + acc_bits[j] = _mm256_add_epi64(acc_bits[j], pop_lo); + acc_bits[j] = _mm256_add_epi64(acc_bits[j], pop_hi); + } + } + + if (i < padded_dim) { + accumulate_warmup_tail(data, query, padded_dim - i, b_query, ip_scalar, ppc_scalar); + } + + for (size_t j = 0; j < b_query; ++j) { + __m128i shift = _mm_cvtsi32_si128(static_cast(j)); + acc_ip = _mm256_add_epi64(acc_ip, _mm256_sll_epi64(acc_bits[j], shift)); + } + + // Standard reduction for a single __m256i + auto mm256_reduce_add_epi64 = [](__m256i v) { + __m128i low = _mm256_castsi256_si128(v); + __m128i high = _mm256_extracti128_si256(v, 1); + __m128i sum = _mm_add_epi64(low, high); + return _mm_extract_epi64(sum, 0) + _mm_extract_epi64(sum, 1); + }; + + ip_scalar += mm256_reduce_add_epi64(acc_ip); + ppc_scalar += mm256_reduce_add_epi64(acc_ppc); + + return (delta * static_cast(ip_scalar)) + (vl * static_cast(ppc_scalar)); +} + +#if defined(__AVX512F__) || defined(_MSC_VER) +template +static inline float warmup_blocks_avx512( + const uint8_t* data, + const uint64_t* query, + float delta, + float vl, + size_t padded_dim, + size_t b_query +) { + __m512i acc_ip = _mm512_setzero_si512(); + __m512i acc_ppc = _mm512_setzero_si512(); + + size_t i = 0; + size_t dim_end_512 = (padded_dim / 512) * 512; + + __m512i acc_bits[MaxQueryBits]; + for (size_t j = 0; j < b_query; ++j) { + acc_bits[j] = _mm512_setzero_si512(); + } + + for (; i < dim_end_512; i += 512) { + __m512i data_vec = _mm512_loadu_si512(data); + data += 64; + + acc_ppc = _mm512_add_epi64(acc_ppc, _mm512_popcnt_epi64(data_vec)); + + for (size_t j = 0; j < b_query; ++j) { + __m512i query_vec = _mm512_loadu_si512(query); + query += 8; + + __m512i pop = _mm512_popcnt_epi64(_mm512_and_si512(data_vec, query_vec)); + acc_bits[j] = _mm512_add_epi64(acc_bits[j], pop); + } + } + + size_t remaining_dim = padded_dim - i; + if (remaining_dim > 0) { + const size_t num_chunks = (remaining_dim + 63) / 64; + const auto valid_mask = static_cast<__mmask8>((1u << num_chunks) - 1u); + // Byte masking zero-extends the partial last word without staging short + // stores through a stack buffer. The query planes already use whole words. + const __m512i data_vec = + _mm512_maskz_loadu_epi8((uint64_t{1} << (remaining_dim / 8)) - 1, data); + acc_ppc = _mm512_add_epi64(acc_ppc, _mm512_popcnt_epi64(data_vec)); + + for (size_t j = 0; j < b_query; ++j) { + __m512i query_vec = _mm512_maskz_loadu_epi64(valid_mask, query); + query += num_chunks; + + __m512i pop = _mm512_popcnt_epi64(_mm512_and_si512(data_vec, query_vec)); + acc_bits[j] = _mm512_add_epi64(acc_bits[j], pop); + } + } + + for (size_t j = 0; j < b_query; ++j) { + __m128i shift = _mm_cvtsi32_si128(static_cast(j)); + acc_ip = _mm512_add_epi64(acc_ip, _mm512_sll_epi64(acc_bits[j], shift)); + } + + const size_t ip_scalar = static_cast(_mm512_reduce_add_epi64(acc_ip)); + const size_t ppc_scalar = static_cast(_mm512_reduce_add_epi64(acc_ppc)); + + return (delta * static_cast(ip_scalar)) + (vl * static_cast(ppc_scalar)); +} +#endif + +} // namespace rabitqlib::simd::detail diff --git a/tests/python/fixtures/README.md b/tests/python/fixtures/README.md index 52c7d3f6..75e1d47f 100644 --- a/tests/python/fixtures/README.md +++ b/tests/python/fixtures/README.md @@ -3,5 +3,18 @@ storage was introduced (legacy, unversioned x86-64 format). It contains three vectors: `[0, 0, 0]`, `[1, 2, 3]`, and `[-1, -2, -3]`, each zero-padded to 64 dimensions, in one cluster with a zero centroid, four-bit L2 quantization, and FHT/Kac rotation. -The zero query has squared distances 0, 14, and 14. This fixture verifies that -adding versioned raw-vector storage preserves legacy loading and serialization. +The zero query has squared distances 0, 14, and 14. This fixture verifies legacy +loading and migration to the current versioned format without changing search results. + +`ivf_legacy_padding65_1bit.index` and `ivf_raw_v1_padding65.index` were saved before +the adaptive-padding change, using the historical unversioned quantized and raw v1 +formats respectively. They have original dimension 65, padded dimension 128, one +zero centroid, and the same three vectors above with zeros in all remaining +coordinates. Their widths are 1 and 32 bits, with L2 distance and FHT/Kac rotation. +They verify that loading does not substitute the new 96-coordinate default. + +`hnsw_legacy_padding65_4bit.index` was written before HNSW switched to 32-coordinate +padding: dimension 65, padded dimension 128, 4 total bits, L2, capacity 8, M=2, +ef_construction=8. Its three vectors and zero centroid match the dimension-65 IVF +fixtures above. A zero query returns IDs 0, 1, 2 with distances 0, 14, 14. Tests +check loading, preserving the old padded domain through save/load, and insertion. diff --git a/tests/python/fixtures/hnsw_legacy_padding65_4bit.index b/tests/python/fixtures/hnsw_legacy_padding65_4bit.index new file mode 100644 index 0000000000000000000000000000000000000000..a5e4f9a047bc09cdc4962f6dccebcd47cf9ddbd8 GIT binary patch literal 1105 zcmd;JfBnguHUuSq-AUG%#> z2b4cbF*G2B0f>zY!<2#KL7)K*gTz3f0ZgLGLxsQ;vOGut#0Oz?&;S$wVRX=dBmiT> zXe22x6C{p~K@uP_FaV_&Py&Rd8;B_FNLXS=w*e$>HYiINbetO|9!aZ1Nrx0((kid(>R2i9xnP5wI^&9vV9!$7aw`AtHLYd^=2vu zOOi?`_vULmpHFg+_7RO&Ot^0OX3L`4yIdkB`lZfDJaDFD7tB6To+Jc#8?Q@w>CV^M z`Ch%LkgNC5r$2F(|2|&s*b>rqsr+=7msS;*^M>jkhLkwLmVG?BvOh0h4m>=eN&ofo K+}1Y%2_FDO289&> literal 0 HcmV?d00001 diff --git a/tests/python/fixtures/ivf_legacy_padding65_1bit.index b/tests/python/fixtures/ivf_legacy_padding65_1bit.index new file mode 100644 index 0000000000000000000000000000000000000000..0a01f6a9a3b4f5234bf3dcb1e6fe0eb449125fbf GIT binary patch literal 1526 zcmZQ(fB;7*%?PDo6e9yOlqq&P+UAd9+}7Ph)|QoNGXrOCz42IR zp+M}1=#Z`Rjq@bJCB2$fc4`!r+}pQ%w$1;@g+|kVR>4dh!E^(X-;voMVO)G%;xIiR zd0c8>a>(+?VlZ)-7y}T%*vK@xI4OLPS`cVpU|?n70yAL*I|B!d38tADcu@FE47@0O z9(Y=W2(vM8qR6B339090U_#N4&gWx5%?;@8MVIGhU_mht-F^lJeiV5|20pxebp7aj zbo(?irc&!pU?TwazG8 #include +#include "rabitqlib/utils/cpu_features.hpp" + namespace rabitqlib::hnsw { struct HnswPruningTestAccess { @@ -1165,3 +1167,64 @@ TEST(HnswRemoveTest, LoadRejectsAnOutOfRangeClusterEvenWithTheMark) { } // namespace } // namespace rabitqlib::hnsw + +TEST(HnswPaddingTest, TailDimensionsMatchEveryAvailableSearchBackend) { + using namespace rabitqlib; + using namespace rabitqlib::hnsw; + constexpr size_t count = 24; + using Search = decltype(&hnsw::detail::search_knn_generic); + std::vector backends{hnsw::detail::search_knn}; +#if defined(__x86_64__) || defined(_M_X64) + if (cpu::has_avx2()) + backends.push_back(hnsw::detail::search_knn_avx2); + if (cpu::has_avx512_core()) + backends.push_back(hnsw::detail::search_knn_avx512_core); + if (cpu::has_avx512_popcnt()) + backends.push_back(hnsw::detail::search_knn_avx512_popcnt); +#endif +#if defined(__aarch64__) || defined(_M_ARM64) + if (cpu::has_neon()) + backends.push_back(hnsw::detail::search_knn_neon); +#endif + for (size_t dim : {64U, 96U, 128U, 160U, 256U, 544U, 576U}) { + std::mt19937 random(642); + std::uniform_real_distribution value(-1, 1); + std::vector data(count * dim), centroid(dim, 0), query(dim); + for (auto& x : data) + x = value(random); + for (auto& x : query) + x = value(random); + for (auto metric : {METRIC_L2, METRIC_IP}) { + for (size_t bits = 1; bits <= 9; ++bits) { + SCOPED_TRACE(dim); + SCOPED_TRACE(bits); + SCOPED_TRACE(metric); + std::vector clusters(count, 0); + HierarchicalNSW index(count, dim, bits, 8, count, 42, metric); + index.construct( + 1, centroid.data(), count - 4, data.data(), clusters.data(), 1, true + ); + // Insertion reconstructs stored sign/extra-bit tails before linking. + index.add(data.data() + (count - 4) * dim, 4, clusters.data(), true, 1); + index.search(query.data(), 1, 5, count, 1); + const auto reference = + hnsw::detail::search_knn_generic(index, query.data(), 5); + for (auto search : backends) { + auto expected = reference; + auto actual = search(index, query.data(), 5); + ASSERT_EQ(actual.size(), expected.size()); + while (!expected.empty()) { + EXPECT_EQ(actual.top().second, expected.top().second); + EXPECT_NEAR( + actual.top().first, + expected.top().first, + 2e-5F * std::max(1.0F, std::abs(expected.top().first)) + ); + actual.pop(); + expected.pop(); + } + } + } + } + } +} diff --git a/tests/unit/rabitqlib/index/ivf_test.cpp b/tests/unit/rabitqlib/index/ivf_test.cpp index c79899db..3bfc8c37 100644 --- a/tests/unit/rabitqlib/index/ivf_test.cpp +++ b/tests/unit/rabitqlib/index/ivf_test.cpp @@ -42,6 +42,7 @@ TEST(IvfSearchTest, BatchSearchMatchesIndividualQueries) { for (auto metric : {METRIC_L2, METRIC_IP}) { for (size_t bits : {1U, 2U, 3U, 4U, 5U, 6U, 7U, 8U, 9U, 32U}) { IVF index(kNum, kDim, kClusters, bits, metric); + EXPECT_EQ(index.padded_dim(), 96U); index.construct(data.data(), centroids.data(), labels.data(), false, 1); std::array expected{}, actual{}, ids_only{}; std::array expected_distances{}, actual_distances{}; @@ -392,12 +393,13 @@ TEST(IvfPersistenceTest, ReloadsRawAndQuantizedStorage) { const std::string path = ::testing::TempDir() + "rabitq_ivf_storage.index"; IVF loaded; for (auto rotator : {RotatorType::FhtKacRotator, RotatorType::MatrixRotator}) { - for (size_t bits : {32UL, 4UL, 32UL, 1UL}) { + for (size_t bits : {32UL, 1UL, 2UL, 3UL, 4UL, 5UL, 6UL, 7UL, 8UL, 9UL}) { IVF index(kNum, kDim, 1, bits, METRIC_L2, rotator); index.construct(data.data(), centroid.data(), clusters.data(), false, 1); index.save(path.c_str()); loaded.load(path.c_str()); EXPECT_EQ(loaded.nbits(), bits); + EXPECT_EQ(loaded.padded_dim(), index.padded_dim()); EXPECT_EQ(loaded.rotator_type(), rotator); std::vector ids(kNum); std::vector loaded_ids(kNum); @@ -505,7 +507,7 @@ TEST(IvfPersistenceTest, RejectsPointIdsUsingTheSearchBufferMarker) { // ── add / remove ──────────────────────────────────────────────────────────────────────── -constexpr size_t kDynDim = 65; // pads to 128, so padded and original dimensions differ +constexpr size_t kDynDim = 65; // Exercise padding and partially occupied FastScan batches. constexpr size_t kDynClusters = 4; std::vector random_rows(size_t rows, uint32_t seed) { diff --git a/tests/unit/rabitqlib/index/qg_test.cpp b/tests/unit/rabitqlib/index/qg_test.cpp index 56bf8d87..d309861d 100644 --- a/tests/unit/rabitqlib/index/qg_test.cpp +++ b/tests/unit/rabitqlib/index/qg_test.cpp @@ -42,6 +42,10 @@ namespace rabitqlib::symqg { struct QGConstructionTestAccess { + static void use_legacy_padding(QuantizedGraph& graph) { + graph.padded_dim_ = round_up_to_multiple(graph.dim_, 64); + graph.initialize(); + } static float build_distance(QGBuilder& builder, PID source_id, PID target_id) { std::vector reconstructed; std::optional prepared; @@ -1339,6 +1343,59 @@ TEST(QuantizedGraphPersistenceTest, RejectsMalformedPayloadWithoutChangingTarget std::remove(path.c_str()); } +TEST(QuantizedGraphPersistenceTest, PreservesNewAndLegacyPaddedDomains) { + constexpr size_t kNum = 33, kDim = 65, kDegree = 32; + std::vector data(kNum * kDim); + for (size_t i = 0; i < data.size(); ++i) { + data[i] = std::sin(static_cast(i) * 0.13F); + } + const std::string path = ::testing::TempDir() + "rabitq_qg_padding.index"; + for (auto metric : {METRIC_L2, METRIC_IP}) { + for (auto rotator : {RotatorType::FhtKacRotator, RotatorType::MatrixRotator}) { + for (size_t bits : {0U, 4U, 8U}) { + for (bool legacy : {false, true}) { + SCOPED_TRACE(::testing::Message() << bits << '/' << legacy); + QuantizedGraph source( + kNum, kDim, kDegree, metric, rotator, bits + ); + EXPECT_EQ(source.padded_dim(), 96U); + if (legacy) { + // The raw and version-1 quantized file layouts are unchanged. + // Recreate their original 64-padded rotation, codes, and rows. + QGConstructionTestAccess::use_legacy_padding(source); + } + QGBuilder builder(source, kDegree, data.data(), 1); + builder.build(2); + source.set_ef(kNum); + std::array expected_ids{}, actual_ids{}; + std::array expected_distances{}, actual_distances{}; + source.search( + data.data(), 5, expected_ids.data(), expected_distances.data() + ); + source.save(path.c_str()); + QuantizedGraph loaded; + loaded.load(path.c_str()); + EXPECT_EQ(loaded.padded_dim(), legacy ? 128U : 96U); + EXPECT_EQ( + QGConstructionTestAccess::rows(source), + QGConstructionTestAccess::rows(loaded) + ); + loaded.set_ef(kNum); + loaded.search( + data.data(), 5, actual_ids.data(), actual_distances.data() + ); + EXPECT_EQ(actual_ids, expected_ids); + EXPECT_EQ(actual_distances, expected_distances); + loaded.save(path.c_str()); + loaded.load(path.c_str()); + EXPECT_EQ(loaded.padded_dim(), source.padded_dim()); + } + } + } + } + std::remove(path.c_str()); +} + TEST(QuantizedGraphLifecycleTest, DestroysConcreteRotatorThroughBasePointer) { QuantizedGraph graph(33, 64, 32, METRIC_L2, RotatorType::MatrixRotator); EXPECT_EQ(graph.num_vertices(), 33U); diff --git a/tests/unit/rabitqlib/quantization/rabitq_test.cpp b/tests/unit/rabitqlib/quantization/rabitq_test.cpp index 8e93b5d9..7a5dca42 100644 --- a/tests/unit/rabitqlib/quantization/rabitq_test.cpp +++ b/tests/unit/rabitqlib/quantization/rabitq_test.cpp @@ -539,3 +539,45 @@ TEST(RabitqSignConventionTest, ZeroResidualDoesNotInflateReconstructionNorm) { } // namespace } // namespace rabitqlib::quant + +TEST(RabitqCompactCodes, WordTailsUseCompactRowStrides) { + using namespace rabitqlib; + for (size_t dim : {96U, 160U}) { + constexpr size_t count = 3; + std::vector data(count * dim), centroid(dim, 0); + for (size_t i = 0; i < data.size(); ++i) + data[i] = static_cast(i % 17) - 8; + for (auto metric : {METRIC_L2, METRIC_IP}) { + std::vector expected(count * dim / 8 + 1, 0xA5), actual = expected; + std::vector add(count), rescale(count), error(count); + std::vector batch_add(count), batch_rescale(count), batch_error(count); + for (size_t i = 0; i < count; ++i) { + quant::rabitq_impl::one_bit::one_bit_compact_code_to_bytes( + data.data() + i * dim, + centroid.data(), + dim, + expected.data() + i * dim / 8, + add[i], + rescale[i], + error[i], + metric + ); + } + quant::rabitq_impl::one_bit::one_bit_compact_codes( + data.data(), + centroid.data(), + count, + dim, + reinterpret_cast(actual.data()), + batch_add.data(), + batch_rescale.data(), + batch_error.data(), + metric + ); + EXPECT_EQ(actual, expected); + EXPECT_EQ(batch_add, add); + EXPECT_EQ(batch_rescale, rescale); + EXPECT_EQ(batch_error, error); + } + } +} diff --git a/tests/unit/rabitqlib/utils/excode_ip_test.cpp b/tests/unit/rabitqlib/utils/excode_ip_test.cpp index e96641c6..b3cf394d 100644 --- a/tests/unit/rabitqlib/utils/excode_ip_test.cpp +++ b/tests/unit/rabitqlib/utils/excode_ip_test.cpp @@ -61,13 +61,9 @@ TEST(ExcodeIp, BackendsMatchScalarAcrossWidthsAndBlockBoundaries) { for (size_t bits = 1; bits <= 8; ++bits) { SCOPED_TRACE(bits); - const size_t block_dim = (bits == 1 || bits == 4 || bits == 8) ? 16 : 64; const auto max_code = static_cast((1U << bits) - 1); - for (size_t dim : std::array{ - 0, 16, 32, 48, 64, 80, 96, 112, 128, 192, 256, 576, 960, 1024, 4096}) { - if (dim % block_dim != 0) { - continue; - } + for (size_t dim : std::array{ + 0, 64, 96, 128, 160, 192, 256, 544, 576, 960, 1024, 4096}) { SCOPED_TRACE(dim); for (size_t pattern = 0; pattern < 4; ++pattern) { SCOPED_TRACE(pattern); diff --git a/tests/unit/rabitqlib/utils/mask_ip_test.cpp b/tests/unit/rabitqlib/utils/mask_ip_test.cpp index 54476d95..52053a68 100644 --- a/tests/unit/rabitqlib/utils/mask_ip_test.cpp +++ b/tests/unit/rabitqlib/utils/mask_ip_test.cpp @@ -34,12 +34,12 @@ TEST(MaskIpX0Q, BackendsMatchScalarAcrossBlocksAndAlignments) { } #endif - for (size_t dim : {0, 64, 128, 192, 256, 448, 576, 1024, 4096}) { + for (size_t dim : {0, 64, 96, 128, 192, 256, 448, 544, 576, 1024, 4096}) { SCOPED_TRACE(dim); for (size_t pattern = 0; pattern < 5; ++pattern) { SCOPED_TRACE(pattern); std::vector query(dim + 1); - std::vector words(dim / 64, 0); + std::vector words((dim + 63) / 64, 0); double expected = 0; double sum_abs = 0; for (size_t i = 0; i < dim; ++i) { @@ -51,18 +51,21 @@ TEST(MaskIpX0Q, BackendsMatchScalarAcrossBlocksAndAlignments) { (pattern == 2 && i % 3 == 0) || (pattern == 3 && (i % 32 == 0 || i % 32 == 31)); if (selected) { - words[i / 64] |= uint64_t{1} << (63 - i % 64); + words[i / 64] |= + uint64_t{1} + << (std::min(size_t{64}, dim - i / 64 * 64) - 1 - i % 64); expected += query[i + 1]; sum_abs += std::abs(static_cast(query[i + 1])); } } for (size_t offset : {0, 1, 3, 7}) { SCOPED_TRACE(offset); - std::vector storage(dim / 8 + 8, 0); + // Vary alignment while ending the code exactly at the allocation boundary. + std::vector storage(8 + offset + dim / 8, 0); + auto* codes = storage.data() + 8 + offset; if (dim != 0) { - std::memcpy(storage.data() + offset, words.data(), dim / 8); + std::memcpy(codes, words.data(), dim / 8); } - const auto* codes = storage.data() + offset; for (auto function : functions) { const float result = function(query.data() + 1, codes, dim); EXPECT_NEAR(result, expected, 2e-6 * std::max(1.0, sum_abs)); @@ -84,7 +87,7 @@ TEST(MaskIpX0Q, DispatchPreservesEveryStoredBitPosition) { } for (size_t bit = 0; bit < dim; ++bit) { SCOPED_TRACE(bit); - std::vector words(dim / 64, 0); + std::vector words((dim + 63) / 64, 0); words[bit / 64] = uint64_t{1} << (63 - bit % 64); EXPECT_EQ(mask_ip_x0_q(query.data(), words.data(), dim), query[bit]); } diff --git a/tests/unit/rabitqlib/utils/neon_hnsw_test.cpp b/tests/unit/rabitqlib/utils/neon_hnsw_test.cpp index ce72ff4a..86ec044c 100644 --- a/tests/unit/rabitqlib/utils/neon_hnsw_test.cpp +++ b/tests/unit/rabitqlib/utils/neon_hnsw_test.cpp @@ -21,21 +21,25 @@ TEST(NeonHnsw, WarmupHandlesUnalignedCodesTailsAndWideCounts) { GTEST_SKIP(); } std::mt19937 random(764); - for (size_t dim : {0U, 64U, 128U, 448U, 512U, 576U, 960U, 1024U, 1088U, 65536U}) { + for (size_t dim : + {0U, 64U, 96U, 128U, 448U, 512U, 544U, 576U, 960U, 1024U, 1088U, 65536U}) { for (size_t bits = 0; bits <= 8; ++bits) { for (bool all_ones : {false, true}) { SCOPED_TRACE(dim); SCOPED_TRACE(bits); SCOPED_TRACE(all_ones); std::vector query(dim), storage(dim / 8 + 1); - std::vector words(dim / 64), transposed(dim / 64 * bits); + std::vector words((dim + 63) / 64), + transposed((dim + 63) / 64 * bits); uint64_t ip = 0, count = 0; for (size_t i = 0; i < dim; ++i) { query[i] = static_cast( (all_ones ? 255 : random()) & ((1U << bits) - 1) ); if (all_ones || random() % 2 != 0) { - words[i / 64] |= uint64_t{1} << (63 - i % 64); + words[i / 64] |= + uint64_t{1} + << (std::min(size_t{64}, dim - i / 64 * 64) - 1 - i % 64); ip += query[i]; ++count; } diff --git a/tests/unit/rabitqlib/utils/portable_kernel_test.cpp b/tests/unit/rabitqlib/utils/portable_kernel_test.cpp index 30deaac4..d0559c09 100644 --- a/tests/unit/rabitqlib/utils/portable_kernel_test.cpp +++ b/tests/unit/rabitqlib/utils/portable_kernel_test.cpp @@ -11,6 +11,7 @@ #include "rabitqlib/simd/pack_excode_dispatch.hpp" #include "rabitqlib/simd/space_dispatch.hpp" #include "rabitqlib/simd/warmup_dispatch.hpp" +#include "rabitqlib/utils/cpu_features.hpp" #include "rabitqlib/utils/space.hpp" #include "rabitqlib/utils/warmup_space.hpp" @@ -34,7 +35,7 @@ TEST(PortableKernels, PackingMatchesDispatchedBytesAndDotProducts) { simd::excode_ipimpl::ip64_fxu7_generic, simd::excode_ipimpl::ip16_fxu8_generic}; for (size_t bits = 1; bits <= 8; ++bits) { - for (size_t dim : {64U, 128U, 192U, 576U}) { + for (size_t dim : {64U, 96U, 128U, 192U, 544U, 576U}) { SCOPED_TRACE(bits); SCOPED_TRACE(dim); std::vector raw(dim), dispatched(dim * bits / 8), @@ -70,18 +71,20 @@ TEST(PortableKernels, PackingMatchesDispatchedBytesAndDotProducts) { } TEST(PortableKernels, QueryTransposeAndWarmupMatchIntegerReference) { - for (size_t dim : {64U, 128U, 448U, 512U, 576U, 960U, 1024U, 1088U}) { + for (size_t dim : {64U, 96U, 128U, 448U, 512U, 544U, 576U, 960U, 1024U, 1088U}) { for (size_t bits = 0; bits <= 8; ++bits) { SCOPED_TRACE(dim); SCOPED_TRACE(bits); std::vector query(dim); - std::vector data(dim / 64), transposed(dim / 64 * bits), + std::vector data((dim + 63) / 64), transposed((dim + 63) / 64 * bits), scalar(transposed.size()); uint64_t expected_ip = 0, expected_count = 0; for (size_t i = 0; i < dim; ++i) { query[i] = static_cast((i * 19 + 7) & ((1U << bits) - 1)); if (i % 3 == 0 || i % 63 == 0) { - data[i / 64] |= uint64_t{1} << (63 - i % 64); + data[i / 64] |= + uint64_t{1} + << (std::min(size_t{64}, dim - i / 64 * 64) - 1 - i % 64); expected_ip += query[i]; ++expected_count; } @@ -91,12 +94,14 @@ TEST(PortableKernels, QueryTransposeAndWarmupMatchIntegerReference) { EXPECT_EQ(scalar, transposed); size_t offset = 0; for (size_t block = 0; block < dim; block += 512) { - const size_t chunks = std::min(size_t{512}, dim - block) / 64; + const size_t chunks = (std::min(size_t{512}, dim - block) + 63) / 64; for (size_t b = 0; b < bits; ++b) { for (size_t c = 0; c < chunks; ++c) { - for (size_t i = 0; i < 64; ++i) { + const size_t width = std::min(size_t{64}, dim - block - c * 64); + for (size_t i = 0; i < width; ++i) { EXPECT_EQ( - (transposed[offset + b * chunks + c] >> (63 - i)) & 1U, + (transposed[offset + b * chunks + c] >> (width - 1 - i)) & + 1U, (query[block + c * 64 + i] >> b) & 1U ); } @@ -145,3 +150,92 @@ TEST(PortableKernels, Uint16TransposePreservesAllBitPlanes) { } } } // namespace rabitqlib + +TEST(PortableKernels, ExtraBitTailsPreserveFullBlockBytesAndUseCompactBitStream) { + using namespace rabitqlib; + for (size_t bits : {2U, 3U, 5U, 6U, 7U}) { + constexpr size_t tail = 32; + const size_t dim = 128 + tail; + std::vector raw(dim), packed(dim * bits / 8 + 2, 0xA5); + std::vector full(128 * bits / 8); + for (size_t i = 0; i < dim; ++i) + raw[i] = (i * 37 + 3) & ((1U << bits) - 1); + quant::rabitq_impl::ex_bits::packing_rabitqplus_code( + raw.data(), full.data(), 128, bits + ); + quant::rabitq_impl::ex_bits::packing_rabitqplus_code( + raw.data(), packed.data() + 1, dim, bits + ); + EXPECT_EQ(packed.front(), 0xA5); + EXPECT_EQ(packed.back(), 0xA5); + EXPECT_TRUE(std::equal(full.begin(), full.end(), packed.begin() + 1)); + // Independent byte-by-byte definition: bit j holds bit j % bits + // of coordinate j / bits, with no gaps or rounded word storage. + for (size_t byte = 0; byte < tail * bits / 8; ++byte) { + uint8_t expected = 0; + for (size_t b = 0; b < 8; ++b) { + const size_t pos = byte * 8 + b; + expected |= ((raw[128 + pos / bits] >> (pos % bits)) & 1U) << b; + } + EXPECT_EQ(packed[1 + full.size() + byte], expected); + } + } +} + +TEST(PortableKernels, PackingMatchesReferenceForMaximumAndIsolatedCodes) { + using namespace rabitqlib; + using Pack = void (*)(const uint8_t*, uint8_t*, size_t); + const std::array generic{ + simd::packing_2bit_excode_generic, + simd::packing_3bit_excode_generic, + simd::packing_4bit_excode_generic, + simd::packing_5bit_excode_generic, + simd::packing_6bit_excode_generic, + simd::packing_7bit_excode_generic}; + std::vector> backends{generic}; +#if defined(__x86_64__) || defined(_M_X64) + if (cpu::has_avx2()) { + backends.push_back({ + simd::packing_2bit_excode_avx2, + simd::packing_3bit_excode_avx2, + simd::packing_4bit_excode_avx2, + simd::packing_5bit_excode_avx2, + simd::packing_6bit_excode_avx2, + simd::packing_7bit_excode_avx2, + }); + } + if (cpu::has_avx512_core()) { + backends.push_back({ + simd::packing_2bit_excode_avx512, + simd::packing_3bit_excode_avx512, + simd::packing_4bit_excode_avx512, + simd::packing_5bit_excode_avx512, + simd::packing_6bit_excode_avx512, + simd::packing_7bit_excode_avx512, + }); + } +#endif + const std::array widths{2, 3, 4, 5, 6, 7}; + for (size_t b = 0; b < widths.size(); ++b) { + const size_t bits = widths[b]; + const auto max_code = static_cast((1U << bits) - 1); + for (size_t dim : {64U, 96U, 128U, 256U}) { + for (size_t pattern = 0; pattern <= dim; ++pattern) { + SCOPED_TRACE(bits); + SCOPED_TRACE(dim); + SCOPED_TRACE(pattern); + std::vector raw(dim + 1); + for (size_t i = 0; i < dim; ++i) { + raw[i + 1] = pattern == dim || i == pattern ? max_code : 0; + } + std::vector expected(dim * bits / 8 + 2, 0xA5); + generic[b](raw.data() + 1, expected.data() + 1, dim); + for (const auto& backend : backends) { + std::vector actual(expected.size(), 0xA5); + backend[b](raw.data() + 1, actual.data() + 1, dim); + EXPECT_EQ(actual, expected); + } + } + } + } +} diff --git a/tests/unit/rabitqlib/utils/query_preparation_test.cpp b/tests/unit/rabitqlib/utils/query_preparation_test.cpp index e1f62bff..f2f8999c 100644 --- a/tests/unit/rabitqlib/utils/query_preparation_test.cpp +++ b/tests/unit/rabitqlib/utils/query_preparation_test.cpp @@ -61,7 +61,7 @@ template void check_transpose( void (*transpose)(const T*, uint64_t*, size_t, size_t), bool blocked512 ) { - for (size_t dim : {0U, 64U, 128U, 448U, 512U, 576U, 960U, 1024U, 1088U}) { + for (size_t dim : {0U, 64U, 96U, 128U, 448U, 512U, 544U, 576U, 960U, 1024U, 1088U}) { std::vector input(dim + 1); for (size_t i = 0; i < dim; ++i) { input[i + 1] = @@ -70,19 +70,20 @@ void check_transpose( for (size_t bits = 0; bits <= sizeof(T) * 8; ++bits) { SCOPED_TRACE(dim); SCOPED_TRACE(bits); - std::vector actual(dim / 64 * bits + 2, UINT64_MAX); + std::vector actual((dim + 63) / 64 * bits + 2, UINT64_MAX); std::vector expected(actual.size(), 0); expected.front() = expected.back() = UINT64_MAX; size_t offset = 1; const size_t block_size = blocked512 ? 512 : 64; for (size_t block = 0; block < dim; block += block_size) { - const size_t chunks = std::min(block_size, dim - block) / 64; + const size_t chunks = (std::min(block_size, dim - block) + 63) / 64; for (size_t b = 0; b < bits; ++b) { for (size_t c = 0; c < chunks; ++c) { - for (size_t i = 0; i < 64; ++i) { + const size_t width = std::min(size_t{64}, dim - block - c * 64); + for (size_t i = 0; i < width; ++i) { expected[offset + b * chunks + c] |= uint64_t{(input[block + c * 64 + i + 1] >> b) & 1U} - << (63 - i); + << (width - 1 - i); } } } diff --git a/tests/unit/rabitqlib/utils/rotator_test.cpp b/tests/unit/rabitqlib/utils/rotator_test.cpp index 576a74ed..506ed671 100644 --- a/tests/unit/rabitqlib/utils/rotator_test.cpp +++ b/tests/unit/rabitqlib/utils/rotator_test.cpp @@ -13,6 +13,7 @@ #include #include "rabitqlib/utils/cpu_features.hpp" +#include "rabitqlib/utils/memory.hpp" #include "test_data.hpp" using namespace rabitqlib; @@ -39,9 +40,9 @@ TEST_F(RotatorTest, DefaultRotatorType) { std::unique_ptr> rotator(choose_rotator(dim)); ASSERT_NE(rotator, nullptr); - // FhtKacRotator pads to multiple of 64 + // FhtKacRotator pads to multiple of 32 size_t padded_dim = rotator->size(); - EXPECT_EQ(padded_dim % 64, 0); + EXPECT_EQ(padded_dim % 32, 0); EXPECT_GE(padded_dim, dim); } @@ -65,8 +66,47 @@ TEST_F(RotatorTest, RejectsZeroDimension) { ); } +TEST(FhtRotatorTest, PadsToThirtyTwoAndAcceptsLegacyPadding) { + for (size_t dim : {64U, 65U, 95U, 96U, 97U, 127U, 128U, 129U}) { + SCOPED_TRACE(dim); + std::unique_ptr> rotator(choose_rotator(dim)); + EXPECT_EQ(rotator->size(), (dim + 31) / 32 * 32); + const size_t legacy_padded = (dim + 63) / 64 * 64; + std::unique_ptr> legacy( + choose_rotator(dim, RotatorType::FhtKacRotator, legacy_padded, 42) + ); + EXPECT_EQ(legacy->size(), legacy_padded); + std::vector state(legacy->dump_bytes()); + legacy->save(state.data()); + std::unique_ptr> loaded( + choose_rotator(dim, RotatorType::FhtKacRotator, legacy_padded, 43) + ); + loaded->load(state.data()); + auto input = TestDataGenerator::GenerateRandomVector(dim, -1.0f, 1.0f, 42); + std::vector expected(legacy_padded), actual(legacy_padded); + legacy->rotate(input.data(), expected.data()); + loaded->rotate(input.data(), actual.data()); + EXPECT_EQ(actual, expected); + } + EXPECT_THROW((rotator_impl::FhtKacRotator(65, 72)), std::invalid_argument); + EXPECT_THROW((rotator_impl::FhtKacRotator(65, 80)), std::invalid_argument); + EXPECT_THROW( + (choose_rotator(65, RotatorType::FhtKacRotator, 80)), std::invalid_argument + ); + EXPECT_THROW( + (choose_rotator(65, RotatorType::FhtKacRotator, 72)), std::invalid_argument + ); + EXPECT_THROW((choose_rotator(63)), std::invalid_argument); + EXPECT_THROW( + (rotator_impl::padding_requirement( + std::numeric_limits::max(), RotatorType::FhtKacRotator + )), + std::invalid_argument + ); +} + TEST(FhtRotatorTest, LargeDimensionsPreserveNormAndSavedRotation) { - for (size_t dim : {4095U, 4096U, 16384U, 16385U, 65535U, 65536U}) { + for (size_t dim : {65U, 97U, 4095U, 4096U, 16384U, 16385U, 65535U, 65536U}) { SCOPED_TRACE(dim); std::unique_ptr> rotator(choose_rotator(dim)); std::vector data(dim), rotated(rotator->size()), restored(rotator->size()); @@ -158,68 +198,71 @@ TEST(FhtDispatchTest, BackendsMatchScalarButterfliesAndPadding) { backends.push_back(simd::fht_rotate_neon); } #endif - for (size_t dim : {64U, 65U, 127U, 128U, 129U, 192U, 193U, 256U, - 512U, 960U, 1024U, 2048U, 2049U, 4095U, 4096U, 4097U, - 8192U, 8193U, 16384U, 16385U, 32768U, 65535U, 65536U}) { + for (size_t dim : {64U, 65U, 95U, 96U, 97U, 127U, 128U, 129U, 192U, + 193U, 256U, 512U, 960U, 1024U, 2048U, 2049U, 4095U, 4096U, + 4097U, 8192U, 8193U, 16384U, 16385U, 32768U, 65535U, 65536U}) { SCOPED_TRACE(dim); - const size_t padded = (dim + 63) / 64 * 64; - size_t trunc = 1; - while (trunc * 2 <= dim) { - trunc *= 2; - } - const float fac = 1.0F / std::sqrt(static_cast(trunc)); - std::vector flip(4 * padded / 8); - for (size_t i = 0; i < flip.size(); ++i) { - flip[i] = static_cast(i * 73 + 13); - } - std::vector data(dim + 1), expected(padded, 0); - for (size_t i = 0; i < dim; ++i) { - data[i + 1] = std::sin(static_cast(i) * 0.13F); - expected[i] = data[i + 1]; - } - for (size_t pass = 0; pass < 4; ++pass) { - for (size_t i = 0; i < padded; ++i) { - if ((flip[pass * padded / 8 + i / 8] >> (i % 8)) & 1) { - expected[i] = -expected[i]; - } + for (size_t multiple : {32U, 64U}) { + const size_t padded = (dim + multiple - 1) / multiple * multiple; + SCOPED_TRACE(padded); + size_t trunc = 1; + while (trunc * 2 <= dim) { + trunc *= 2; + } + const float fac = 1.0F / std::sqrt(static_cast(trunc)); + std::vector flip(4 * padded / 8); + for (size_t i = 0; i < flip.size(); ++i) { + flip[i] = static_cast(i * 73 + 13); + } + std::vector data(dim + 1), expected(padded, 0); + for (size_t i = 0; i < dim; ++i) { + data[i + 1] = std::sin(static_cast(i) * 0.13F); + expected[i] = data[i + 1]; } - const size_t start = pass % 2 == 0 ? 0 : padded - trunc; - for (size_t width = 1; width < trunc; width *= 2) { - for (size_t block = 0; block < trunc; block += width * 2) { - for (size_t i = 0; i < width; ++i) { - const size_t x = start + block + i, y = x + width; - const float a = expected[x], b = expected[y]; - expected[x] = a + b; - expected[y] = a - b; + for (size_t pass = 0; pass < 4; ++pass) { + for (size_t i = 0; i < padded; ++i) { + if ((flip[pass * padded / 8 + i / 8] >> (i % 8)) & 1) { + expected[i] = -expected[i]; + } + } + const size_t start = pass % 2 == 0 ? 0 : padded - trunc; + for (size_t width = 1; width < trunc; width *= 2) { + for (size_t block = 0; block < trunc; block += width * 2) { + for (size_t i = 0; i < width; ++i) { + const size_t x = start + block + i, y = x + width; + const float a = expected[x], b = expected[y]; + expected[x] = a + b; + expected[y] = a - b; + } + } + } + for (size_t i = start; i < start + trunc; ++i) { + expected[i] *= fac; + } + if (padded != trunc) { + for (size_t i = 0; i < padded / 2; ++i) { + const float a = expected[i], b = expected[i + padded / 2]; + expected[i] = a + b; + expected[i + padded / 2] = a - b; } } - } - for (size_t i = start; i < start + trunc; ++i) { - expected[i] *= fac; } if (padded != trunc) { - for (size_t i = 0; i < padded / 2; ++i) { - const float a = expected[i], b = expected[i + padded / 2]; - expected[i] = a + b; - expected[i + padded / 2] = a - b; + for (float& value : expected) { + value *= 0.25F; } } - } - if (padded != trunc) { - for (float& value : expected) { - value *= 0.25F; - } - } - for (auto backend : backends) { - std::vector actual(padded + 2, 12345); - backend( - data.data() + 1, actual.data() + 1, dim, padded, trunc, fac, flip.data() - ); - for (size_t i = 0; i < padded; ++i) { - EXPECT_NEAR(actual[i + 1], expected[i], 2e-5F); + for (auto backend : backends) { + std::vector actual(padded + 2, 12345); + backend( + data.data() + 1, actual.data() + 1, dim, padded, trunc, fac, flip.data() + ); + for (size_t i = 0; i < padded; ++i) { + EXPECT_NEAR(actual[i + 1], expected[i], 2e-5F); + } + EXPECT_EQ(actual.front(), 12345); + EXPECT_EQ(actual.back(), 12345); } - EXPECT_EQ(actual.front(), 12345); - EXPECT_EQ(actual.back(), 12345); } } } @@ -242,36 +285,41 @@ TEST(FhtDispatchTest, PreservesZeroNormAndInnerProduct) { backends.push_back(simd::fht_rotate_neon); } #endif - for (size_t dim : {64U, 65U, 128U, 192U, 256U, 512U, 4097U, 65536U}) { + for (size_t dim : {64U, 65U, 97U, 128U, 192U, 256U, 512U, 4097U, 65536U}) { SCOPED_TRACE(dim); - const size_t padded = (dim + 63) / 64 * 64; - size_t trunc = 1; - while (trunc * 2 <= dim) { - trunc *= 2; - } - const float fac = 1.0F / std::sqrt(static_cast(trunc)); - std::vector flip(padded / 2, 0xA5); - std::vector a(dim, 0), b(dim), rotated_a(padded), rotated_b(padded); - a[dim - 1] = 1; // Exercise the padded tail with an impulse. - for (size_t i = 0; i < dim; ++i) { - b[i] = i % 2 == 0 ? 1.0F : -1.0F; - } - for (auto backend : backends) { - backend(a.data(), rotated_a.data(), dim, padded, trunc, fac, flip.data()); - backend(b.data(), rotated_b.data(), dim, padded, trunc, fac, flip.data()); - double norm_a = 0, norm_b = 0, dot = 0; - for (size_t i = 0; i < padded; ++i) { - norm_a += static_cast(rotated_a[i]) * rotated_a[i]; - norm_b += static_cast(rotated_b[i]) * rotated_b[i]; - dot += static_cast(rotated_a[i]) * rotated_b[i]; + for (size_t multiple : {32U, 64U}) { + const size_t padded = (dim + multiple - 1) / multiple * multiple; + SCOPED_TRACE(padded); + size_t trunc = 1; + while (trunc * 2 <= dim) { + trunc *= 2; } - EXPECT_NEAR(norm_a, 1.0, 2e-6); - EXPECT_NEAR(norm_b, static_cast(dim), 2e-6 * dim); - EXPECT_NEAR(dot, b.back(), 2e-6); - std::vector zero(dim, 0); - backend(zero.data(), rotated_a.data(), dim, padded, trunc, fac, flip.data()); - for (float value : rotated_a) { - EXPECT_EQ(value, 0.0F); + const float fac = 1.0F / std::sqrt(static_cast(trunc)); + std::vector flip(padded / 2, 0xA5); + std::vector a(dim, 0), b(dim), rotated_a(padded), rotated_b(padded); + a[dim - 1] = 1; // Exercise the padded tail with an impulse. + for (size_t i = 0; i < dim; ++i) { + b[i] = i % 2 == 0 ? 1.0F : -1.0F; + } + for (auto backend : backends) { + backend(a.data(), rotated_a.data(), dim, padded, trunc, fac, flip.data()); + backend(b.data(), rotated_b.data(), dim, padded, trunc, fac, flip.data()); + double norm_a = 0, norm_b = 0, dot = 0; + for (size_t i = 0; i < padded; ++i) { + norm_a += static_cast(rotated_a[i]) * rotated_a[i]; + norm_b += static_cast(rotated_b[i]) * rotated_b[i]; + dot += static_cast(rotated_a[i]) * rotated_b[i]; + } + EXPECT_NEAR(norm_a, 1.0, 2e-6); + EXPECT_NEAR(norm_b, static_cast(dim), 2e-6 * dim); + EXPECT_NEAR(dot, b.back(), 2e-6); + std::vector zero(dim, 0); + backend( + zero.data(), rotated_a.data(), dim, padded, trunc, fac, flip.data() + ); + for (float value : rotated_a) { + EXPECT_EQ(value, 0.0F); + } } } } @@ -307,23 +355,25 @@ TEST(FlipSignTest, BackendsPreserveBitsAndUnalignedBuffers) { // Signed zeros, infinities, quiet NaN payloads, subnormals and finite values. const std::array values{ 0, 0x80000000U, 0x7f800000U, 0xff800000U, 0x7fc12345U, 0xffc54321U, 1, 0x3f812345U}; - for (unsigned mask = 0; mask < 256; ++mask) { - std::array flip{}; - std::fill(flip.begin(), flip.end(), static_cast(mask)); - for (auto backend : backends) { - std::array data{}; - data.front() = data.back() = 12345; - for (size_t i = 0; i < 64; ++i) { - std::memcpy(&data[i + 1], &values[i % 8], sizeof(float)); - } - backend(flip.data() + 1, data.data() + 1, 64); - for (size_t i = 0; i < 64; ++i) { - uint32_t actual; - std::memcpy(&actual, &data[i + 1], sizeof(actual)); - EXPECT_EQ(actual, values[i % 8] ^ (((mask >> (i % 8)) & 1U) << 31)); + for (size_t dim : {64U, 96U, 128U, 160U}) { + for (unsigned mask = 0; mask < 256; ++mask) { + std::vector flip(dim / 8 + 1); + std::fill(flip.begin(), flip.end(), static_cast(mask)); + for (auto backend : backends) { + std::vector data(dim + 2); + data.front() = data.back() = 12345; + for (size_t i = 0; i < dim; ++i) { + std::memcpy(&data[i + 1], &values[i % 8], sizeof(float)); + } + backend(flip.data() + 1, data.data() + 1, dim); + for (size_t i = 0; i < dim; ++i) { + uint32_t actual; + std::memcpy(&actual, &data[i + 1], sizeof(actual)); + EXPECT_EQ(actual, values[i % 8] ^ (((mask >> (i % 8)) & 1U) << 31)); + } + EXPECT_EQ(data.front(), 12345); + EXPECT_EQ(data.back(), 12345); } - EXPECT_EQ(data.front(), 12345); - EXPECT_EQ(data.back(), 12345); } } } @@ -340,16 +390,40 @@ TEST(KacsWalkTest, BackendsMatchScalarAndRespectBufferBounds) { if (cpu::has_avx512_core()) backends.push_back(simd::kacs_walk_avx512); #endif - for (size_t dim : {64U, 128U, 192U, 960U, 1024U, 65536U}) { - std::vector input(dim + 2, 12345); - for (size_t i = 0; i < dim; ++i) - input[i + 1] = std::sin(static_cast(i)); - auto expected = input; - simd::kacs_walk_generic(expected.data() + 1, dim); - for (auto backend : backends) { - auto actual = input; - backend(actual.data() + 1, dim); - EXPECT_EQ(actual, expected); + for (size_t dim : + {64U, + 96U, + 128U, + 192U, + 224U, + 256U, + 288U, + 320U, + 544U, + 576U, + 960U, + 1024U, + 1056U, + 1088U, + 4128U, + 4160U, + 65504U, + 65536U}) { + for (size_t offset : {16U, 17U, 24U}) { + SCOPED_TRACE(dim); + SCOPED_TRACE(offset); + std::vector> input( + dim + offset + 1, 12345 + ); + for (size_t i = 0; i < dim; ++i) + input[i + offset] = std::sin(static_cast(i)); + auto expected = input; + simd::kacs_walk_generic(expected.data() + offset, dim); + for (auto backend : backends) { + auto actual = input; + backend(actual.data() + offset, dim); + EXPECT_EQ(actual, expected); + } } } } @@ -359,8 +433,8 @@ TEST(FhtNeonTest, MatchesSavedScalarRotationInPlaceAndWithExtraPadding) { if (!cpu::has_neon()) GTEST_SKIP(); for (size_t dim : {64U, 65U, 127U, 128U, 193U, 960U, 1025U, 4097U, 65536U}) { - for (size_t extra : {0U, 64U}) { - const size_t padded = (dim + 63) / 64 * 64 + extra; + for (size_t extra : {0U, 32U, 64U}) { + const size_t padded = (dim + 31) / 32 * 32 + extra; if (padded > 65536) continue; SCOPED_TRACE(dim); @@ -432,3 +506,80 @@ TEST(NeonRotationTest, StandaloneKernelsHandleShortTails) { } } #endif + +#if defined(__x86_64__) || defined(_M_X64) +TEST(FhtDispatchTest, AvxBackendsMatchBitwiseForCompactTails) { + if (!cpu::has_avx2() || !cpu::has_avx512_core()) { + GTEST_SKIP() << "AVX2 and AVX-512 are required for direct backend comparison"; + } + using AlignedVector = std::vector>; + for (size_t dim : + {65U, + 97U, + 129U, + 161U, + 257U, + 289U, + 513U, + 545U, + 769U, + 801U, + 1025U, + 1057U, + 4097U, + 4129U, + 65535U}) { + for (size_t multiple : {32U, 64U}) { + const size_t padded = (dim + multiple - 1) / multiple * multiple; + size_t trunc = 1; + while (trunc * 2 <= dim) { + trunc *= 2; + } + const float fac = 1.0F / std::sqrt(static_cast(trunc)); + std::vector flip(padded / 2 + 1); + for (size_t i = 0; i < flip.size(); ++i) { + flip[i] = static_cast(i * 73 + 13); + } + for (size_t offset : {0U, 1U, 8U}) { + SCOPED_TRACE(dim); + SCOPED_TRACE(padded); + SCOPED_TRACE(offset); + AlignedVector input(dim + offset); + AlignedVector expected(padded + offset + 1, 12345); + AlignedVector actual(expected); + for (size_t i = 0; i < dim; ++i) { + input[offset + i] = std::sin(static_cast(i) * 0.13F); + } + simd::fht_rotate_avx2( + input.data() + offset, + expected.data() + offset, + dim, + padded, + trunc, + fac, + flip.data() + 1 + ); + simd::fht_rotate_avx512( + input.data() + offset, + actual.data() + offset, + dim, + padded, + trunc, + fac, + flip.data() + 1 + ); + EXPECT_EQ( + std::memcmp( + actual.data(), expected.data(), actual.size() * sizeof(float) + ), + 0 + ); + for (size_t i = 0; i < offset; ++i) { + EXPECT_EQ(actual[i], 12345); + } + EXPECT_EQ(actual.back(), 12345); + } + } + } +} +#endif diff --git a/tests/unit/rabitqlib/utils/space_test.cpp b/tests/unit/rabitqlib/utils/space_test.cpp index e962653c..aa7ec38a 100644 --- a/tests/unit/rabitqlib/utils/space_test.cpp +++ b/tests/unit/rabitqlib/utils/space_test.cpp @@ -12,11 +12,50 @@ #include "rabitqlib/simd/pack_excode_dispatch.hpp" #include "rabitqlib/simd/space_dispatch.hpp" #include "rabitqlib/simd/warmup_dispatch.hpp" +#include "rabitqlib/utils/bitops.hpp" #include "rabitqlib/utils/cpu_features.hpp" #include "rabitqlib/utils/warmup_space.hpp" using namespace rabitqlib; +TEST(BitOps, Popcount64MatchesBitReference) { + const auto check = [](uint64_t value) { + unsigned expected = 0; + for (size_t bit = 0; bit < 64; ++bit) { + expected += static_cast((value >> bit) & 1U); + } + EXPECT_EQ(bitops::popcount64(value), expected) << value; + }; + check(0); + check(std::numeric_limits::max()); + for (size_t bit = 0; bit < 64; ++bit) { + check(uint64_t{1} << bit); + check(~(uint64_t{1} << bit)); + } + uint64_t value = 17; + for (size_t i = 0; i < 1024; ++i) { + value = value * 6364136223846793005ULL + 1442695040888963407ULL; + check(value); + } +} + +TEST(BitOps, CompactWordLoadsZeroExtendWithoutReadingPastInput) { + for (size_t bytes = 1; bytes <= 8; ++bytes) { + for (size_t offset : {0U, 1U}) { + SCOPED_TRACE(bytes); + SCOPED_TRACE(offset); + std::vector storage(bytes + offset); + uint64_t expected = 0; + for (size_t i = 0; i < bytes; ++i) { + const auto value = static_cast(0x80U + i * 17U); + storage[offset + i] = value; + expected |= uint64_t{value} << (8 * i); + } + EXPECT_EQ(bitops::load_word(storage.data() + offset, bytes), expected); + } + } +} + TEST(PackBinary, SupportsUnalignedOutput) { constexpr size_t dim = 128; std::array binary_code{}; @@ -40,6 +79,40 @@ TEST(PackBinary, SupportsUnalignedOutput) { } } +TEST(PackBinary, AllWordSizesPreserveFullAndPartialWordLayout) { + const auto check = [](auto word_type) { + using Word = decltype(word_type); + constexpr size_t kWordBits = sizeof(Word) * 8; + for (size_t dim : {8U, 16U, 24U, 32U, 48U, 64U, 80U, 96U, 112U, 128U}) { + SCOPED_TRACE(kWordBits); + SCOPED_TRACE(dim); + std::vector signs(dim); + // All zeros, all ones, and each isolated bit catch order and tail errors. + for (size_t pattern = 0; pattern < dim + 2; ++pattern) { + for (size_t i = 0; i < dim; ++i) { + signs[i] = pattern == dim + 1 || i == pattern; + } + std::vector output(dim / 8 + 2, 0xA5); + pack_binary_to_bytes(signs.data(), output.data() + 1, dim); + EXPECT_EQ(output.front(), 0xA5); + EXPECT_EQ(output.back(), 0xA5); + for (size_t start = 0; start < dim; start += kWordBits) { + const size_t width = std::min(kWordBits, dim - start); + Word actual = 0; + std::memcpy(&actual, output.data() + 1 + start / 8, width / 8); + for (size_t bit = 0; bit < width; ++bit) { + EXPECT_EQ((actual >> (width - 1 - bit)) & 1U, signs[start + bit]); + } + } + } + } + }; + check(uint8_t{}); + check(uint16_t{}); + check(uint32_t{}); + check(uint64_t{}); +} + TEST(MaskIpX0Q, SupportsUnalignedCodes) { constexpr size_t dim = 128; std::array binary_code{}; @@ -72,7 +145,20 @@ TEST(MaskIpX0Q, SupportsUnalignedCodes) { TEST(WarmupIpX0Q, SupportsUnalignedCodes) { constexpr float delta = 0.5F; constexpr float vl = -0.25F; - for (size_t dim : {64UL, 128UL, 192UL, 448UL, 512UL, 576UL}) { + for (size_t dim : + {0UL, + 64UL, + 96UL, + 128UL, + 160UL, + 192UL, + 224UL, + 256UL, + 288UL, + 448UL, + 512UL, + 544UL, + 576UL}) { SCOPED_TRACE(dim); std::vector data_bits(dim); size_t data_popcount = 0; @@ -86,7 +172,7 @@ TEST(WarmupIpX0Q, SupportsUnalignedCodes) { ASSERT_NE(reinterpret_cast(codes) % alignof(uint64_t), 0U); pack_binary_to_bytes(data_bits.data(), codes, dim); - for (size_t b_query : {1UL, 4UL, 8UL}) { + for (size_t b_query = 0; b_query <= 8; ++b_query) { SCOPED_TRACE(b_query); std::vector query_values(dim); size_t weighted_intersection = 0; @@ -98,18 +184,19 @@ TEST(WarmupIpX0Q, SupportsUnalignedCodes) { } } - std::vector query((dim / 64) * b_query, 0); + std::vector query(((dim + 63) / 64) * b_query, 0); size_t query_offset = 0; for (size_t block = 0; block < dim; block += 512) { const size_t block_dim = (dim - block < 512) ? dim - block : 512; - const size_t chunks = block_dim / 64; + const size_t chunks = (block_dim + 63) / 64; for (size_t bit = 0; bit < b_query; ++bit) { for (size_t chunk = 0; chunk < chunks; ++chunk) { uint64_t packed = 0; - for (size_t k = 0; k < 64; ++k) { + const size_t width = std::min(size_t{64}, block_dim - chunk * 64); + for (size_t k = 0; k < width; ++k) { const size_t i = block + chunk * 64 + k; packed |= static_cast((query_values[i] >> bit) & 1U) - << (63 - k); + << (width - 1 - k); } query[query_offset + bit * chunks + chunk] = packed; } @@ -464,3 +551,29 @@ TEST(ip_fxu8_avx, ip_works) { } #endif } + +TEST(PackBinary, CompactTailWordsMatchBinaryArithmetic) { + for (size_t dim : {64U, 96U, 128U, 160U}) { + std::vector signs(dim); + std::vector values(dim); + std::vector code(dim / 8 + 1, 0xA5); + std::vector transposed((dim + 63) / 64 * 4); + unsigned expected_count = 0, expected_ip = 0; + for (size_t i = 0; i < dim; ++i) { + signs[i] = i % 3 != 0; + values[i] = (i * 7) % 16; + expected_count += signs[i]; + expected_ip += signs[i] * values[i]; + } + pack_binary_to_bytes(signs.data(), code.data(), dim); + EXPECT_EQ(code.back(), 0xA5); + const auto* words = reinterpret_cast(code.data()); + EXPECT_EQ(popcount(words, dim), expected_count); + EXPECT_EQ(ip_bin_bin(words, words, dim), expected_count); + new_transpose_bin(values.data(), transposed.data(), dim, 4); + EXPECT_EQ(ip_x0_q(words, transposed.data(), 1, 0, dim, 4), expected_ip); + EXPECT_EQ( + rabitqlib::warmup_ip_x0_q<4>(words, transposed.data(), 1, 0, dim), expected_ip + ); + } +} From 48e9779d1f68522a62f13e22a8a81d92d09aff92 Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Thu, 8 Oct 2026 12:22:29 +0800 Subject: [PATCH 2/3] update review --- include/rabitqlib/utils/bitops.hpp | 2 +- src/simd/pack_excode_kernels.hpp | 13 +++++++++---- src/simd/packed_tail_avx2.hpp | 2 ++ tests/python/fixtures/README.md | 5 +++++ tests/python/test_ivf.py | 27 +++++++++++++++++++-------- 5 files changed, 36 insertions(+), 13 deletions(-) diff --git a/include/rabitqlib/utils/bitops.hpp b/include/rabitqlib/utils/bitops.hpp index 0702a666..f75d4356 100644 --- a/include/rabitqlib/utils/bitops.hpp +++ b/include/rabitqlib/utils/bitops.hpp @@ -37,7 +37,7 @@ inline unsigned popcount32(uint32_t value) { #endif } -inline unsigned popcount64(uint64_t value) { +static inline unsigned popcount64(uint64_t value) { #if defined(_MSC_VER) || \ ((defined(__x86_64__) || defined(__i386__)) && !defined(__POPCNT__)) // Inline the software fallback on generic x86 instead of calling libgcc diff --git a/src/simd/pack_excode_kernels.hpp b/src/simd/pack_excode_kernels.hpp index 2a250f23..cbaa6316 100644 --- a/src/simd/pack_excode_kernels.hpp +++ b/src/simd/pack_excode_kernels.hpp @@ -3,6 +3,7 @@ #include #include +#include #include #include #include @@ -12,6 +13,7 @@ namespace rabitqlib::simd::detail { template inline void pack_excode_tail_intrinsics(const uint8_t* raw, uint8_t* compact, size_t dim) { static_assert(Bits == 2 || Bits == 3 || Bits == 5 || Bits == 6 || Bits == 7); + assert(dim == 0 || dim == 32); if (dim == 0) return; // A partial block is always 32 coordinates. @@ -22,10 +24,13 @@ inline void pack_excode_tail_intrinsics(const uint8_t* raw, uint8_t* compact, si alignas(16) static constexpr auto kPositions = [] { std::array positions{}; for (size_t i = 0; i < 16; ++i) { - positions[i] = static_cast( - Bits % 2 == 0 ? (i < 2 * Bits ? (i / (Bits / 2)) * 4 + i % (Bits / 2) : 128) - : (i < Bits ? i : (i < 2 * Bits ? 8 + i - Bits : 128)) - ); + if (i >= 2 * Bits) { + positions[i] = 128; // Zero this byte in the shuffle. + } else if constexpr (Bits % 2 == 0) { + positions[i] = static_cast((i / (Bits / 2)) * 4 + i % (Bits / 2)); + } else { + positions[i] = static_cast(i < Bits ? i : 8 + i - Bits); + } } return positions; }(); diff --git a/src/simd/packed_tail_avx2.hpp b/src/simd/packed_tail_avx2.hpp index 7994e12f..dda28af7 100644 --- a/src/simd/packed_tail_avx2.hpp +++ b/src/simd/packed_tail_avx2.hpp @@ -2,6 +2,7 @@ #include +#include #include #include #include @@ -36,6 +37,7 @@ template inline void accumulate_excode_tail_avx2( const float* query, const uint8_t* compact, size_t dim, __m256& sum, __m256& sum_hi ) { + assert(dim == 0 || dim == 32); if (dim == 0) return; // The fixed 32-coordinate tail contains two 16-code groups. diff --git a/tests/python/fixtures/README.md b/tests/python/fixtures/README.md index 75e1d47f..56972902 100644 --- a/tests/python/fixtures/README.md +++ b/tests/python/fixtures/README.md @@ -13,6 +13,11 @@ zero centroid, and the same three vectors above with zeros in all remaining coordinates. Their widths are 1 and 32 bits, with L2 distance and FHT/Kac rotation. They verify that loading does not substitute the new 96-coordinate default. +The IVF tests also query `q[i] = (i % 9 - 4) / 4` before and after migration. +The expected IDs and distances were recorded using the pad64 implementation at +commit `a065785`, covering nonzero query rotation and quantized distance estimation +as well as raw-vector reranking. + `hnsw_legacy_padding65_4bit.index` was written before HNSW switched to 32-coordinate padding: dimension 65, padded dimension 128, 4 total bits, L2, capacity 8, M=2, ef_construction=8. Its three vectors and zero centroid match the dimension-65 IVF diff --git a/tests/python/test_ivf.py b/tests/python/test_ivf.py index 5f0f83ae..da21d258 100644 --- a/tests/python/test_ivf.py +++ b/tests/python/test_ivf.py @@ -347,14 +347,20 @@ def test_raw_reranking(dim, metric, high_accuracy, fast_quantization, tmp_path): @pytest.mark.parametrize( - "filename,dim,bits,padded", + "filename,dim,bits,padded,nonzero_distances", [ - ("ivf_legacy_4bit.index", 64, 4, 64), - ("ivf_legacy_padding65_1bit.index", 65, 1, 128), - ("ivf_raw_v1_padding65.index", 65, 32, 128), + ("ivf_legacy_4bit.index", 64, 4, 64, [27.25, 32.31400681, 50.34958267]), + ( + "ivf_legacy_padding65_1bit.index", + 65, + 1, + 128, + [27.81249809, 36.12752914, 47.54097748], + ), + ("ivf_raw_v1_padding65.index", 65, 32, 128, [27.8125, 33.8125, 49.8125]), ], ) -def test_legacy_ivf_fixture(tmp_path, filename, dim, bits, padded): +def test_legacy_ivf_fixture(tmp_path, filename, dim, bits, padded, nonzero_distances): path = Path(__file__).parent / "fixtures" / filename idx = IvfIndex.load(str(path)) assert (idx.dim, idx.max_elements, idx.num_clusters, idx.nbits, idx.metric) == ( @@ -364,10 +370,15 @@ def test_legacy_ivf_fixture(tmp_path, filename, dim, bits, padded): bits, "l2", ) - queries = np.zeros((1, dim), np.float32) + queries = np.zeros((2, dim), np.float32) + # Expected nonzero-query results were recorded with the pad64 implementation + # at a065785, before the padding changes. + queries[1] = (np.arange(dim) % 9 - 4) / 4 ids, distances = idx.search(queries, 3, 1) - np.testing.assert_array_equal(np.sort(ids), [[0, 1, 2]]) - np.testing.assert_allclose(distances, [[0, 14, 14]], atol=2e-5) + np.testing.assert_array_equal(np.sort(ids[0]), [0, 1, 2]) + np.testing.assert_allclose(distances[0], [0, 14, 14], atol=2e-5) + np.testing.assert_array_equal(ids[1], [0, 2, 1]) + np.testing.assert_allclose(distances[1], nonzero_distances, rtol=2e-6) saved = tmp_path / "legacy.index" idx.save(str(saved)) payload = saved.read_bytes() From 024e6135806ab1dbfff0c7bda53c022056e1e4cd Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Thu, 8 Oct 2026 13:46:18 +0800 Subject: [PATCH 3/3] Opt for neon --- src/simd/query_neon.cpp | 21 +++++++++++++------ .../utils/query_preparation_test.cpp | 16 +++++++++++++- 2 files changed, 30 insertions(+), 7 deletions(-) diff --git a/src/simd/query_neon.cpp b/src/simd/query_neon.cpp index 664692b1..914dbffc 100644 --- a/src/simd/query_neon.cpp +++ b/src/simd/query_neon.cpp @@ -92,12 +92,10 @@ void new_transpose_bin_512_neon( const uint8_t* query, uint64_t* transposed, size_t dim, size_t bits ) { for (size_t block = 0; block < dim; block += 512) { - if ((dim - block) < 512 && dim % 64 != 0) { - new_transpose_bin_512_generic(query + block, transposed, dim - block, bits); - break; - } - const size_t chunks = std::min(size_t{512}, dim - block) / 64; - for (size_t c = 0; c < chunks; ++c) { + const size_t block_dim = std::min(size_t{512}, dim - block); + const size_t full_chunks = block_dim / 64; + const size_t chunks = (block_dim + 63) / 64; + for (size_t c = 0; c < full_chunks; ++c) { const uint8_t* row = query + block + c * 64; const uint8x16_t a = vld1q_u8(row), b = vld1q_u8(row + 16), d = vld1q_u8(row + 32), e = vld1q_u8(row + 48); @@ -110,6 +108,17 @@ void new_transpose_bin_512_neon( byte_mask(vtstq_u8(e, bit)); } } + const size_t tail_dim = block_dim % 64; + if (tail_dim != 0) { + const uint8_t* tail = query + block + full_chunks * 64; + for (size_t plane = 0; plane < bits; ++plane) { + uint64_t word = 0; + for (size_t i = 0; i < tail_dim; ++i) { + word |= uint64_t{(tail[i] >> plane) & 1U} << (tail_dim - 1 - i); + } + transposed[plane * chunks + full_chunks] = word; + } + } transposed += chunks * bits; } } diff --git a/tests/unit/rabitqlib/utils/query_preparation_test.cpp b/tests/unit/rabitqlib/utils/query_preparation_test.cpp index f2f8999c..6904ebe3 100644 --- a/tests/unit/rabitqlib/utils/query_preparation_test.cpp +++ b/tests/unit/rabitqlib/utils/query_preparation_test.cpp @@ -6,6 +6,7 @@ #include #include #include +#include #include #include "rabitqlib/simd/space_dispatch.hpp" @@ -61,7 +62,8 @@ template void check_transpose( void (*transpose)(const T*, uint64_t*, size_t, size_t), bool blocked512 ) { - for (size_t dim : {0U, 64U, 96U, 128U, 448U, 512U, 544U, 576U, 960U, 1024U, 1088U}) { + for (size_t dim : {0U, 32U, 64U, 96U, 128U, 160U, 256U, 288U, 448U, 480U, 512U, + 544U, 576U, 768U, 960U, 992U, 1024U, 1056U, 1088U, 4096U, 4128U}) { std::vector input(dim + 1); for (size_t i = 0; i < dim; ++i) { input[i + 1] = @@ -91,6 +93,18 @@ void check_transpose( } transpose(input.data() + 1, actual.data() + 1, dim, bits); EXPECT_EQ(actual, expected); + + // Keep both buffers exact at the end for sanitizer overread/overwrite + // coverage, while offsetting SIMD accesses by one natural element. + auto exact_input = std::make_unique(dim + 1); + std::copy(input.begin(), input.end(), exact_input.get()); + auto exact_output = std::make_unique(actual.size() - 1); + exact_output[0] = UINT64_MAX; + transpose(exact_input.get() + 1, exact_output.get() + 1, dim, bits); + EXPECT_EQ(exact_output[0], UINT64_MAX); + EXPECT_TRUE( + std::equal(expected.begin() + 1, expected.end() - 1, exact_output.get() + 1) + ); } } }