From 42fdff7e8f51e2fc4f906d45ed2bc3cea1300224 Mon Sep 17 00:00:00 2001 From: menghuaban520 <270288916+menghuaban520@users.noreply.github.com> Date: Thu, 3 Sep 2026 14:45:23 +0800 Subject: [PATCH] Add AArch64 C++ backend with native NEON FastScan --- CMakeLists.txt | 90 +++++++++++++------ README.md | 30 ++++++- include/rabitqlib/index/hnsw/hnsw.hpp | 7 +- include/rabitqlib/index/ivf/ivf.hpp | 3 +- include/rabitqlib/index/symqg/qg.hpp | 3 +- include/rabitqlib/index/symqg/qg_builder.hpp | 4 +- include/rabitqlib/quantization/rabitq.hpp | 3 +- .../rabitqlib/quantization/rabitq_impl.hpp | 12 ++- include/rabitqlib/simd/fastscan_dispatch.hpp | 4 + include/rabitqlib/simd/x86_compat.hpp | 9 ++ include/rabitqlib/utils/fht_portable.hpp | 23 +++++ include/rabitqlib/utils/memory.hpp | 16 ++++ include/rabitqlib/utils/parallel.hpp | 23 +++++ include/rabitqlib/utils/rotator.hpp | 4 + include/rabitqlib/utils/space.hpp | 7 +- src/index/hnsw_search_avx2_kernels.hpp | 3 +- src/simd/dispatch.cpp | 71 +++++++++++++++ src/simd/fastscan_avx2.cpp | 3 +- src/simd/fastscan_neon.cpp | 37 ++++++++ src/simd/pack_excode_kernels.hpp | 4 +- src/simd/rotator_avx2.cpp | 3 +- src/simd/space_avx2.cpp | 15 ++-- src/simd/space_excode_avx2.cpp | 3 +- src/simd/warmup_avx2.cpp | 4 +- tests/integration/bit_pack_unpack_test.cpp | 10 +-- .../unit/rabitqlib/utils/arm_backend_test.cpp | 69 ++++++++++++++ tests/unit/rabitqlib/utils/memory_test.cpp | 15 ++++ tests/unit/rabitqlib/utils/space_test.cpp | 64 ++++++++++--- 28 files changed, 455 insertions(+), 84 deletions(-) create mode 100644 include/rabitqlib/simd/x86_compat.hpp create mode 100644 include/rabitqlib/utils/fht_portable.hpp create mode 100644 include/rabitqlib/utils/parallel.hpp create mode 100644 src/simd/fastscan_neon.cpp create mode 100644 tests/unit/rabitqlib/utils/arm_backend_test.cpp create mode 100644 tests/unit/rabitqlib/utils/memory_test.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index d70f2da..8e97998 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -16,7 +16,10 @@ 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) @@ -24,7 +27,13 @@ target_compile_options(rabitq_compile_options INTERFACE $<$,$>:-Ofast> $<$,$>:-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 + $<$:-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) diff --git a/README.md b/README.md index f38d0fb..7c5d136 100644 --- a/README.md +++ b/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
CPU dispatch details @@ -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 diff --git a/include/rabitqlib/index/hnsw/hnsw.hpp b/include/rabitqlib/index/hnsw/hnsw.hpp index 6e5b0af..fdd1917 100644 --- a/include/rabitqlib/index/hnsw/hnsw.hpp +++ b/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 - #include #include #include @@ -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>> HierarchicalNSW::search( inline maxheap> 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> HierarchicalNSW::search_knn( } throw std::runtime_error("HNSW search requires AVX2/FMA or AVX512 support"); +#endif } template diff --git a/include/rabitqlib/index/ivf/ivf.hpp b/include/rabitqlib/index/ivf/ivf.hpp index dc8600b..9cb3f6f 100644 --- a/include/rabitqlib/index/ivf/ivf.hpp +++ b/include/rabitqlib/index/ivf/ivf.hpp @@ -1,7 +1,5 @@ #pragma once -#include - #include #include #include @@ -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" diff --git a/include/rabitqlib/index/symqg/qg.hpp b/include/rabitqlib/index/symqg/qg.hpp index 4cd9dc1..f4a75f6 100644 --- a/include/rabitqlib/index/symqg/qg.hpp +++ b/include/rabitqlib/index/symqg/qg.hpp @@ -1,7 +1,5 @@ #pragma once -#include - #include #include #include @@ -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" diff --git a/include/rabitqlib/index/symqg/qg_builder.hpp b/include/rabitqlib/index/symqg/qg_builder.hpp index 8b75bfa..4d88235 100644 --- a/include/rabitqlib/index/symqg/qg_builder.hpp +++ b/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(num_threads_)); + parallel::set_thread_count(static_cast(num_threads_)); std::vector 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); diff --git a/include/rabitqlib/quantization/rabitq.hpp b/include/rabitqlib/quantization/rabitq.hpp index 229b184..1bf5657 100644 --- a/include/rabitqlib/quantization/rabitq.hpp +++ b/include/rabitqlib/quantization/rabitq.hpp @@ -1,7 +1,5 @@ #pragma once -#include - #include #include #include @@ -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 { diff --git a/include/rabitqlib/quantization/rabitq_impl.hpp b/include/rabitqlib/quantization/rabitq_impl.hpp index bd05ee3..c9ddb55 100644 --- a/include/rabitqlib/quantization/rabitq_impl.hpp +++ b/include/rabitqlib/quantization/rabitq_impl.hpp @@ -1,7 +1,5 @@ #pragma once -#include - #include #include #include @@ -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(((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 res_arr(residual, 1, dim); // get normalized abs residual for plus code - RowMajorArray abs_res = res_arr.rowwise().normalized().abs(); + RowMajorArray abs_res = res_arr.abs(); + const T norm = abs_res.matrix().norm(); + if (norm > 0) { + abs_res /= norm; + } // quantize data T ipnorm_inv = 1; diff --git a/include/rabitqlib/simd/fastscan_dispatch.hpp b/include/rabitqlib/simd/fastscan_dispatch.hpp index 8bf44f5..cf1ff9b 100644 --- a/include/rabitqlib/simd/fastscan_dispatch.hpp +++ b/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, diff --git a/include/rabitqlib/simd/x86_compat.hpp b/include/rabitqlib/simd/x86_compat.hpp new file mode 100644 index 0000000..ff42c8a --- /dev/null +++ b/include/rabitqlib/simd/x86_compat.hpp @@ -0,0 +1,9 @@ +#pragma once + +#if defined(__aarch64__) +#define SIMDE_ENABLE_NATIVE_ALIASES +#include +#include +#else +#include +#endif diff --git a/include/rabitqlib/utils/fht_portable.hpp b/include/rabitqlib/utils/fht_portable.hpp new file mode 100644 index 0000000..7fe1e6f --- /dev/null +++ b/include/rabitqlib/utils/fht_portable.hpp @@ -0,0 +1,23 @@ +#pragma once +#include + +// Same unnormalized Sylvester Hadamard transform as upstream's x86 assembly. +template +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); } diff --git a/include/rabitqlib/utils/memory.hpp b/include/rabitqlib/utils/memory.hpp index 8dced55..47c5642 100644 --- a/include/rabitqlib/utils/memory.hpp +++ b/include/rabitqlib/utils/memory.hpp @@ -45,12 +45,26 @@ class AlignedAllocator { auto nbytes = round_up_to_multiple_of(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(ptr); } void deallocate(T* ptr, [[maybe_unused]] std::size_t n) { std::free(ptr); } + + template + constexpr bool operator==(const AlignedAllocator&) + const noexcept { + return true; + } + + template + constexpr bool operator!=(const AlignedAllocator&) + const noexcept { + return false; + } }; template @@ -81,7 +95,9 @@ inline T* align_allocate(size_t nbytes) { auto size = round_up_to_multiple_of(nbytes, Alignment); void* ptr = std::aligned_alloc(Alignment, size); if (HugePage) { +#if defined(__linux__) madvise(ptr, size, MADV_HUGEPAGE); +#endif } return static_cast(ptr); } diff --git a/include/rabitqlib/utils/parallel.hpp b/include/rabitqlib/utils/parallel.hpp new file mode 100644 index 0000000..1dd4633 --- /dev/null +++ b/include/rabitqlib/utils/parallel.hpp @@ -0,0 +1,23 @@ +#pragma once + +#ifdef _OPENMP +#include +#endif + +namespace rabitqlib::parallel { +inline int thread_index() { +#ifdef _OPENMP + return omp_get_thread_num(); +#else + return 0; +#endif +} + +inline void set_thread_count(int count) { +#ifdef _OPENMP + omp_set_num_threads(count); +#else + (void)count; +#endif +} +} // namespace rabitqlib::parallel diff --git a/include/rabitqlib/utils/rotator.hpp b/include/rabitqlib/utils/rotator.hpp index 1611e7b..366ec0c 100644 --- a/include/rabitqlib/utils/rotator.hpp +++ b/include/rabitqlib/utils/rotator.hpp @@ -12,7 +12,11 @@ #include "rabitqlib/defines.hpp" #include "rabitqlib/simd/rotator_dispatch.hpp" +#if defined(__aarch64__) +#include "rabitqlib/utils/fht_portable.hpp" +#else #include "rabitqlib/utils/fht_avx.hpp" +#endif #include "rabitqlib/utils/space.hpp" #include "rabitqlib/utils/tools.hpp" diff --git a/include/rabitqlib/utils/space.hpp b/include/rabitqlib/utils/space.hpp index 4ffe035..ed800ec 100644 --- a/include/rabitqlib/utils/space.hpp +++ b/include/rabitqlib/utils/space.hpp @@ -1,7 +1,5 @@ #pragma once -#include - #include #include #include @@ -13,6 +11,7 @@ #include "rabitqlib/defines.hpp" #include "rabitqlib/simd/space_dispatch.hpp" +#include "rabitqlib/utils/parallel.hpp" #include "rabitqlib/utils/tools.hpp" namespace rabitqlib { @@ -175,7 +174,7 @@ inline std::vector compute_centroid( #pragma omp parallel for schedule(dynamic) num_threads(thread_count) for (size_t i = 0; i < num_points; ++i) { - auto tid = omp_get_thread_num(); + auto tid = parallel::thread_index(); std::vector& cur_results = all_results[tid]; const T* cur_data = data + (dim * i); for (size_t k = 0; k < dim; ++k) { @@ -219,7 +218,7 @@ inline PID exact_nn( #pragma omp parallel for schedule(dynamic) num_threads(thread_count) for (size_t i = 0; i < num_points; ++i) { - auto tid = omp_get_thread_num(); + auto tid = parallel::thread_index(); AnnCandidate& cur_entry = best_entries[tid]; const T* cur_data = data + (dim * i); diff --git a/src/index/hnsw_search_avx2_kernels.hpp b/src/index/hnsw_search_avx2_kernels.hpp index bf20fc0..96ce564 100644 --- a/src/index/hnsw_search_avx2_kernels.hpp +++ b/src/index/hnsw_search_avx2_kernels.hpp @@ -1,12 +1,11 @@ #pragma once -#include - #include #include #include "rabitqlib/index/query.hpp" #include "rabitqlib/simd/space_dispatch.hpp" +#include "rabitqlib/simd/x86_compat.hpp" namespace rabitqlib::hnsw::detail { diff --git a/src/simd/dispatch.cpp b/src/simd/dispatch.cpp index 1a3eed6..a1bf213 100644 --- a/src/simd/dispatch.cpp +++ b/src/simd/dispatch.cpp @@ -14,11 +14,13 @@ namespace rabitqlib::simd { +#if !defined(__aarch64__) [[noreturn]] static void missing_feature(const char* feature_name) { throw std::runtime_error( std::string(feature_name) + " requires AVX2/FMA or AVX512 support" ); } +#endif // With zero extra bits there is no extra code to contribute to the inner // product, so the ex_bits == 0 slot must be a constant-zero stub rather than @@ -29,6 +31,7 @@ static float ip_fxu0( return 0.0F; } +#if !defined(__aarch64__) static float missing_excode_ip(const float*, const uint8_t*, size_t) { missing_feature("excode ip functions"); } @@ -82,8 +85,22 @@ static float missing_warmup_ip_x0_q_512( ) { missing_feature("warmup_ip_x0_q_512"); } +#endif ExcodeIpTable resolve_excode_ip_table() { +#if defined(__aarch64__) + return { + ip_fxu0, + excode_ipimpl::ip16_fxu1_avx2, + excode_ipimpl::ip64_fxu2_avx2, + excode_ipimpl::ip64_fxu3_avx2, + excode_ipimpl::ip16_fxu4_avx2, + excode_ipimpl::ip64_fxu5_avx2, + excode_ipimpl::ip64_fxu6_avx2, + excode_ipimpl::ip64_fxu7_avx2, + excode_ipimpl::ip16_fxu8_avx2, + }; +#else if (cpu::has_avx512_core()) { return { ip_fxu0, @@ -121,10 +138,14 @@ ExcodeIpTable resolve_excode_ip_table() { missing_excode_ip, }; } +#endif } using FlipSignFn = void (*)(const uint8_t*, float*, size_t); const FlipSignFn kFlipSignFn = [] { +#if defined(__aarch64__) + return flip_sign_avx2; +#else if (cpu::has_avx512_core()) { return flip_sign_avx512; } else if (cpu::has_avx2()) { @@ -132,10 +153,14 @@ const FlipSignFn kFlipSignFn = [] { } else { return missing_flip_sign; } +#endif }(); using KacsWalkFn = void (*)(float*, size_t); const KacsWalkFn kKacsWalkFn = [] { +#if defined(__aarch64__) + return kacs_walk_avx2; +#else if (cpu::has_avx512_core()) { return kacs_walk_avx512; } else if (cpu::has_avx2()) { @@ -143,10 +168,14 @@ const KacsWalkFn kKacsWalkFn = [] { } else { return missing_kacs_walk; } +#endif }(); using ScalarQuantizeUint8Fn = void (*)(uint8_t*, const float*, size_t, float, float); const ScalarQuantizeUint8Fn kScalarQuantizeUint8Fn = [] { +#if defined(__aarch64__) + return scalar_quantize_uint8_avx2; +#else if (cpu::has_avx512_core()) { return scalar_quantize_uint8_avx512; } else if (cpu::has_avx2()) { @@ -154,10 +183,14 @@ const ScalarQuantizeUint8Fn kScalarQuantizeUint8Fn = [] { } else { return missing_scalar_quantize_uint8; } +#endif }(); using ScalarQuantizeUint16Fn = void (*)(uint16_t*, const float*, size_t, float, float); const ScalarQuantizeUint16Fn kScalarQuantizeUint16Fn = [] { +#if defined(__aarch64__) + return scalar_quantize_uint16_avx2; +#else if (cpu::has_avx512_core()) { return scalar_quantize_uint16_avx512; } else if (cpu::has_avx2()) { @@ -165,10 +198,19 @@ const ScalarQuantizeUint16Fn kScalarQuantizeUint16Fn = [] { } else { return missing_scalar_quantize_uint16; } +#endif }(); using PackExcodeFn = void (*)(const uint8_t*, uint8_t*, size_t); +#if defined(__aarch64__) +const PackExcodeFn kPacking2BitExcodeFn = packing_2bit_excode_avx2; +const PackExcodeFn kPacking3BitExcodeFn = packing_3bit_excode_avx2; +const PackExcodeFn kPacking4BitExcodeFn = packing_4bit_excode_avx2; +const PackExcodeFn kPacking5BitExcodeFn = packing_5bit_excode_avx2; +const PackExcodeFn kPacking6BitExcodeFn = packing_6bit_excode_avx2; +const PackExcodeFn kPacking7BitExcodeFn = packing_7bit_excode_avx2; +#else static PackExcodeFn resolve_pack_excode_fn(PackExcodeFn avx512_fn, PackExcodeFn avx2_fn) { if (cpu::has_avx512_core()) { return avx512_fn; @@ -191,6 +233,7 @@ const PackExcodeFn kPacking6BitExcodeFn = resolve_pack_excode_fn(packing_6bit_excode_avx512, packing_6bit_excode_avx2); const PackExcodeFn kPacking7BitExcodeFn = resolve_pack_excode_fn(packing_7bit_excode_avx512, packing_7bit_excode_avx2); +#endif void flip_sign(const uint8_t* flip, float* data, size_t dim) { kFlipSignFn(flip, data, dim); @@ -250,6 +293,9 @@ const ex_ipfunc kIp64Fxu7AvxFn = kExcodeIpTable[7]; using NewTransposeBinFn = void (*)(const uint16_t*, uint64_t*, size_t, size_t); const NewTransposeBinFn kNewTransposeBinFn = [] { +#if defined(__aarch64__) + return simd::new_transpose_bin_avx2; +#else if (cpu::has_avx512_core()) { return simd::new_transpose_bin_avx512; } else if (cpu::has_avx2()) { @@ -257,10 +303,14 @@ const NewTransposeBinFn kNewTransposeBinFn = [] { } else { return simd::missing_new_transpose_bin; } +#endif }(); using NewTransposeBin512Fn = void (*)(const uint8_t*, uint64_t*, size_t, size_t); const NewTransposeBin512Fn kNewTransposeBin512Fn = [] { +#if defined(__aarch64__) + return simd::new_transpose_bin_512_avx2; +#else if (cpu::has_avx512_core()) { return simd::new_transpose_bin_512_avx512; } else if (cpu::has_avx2()) { @@ -268,10 +318,14 @@ const NewTransposeBin512Fn kNewTransposeBin512Fn = [] { } else { return simd::missing_new_transpose_bin_512; } +#endif }(); using MaskIpX0QFn = float (*)(const float*, const uint64_t*, size_t); const MaskIpX0QFn kMaskIpX0QFn = [] { +#if defined(__aarch64__) + return simd::mask_ip_x0_q_avx2; +#else if (cpu::has_avx512_core()) { return simd::mask_ip_x0_q_avx512; } else if (cpu::has_avx2()) { @@ -279,6 +333,7 @@ const MaskIpX0QFn kMaskIpX0QFn = [] { } else { return simd::missing_mask_ip_x0_q; } +#endif }(); ex_ipfunc select_excode_ipfunc(size_t ex_bits) { @@ -351,6 +406,9 @@ namespace rabitqlib::fastscan { using AccumulateFn = void (*)(const uint8_t*, const uint8_t*, uint16_t*, size_t); const AccumulateFn kAccumulateFn = [] { +#if defined(__aarch64__) + return simd::accumulate_neon; +#else if (cpu::has_avx512_core()) { return simd::accumulate_avx512; } else if (cpu::has_avx2()) { @@ -358,10 +416,14 @@ const AccumulateFn kAccumulateFn = [] { } else { return rabitqlib::simd::missing_fastscan_accumulate; } +#endif }(); using TransferLutHaccFn = void (*)(const uint16_t*, size_t, uint8_t*); const TransferLutHaccFn kTransferLutHaccFn = [] { +#if defined(__aarch64__) + return simd::transfer_lut_hacc_avx2; +#else if (cpu::has_avx512_core()) { return simd::transfer_lut_hacc_avx512; } else if (cpu::has_avx2()) { @@ -369,10 +431,14 @@ const TransferLutHaccFn kTransferLutHaccFn = [] { } else { return rabitqlib::simd::missing_fastscan_transfer_lut_hacc; } +#endif }(); using AccumulateHaccFn = void (*)(const uint8_t*, const uint8_t*, int32_t*, size_t); const AccumulateHaccFn kAccumulateHaccFn = [] { +#if defined(__aarch64__) + return simd::accumulate_hacc_avx2; +#else if (cpu::has_avx512_core()) { return simd::accumulate_hacc_avx512; } else if (cpu::has_avx2()) { @@ -380,6 +446,7 @@ const AccumulateHaccFn kAccumulateHaccFn = [] { } else { return rabitqlib::simd::missing_fastscan_accumulate_hacc; } +#endif }(); void accumulate( @@ -411,6 +478,9 @@ namespace rabitqlib { using WarmupIpX0Q512Fn = float (*)(const uint64_t*, const uint64_t*, float, float, size_t, size_t); const WarmupIpX0Q512Fn kWarmupIpX0Q512Fn = [] { +#if defined(__aarch64__) + return rabitqlib::simd::warmup_ip_x0_q_512_avx2; +#else if (rabitqlib::cpu::has_avx512_popcnt()) { return rabitqlib::simd::warmup_ip_x0_q_512_avx512; } else if (rabitqlib::cpu::has_avx2()) { @@ -418,6 +488,7 @@ const WarmupIpX0Q512Fn kWarmupIpX0Q512Fn = [] { } else { return rabitqlib::simd::missing_warmup_ip_x0_q_512; } +#endif }(); float warmup_ip_x0_q_512( diff --git a/src/simd/fastscan_avx2.cpp b/src/simd/fastscan_avx2.cpp index 5d2e5c0..6b4dbfe 100644 --- a/src/simd/fastscan_avx2.cpp +++ b/src/simd/fastscan_avx2.cpp @@ -1,9 +1,8 @@ -#include - #include #include "rabitqlib/fastscan/fastscan.hpp" #include "rabitqlib/fastscan/highacc_fastscan.hpp" +#include "rabitqlib/simd/x86_compat.hpp" namespace rabitqlib::fastscan::simd { diff --git a/src/simd/fastscan_neon.cpp b/src/simd/fastscan_neon.cpp new file mode 100644 index 0000000..6f65cbe --- /dev/null +++ b/src/simd/fastscan_neon.cpp @@ -0,0 +1,37 @@ +#include + +#include "rabitqlib/simd/fastscan_dispatch.hpp" + +namespace rabitqlib::fastscan::simd { + +void accumulate_neon( + const std::uint8_t* codes, + const std::uint8_t* lut, + std::uint16_t* result, + std::size_t dim +) { + auto lo0 = vdupq_n_u16(0); + auto lo1 = vdupq_n_u16(0); + auto hi0 = vdupq_n_u16(0); + auto hi1 = vdupq_n_u16(0); + const auto mask = vdupq_n_u8(15); + + for (std::size_t offset = 0; offset < dim * 4; offset += 16) { + const auto packed = vld1q_u8(codes + offset); + const auto table = vld1q_u8(lut + offset); + const auto low = vqtbl1q_u8(table, vandq_u8(packed, mask)); + const auto high = vqtbl1q_u8(table, vshrq_n_u8(packed, 4)); + lo0 = vaddw_u8(lo0, vget_low_u8(low)); + lo1 = vaddw_high_u8(lo1, low); + hi0 = vaddw_u8(hi0, vget_low_u8(high)); + hi1 = vaddw_high_u8(hi1, high); + } + + // Upstream packs lanes as 0,8,1,9,...,7,15; unzip restores vector order. + vst1q_u16(result, vuzp1q_u16(lo0, lo1)); + vst1q_u16(result + 8, vuzp2q_u16(lo0, lo1)); + vst1q_u16(result + 16, vuzp1q_u16(hi0, hi1)); + vst1q_u16(result + 24, vuzp2q_u16(hi0, hi1)); +} + +} // namespace rabitqlib::fastscan::simd diff --git a/src/simd/pack_excode_kernels.hpp b/src/simd/pack_excode_kernels.hpp index b00ac5c..e827ff0 100644 --- a/src/simd/pack_excode_kernels.hpp +++ b/src/simd/pack_excode_kernels.hpp @@ -1,11 +1,11 @@ #pragma once -#include - #include #include #include +#include "rabitqlib/simd/x86_compat.hpp" + namespace rabitqlib::simd::detail { inline void packing_2bit_excode_intrinsics( diff --git a/src/simd/rotator_avx2.cpp b/src/simd/rotator_avx2.cpp index d21d576..aa3773a 100644 --- a/src/simd/rotator_avx2.cpp +++ b/src/simd/rotator_avx2.cpp @@ -1,9 +1,8 @@ -#include - #include #include #include "rabitqlib/simd/rotator_dispatch.hpp" +#include "rabitqlib/simd/x86_compat.hpp" namespace rabitqlib::simd { diff --git a/src/simd/space_avx2.cpp b/src/simd/space_avx2.cpp index a618569..33edd97 100644 --- a/src/simd/space_avx2.cpp +++ b/src/simd/space_avx2.cpp @@ -1,8 +1,8 @@ -#include - #include #include +#include +#include "rabitqlib/simd/x86_compat.hpp" #include "rabitqlib/utils/space.hpp" namespace rabitqlib::simd { @@ -25,7 +25,8 @@ void scalar_quantize_uint8_avx2( __m128i hi32 = _mm256_extracti128_si256(i32, 1); __m128i i16 = _mm_packus_epi32(lo32, hi32); __m128i i8 = _mm_packus_epi16(i16, zero); - _mm_storel_epi64(reinterpret_cast<__m128i*>(&result[i]), i8); + const int64_t packed = _mm_cvtsi128_si64(i8); + std::memcpy(result + i, &packed, sizeof(packed)); } for (; i < dim; ++i) { result[i] = static_cast(std::round((vec0[i] - lo) * one_over_delta)); @@ -65,10 +66,10 @@ void new_transpose_bin_avx2( // 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); + vec_00_to_15 = _mm256_sll_epi32(vec_00_to_15, _mm_cvtsi32_si128(shift)); + vec_16_to_31 = _mm256_sll_epi32(vec_16_to_31, _mm_cvtsi32_si128(shift)); + vec_32_to_47 = _mm256_sll_epi32(vec_32_to_47, _mm_cvtsi32_si128(shift)); + vec_48_to_63 = _mm256_sll_epi32(vec_48_to_63, _mm_cvtsi32_si128(shift)); for (size_t j = 0; j < b_query; ++j) { // pack two 16-bit vectors to 8-bit interleaved vectors diff --git a/src/simd/space_excode_avx2.cpp b/src/simd/space_excode_avx2.cpp index 6a77fbf..d2c9b24 100644 --- a/src/simd/space_excode_avx2.cpp +++ b/src/simd/space_excode_avx2.cpp @@ -1,10 +1,9 @@ -#include - #include #include #include #include "rabitqlib/simd/space_dispatch.hpp" +#include "rabitqlib/simd/x86_compat.hpp" namespace rabitqlib::simd::excode_ipimpl { diff --git a/src/simd/warmup_avx2.cpp b/src/simd/warmup_avx2.cpp index d31e381..12b95a2 100644 --- a/src/simd/warmup_avx2.cpp +++ b/src/simd/warmup_avx2.cpp @@ -1,10 +1,10 @@ -#include - #include #include #include #include +#include "rabitqlib/simd/x86_compat.hpp" + namespace rabitqlib::simd { // Helper: AVX2 64-bit Popcount; Mula's method diff --git a/tests/integration/bit_pack_unpack_test.cpp b/tests/integration/bit_pack_unpack_test.cpp index 5401ae8..5632a71 100644 --- a/tests/integration/bit_pack_unpack_test.cpp +++ b/tests/integration/bit_pack_unpack_test.cpp @@ -48,17 +48,17 @@ class BitPackUnpackTest : public ::testing::Test { } // 4. Helper for Ground Truth Calculation - float CalculateExpected() const { - float expected_result = 0.0f; + double CalculateExpected() const { + double expected_result = 0.0; for (size_t i = 0; i < dim; ++i) { - expected_result += query[i] * static_cast(code[i]); + expected_result += static_cast(query[i]) * code[i]; } return expected_result; } void ExpectIpNear(float result) const { - float expected = CalculateExpected(); - ASSERT_NEAR(expected, result, std::max(0.1f, std::abs(expected) * 1e-6f)); + double expected = CalculateExpected(); + ASSERT_NEAR(expected, result, std::max(0.1, std::abs(expected) * 1e-6)); } }; diff --git a/tests/unit/rabitqlib/utils/arm_backend_test.cpp b/tests/unit/rabitqlib/utils/arm_backend_test.cpp new file mode 100644 index 0000000..2981021 --- /dev/null +++ b/tests/unit/rabitqlib/utils/arm_backend_test.cpp @@ -0,0 +1,69 @@ +#include + +#include +#include +#include +#include + +#include "rabitqlib/fastscan/fastscan.hpp" +#include "rabitqlib/utils/rotator.hpp" + +TEST(FastScanDispatch, MatchesUnpackedRowsWithTailsAndWrapping) { + std::mt19937 gen(42); + for (size_t dim : {16, 64, 512, 1040, 4096}) { + for (size_t count : {1, 31, 32, 33, 97}) { + SCOPED_TRACE(::testing::Message() << "dim=" << dim << " count=" << count); + const size_t batches = (count + 31) / 32; + std::vector rows(count * dim / 8), lut(dim * 4 + 1), + packed(batches * dim * 4 + 1); + for (auto& value : rows) + value = static_cast(gen()); + for (auto& value : lut) + value = static_cast(gen()); + rabitqlib::fastscan::pack_codes(dim, rows.data(), count, packed.data() + 1); + for (size_t batch = 0; batch < batches; ++batch) { + std::array actual{}; + actual.front() = actual.back() = 0xcafe; + rabitqlib::fastscan::accumulate( + packed.data() + 1 + batch * dim * 4, + lut.data() + 1, + actual.data() + 1, + dim + ); + for (size_t lane = 0; lane < 32; ++lane) { + const size_t row = batch * 32 + lane; + uint32_t expected = 0; + for (size_t col = 0; col < dim / 8; ++col) { + const uint8_t code = row < count ? rows[row * dim / 8 + col] : 0; + expected += lut[1 + col * 32 + (code >> 4)] + + lut[1 + col * 32 + 16 + (code & 15)]; + } + EXPECT_EQ(actual[lane + 1], static_cast(expected)); + } + EXPECT_EQ(actual.front(), 0xcafe); + EXPECT_EQ(actual.back(), 0xcafe); + } + } + } +} + +TEST(HadamardBackend, MatchesIndependentSylvesterMatrix) { + std::array original{}, actual{}; + for (size_t i = 0; i < original.size(); ++i) + original[i] = static_cast(i % 17) - 8; + actual = original; + helper_float_9(actual.data()); + for (size_t row = 0; row < actual.size(); ++row) { + float expected = 0; + for (size_t col = 0; col < original.size(); ++col) { + size_t overlap = row & col; + unsigned parity = 0; + while (overlap) { + parity ^= (overlap & 1); + overlap >>= 1; + } + expected += (parity ? -1.0F : 1.0F) * original[col]; + } + EXPECT_EQ(actual[row], expected); + } +} diff --git a/tests/unit/rabitqlib/utils/memory_test.cpp b/tests/unit/rabitqlib/utils/memory_test.cpp new file mode 100644 index 0000000..4545a0d --- /dev/null +++ b/tests/unit/rabitqlib/utils/memory_test.cpp @@ -0,0 +1,15 @@ +#include "rabitqlib/utils/memory.hpp" + +#include + +#include +#include + +TEST(AlignedAllocator, SupportsVectorMoveAssignment) { + using Allocator = rabitqlib::memory::AlignedAllocator; + std::vector source{1, 2, 3}, destination; + destination = std::move(source); + ASSERT_EQ(destination.size(), 3U); + EXPECT_EQ(destination[0], 1); + EXPECT_EQ(destination[2], 3); +} diff --git a/tests/unit/rabitqlib/utils/space_test.cpp b/tests/unit/rabitqlib/utils/space_test.cpp index 9044ed8..c107e17 100644 --- a/tests/unit/rabitqlib/utils/space_test.cpp +++ b/tests/unit/rabitqlib/utils/space_test.cpp @@ -13,6 +13,32 @@ using namespace rabitqlib; +namespace { +bool has_avx2_backend() { +#if defined(__aarch64__) + return true; // The AVX2 implementation is compiled through SIMDe. +#else + return cpu::has_avx2(); +#endif +} +} // namespace + +TEST(ScalarQuantizeUint8, SupportsUnalignedOutputAcrossVectorBlocks) { + constexpr size_t dim = 137; + std::array query{}; + std::array output{}; + output.front() = output.back() = 0xa5; + for (size_t i = 0; i < dim; ++i) { + query[i] = static_cast((i * 13) % 256); + } + simd::scalar_quantize_uint8(output.data() + 1, query.data(), dim, 0.0F, 1.0F); + for (size_t i = 0; i < dim; ++i) { + EXPECT_EQ(output[i + 1], static_cast(query[i])); + } + EXPECT_EQ(output.front(), 0xa5); + EXPECT_EQ(output.back(), 0xa5); +} + TEST(PackBinary, SupportsUnalignedOutput) { constexpr size_t dim = 128; std::array binary_code{}; @@ -52,12 +78,14 @@ TEST(MaskIpX0Q, SupportsUnalignedCodes) { ASSERT_NE(reinterpret_cast(codes) % alignof(uint64_t), 0U); pack_binary(binary_code.data(), codes, dim); - if (cpu::has_avx2()) { + if (has_avx2_backend()) { EXPECT_FLOAT_EQ(simd::mask_ip_x0_q_avx2(query.data(), codes, dim), expected); } +#if !defined(__aarch64__) if (cpu::has_avx512_core()) { EXPECT_FLOAT_EQ(simd::mask_ip_x0_q_avx512(query.data(), codes, dim), expected); } +#endif EXPECT_FLOAT_EQ(mask_ip_x0_q(query.data(), codes, dim), expected); } @@ -96,11 +124,15 @@ TEST(Select_IP_Func, returns_stable_function_pointer) { ip_func = select_excode_ipfunc(8); ASSERT_NE(ip_func, nullptr); +#if !defined(__aarch64__) if (cpu::has_avx512_core()) { ASSERT_EQ(ip_func, simd::excode_ipimpl::ip16_fxu8_avx512); } else { ASSERT_EQ(ip_func, simd::excode_ipimpl::ip16_fxu8_avx2); } +#else + ASSERT_EQ(ip_func, simd::excode_ipimpl::ip16_fxu8_avx2); +#endif } TEST(Select_IP_Func, zero_ex_bits_contributes_nothing) { @@ -167,16 +199,18 @@ TEST(ip16_fxu1_avx, ip_works) { uint8_t codes[dim / 8]; for (size_t i = 0; i < dim; ++i) { - query[i] = static_cast(rand()) / RAND_MAX * 1000.0f; + query[i] = static_cast(rand()) / static_cast(RAND_MAX) * 1000.0f; } for (size_t i = 0; i < dim / 8; ++i) { codes[i] = static_cast(rand() % 256); } - ASSERT_NEAR( - rabitqlib::excode_ipimpl::ip16_fxu1_avx(query, codes, dim), 15055.81f, 0.1f - ); + double expected = 0; + for (size_t i = 0; i < dim; ++i) { + expected += ((codes[i / 8] >> (i % 8)) & 1U) * static_cast(query[i]); + } + ASSERT_NEAR(rabitqlib::excode_ipimpl::ip16_fxu1_avx(query, codes, dim), expected, 0.1f); } TEST(ip64_fxu2_avx, ip_works) { @@ -186,15 +220,19 @@ TEST(ip64_fxu2_avx, ip_works) { uint8_t codes[dim / 4]; for (size_t i = 0; i < dim; ++i) { - query[i] = static_cast(rand()) / RAND_MAX * 1000.0f; + query[i] = static_cast(rand()) / static_cast(RAND_MAX) * 1000.0f; } for (size_t i = 0; i < dim / 4; ++i) { codes[i] = static_cast(rand() % 256); } - ASSERT_NEAR( - rabitqlib::excode_ipimpl::ip64_fxu2_avx(query, codes, dim), 217584.15f, 0.1f - ); + double expected = 0; + for (size_t i = 0; i < dim; ++i) { + const size_t block = i / 64, lane = i % 64; + const uint8_t code = (codes[block * 16 + lane % 16] >> (2 * (lane / 16))) & 3U; + expected += static_cast(query[i]) * code; + } + ASSERT_NEAR(rabitqlib::excode_ipimpl::ip64_fxu2_avx(query, codes, dim), expected, 0.1f); } TEST(OddBitExcodeIp, MatchesScalarInnerProduct) { @@ -229,7 +267,7 @@ TEST(OddBitExcodeIp, MatchesScalarInnerProduct) { } const float expected_float = static_cast(expected); - if (cpu::has_avx2()) { + if (has_avx2_backend()) { const std::array avx2_functions{ nullptr, simd::excode_ipimpl::ip16_fxu1_avx2, @@ -246,6 +284,7 @@ TEST(OddBitExcodeIp, MatchesScalarInnerProduct) { 0.1F ); } +#if !defined(__aarch64__) if (cpu::has_avx512_core()) { const std::array avx512_functions{ nullptr, @@ -263,6 +302,7 @@ TEST(OddBitExcodeIp, MatchesScalarInnerProduct) { 0.1F ); } +#endif } } @@ -281,13 +321,14 @@ TEST(ip_fxu8_avx, ip_works) { const float expected_float = static_cast(expected); ex_ipfunc ip_func = select_excode_ipfunc(8); ASSERT_NEAR(ip_func(query.data(), codes.data(), dim), expected_float, 0.1F); - if (cpu::has_avx2()) { + if (has_avx2_backend()) { ASSERT_NEAR( simd::excode_ipimpl::ip16_fxu8_avx2(query.data(), codes.data(), dim), expected_float, 0.1F ); } +#if !defined(__aarch64__) if (cpu::has_avx512_core()) { ASSERT_NEAR( simd::excode_ipimpl::ip16_fxu8_avx512(query.data(), codes.data(), dim), @@ -295,4 +336,5 @@ TEST(ip_fxu8_avx, ip_works) { 0.1F ); } +#endif }