Fix unaligned packed binary code access

include/rabitqlib/utils/space.hpp

@@ -5,6 +5,7 @@
 #include <array>
 #include <cassert>
 #include <cstdint>
+#include <cstring>
 #include <iostream>
 #include <limits>
 #include <optional>
@@ -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 <typename T>
 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<uint8_t*>(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<T>(binary_code[i + j]) << (kTypeBits - 1 - j));
         }
-        *compact_code = cur;
-        ++compact_code;
+        std::memcpy(output, &cur, sizeof(cur));
+        output += sizeof(cur);
     }
 }
 

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<const uint8_t*>(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];

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<const uint8_t*>(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<const char*>(it_query + 128), _MM_HINT_T1);
 
-        ++it_data;
+        it_data += sizeof(uint64_t);
         it_query += 64;
     }
 

src/simd/space_avx2.cpp

@@ -153,15 +153,15 @@ 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<const uint8_t*>(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 = 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;

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<const uint8_t*>(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<const char*>(it_query + 128), _MM_HINT_T1);
 
-        ++it_data;
+        it_data += sizeof(uint64_t);
         it_query += 64;
     }
 

tests/unit/rabitqlib/utils/space_test.cpp

@@ -13,6 +13,54 @@
 
 using namespace rabitqlib;
 
+TEST(PackBinary, SupportsUnalignedOutput) {
+    constexpr size_t dim = 128;
+    std::array<int, dim> binary_code{};
+    for (size_t i = 0; i < dim; ++i) {
+        binary_code[i] = (i % 3) == 0;
+    }
+
+    alignas(uint64_t) std::array<uint8_t, (2 * sizeof(uint64_t)) + 1> storage{};
+    auto* output = reinterpret_cast<uint64_t*>(storage.data() + 1);
+    ASSERT_NE(reinterpret_cast<uintptr_t>(output) % alignof(uint64_t), 0U);
+    pack_binary(binary_code.data(), output, dim);
+
+    const std::array<uint64_t, 2> 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<int>((packed[word] >> (63 - (i % 64))) & 1U);
+        EXPECT_EQ(bit, binary_code[i]) << "bit " << i;
+    }
+}
+
+TEST(MaskIpX0Q, SupportsUnalignedCodes) {
+    constexpr size_t dim = 128;
+    std::array<int, dim> binary_code{};
+    std::array<float, dim> query{};
+    float expected = 0;
+    for (size_t i = 0; i < dim; ++i) {
+        binary_code[i] = (i % 3) == 0;
+        query[i] = static_cast<float>(i + 1);
+        expected += binary_code[i] != 0 ? query[i] : 0;
+    }
+
+    alignas(uint64_t) std::array<uint8_t, (2 * sizeof(uint64_t)) + 1> storage{};
+    auto* codes = reinterpret_cast<uint64_t*>(storage.data() + 1);
+    ASSERT_NE(reinterpret_cast<uintptr_t>(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);