feat(rabitq): enable runtime AVX2 and AVX512 dispatch

CMakeLists.txt

@@ -89,27 +89,26 @@ message(STATUS "BUILD_PYTHON_BINDINGS:${BUILD_PYTHON_BINDINGS}")
 message(STATUS "BUILD_C_BINDINGS:${BUILD_C_BINDINGS}")
 message(STATUS "BUILD_TOOLS:${BUILD_TOOLS}")
 
-option(RABITQ_ENABLE_AVX512 "Compile RaBitQ with AVX-512 support" OFF)
-
 if(CMAKE_SYSTEM_NAME STREQUAL "Linux" AND CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|amd64|AMD64" AND NOT ANDROID AND NOT IOS)
-  include(CheckCCompilerFlag)
-
-  check_c_compiler_flag("-mavx2" COMPILER_SUPPORTS_AVX2)
-  check_c_compiler_flag("-mavx512f -mavx512bw -mavx512vl" COMPILER_SUPPORTS_AVX512)
-
-  if(COMPILER_SUPPORTS_AVX2 OR COMPILER_SUPPORTS_AVX512)
+  include(CheckCXXCompilerFlag)
+
+  check_cxx_compiler_flag("-mavx2 -mfma" COMPILER_SUPPORTS_RABITQ_AVX2)
+  check_cxx_compiler_flag(
+      "-mavx512f -mavx512bw -mavx512dq -mfma"
+      COMPILER_SUPPORTS_RABITQ_AVX512)
+  check_cxx_compiler_flag(
+      "-mavx512f -mavx512bw -mavx512dq -mavx512vpopcntdq -mfma"
+      COMPILER_SUPPORTS_RABITQ_AVX512_POPCNT)
+
+  if(COMPILER_SUPPORTS_RABITQ_AVX2 AND
+     COMPILER_SUPPORTS_RABITQ_AVX512 AND
+     COMPILER_SUPPORTS_RABITQ_AVX512_POPCNT)
     set(RABITQ_SUPPORTED ON)
     add_definitions(-DRABITQ_SUPPORTED=1)
-    if(RABITQ_ENABLE_AVX512 AND COMPILER_SUPPORTS_AVX512)
-      add_definitions(-DRABITQ_COMPILED_AVX512=1)
-      set(RABITQ_ARCH_FLAG "${MATH_MARCH_FLAG_AVX512}")
-    else()
-      set(RABITQ_ARCH_FLAG "${MATH_MARCH_FLAG_AVX2}")
-    endif()
   else()
     set(RABITQ_SUPPORTED OFF)
     add_definitions(-DRABITQ_SUPPORTED=0)
-    message(STATUS "RaBitQ support disabled - compiler does not support AVX2 or AVX-512")
+    message(STATUS "RaBitQ support disabled - compiler cannot build all runtime dispatch variants")
   endif()
 elseif(IOS)
   set(RABITQ_SUPPORTED OFF)
@@ -120,7 +119,7 @@ else()
   add_definitions(-DRABITQ_SUPPORTED=0)
   message(STATUS "RaBitQ support disabled - only supported on Linux x86_64")
 endif()
-message(STATUS "RABITQ_ARCH_FLAG: ${RABITQ_ARCH_FLAG}")
+message(STATUS "RABITQ_SUPPORTED: ${RABITQ_SUPPORTED}")
 
 # DiskAnn support (Linux x86_64 only; libaio loaded at runtime via dlopen)
 if(CMAKE_SYSTEM_NAME STREQUAL "Linux" AND CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|i686|i386" AND NOT ANDROID AND NOT IOS)

src/core/CMakeLists.txt

@@ -1,26 +1,6 @@
 include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake)
 include(${PROJECT_ROOT_DIR}/cmake/option.cmake)
 
-if(RABITQ_SUPPORTED AND AUTO_DETECT_ARCH)
-  set(HNSW_RABITQ_FILES
-      hnsw_rabitq_query_algorithm.cc
-      hnsw_rabitq_streamer.cc
-      hnsw_rabitq_searcher.cc
-      hnsw_rabitq_entity.cc
-      rabitq_reformer.cc
-      rabitq_converter.cc
-  )
-  set(HNSW_RABITQ_FILES_FULL ${HNSW_RABITQ_FILES})
-  list(TRANSFORM HNSW_RABITQ_FILES_FULL PREPEND "algorithm/hnsw_rabitq/")
-  foreach(FILE ${HNSW_RABITQ_FILES_FULL})
-      set_source_files_properties(
-          ${FILE}
-          PROPERTIES
-          COMPILE_FLAGS "${RABITQ_ARCH_FLAG}"
-      )
-  endforeach()
-endif()
-
 # utility/block_heap.cc uses AVX2 intrinsics guarded by __AVX2__. When the
 # host toolchain supports it, compile this source with an AVX2-capable
 # -march so AVX2 codegen is emitted. zvec_core glob-collects this source

src/core/algorithm/hnsw_rabitq/CMakeLists.txt

@@ -1,16 +1,6 @@
 include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake)
 include(${PROJECT_ROOT_DIR}/cmake/option.cmake)
 
