diff --git a/CMakeLists.txt b/CMakeLists.txt index 956ee5599..28fbd0866 100644 --- a/CMakeLists.txt +++ b/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) diff --git a/src/core/CMakeLists.txt b/src/core/CMakeLists.txt index d875176f2..9cf9877cb 100644 --- a/src/core/CMakeLists.txt +++ b/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 diff --git a/src/core/algorithm/hnsw_rabitq/CMakeLists.txt b/src/core/algorithm/hnsw_rabitq/CMakeLists.txt index 09ce72f55..aaeaa9467 100644 --- a/src/core/algorithm/hnsw_rabitq/CMakeLists.txt +++ b/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}" - ) \ No newline at end of file + ) diff --git a/src/db/index/common/schema.cc b/src/db/index/common/schema.cc index 532696803..69da579a7 100644 --- a/src/db/index/common/schema.cc +++ b/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> 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"); } } 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 1e5827648..2501a734e 100644 --- a/tests/core/algorithm/hnsw_rabitq/hnsw_rabitq_streamer_test.cc +++ b/tests/core/algorithm/hnsw_rabitq/hnsw_rabitq_streamer_test.cc @@ -16,6 +16,7 @@ #include #include #include +#include #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 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()); diff --git a/tests/core/algorithm/hnsw_rabitq/rabitq_dispatch_test.cc b/tests/core/algorithm/hnsw_rabitq/rabitq_dispatch_test.cc new file mode 100644 index 000000000..53df68837 --- /dev/null +++ b/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 +#include +#include +#include + +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 diff --git a/thirdparty/RaBitQ-Library/CMakeLists.txt b/thirdparty/RaBitQ-Library/CMakeLists.txt index 6f0ee938c..ec152a2a2 100644 --- a/thirdparty/RaBitQ-Library/CMakeLists.txt +++ b/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() diff --git a/thirdparty/RaBitQ-Library/RaBitQ-Library-0.1 b/thirdparty/RaBitQ-Library/RaBitQ-Library-0.1 index 858b0d6c4..540242ea0 160000 --- a/thirdparty/RaBitQ-Library/RaBitQ-Library-0.1 +++ b/thirdparty/RaBitQ-Library/RaBitQ-Library-0.1 @@ -1 +1 @@ -Subproject commit 858b0d6c480766d0e4f08fc5e02f34b53d698fad +Subproject commit 540242ea0a68926f1b827bf1f9add844f07a427b