Add an AArch64 C++ backend with native NEON FastScan

CMakeLists.txt

@@ -16,15 +16,24 @@ option(
     ON
 )
 
-find_package(OpenMP REQUIRED)
+option(RABITQ_ENABLE_OPENMP "Enable OpenMP parallel loops" ON)
+if(RABITQ_ENABLE_OPENMP)
+    find_package(OpenMP REQUIRED)
+endif()
 
 add_library(rabitq_compile_options INTERFACE)
 target_compile_features(rabitq_compile_options INTERFACE cxx_std_17)
 target_compile_options(rabitq_compile_options INTERFACE
     $<$<AND:$<CONFIG:Release>,$<COMPILE_LANG_AND_ID:CXX,Clang,GNU>>:-Ofast>
     $<$<AND:$<BOOL:${RABITQ_ENABLE_NATIVE_OPTIMIZATION}>,$<COMPILE_LANG_AND_ID:CXX,Clang,GNU>>:-march=native>
 )
-target_link_libraries(rabitq_compile_options INTERFACE OpenMP::OpenMP_CXX)
+if(RABITQ_ENABLE_OPENMP)
+    target_link_libraries(rabitq_compile_options INTERFACE OpenMP::OpenMP_CXX)
+else()
+    target_compile_options(rabitq_compile_options INTERFACE
+        $<$<COMPILE_LANG_AND_ID:CXX,AppleClang,Clang,GNU>:-Wno-unknown-pragmas>
+    )
+endif()
 
 function(rabitq_enable_warnings target)
     target_compile_options(${target} PRIVATE
@@ -70,39 +79,62 @@ set(RABITQ_HNSW_AVX512_POPCNT_SOURCES
     src/index/hnsw_search_avx512_popcnt.cpp
 )
 
-add_library(rabitq_core STATIC
-    ${RABITQ_COMMON_SOURCES}
-    ${RABITQ_AVX2_SOURCES}
-    ${RABITQ_AVX512_SOURCES}
-    ${RABITQ_AVX512_POPCNT_SOURCES}
-    ${RABITQ_HNSW_AVX2_SOURCES}
-    ${RABITQ_HNSW_AVX512_CORE_SOURCES}
-    ${RABITQ_HNSW_AVX512_POPCNT_SOURCES}
-)
+include(CheckCXXSourceCompiles)
+check_cxx_source_compiles("
+#if !defined(__aarch64__)
+#error Not an AArch64 target
+#endif
+int main() { return 0; }" RABITQ_TARGET_AARCH64)
+
+if(RABITQ_TARGET_AARCH64)
+    find_path(SIMDE_INCLUDE_DIR simde/x86/avx2.h)
+    if(NOT SIMDE_INCLUDE_DIR)
+        message(FATAL_ERROR "AArch64 builds require SIMDe headers. Set SIMDE_INCLUDE_DIR to the directory containing simde/.")
+    endif()
+    add_library(rabitq_core STATIC
+        ${RABITQ_COMMON_SOURCES}
+        ${RABITQ_AVX2_SOURCES}
+        ${RABITQ_HNSW_AVX2_SOURCES}
+        src/simd/fastscan_neon.cpp
+    )
+    target_include_directories(rabitq_core SYSTEM PRIVATE "${SIMDE_INCLUDE_DIR}")
+else()
+    add_library(rabitq_core STATIC
+        ${RABITQ_COMMON_SOURCES}
+        ${RABITQ_AVX2_SOURCES}
+        ${RABITQ_AVX512_SOURCES}
+        ${RABITQ_AVX512_POPCNT_SOURCES}
+        ${RABITQ_HNSW_AVX2_SOURCES}
+        ${RABITQ_HNSW_AVX512_CORE_SOURCES}
+        ${RABITQ_HNSW_AVX512_POPCNT_SOURCES}
+    )
+endif()
 
 target_include_directories(rabitq_core PUBLIC ${PROJECT_SOURCE_DIR}/include)
 target_link_libraries(rabitq_core PUBLIC rabitq_compile_options)
 set_target_properties(rabitq_core PROPERTIES POSITION_INDEPENDENT_CODE ON)
 rabitq_enable_warnings(rabitq_core)
 
-set_source_files_properties(${RABITQ_AVX2_SOURCES} PROPERTIES
-    COMPILE_OPTIONS "-mavx2;-mfma"
-)
-set_source_files_properties(${RABITQ_AVX512_SOURCES} PROPERTIES
-    COMPILE_OPTIONS "-mavx512f;-mavx512bw;-mavx512dq;-mfma"
-)
-set_source_files_properties(${RABITQ_AVX512_POPCNT_SOURCES} PROPERTIES
-    COMPILE_OPTIONS "-mavx512f;-mavx512bw;-mavx512dq;-mavx512vpopcntdq;-mfma"
-)
-set_source_files_properties(${RABITQ_HNSW_AVX2_SOURCES} PROPERTIES
-    COMPILE_OPTIONS "-mavx2;-mfma"
-)
-set_source_files_properties(${RABITQ_HNSW_AVX512_CORE_SOURCES} PROPERTIES
-    COMPILE_OPTIONS "-mavx512f;-mavx512bw;-mavx512dq;-mavx2;-mfma"
-)
-set_source_files_properties(${RABITQ_HNSW_AVX512_POPCNT_SOURCES} PROPERTIES
-    COMPILE_OPTIONS "-mavx512f;-mavx512bw;-mavx512dq;-mavx512vpopcntdq;-mfma"
-)
+if(NOT RABITQ_TARGET_AARCH64)
+    set_source_files_properties(${RABITQ_AVX2_SOURCES} PROPERTIES
+        COMPILE_OPTIONS "-mavx2;-mfma"
+    )
+    set_source_files_properties(${RABITQ_AVX512_SOURCES} PROPERTIES
+        COMPILE_OPTIONS "-mavx512f;-mavx512bw;-mavx512dq;-mfma"
+    )
+    set_source_files_properties(${RABITQ_AVX512_POPCNT_SOURCES} PROPERTIES
+        COMPILE_OPTIONS "-mavx512f;-mavx512bw;-mavx512dq;-mavx512vpopcntdq;-mfma"
+    )
+    set_source_files_properties(${RABITQ_HNSW_AVX2_SOURCES} PROPERTIES
+        COMPILE_OPTIONS "-mavx2;-mfma"
+    )
+    set_source_files_properties(${RABITQ_HNSW_AVX512_CORE_SOURCES} PROPERTIES
+        COMPILE_OPTIONS "-mavx512f;-mavx512bw;-mavx512dq;-mavx2;-mfma"
+    )
+    set_source_files_properties(${RABITQ_HNSW_AVX512_POPCNT_SOURCES} PROPERTIES
+        COMPILE_OPTIONS "-mavx512f;-mavx512bw;-mavx512dq;-mavx512vpopcntdq;-mfma"
+    )
+endif()
 
 add_library(rabitq_headers INTERFACE)
 target_include_directories(rabitq_headers INTERFACE ${PROJECT_SOURCE_DIR}/include)

README.md

@@ -162,9 +162,9 @@ python -m pip install .
 ### Requirements
 
 - CMake 3.15 or newer
-- a C++17 compiler with OpenMP support
-- an x86-64 CPU supported by the selected kernels: most paths accept either
-  AVX2 with FMA or AVX-512F/BW/DQ with FMA
+- a C++17 compiler; OpenMP is enabled by default and can be disabled for serial builds
+- an x86-64 CPU with AVX2/FMA or AVX-512F/BW/DQ/FMA, or an AArch64 CPU with
+  [SIMDe](https://github.com/simd-everywhere/simde) headers available at build time
 
 <details>
 <summary>CPU dispatch details</summary>
@@ -187,6 +187,30 @@ cmake -S . -B build -DCMAKE_BUILD_TYPE=Release
 cmake --build build --parallel
 ```
 
+For an AArch64 source build, supply the directory containing `simde/`:
+
+```bash
+cmake -S . -B build-arm -DCMAKE_BUILD_TYPE=Release \
+  -DSIMDE_INCLUDE_DIR=/path/to/simde-checkout \
+  -DRABITQ_ENABLE_OPENMP=OFF -DRABITQ_BUILD_TESTS=ON
+cmake --build build-arm --parallel
+ctest --test-dir build-arm --output-on-failure
+```
+
+The ARM backend uses a native NEON ordinary FastScan accumulator, SIMDe for
+the remaining AVX2 kernels, and a portable Hadamard transform. SIMDe is a
+private build dependency of the compiled library; consumers link the existing
+`rabitq_headers` target without including SIMDe themselves. The public API,
+saved-index layout, and production rotation seed policy are unchanged.
+The implementation was tested with SIMDe revision
+`b49b253a24ba1f26c7e619ed88330cae53135950` on Apple Silicon.
+
+`RABITQ_ENABLE_OPENMP=OFF` selects serial execution. With an OpenMP compiler
+and runtime installed, omit that option to retain parallel construction and
+batch operations. ARM support here covers C++ source builds; prebuilt Python
+wheels remain Linux x86-64 only. Multi-architecture macOS universal binaries
+require separate per-architecture builds and are not supported by this configuration.
+
 Release builds enable native CPU tuning by default. To build a binary that can
 be moved between AVX2- and AVX-512-capable machines, configure with
 `-DRABITQ_ENABLE_NATIVE_OPTIMIZATION=OFF`; the ISA-specific kernels will still

include/rabitqlib/index/hnsw/hnsw.hpp

@@ -1,8 +1,6 @@
 // HNSW is developed from the [HNSW library](https://github.com/nmslib/hnswlib)
 #pragma once
 
-#include <omp.h>
-
 #include <atomic>
 #include <cassert>
 #include <cstddef>
@@ -22,6 +20,7 @@
 #include "rabitqlib/quantization/rabitq.hpp"
 #include "rabitqlib/utils/buffer.hpp"
 #include "rabitqlib/utils/cpu_features.hpp"
+#include "rabitqlib/utils/parallel.hpp"
 #include "rabitqlib/utils/rotator.hpp"
 #include "rabitqlib/utils/space.hpp"
 #include "rabitqlib/utils/tools.hpp"
@@ -1111,6 +1110,9 @@ inline std::vector<std::vector<std::pair<float, PID>>> HierarchicalNSW::search(
 inline maxheap<std::pair<float, PID>> HierarchicalNSW::search_knn(
     const float* rotated_query, size_t TOPK
 ) {
+#if defined(__aarch64__)
+    return detail::search_knn_avx2(*this, rotated_query, TOPK);
+#else
     if (rabitqlib::cpu::has_avx512_popcnt()) {
         return detail::search_knn_avx512_popcnt(*this, rotated_query, TOPK);
     }
@@ -1122,6 +1124,7 @@ inline maxheap<std::pair<float, PID>> HierarchicalNSW::search_knn(
     }
 
     throw std::runtime_error("HNSW search requires AVX2/FMA or AVX512 support");
+#endif
 }
 
 template <class Kernel>

include/rabitqlib/index/ivf/ivf.hpp

@@ -1,7 +1,5 @@
 #pragma once
 
-#include <omp.h>
-
 #include <algorithm>
 #include <cassert>
 #include <cmath>
@@ -22,6 +20,7 @@
 #include "rabitqlib/quantization/rabitq.hpp"
 #include "rabitqlib/utils/buffer.hpp"
 #include "rabitqlib/utils/memory.hpp"
+#include "rabitqlib/utils/parallel.hpp"
 #include "rabitqlib/utils/rotator.hpp"
 #include "rabitqlib/utils/space.hpp"
 #include "rabitqlib/utils/tools.hpp"

include/rabitqlib/index/symqg/qg.hpp

@@ -1,7 +1,5 @@
 #pragma once
 
-#include <omp.h>
-
 #include <cassert>
 #include <cstddef>
 #include <cstdint>
@@ -23,6 +21,7 @@
 #include "rabitqlib/utils/buffer.hpp"
 #include "rabitqlib/utils/io.hpp"
 #include "rabitqlib/utils/memory.hpp"
+#include "rabitqlib/utils/parallel.hpp"
 #include "rabitqlib/utils/rotator.hpp"
 #include "rabitqlib/utils/space.hpp"
 #include "rabitqlib/utils/visited_pool.hpp"

include/rabitqlib/index/symqg/qg_builder.hpp

@@ -70,7 +70,7 @@ class QGBuilder {
               VisitedSet(num_nodes_, std::min(ef_build_ * ef_build_, num_nodes_ / 10))
           )
         , degrees_(qg_.num_vertices(), degree_bound_) {
-        omp_set_num_threads(static_cast<int>(num_threads_));
+        parallel::set_thread_count(static_cast<int>(num_threads_));
 
         std::vector<float> centroid =
             compute_centroid(data, num_nodes_, dim_, num_threads_);
@@ -233,7 +233,7 @@ inline void QGBuilder::search_new_neighbors(bool refine) {
 #pragma omp parallel for schedule(dynamic)
     for (size_t i = 0; i < num_nodes_; ++i) {
         PID cur_id = i;
-        auto tid = omp_get_thread_num();
+        auto tid = parallel::thread_index();
         CandidateList candidates;
         VisitedSet& vis = visited_list_[tid];
         candidates.reserve(2 * kMaxCandidatePoolSize);

include/rabitqlib/quantization/rabitq.hpp

@@ -1,7 +1,5 @@
 #pragma once
 
-#include <omp.h>
-
 #include <cmath>
 #include <cstddef>
 #include <vector>
@@ -10,6 +8,7 @@
 #include "rabitqlib/fastscan/fastscan.hpp"
 #include "rabitqlib/quantization/data_layout.hpp"
 #include "rabitqlib/quantization/rabitq_impl.hpp"
+#include "rabitqlib/utils/parallel.hpp"
 
 namespace rabitqlib::quant {
 

include/rabitqlib/quantization/rabitq_impl.hpp

@@ -1,7 +1,5 @@
 #pragma once
 
-#include <omp.h>
-
 #include <algorithm>
 #include <cassert>
 #include <cmath>
@@ -13,6 +11,7 @@
 #include "rabitqlib/fastscan/fastscan.hpp"
 #include "rabitqlib/quantization/data_layout.hpp"
 #include "rabitqlib/quantization/pack_excode.hpp"
+#include "rabitqlib/utils/parallel.hpp"
 #include "rabitqlib/utils/space.hpp"
 
 namespace rabitqlib::quant::rabitq_impl {
@@ -291,6 +290,9 @@ inline double best_rescale_factor(const T* o_abs, size_t dim, size_t ex_bits) {
     constexpr double kEps = 1e-5;
     constexpr int kNEnum = 10;
     double max_o = *std::max_element(o_abs, o_abs + dim);
+    if (max_o == 0) {
+        return 0;
+    }
 
     double t_end = static_cast<double>(((1 << ex_bits) - 1) + kNEnum) / max_o;
     double t_start = t_end * kTightStart[ex_bits];
@@ -424,7 +426,11 @@ inline T ex_bits_code(
     ConstRowMajorArrayMap<T> res_arr(residual, 1, dim);
 
     // get normalized abs residual for plus code
-    RowMajorArray<T> abs_res = res_arr.rowwise().normalized().abs();
+    RowMajorArray<T> abs_res = res_arr.abs();
+    const T norm = abs_res.matrix().norm();
+    if (norm > 0) {
+        abs_res /= norm;
+    }
 
     // quantize data
     T ipnorm_inv = 1;

include/rabitqlib/simd/fastscan_dispatch.hpp

@@ -5,6 +5,10 @@
 
 namespace rabitqlib::fastscan::simd {
 
+void accumulate_neon(
+    const uint8_t* codes, const uint8_t* lut, uint16_t* result, size_t dim
+);
+
 void accumulate_avx2(
     const uint8_t* __restrict__ codes,
     const uint8_t* __restrict__ lp_table,

include/rabitqlib/simd/x86_compat.hpp

@@ -0,0 +1,9 @@
+#pragma once
+
+#if defined(__aarch64__)
+#define SIMDE_ENABLE_NATIVE_ALIASES
+#include <simde/x86/avx2.h>
+#include <simde/x86/fma.h>
+#else
+#include <immintrin.h>
+#endif

include/rabitqlib/utils/fht_portable.hpp

@@ -0,0 +1,23 @@
+#pragma once
+#include <cstddef>
+
+// Same unnormalized Sylvester Hadamard transform as upstream's x86 assembly.
+template <std::size_t LogN>
+inline void portable_hadamard(float* data) {
+    constexpr std::size_t n = 1UL << LogN;
+    for (std::size_t half = 1; half < n; half *= 2) {
+        for (std::size_t base = 0; base < n; base += 2 * half) {
+            for (std::size_t i = 0; i < half; ++i) {
+                const float a = data[base + i], b = data[base + half + i];
+                data[base + i] = a + b;
+                data[base + half + i] = a - b;
+            }
+        }
+    }
+}
+inline void helper_float_6(float* p) { portable_hadamard<6>(p); }
+inline void helper_float_7(float* p) { portable_hadamard<7>(p); }
+inline void helper_float_8(float* p) { portable_hadamard<8>(p); }
+inline void helper_float_9(float* p) { portable_hadamard<9>(p); }
+inline void helper_float_10(float* p) { portable_hadamard<10>(p); }
+inline void helper_float_11(float* p) { portable_hadamard<11>(p); }

include/rabitqlib/utils/memory.hpp

@@ -45,12 +45,26 @@ class AlignedAllocator {
         auto nbytes = round_up_to_multiple_of<size_t>(n * sizeof(T), Alignment);
         auto* ptr = std::aligned_alloc(Alignment, nbytes);
         if (HugePage) {
+#if defined(__linux__)
             madvise(ptr, nbytes, MADV_HUGEPAGE);
+#endif
         }
         return reinterpret_cast<T*>(ptr);
     }
 
     void deallocate(T* ptr, [[maybe_unused]] std::size_t n) { std::free(ptr); }
+
+    template <typename U>
+    constexpr bool operator==(const AlignedAllocator<U, Alignment, HugePage>&)
+        const noexcept {
+        return true;
+    }
+
+    template <typename U>
+    constexpr bool operator!=(const AlignedAllocator<U, Alignment, HugePage>&)
+        const noexcept {
+        return false;
+    }
 };
 
 template <typename T>
@@ -81,7 +95,9 @@ inline T* align_allocate(size_t nbytes) {
     auto size = round_up_to_multiple_of<size_t>(nbytes, Alignment);
     void* ptr = std::aligned_alloc(Alignment, size);
     if (HugePage) {
+#if defined(__linux__)
         madvise(ptr, size, MADV_HUGEPAGE);
+#endif
     }
     return static_cast<T*>(ptr);
 }