-if(AUTO_DETECT_ARCH)
-  foreach(FILE ${HNSW_RABITQ_FILES})
-      set_source_files_properties(
-          ${FILE}
-          PROPERTIES
-          COMPILE_FLAGS "${RABITQ_ARCH_FLAG}"
-      )
-  endforeach()
-endif()
-
 if(NOT APPLE)
   set(CORE_KNN_HNSW_RABITQ_LDFLAGS
       "-Wl,--exclude-libs,libparquet.a:libarrow.a:libarrow_bundled_dependencies.a")
@@ -24,4 +14,4 @@ cc_library(
     INCS . ${PROJECT_ROOT_DIR}/src ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm
     LDFLAGS "${CORE_KNN_HNSW_RABITQ_LDFLAGS}"
     VERSION "${PROXIMA_ZVEC_VERSION}"
-  )
+  )

src/db/index/common/schema.cc

@@ -30,12 +30,6 @@
 
 namespace zvec {
 
-#if defined(RABITQ_COMPILED_AVX512)
-constexpr const int kRabitqCompiledAvx512 = RABITQ_COMPILED_AVX512;
-#else
-constexpr const int kRabitqCompiledAvx512 = 0;
-#endif
-
 std::unordered_map<DataType, std::set<QuantizeType>> quantize_type_map = {
     {DataType::VECTOR_FP32,
      {QuantizeType::FP16, QuantizeType::INT4, QuantizeType::INT8,
@@ -184,16 +178,12 @@ Status FieldSchema::validate() const {
             "RabitQ is not supported on this platform (Linux x86_64 only)");
 #endif
         auto &flags = zvec::ailego::internal::CpuFeatures::static_flags_;
-        if (!flags.AVX2 && !flags.AVX512F) {
+        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) {
           return Status::NotSupported(
-              "RabitQ requires AVX2/AVX512F to be supported");
-        }
-
-        if constexpr (kRabitqCompiledAvx512) {
-          if (!flags.AVX512F) {
-            return Status::NotSupported(
-                "RabitQ compiled with AVX512F while runtime does not support");
-          }
+              "RabitQ requires AVX2/FMA or AVX512F/BW/DQ to be supported");
         }
       }
 

tests/core/algorithm/hnsw_rabitq/hnsw_rabitq_streamer_test.cc

@@ -16,6 +16,7 @@
 #include <memory>
 #include <string>
 #include <gtest/gtest.h>
+#include <rabitqlib/utils/cpu_features.hpp>
 #include "zvec/ailego/container/params.h"
 #include "zvec/ailego/utility/file_helper.h"
 #include "zvec/core/framework/index_holder.h"
@@ -45,6 +46,10 @@ std::string HnswRabitqStreamerTest::dir_("hnswRabitqStreamerTest");
 shared_ptr<IndexMeta> HnswRabitqStreamerTest::index_meta_ptr_;
 
 void HnswRabitqStreamerTest::SetUp(void) {
+  if (!rabitqlib::cpu::has_avx512_core() && !rabitqlib::cpu::has_avx2()) {
+    GTEST_SKIP() << "CPU does not support AVX2/FMA or AVX512F/BW/DQ";
+  }
+
   index_meta_ptr_.reset(new (nothrow)
                             IndexMeta(IndexMeta::DataType::DT_FP32, dim));
   index_meta_ptr_->set_metric("SquaredEuclidean", 0, ailego::Params());

tests/core/algorithm/hnsw_rabitq/rabitq_dispatch_test.cc

@@ -0,0 +1,41 @@
+// Copyright 2025-present the zvec project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <gtest/gtest.h>
+#include <rabitqlib/simd/dispatch.hpp>
+#include <rabitqlib/simd/space_dispatch.hpp>
+#include <rabitqlib/utils/cpu_features.hpp>
+
+namespace zvec {
+namespace core {
+
+TEST(RabitqDispatchTest, SelectsBestAvailableExcodeImplementation) {
+  if (!rabitqlib::cpu::has_avx512_core() && !rabitqlib::cpu::has_avx2()) {
+    GTEST_SKIP() << "CPU does not support AVX2/FMA or AVX512F/BW/DQ";
+  }
+
+  constexpr size_t kExBits = 2;
+  const auto table = rabitqlib::simd::resolve_excode_ip_table();
+
+  if (rabitqlib::cpu::has_avx512_core()) {
+    EXPECT_EQ(table[kExBits],
+              &rabitqlib::simd::excode_ipimpl::ip64_fxu2_avx512);
+  } else {
+    ASSERT_TRUE(rabitqlib::cpu::has_avx2());
+    EXPECT_EQ(table[kExBits], &rabitqlib::simd::excode_ipimpl::ip64_fxu2_avx2);
+  }
+}
+
+}  // namespace core
+}  // namespace zvec

thirdparty/RaBitQ-Library/CMakeLists.txt

@@ -2,3 +2,79 @@ add_library(rabitqlib INTERFACE)
 target_include_directories(
     rabitqlib SYSTEM INTERFACE "${CMAKE_CURRENT_SOURCE_DIR}/RaBitQ-Library-0.1/include"
   )
+
+if(RABITQ_SUPPORTED)
+  set(RABITQ_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/RaBitQ-Library-0.1")
+
+  set(RABITQ_COMMON_SOURCES
+      "${RABITQ_SOURCE_DIR}/src/utils/cpu_features.cpp"
+      "${RABITQ_SOURCE_DIR}/src/simd/dispatch.cpp"
+  )
+  set(RABITQ_AVX2_SOURCES
+      "${RABITQ_SOURCE_DIR}/src/simd/pack_excode_avx2.cpp"
+      "${RABITQ_SOURCE_DIR}/src/simd/space_excode_avx2.cpp"
+      "${RABITQ_SOURCE_DIR}/src/simd/space_avx2.cpp"
+      "${RABITQ_SOURCE_DIR}/src/simd/fastscan_avx2.cpp"
+      "${RABITQ_SOURCE_DIR}/src/simd/warmup_avx2.cpp"
+      "${RABITQ_SOURCE_DIR}/src/simd/rotator_avx2.cpp"
+      "${RABITQ_SOURCE_DIR}/src/index/hnsw_search_avx2.cpp"
+  )
+  set(RABITQ_AVX512_SOURCES
+      "${RABITQ_SOURCE_DIR}/src/simd/pack_excode_avx512.cpp"
+      "${RABITQ_SOURCE_DIR}/src/simd/space_excode_avx512.cpp"
+      "${RABITQ_SOURCE_DIR}/src/simd/space_avx512.cpp"
+      "${RABITQ_SOURCE_DIR}/src/simd/fastscan_avx512.cpp"
+      "${RABITQ_SOURCE_DIR}/src/simd/rotator_avx512.cpp"
+  )
+  set(RABITQ_AVX512_POPCNT_SOURCES
+      "${RABITQ_SOURCE_DIR}/src/simd/warmup_avx512.cpp"
+  )
+  set(RABITQ_HNSW_AVX512_SOURCES
+      "${RABITQ_SOURCE_DIR}/src/index/hnsw_search_avx512_core.cpp"
+  )
+  set(RABITQ_HNSW_AVX512_POPCNT_SOURCES
+      "${RABITQ_SOURCE_DIR}/src/index/hnsw_search_avx512_popcnt.cpp"
+  )
+
+  add_library(
+      rabitqlib_runtime STATIC
+      ${RABITQ_COMMON_SOURCES}
+      ${RABITQ_AVX2_SOURCES}
+      ${RABITQ_AVX512_SOURCES}
+      ${RABITQ_AVX512_POPCNT_SOURCES}
+      ${RABITQ_HNSW_AVX512_SOURCES}
+      ${RABITQ_HNSW_AVX512_POPCNT_SOURCES}
+  )
+  target_include_directories(
+      rabitqlib_runtime SYSTEM PUBLIC "${RABITQ_SOURCE_DIR}/include"
+  )
+  set_target_properties(
+      rabitqlib_runtime PROPERTIES POSITION_INDEPENDENT_CODE ON
+  )
+
+  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_AVX512_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"
+  )
+
+  target_link_libraries(rabitqlib INTERFACE rabitqlib_runtime)
+endif()