From 801ffc39ada9ea5007e4f92d55dd943e4adf3d0a Mon Sep 17 00:00:00 2001 From: egolearner Date: Tue, 20 Oct 2026 17:22:04 +0800 Subject: [PATCH] Enable RaBitQ on ARM64 targets and unify DB test linkage --- .github/workflows/04-android-build.yml | 5 ++- .github/workflows/06-ios-build.yml | 8 ++-- CMakeLists.txt | 40 +++++++++++++------ src/core/algorithm/CMakeLists.txt | 10 ++--- .../algorithm/hnsw_rabitq/rabitq_params.h | 5 ++- src/db/index/common/schema.cc | 5 ++- .../hnsw_rabitq/hnsw_rabitq_streamer_test.cc | 5 ++- tests/db/CMakeLists.txt | 29 ++++++++------ tests/db/crash_recovery/CMakeLists.txt | 14 ++----- tests/db/index/CMakeLists.txt | 10 ++--- tests/db/sqlengine/CMakeLists.txt | 13 ++---- tests/db/utf8_acp/CMakeLists.txt | 13 +----- 12 files changed, 78 insertions(+), 79 deletions(-) diff --git a/.github/workflows/04-android-build.yml b/.github/workflows/04-android-build.yml index 5b9d075806..b38300035c 100644 --- a/.github/workflows/04-android-build.yml +++ b/.github/workflows/04-android-build.yml @@ -116,9 +116,10 @@ jobs: -DCMAKE_CXX_COMPILER_LAUNCHER=ccache if [ "${{ matrix.abi }}" = "arm64-v8a" ]; then - echo "Building focused DiskAnn tests and public C++ examples for arm64-v8a..." + echo "Building focused ARM64 index tests and public C++ examples..." cmake --build "$BUILD_DIR" \ - --target diskann_mobile_compat_test diskann_mobile_collection_test zvec_cpp_examples \ + --target diskann_mobile_compat_test diskann_mobile_collection_test \ + hnsw_rabitq_streamer_test ivf_rabitq_streamer_test zvec_cpp_examples \ --parallel exit 0 fi diff --git a/.github/workflows/06-ios-build.yml b/.github/workflows/06-ios-build.yml index 0d4bc90e48..c730c6b202 100644 --- a/.github/workflows/06-ios-build.yml +++ b/.github/workflows/06-ios-build.yml @@ -80,18 +80,20 @@ jobs: fi done - - name: Build required DiskAnn test targets + - name: Build required ARM64 index test targets run: | NPROC=$(sysctl -n hw.ncpu) BUILD_DIR=build_ios_${{ matrix.platform }} cmake --build "$BUILD_DIR" \ --target diskann_mobile_compat_test diskann_mobile_collection_test \ + hnsw_rabitq_streamer_test ivf_rabitq_streamer_test \ --parallel "$NPROC" - for test_name in diskann_mobile_compat_test diskann_mobile_collection_test; do + for test_name in diskann_mobile_compat_test diskann_mobile_collection_test \ + hnsw_rabitq_streamer_test ivf_rabitq_streamer_test; do test_app="$BUILD_DIR/bin/${test_name}.app" if [ ! -d "$test_app" ]; then - echo "::error::Missing required iOS DiskAnn test app: $test_app" + echo "::error::Missing required iOS index test app: $test_app" exit 1 fi done diff --git a/CMakeLists.txt b/CMakeLists.txt index 4999c64552..2bb5f5c748 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -145,14 +145,30 @@ message(STATUS "BUILD_C_BINDINGS:${BUILD_C_BINDINGS}") message(STATUS "BUILD_TOOLS:${BUILD_TOOLS}") message(STATUS "BUILD_CPP_EXAMPLES:${BUILD_CPP_EXAMPLES}") -# RaBitQ support (rabitqlib ships x86-64 AVX2/AVX-512 kernels with runtime dispatch): -# - Linux x86_64 with GCC/Clang -# - Windows x86_64 with MSVC (added upstream in v0.3.8) -# Everything else (ARM, Android, iOS, 32-bit) falls back to the stub libraries. +# RaBitQ support: +# - Linux x86_64 and ARM64 with GCC/Clang +# - macOS x86_64 and ARM64 with AppleClang +# - Windows x86_64 with MSVC +# - Android ARM64 and iOS ARM64 with Clang +# ARM64 uses NEON; x86-64 uses runtime AVX2/AVX-512 dispatch. +# Other targets, including Android x86_64, iOS x86_64, Windows ARM64, and +# 32-bit platforms, use stub libraries. +set(RABITQ_TARGET_PROCESSOR "${CMAKE_SYSTEM_PROCESSOR}") +if(APPLE AND CMAKE_OSX_ARCHITECTURES) + set(RABITQ_TARGET_PROCESSOR "${CMAKE_OSX_ARCHITECTURES}") +endif() +string(TOLOWER "${RABITQ_TARGET_PROCESSOR}" RABITQ_TARGET_PROCESSOR) set(RABITQ_SUPPORTED OFF) -if(NOT ANDROID AND NOT IOS AND CMAKE_SIZEOF_VOID_P EQUAL 8 - AND CMAKE_SYSTEM_PROCESSOR MATCHES "^(x86_64|amd64|AMD64)$" - AND (CMAKE_SYSTEM_NAME STREQUAL "Linux" OR (WIN32 AND MSVC))) +if(CMAKE_SIZEOF_VOID_P EQUAL 8 + AND RABITQ_TARGET_PROCESSOR MATCHES "^(arm64|aarch64)$" + AND (CMAKE_SYSTEM_NAME STREQUAL "Linux" OR + CMAKE_SYSTEM_NAME STREQUAL "Darwin" OR + (ANDROID AND ANDROID_ABI STREQUAL "arm64-v8a") OR IOS)) + set(RABITQ_SUPPORTED ON) +elseif(NOT ANDROID AND NOT IOS AND CMAKE_SIZEOF_VOID_P EQUAL 8 + AND RABITQ_TARGET_PROCESSOR MATCHES "^(x86_64|amd64)$" + AND (CMAKE_SYSTEM_NAME STREQUAL "Linux" OR + CMAKE_SYSTEM_NAME STREQUAL "Darwin" OR (WIN32 AND MSVC))) if(MSVC) # /arch:AVX512 turns on the complete F/CD/VL/BW/DQ group. MSVC has no # separate switch for VPOPCNTDQ, so the popcount backend is compiled with @@ -181,15 +197,15 @@ if(NOT ANDROID AND NOT IOS AND CMAKE_SIZEOF_VOID_P EQUAL 8 endforeach() unset(_rabitq_probe) - if(RABITQ_SUPPORTED) - add_definitions(-DRABITQ_SUPPORTED=1) - else() - add_definitions(-DRABITQ_SUPPORTED=0) + if(NOT RABITQ_SUPPORTED) message(STATUS "RaBitQ support disabled - compiler cannot build all runtime dispatch variants") endif() +endif() +if(RABITQ_SUPPORTED) + add_definitions(-DRABITQ_SUPPORTED=1) else() add_definitions(-DRABITQ_SUPPORTED=0) - message(STATUS "RaBitQ support disabled - only supported on Linux x86_64 and Windows x86_64 (MSVC)") + message(STATUS "RaBitQ support disabled - unsupported platform or compiler") endif() message(STATUS "RABITQ_SUPPORTED: ${RABITQ_SUPPORTED}") diff --git a/src/core/algorithm/CMakeLists.txt b/src/core/algorithm/CMakeLists.txt index bb2560e4ba..a386180893 100644 --- a/src/core/algorithm/CMakeLists.txt +++ b/src/core/algorithm/CMakeLists.txt @@ -56,10 +56,10 @@ if(RABITQ_SUPPORTED) cc_directory(ivf_rabitq) else() message(STATUS "NOT BUILD RABITQ") - # Empty stub library for unsupported platforms + # Empty stub library when RaBitQ is disabled for this build file(WRITE ${CMAKE_CURRENT_BINARY_DIR}/rabitq_stub.cc - "// Stub implementation for unsupported platforms\n" - "// RaBitQ only supports Linux x86_64 and Windows x86_64\n" + "// Stub implementation for build compatibility\n" + "// RaBitQ is disabled for this build\n" "namespace zvec { namespace core { /* empty namespace for compatibility */ } }\n" ) @@ -73,8 +73,8 @@ else() ) file(WRITE ${CMAKE_CURRENT_BINARY_DIR}/ivf_rabitq_stub.cc - "// Stub implementation for unsupported platforms\n" - "// IVF RaBitQ only supports Linux x86_64 and Windows x86_64\n" + "// Stub implementation for build compatibility\n" + "// IVF RaBitQ is disabled for this build\n" "namespace zvec { namespace core { /* empty namespace for compatibility */ } }\n" ) diff --git a/src/core/algorithm/hnsw_rabitq/rabitq_params.h b/src/core/algorithm/hnsw_rabitq/rabitq_params.h index b4a6900eca..515f938d91 100644 --- a/src/core/algorithm/hnsw_rabitq/rabitq_params.h +++ b/src/core/algorithm/hnsw_rabitq/rabitq_params.h @@ -45,9 +45,10 @@ constexpr int kMinRabitqDimSize = 64; constexpr int kMaxRabitqDimSize = 4095; // Reported wherever RaBitQ is compiled out. Keep in sync with RABITQ_SUPPORTED -// in the top-level CMakeLists.txt; this is the only place naming the platforms. +// in the top-level CMakeLists.txt. constexpr char kRabitqUnsupportedPlatform[] = - "RaBitQ is not supported on this platform (Linux/Windows x86_64)"; + "RaBitQ is not supported on this platform (supported: Linux " + "x86_64/ARM64, macOS x86_64/ARM64, Windows x86_64, Android/iOS ARM64)"; // Original vector dimension before converter (e.g., CosineNormalizeConverter) // modifies it. Used by IVF-RaBitQ to ignore extra dimensions from converter. diff --git a/src/db/index/common/schema.cc b/src/db/index/common/schema.cc index 3b48d38658..bea3f86328 100644 --- a/src/db/index/common/schema.cc +++ b/src/db/index/common/schema.cc @@ -180,9 +180,10 @@ Status FieldSchema::validate() const { const bool supports_rabitq_avx2 = flags.AVX2 && flags.FMA; const bool supports_rabitq_avx512 = flags.AVX512F && flags.AVX512BW && flags.AVX512DQ; - if (!supports_rabitq_avx2 && !supports_rabitq_avx512) { + if (!supports_rabitq_avx2 && !supports_rabitq_avx512 && !flags.NEON) { return Status::NotSupported( - "RabitQ requires AVX2/FMA or AVX512F/BW/DQ to be supported"); + "RabitQ requires AVX2/FMA, AVX512F/BW/DQ, or NEON to be " + "supported"); } } diff --git a/tests/core/algorithm/hnsw_rabitq/hnsw_rabitq_streamer_test.cc b/tests/core/algorithm/hnsw_rabitq/hnsw_rabitq_streamer_test.cc index 1b8552ea82..deace54cd1 100644 --- a/tests/core/algorithm/hnsw_rabitq/hnsw_rabitq_streamer_test.cc +++ b/tests/core/algorithm/hnsw_rabitq/hnsw_rabitq_streamer_test.cc @@ -46,8 +46,9 @@ std::string HnswRabitqStreamerTest::dir_("hnswRabitqStreamerTest"); shared_ptr HnswRabitqStreamerTest::index_meta_ptr_; void HnswRabitqStreamerTest::SetUp() { - if (!rabitqlib::cpu::has_avx512_core() && !rabitqlib::cpu::has_avx2()) { - GTEST_SKIP() << "CPU does not support AVX2/FMA or AVX512F/BW/DQ"; + if (!rabitqlib::cpu::has_avx512_core() && !rabitqlib::cpu::has_avx2() && + !rabitqlib::cpu::has_neon()) { + GTEST_SKIP() << "CPU does not support AVX2/FMA, AVX512F/BW/DQ, or NEON"; } index_meta_ptr_.reset(new (nothrow) diff --git a/tests/db/CMakeLists.txt b/tests/db/CMakeLists.txt index 219d979d21..2bee05936d 100644 --- a/tests/db/CMakeLists.txt +++ b/tests/db/CMakeLists.txt @@ -1,6 +1,20 @@ include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) include(${PROJECT_ROOT_DIR}/cmake/option.cmake) +function(db_test_force_load_core TARGET) + # Include core registries and keep their RTTI in the test executable. + if(MSVC) + target_link_options(${TARGET} PRIVATE + "/WHOLEARCHIVE:$") + elseif(APPLE) + target_link_options(${TARGET} PRIVATE + "-Wl,-force_load,$") + else() + target_link_options(${TARGET} PRIVATE + "-Wl,--whole-archive,$,--no-whole-archive") + endif() +endfunction() + cc_directory(common) cc_directories(crash_recovery) cc_directory(sqlengine) @@ -38,24 +52,13 @@ foreach(CC_SRCS ${ALL_TEST_SRCS}) cc_gmock( NAME ${CC_TARGET} STRICT DEFS DISKANN_STRESS_TESTS=${DISKANN_STRESS_TESTS} - LIBS zvec - core_knn_flat - core_knn_flat_sparse - core_knn_hnsw - core_knn_hnsw_rabitq - core_knn_ivf_rabitq - core_knn_hnsw_sparse - core_knn_ivf - core_knn_diskann - core_mix_reducer - core_metric - core_utility - core_quantizer + LIBS zvec zvec_core ${CMAKE_THREAD_LIBS_INIT} ${CMAKE_DL_LIBS} SRCS ${CC_SRCS} index/utils/utils.cc INCS . .. ../../src LDFLAGS ${APPLE_FRAMEWORK_LIBS} ) + db_test_force_load_core(${CC_TARGET}) cc_test_suite(zvec ${CC_TARGET}) endforeach() diff --git a/tests/db/crash_recovery/CMakeLists.txt b/tests/db/crash_recovery/CMakeLists.txt index 4e28c277cf..3e3bf9fa16 100644 --- a/tests/db/crash_recovery/CMakeLists.txt +++ b/tests/db/crash_recovery/CMakeLists.txt @@ -27,16 +27,7 @@ endif() # Common libraries set(CRASH_RECOVERY_COMMON_LIBS zvec - core_knn_flat - core_knn_flat_sparse - core_knn_hnsw - core_knn_hnsw_sparse - core_knn_ivf - core_knn_hnsw_rabitq - core_mix_reducer - core_metric - core_utility - core_quantizer + zvec_core ${CMAKE_THREAD_LIBS_INIT} ${CMAKE_DL_LIBS} ) @@ -51,6 +42,7 @@ cc_binary( LDFLAGS ${APPLE_FRAMEWORK_LIBS} ) set_target_properties(data_generator PROPERTIES EXCLUDE_FROM_ALL TRUE) +db_test_force_load_core(data_generator) # Build collection_optimizer executable @@ -62,6 +54,7 @@ cc_binary( LDFLAGS ${APPLE_FRAMEWORK_LIBS} ) set_target_properties(collection_optimizer PROPERTIES EXCLUDE_FROM_ALL TRUE) +db_test_force_load_core(collection_optimizer) # Build test executables @@ -75,6 +68,7 @@ foreach(CC_SRCS ${ALL_TEST_SRCS}) INCS .. ../../src LDFLAGS ${APPLE_FRAMEWORK_LIBS} ) + db_test_force_load_core(${CC_TARGET}) add_dependencies(${CC_TARGET} data_generator) add_dependencies(${CC_TARGET} collection_optimizer) set_tests_properties(${CC_TARGET} PROPERTIES diff --git a/tests/db/index/CMakeLists.txt b/tests/db/index/CMakeLists.txt index ccf9663f25..25cdecdb28 100644 --- a/tests/db/index/CMakeLists.txt +++ b/tests/db/index/CMakeLists.txt @@ -38,19 +38,15 @@ foreach(CC_SRCS ${ALL_TEST_SRCS}) NAME ${CC_TARGET} STRICT LIBS zvec zvec_ailego - core_knn_diskann - core_metric_static - core_utility_static - core_quantizer_static - core_knn_hnsw core_knn_hnsw_sparse sparsehash - core_knn_flat core_knn_flat_sparse core_knn_ivf - core_knn_hnsw_rabitq core_knn_ivf_rabitq core_mix_reducer + zvec_core + sparsehash ${CMAKE_THREAD_LIBS_INIT} ${CMAKE_DL_LIBS} SRCS ${CC_SRCS} utils/utils.cc INCS . .. ../../src LDFLAGS ${APPLE_FRAMEWORK_LIBS} ) + db_test_force_load_core(${CC_TARGET}) cc_test_suite(zvec_index ${CC_TARGET}) endforeach() diff --git a/tests/db/sqlengine/CMakeLists.txt b/tests/db/sqlengine/CMakeLists.txt index 42a3ee260c..8543679ac3 100644 --- a/tests/db/sqlengine/CMakeLists.txt +++ b/tests/db/sqlengine/CMakeLists.txt @@ -18,25 +18,18 @@ if(APPLE) endif() file(GLOB ALL_TEST_SRCS *_test.cc) + foreach(CC_SRCS ${ALL_TEST_SRCS}) get_filename_component(CC_TARGET ${CC_SRCS} NAME_WE) cc_gmock( NAME ${CC_TARGET} STRICT - LIBS zvec_common - zvec_sqlengine - zvec - zvec_ailego - core_metric - core_utility - core_quantizer - core_knn_hnsw core_knn_hnsw_sparse sparsehash - core_knn_flat core_knn_flat_sparse core_knn_ivf - core_knn_hnsw_rabitq core_mix_reducer + LIBS zvec zvec_core zvec_ailego sparsehash ${CMAKE_THREAD_LIBS_INIT} ${CMAKE_DL_LIBS} SRCS ${CC_SRCS} INCS . ${PROJECT_ROOT_DIR}/src ${PROJECT_ROOT_DIR}/src/db ${PROJECT_ROOT_DIR}/src/db/common LDFLAGS ${APPLE_FRAMEWORK_LIBS} ) + db_test_force_load_core(${CC_TARGET}) cc_test_suite(zvec_sqlengine ${CC_TARGET}) endforeach() diff --git a/tests/db/utf8_acp/CMakeLists.txt b/tests/db/utf8_acp/CMakeLists.txt index 4f669de014..cc15d1101c 100644 --- a/tests/db/utf8_acp/CMakeLists.txt +++ b/tests/db/utf8_acp/CMakeLists.txt @@ -9,17 +9,7 @@ endif() set(UTF8_ACP_LIBS zvec - core_knn_flat - core_knn_flat_sparse - core_knn_hnsw - core_knn_hnsw_rabitq - core_knn_hnsw_sparse - core_knn_ivf - core_knn_diskann - core_mix_reducer - core_metric - core_utility - core_quantizer + zvec_core ${CMAKE_THREAD_LIBS_INIT} ${CMAKE_DL_LIBS} ) @@ -50,6 +40,7 @@ foreach(UTF8_ACP_VARIANT ${UTF8_ACP_VARIANTS}) INCS ${CMAKE_CURRENT_SOURCE_DIR}/.. DEFS ZVEC_TEST_EXPECTED_ACP=${ZVEC_ACP_CODE} ) + db_test_force_load_core(${ZVEC_ACP_TARGET}) # CMake otherwise performs another mt.exe pass for incremental linking. The # manifest is already embedded explicitly, so disable that redundant pass.