fix: fix import wrong package in ARM architecture (#49)

huanglune · web-flow · commit 8f1284fcbee5 · 2026-01-22T16:14:51.000+08:00
diff --git a/.clang-format b/.clang-format
@@ -4,6 +4,9 @@ DerivePointerAlignment                   : false
 PointerAlignment                         : Right
 ColumnLimit                              : 100
 
+# Preprocessor directive indentation
+IndentPPDirectives                       : BeforeHash
+
 # Default for clang-8, changed in later clangs. Set explicitly for forwards
 # compatibility for students with modern clangs
 IncludeBlocks                            : Preserve
diff --git a/include/executor/jobs/graph_search_job.hpp b/include/executor/jobs/graph_search_job.hpp
@@ -33,7 +33,7 @@
 #include "utils/rabitq_utils/search_utils/hashset.hpp"
 
 #if defined(__linux__)
-#include "coro/task.hpp"
+  #include "coro/task.hpp"
 #endif
 
 namespace alaya {
diff --git a/include/index/graph/qg/qg_builder.hpp b/include/index/graph/qg/qg_builder.hpp
@@ -130,7 +130,7 @@ class QGBuilder {
   void search_new_neighbors(bool sup) {
     LOG_INFO("Searching for new neighbor candidates...");
 #if defined(__AVX512F__)
-#pragma omp parallel for schedule(dynamic)
+  #pragma omp parallel for schedule(dynamic)
     for (size_t i = 0; i < num_nodes_; ++i) {
       IDType cur_id = i;
       auto tid = omp_get_thread_num();
@@ -165,7 +165,7 @@ class QGBuilder {
 #if defined(__AVX512F__)
     std::vector<std::mutex> locks(num_nodes_);
     std::vector<CandidateList> reverse_buffer(num_nodes_);
-#pragma omp parallel for schedule(dynamic)
+  #pragma omp parallel for schedule(dynamic)
     for (IDType data_id = 0; data_id < num_nodes_; ++data_id) {  // for every vertex
       for (const auto &nei : new_neighbors_[data_id]) {
         auto dst = nei.id_;
@@ -191,7 +191,7 @@ class QGBuilder {
         }
       }
     }
-#pragma omp parallel for schedule(dynamic)
+  #pragma omp parallel for schedule(dynamic)
     for (IDType data_id = 0; data_id < num_nodes_; ++data_id) {  // prune for every vertex
       CandidateList &tmp_pool = reverse_buffer[data_id];
       tmp_pool.reserve(tmp_pool.size() + degree_bound_);
@@ -214,7 +214,7 @@ class QGBuilder {
   void angle_based_supplement() {
     LOG_INFO("Supplementing edges...");
 #if defined(__AVX512F__)
-#pragma omp parallel for schedule(dynamic)
+  #pragma omp parallel for schedule(dynamic)
     for (size_t i = 0; i < num_nodes_; ++i) {
       CandidateList &cur_neighbors = new_neighbors_[i];
       size_t cur_degree = cur_neighbors.size();
diff --git a/include/index/index_type.hpp b/include/index/index_type.hpp
@@ -21,11 +21,11 @@
 #include <tuple>
 
 #if defined(_MSC_VER)
-#define ALAYA_UNREACHABLE __assume(0)
+  #define ALAYA_UNREACHABLE __assume(0)
 #elif defined(__GNUC__) || defined(__clang__)
-#define ALAYA_UNREACHABLE __builtin_unreachable()
+  #define ALAYA_UNREACHABLE __builtin_unreachable()
 #else
-#define ALAYA_UNREACHABLE
+  #define ALAYA_UNREACHABLE
 #endif
 
 namespace alaya {
diff --git a/include/space/distance/dist_config.hpp b/include/space/distance/dist_config.hpp
@@ -18,47 +18,59 @@
 #include <sys/types.h>
 
 #ifndef NO_MANUAL_VECTORIZATION
-#if (defined(__SSE__) || _M_IX86_FP > 0 || defined(_M_AMD64) || defined(_M_X64))
-#define USE_SSE
-#ifdef __AVX__
-#define USE_AVX
-#ifdef __AVX512F__
-#define USE_AVX512
-#endif
-#endif
-#endif
+  #if (defined(__SSE__) || _M_IX86_FP > 0 || defined(_M_AMD64) || defined(_M_X64))
+    #define USE_SSE
+    #ifdef __AVX__
+      #define USE_AVX
+      #ifdef __AVX512F__
+        #define USE_AVX512
+      #endif
+    #endif
+  #endif
+
+  // ARM NEON detection
+  #if defined(__aarch64__) || defined(_M_ARM64)
+    #define USE_NEON
+  #endif
 #endif
 
 #if defined(USE_AVX) || defined(USE_SSE)
-#ifdef _MSC_VER
-#include <intrin.h>
-#include <stdexcept>
-#else
-#include <cpuid.h>
-#include <x86intrin.h>
+  #ifdef _MSC_VER
+    #include <intrin.h>
+    #include <stdexcept>
+  #else
+    #include <cpuid.h>
+    #include <x86intrin.h>
+
+  #endif
 
+  #if defined(USE_AVX512)
+    #include <immintrin.h>
+  #endif
 #endif
 
-#if defined(USE_AVX512)
-#include <immintrin.h>
+// ARM NEON headers
+#if defined(USE_NEON)
+  #include <arm_neon.h>
 #endif
 
+// Alignment macros (available on all platforms)
 #if defined(__GNUC__) || defined(__clang__)
-#define PORTABLE_ALIGN32 __attribute__((aligned(32)))
-#define PORTABLE_ALIGN64 __attribute__((aligned(64)))
+  #define PORTABLE_ALIGN32 __attribute__((aligned(32)))
+  #define PORTABLE_ALIGN64 __attribute__((aligned(64)))
 #else
-#define PORTABLE_ALIGN32 __declspec(align(32))
-#define PORTABLE_ALIGN64 __declspec(align(64))
-#endif
+  #define PORTABLE_ALIGN32 __declspec(align(32))
+  #define PORTABLE_ALIGN64 __declspec(align(64))
 #endif
 
 #if defined(__GNUC__) && !defined(__clang__)
-#define FAST_BEGIN _Pragma("GCC push_options") _Pragma("GCC optimize (\"unroll-loops,fast-math\")")
-#define FAST_END _Pragma("GCC pop_options")
+  #define FAST_BEGIN \
+    _Pragma("GCC push_options") _Pragma("GCC optimize (\"unroll-loops,fast-math\")")
+  #define FAST_END _Pragma("GCC pop_options")
 #else
-// Clang / MSVC / others: no-op
-#define FAST_BEGIN
-#define FAST_END
+  // Clang / MSVC / others: no-op
+  #define FAST_BEGIN
+  #define FAST_END
 #endif
 
 #define CAST(Type, Ptr) (*reinterpret_cast<const Type *>(Ptr))
diff --git a/include/space/raw_space.hpp b/include/space/raw_space.hpp
@@ -33,7 +33,7 @@
 #include "utils/metric_type.hpp"
 
 #ifdef _MSC_VER
-#include <malloc.h>
+  #include <malloc.h>
 #endif
 
 namespace alaya {
diff --git a/include/space/sq4_space.hpp b/include/space/sq4_space.hpp
@@ -39,7 +39,7 @@
 #include "utils/types.hpp"
 
 #ifdef _MSC_VER
-#include <malloc.h>
+  #include <malloc.h>
 #endif
 
 namespace alaya {
diff --git a/include/space/sq8_space.hpp b/include/space/sq8_space.hpp
@@ -31,7 +31,7 @@
 #include "utils/prefetch.hpp"
 
 #ifdef _MSC_VER
-#include <malloc.h>
+  #include <malloc.h>
 #endif
 
 namespace alaya {
diff --git a/include/storage/sequential_storage.hpp b/include/storage/sequential_storage.hpp
@@ -30,7 +30,7 @@
 #include "utils/types.hpp"
 
 #ifdef _MSC_VER
-#include <malloc.h>
+  #include <malloc.h>
 #endif
 
 namespace alaya {
diff --git a/include/utils/log.hpp b/include/utils/log.hpp
@@ -22,15 +22,15 @@
 #include <string>
 
 #ifdef PROJECT_ROOT
-#define RELATIVE_FILE get_relative_path(__FILE__, PROJECT_ROOT)
+  #define RELATIVE_FILE get_relative_path(__FILE__, PROJECT_ROOT)
 auto inline get_relative_path(const std::string &full_path, const std::string &base_path)
     -> std::string {
   std::filesystem::path full(full_path);
   std::filesystem::path base(base_path);
   return std::filesystem::relative(full, base).string();
 }
 #else
-#define RELATIVE_FILE __FILE__
+  #define RELATIVE_FILE __FILE__
 #endif
 
 #define CONCATENATE_STRINGS(a, b) a b
diff --git a/include/utils/memory.hpp b/include/utils/memory.hpp
@@ -17,7 +17,7 @@
 #pragma once
 
 #ifdef __linux__
-#include <sys/mman.h>
+  #include <sys/mman.h>
 #endif
 
 #include <cstdio>
@@ -26,7 +26,7 @@
 #include <type_traits>
 
 #ifdef _MSC_VER
-#include <malloc.h>
+  #include <malloc.h>
 #endif
 
 namespace alaya {
diff --git a/include/utils/prefetch.hpp b/include/utils/prefetch.hpp
@@ -20,11 +20,11 @@
 #include <cstdio>
 
 #if defined(_MSC_VER)
-#include <intrin.h>
+  #include <intrin.h>
 #elif defined(__SSE2__)
-#include <immintrin.h>
+  #include <immintrin.h>
 #elif defined(__aarch64__)
-#include <arm_neon.h>
+  #include <arm_neon.h>
 #endif
 
 namespace alaya {
diff --git a/include/utils/quantization_type.hpp b/include/utils/quantization_type.hpp
@@ -21,11 +21,11 @@
 #include <tuple>
 
 #if defined(_MSC_VER)
-#define ALAYA_UNREACHABLE __assume(0)
+  #define ALAYA_UNREACHABLE __assume(0)
 #elif defined(__GNUC__) || defined(__clang__)
-#define ALAYA_UNREACHABLE __builtin_unreachable()
+  #define ALAYA_UNREACHABLE __builtin_unreachable()
 #else
-#define ALAYA_UNREACHABLE
+  #define ALAYA_UNREACHABLE
 #endif
 
 namespace alaya {
diff --git a/include/utils/query_utils.hpp b/include/utils/query_utils.hpp
@@ -25,7 +25,7 @@
 #include "memory.hpp"
 
 #if defined(_MSC_VER)
-#include <intrin.h>
+  #include <intrin.h>
 #endif
 
 namespace alaya {
diff --git a/include/utils/rabitq_utils/defines.hpp b/include/utils/rabitq_utils/defines.hpp
@@ -21,11 +21,11 @@
 #define LOWBIT(x) ((x) & (-(x)))
 
 #if defined(_MSC_VER)
-#define ALAYA_UNREACHABLE __assume(0)
+  #define ALAYA_UNREACHABLE __assume(0)
 #elif defined(__GNUC__) || defined(__clang__)
-#define ALAYA_UNREACHABLE __builtin_unreachable()
+  #define ALAYA_UNREACHABLE __builtin_unreachable()
 #else
-#define ALAYA_UNREACHABLE
+  #define ALAYA_UNREACHABLE
 #endif
 
 namespace alaya {
diff --git a/include/utils/rabitq_utils/fastscan.hpp b/include/utils/rabitq_utils/fastscan.hpp
@@ -16,7 +16,14 @@
 
 #pragma once
 
-#include <immintrin.h>
+// Architecture-specific SIMD headers
+#if defined(__x86_64__) || defined(_M_X64) || defined(__i386__) || defined(_M_IX86)
+  #define ALAYA_X86
+  #include <immintrin.h>
+#elif defined(__aarch64__) || defined(_M_ARM64)
+  #define ALAYA_ARM64
+  #include <arm_neon.h>
+#endif
 
 #include <array>
 #include <cassert>
@@ -218,6 +225,40 @@ inline void accumulate(const uint8_t *__restrict__ codes,
   __m256i dis1 = _mm256_add_epi16(_mm256_permute2f128_si256(accu2, accu3, 0x21),
                                   _mm256_blend_epi32(accu2, accu3, 0xF0));
   _mm256_storeu_si256((__m256i *)&result[16], dis1);
+#else
+  // Scalar fallback for non-x86 architectures (ARM, etc.)
+  // This implementation is slower but functionally equivalent to the SIMD versions
+  std::memset(result, 0, 32 * sizeof(uint16_t));
+
+  for (size_t i = 0; i < code_length; i += 32) {
+    // First 16 bytes: upper 4-bit codes (col_0 from pack_codes)
+    for (size_t j = 0; j < 16; ++j) {
+      uint8_t c = codes[i + j];
+      uint8_t lo_code = c & 0x0f;
+      uint8_t hi_code = c >> 4;
+
+      // LUT lookup for upper 4 bits - first 16 bytes of LUT block
+      int8_t val_lo = static_cast<int8_t>(lp_table[i + lo_code]);
+      int8_t val_hi = static_cast<int8_t>(lp_table[i + hi_code]);
+
+      result[kPerm0[j]] += val_lo;       // Vector at kPerm0[j]
+      result[kPerm0[j] + 16] += val_hi;  // Vector at kPerm0[j] + 16
+    }
+
+    // Second 16 bytes: lower 4-bit codes (col_1 from pack_codes)
+    for (size_t j = 0; j < 16; ++j) {
+      uint8_t c = codes[i + 16 + j];
+      uint8_t lo_code = c & 0x0f;
+      uint8_t hi_code = c >> 4;
+
+      // LUT lookup for lower 4 bits - second 16 bytes of LUT block
+      int8_t val_lo = static_cast<int8_t>(lp_table[i + 16 + lo_code]);
+      int8_t val_hi = static_cast<int8_t>(lp_table[i + 16 + hi_code]);
+
+      result[kPerm0[j]] += val_lo;
+      result[kPerm0[j] + 16] += val_hi;
+    }
+  }
 #endif
 }
 // NOLINTEND
diff --git a/include/utils/rabitq_utils/rotator.hpp b/include/utils/rabitq_utils/rotator.hpp
@@ -321,6 +321,13 @@ template <typename T>
 std::unique_ptr<Rotator<T>> choose_rotator(size_t dim,
                                            RotatorType type = RotatorType::FhtKacRotator,
                                            size_t padded_dim = 0) {
+// RaBitQ requires AVX512 for FhtKacRotator, not supported on ARM
+#if defined(__aarch64__) || defined(_M_ARM64)
+  throw std::runtime_error(
+      "RaBitQ is not supported on ARM architecture. "
+      "AVX512 instructions are required for FhtKacRotator.");
+#endif
+
   if (padded_dim == 0) {
     padded_dim = rotator_impl::padding_requirement(dim, type);
     if (padded_dim != dim) {
