diff --git a/include/rabitqlib/utils/space.hpp b/include/rabitqlib/utils/space.hpp index f6cc24e..4ffe035 100644 --- a/include/rabitqlib/utils/space.hpp +++ b/include/rabitqlib/utils/space.hpp @@ -5,6 +5,7 @@ #include #include #include +#include #include #include #include @@ -15,6 +16,13 @@ #include "rabitqlib/utils/tools.hpp" namespace rabitqlib { + +inline uint64_t load_unaligned_u64(const void* data) { + uint64_t value = 0; + std::memcpy(&value, data, sizeof(value)); + return value; +} + namespace scalar_impl { template void scalar_quantize_normal( @@ -124,14 +132,15 @@ inline void pack_binary( const int* __restrict__ binary_code, T* __restrict__ compact_code, size_t length ) { constexpr size_t kTypeBits = sizeof(T) * 8; + auto* output = reinterpret_cast(compact_code); for (size_t i = 0; i < length; i += kTypeBits) { T cur = 0; for (size_t j = 0; j < kTypeBits; ++j) { cur |= (static_cast(binary_code[i + j]) << (kTypeBits - 1 - j)); } - *compact_code = cur; - ++compact_code; + std::memcpy(output, &cur, sizeof(cur)); + output += sizeof(cur); } } diff --git a/src/index/hnsw_search_avx2_kernels.hpp b/src/index/hnsw_search_avx2_kernels.hpp index 4fbe015..bf20fc0 100644 --- a/src/index/hnsw_search_avx2_kernels.hpp +++ b/src/index/hnsw_search_avx2_kernels.hpp @@ -68,14 +68,14 @@ 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 uint64_t* it_data = data; + const uint8_t* it_data = reinterpret_cast(data); const float* it_query = query; __m256 sum = _mm256_setzero_ps(); __m256i bit_checker = _mm256_set_epi32(0x80, 0x40, 0x20, 0x10, 0x08, 0x04, 0x02, 0x01); for (size_t i = 0; i < num_blk; ++i) { - uint64_t bits = rabitqlib::reverse_bits_u64(*it_data); + uint64_t bits = rabitqlib::reverse_bits_u64(rabitqlib::load_unaligned_u64(it_data)); // 64 bits / 8 floats = 8 iterations for (int j = 0; j < 8; ++j) { @@ -91,7 +91,7 @@ static inline float hnsw_mask_ip_x0_q_avx2( it_query += 8; } - ++it_data; + it_data += sizeof(uint64_t); } alignas(32) float lanes[8]; diff --git a/src/index/hnsw_search_avx512_kernels.hpp b/src/index/hnsw_search_avx512_kernels.hpp index 5c8893b..f6282ba 100644 --- a/src/index/hnsw_search_avx512_kernels.hpp +++ b/src/index/hnsw_search_avx512_kernels.hpp @@ -14,7 +14,7 @@ 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 uint64_t* it_data = data; + const uint8_t* it_data = reinterpret_cast(data); const float* it_query = query; // __m512 sum0 = _mm512_setzero_ps(); @@ -24,7 +24,7 @@ static inline float hnsw_mask_ip_x0_q_avx512( __m512 sum = _mm512_setzero_ps(); for (size_t i = 0; i < num_blk; ++i) { - uint64_t bits = rabitqlib::reverse_bits_u64(*it_data); + 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); @@ -43,7 +43,7 @@ static inline float hnsw_mask_ip_x0_q_avx512( // _mm_prefetch(reinterpret_cast(it_query + 128), _MM_HINT_T1); - ++it_data; + it_data += sizeof(uint64_t); it_query += 64; } diff --git a/src/simd/space_avx2.cpp b/src/simd/space_avx2.cpp index d533964..a618569 100644 --- a/src/simd/space_avx2.cpp +++ b/src/simd/space_avx2.cpp @@ -153,7 +153,7 @@ void new_transpose_bin_512_avx2( float 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 uint64_t* it_data = data; + const uint8_t* it_data = reinterpret_cast(data); const float* it_query = query; __m256 sum = _mm256_setzero_ps(); @@ -161,7 +161,7 @@ float mask_ip_x0_q_avx2(const float* query, const uint64_t* data, size_t padded_ __m256i bit_checker = _mm256_set_epi32(0x80, 0x40, 0x20, 0x10, 0x08, 0x04, 0x02, 0x01); for (size_t i = 0; i < num_blk; ++i) { - uint64_t bits = reverse_bits_u64(*it_data); + uint64_t bits = reverse_bits_u64(load_unaligned_u64(it_data)); // 64 bits / 8 floats = 8 iterations for (int j = 0; j < 8; ++j) { @@ -177,7 +177,7 @@ float mask_ip_x0_q_avx2(const float* query, const uint64_t* data, size_t padded_ it_query += 8; } - ++it_data; + it_data += sizeof(uint64_t); } float result = 0.0f; diff --git a/src/simd/space_avx512.cpp b/src/simd/space_avx512.cpp index e164774..349cc66 100644 --- a/src/simd/space_avx512.cpp +++ b/src/simd/space_avx512.cpp @@ -110,7 +110,7 @@ void new_transpose_bin_512_avx512( float 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 uint64_t* it_data = data; + const uint8_t* it_data = reinterpret_cast(data); const float* it_query = query; // __m512 sum0 = _mm512_setzero_ps(); @@ -120,7 +120,7 @@ float mask_ip_x0_q_avx512(const float* query, const uint64_t* data, size_t padde __m512 sum = _mm512_setzero_ps(); for (size_t i = 0; i < num_blk; ++i) { - uint64_t bits = reverse_bits_u64(*it_data); + uint64_t bits = reverse_bits_u64(load_unaligned_u64(it_data)); auto mask0 = static_cast<__mmask16>(bits); auto mask1 = static_cast<__mmask16>(bits >> 16); @@ -139,7 +139,7 @@ float mask_ip_x0_q_avx512(const float* query, const uint64_t* data, size_t padde // _mm_prefetch(reinterpret_cast(it_query + 128), _MM_HINT_T1); - ++it_data; + it_data += sizeof(uint64_t); it_query += 64; } diff --git a/tests/unit/rabitqlib/utils/space_test.cpp b/tests/unit/rabitqlib/utils/space_test.cpp index fcbc4c8..9044ed8 100644 --- a/tests/unit/rabitqlib/utils/space_test.cpp +++ b/tests/unit/rabitqlib/utils/space_test.cpp @@ -13,6 +13,54 @@ using namespace rabitqlib; +TEST(PackBinary, SupportsUnalignedOutput) { + constexpr size_t dim = 128; + std::array binary_code{}; + for (size_t i = 0; i < dim; ++i) { + binary_code[i] = (i % 3) == 0; + } + + alignas(uint64_t) std::array storage{}; + auto* output = reinterpret_cast(storage.data() + 1); + ASSERT_NE(reinterpret_cast(output) % alignof(uint64_t), 0U); + pack_binary(binary_code.data(), output, dim); + + const std::array packed{ + load_unaligned_u64(storage.data() + 1), + load_unaligned_u64(storage.data() + 1 + sizeof(uint64_t)), + }; + for (size_t i = 0; i < dim; ++i) { + const size_t word = i / 64; + const auto bit = static_cast((packed[word] >> (63 - (i % 64))) & 1U); + EXPECT_EQ(bit, binary_code[i]) << "bit " << i; + } +} + +TEST(MaskIpX0Q, SupportsUnalignedCodes) { + constexpr size_t dim = 128; + std::array binary_code{}; + std::array query{}; + float expected = 0; + for (size_t i = 0; i < dim; ++i) { + binary_code[i] = (i % 3) == 0; + query[i] = static_cast(i + 1); + expected += binary_code[i] != 0 ? query[i] : 0; + } + + alignas(uint64_t) std::array storage{}; + auto* codes = reinterpret_cast(storage.data() + 1); + ASSERT_NE(reinterpret_cast(codes) % alignof(uint64_t), 0U); + pack_binary(binary_code.data(), codes, dim); + + if (cpu::has_avx2()) { + EXPECT_FLOAT_EQ(simd::mask_ip_x0_q_avx2(query.data(), codes, dim), expected); + } + if (cpu::has_avx512_core()) { + EXPECT_FLOAT_EQ(simd::mask_ip_x0_q_avx512(query.data(), codes, dim), expected); + } + EXPECT_FLOAT_EQ(mask_ip_x0_q(query.data(), codes, dim), expected); +} + TEST(Select_IP_Func, returns_stable_function_pointer) { auto ip_func = select_excode_ipfunc(0); ASSERT_NE(ip_func, nullptr);