Make knowhere run on old cpus which only support sse4.2 (#144)

Signed-off-by: chasingegg <[email protected]>
zilliztech · Oct 13, 2023 · 32f8308 · 32f8308
1 parent 6ed164e
commit 32f8308
Show file tree

Hide file tree

Showing 9 changed files with 69 additions and 15 deletions.
diff --git a/cmake/libs/libdiskann.cmake b/cmake/libs/libdiskann.cmake
@@ -32,6 +32,6 @@ target_link_libraries(diskann PUBLIC ${AIO_LIBRARIES}
 if(__X86_64)
   target_compile_options(
     diskann PRIVATE -fno-builtin-malloc -fno-builtin-calloc
-                    -fno-builtin-realloc -fno-builtin-free -mavx2 -DUSE_AVX2)
+                    -fno-builtin-realloc -fno-builtin-free)
 endif()
 list(APPEND KNOWHERE_LINKER_LIBS diskann)
diff --git a/cmake/libs/libfaiss.cmake b/cmake/libs/libfaiss.cmake
@@ -6,6 +6,11 @@ knowhere_file_glob(
 knowhere_file_glob(GLOB FAISS_AVX512_SRCS
                    thirdparty/faiss/faiss/impl/*avx512.cpp)
 
+knowhere_file_glob(GLOB FAISS_AVX2_SRCS
+                   thirdparty/faiss/faiss/impl/*avx.cpp
+                   thirdparty/faiss/faiss/impl/pq4_fast_scan_search_qbs.cpp
+                   thirdparty/faiss/faiss/IndexIVFPQFastScan.cpp)
+
 list(REMOVE_ITEM FAISS_SRCS ${FAISS_AVX512_SRCS})
 
 if(__X86_64)
@@ -18,10 +23,10 @@ if(__X86_64)
   add_library(utils_avx OBJECT ${UTILS_AVX_SRC})
   add_library(utils_avx512 OBJECT ${UTILS_AVX512_SRC})
 
-  target_compile_options(utils_sse PRIVATE -msse4.2)
-  target_compile_options(utils_avx PRIVATE -mf16c -mavx2)
-  target_compile_options(utils_avx512 PRIVATE -mf16c -mavx512f -mavx512dq
-                                              -mavx512bw)
+  target_compile_options(utils_sse PRIVATE -msse4.2 -mpopcnt)
+  target_compile_options(utils_avx PRIVATE -mfma -mf16c -mavx2 -mpopcnt)
+  target_compile_options(utils_avx512 PRIVATE -mfma -mf16c -mavx512f -mavx512dq
+                                              -mavx512bw -mpopcnt)
 
   add_library(
     knowhere_utils STATIC
@@ -52,6 +57,17 @@ endif()
 find_package(LAPACK REQUIRED)
 
 if(__X86_64)
+  list(REMOVE_ITEM FAISS_SRCS ${FAISS_AVX2_SRCS})
+
+  add_library(faiss_avx2 OBJECT ${FAISS_AVX2_SRCS})
+  target_compile_options(
+    faiss_avx2
+    PRIVATE $<$<COMPILE_LANGUAGE:CXX>:
+            -msse4.2
+            -mavx2
+            -mfma
+            -mf16c
+            -mpopcnt>)
   add_library(faiss_avx512 OBJECT ${FAISS_AVX512_SRCS})
   target_compile_options(
     faiss_avx512
@@ -60,19 +76,19 @@ if(__X86_64)
             -mavx2
             -mfma
             -mf16c
+            -mavx512f
             -mavx512dq
-            -mavx512bw>)
+            -mavx512bw
+            -mpopcnt>)
 
   add_library(faiss STATIC ${FAISS_SRCS})
 
-  add_dependencies(faiss faiss_avx512 knowhere_utils)
+  add_dependencies(faiss faiss_avx2 faiss_avx512 knowhere_utils)
   target_compile_options(
     faiss
     PRIVATE $<$<COMPILE_LANGUAGE:CXX>:
             -msse4.2
-            -mavx2
-            -mfma
-            -mf16c
+            -mpopcnt
             -Wno-sign-compare
             -Wno-unused-variable
             -Wno-reorder
@@ -81,7 +97,7 @@ if(__X86_64)
             -Wno-strict-aliasing>)
   target_link_libraries(
     faiss PUBLIC OpenMP::OpenMP_CXX ${BLAS_LIBRARIES} ${LAPACK_LIBRARIES}
-                 faiss_avx512 knowhere_utils)
+                 faiss_avx2 faiss_avx512 knowhere_utils)
   target_compile_definitions(faiss PRIVATE FINTEGER=int)
 endif()
 

diff --git a/include/knowhere/config.h b/include/knowhere/config.h
@@ -564,7 +564,7 @@ class BaseConfig : public Config {
     }
 
     virtual Status
-    CheckAndAdjustForRangeSearch() {
+    CheckAndAdjustForRangeSearch(std::string* err_msg) {
         return Status::success;
     }
 

diff --git a/include/knowhere/expected.h b/include/knowhere/expected.h
@@ -39,6 +39,7 @@ enum class Status {
     arithmetic_overflow = 17,
     raft_inner_error = 18,
     invalid_binary_set = 19,
+    invalid_instruction_set = 20,
 };
 
 inline std::string
@@ -80,6 +81,8 @@ Status2String(knowhere::Status status) {
             return "raft inner error";
         case knowhere::Status::invalid_binary_set:
             return "invalid binary set";
+        case knowhere::Status::invalid_instruction_set:
+            return "the current index is not supported on the current CPU model";
         default:
             return "unexpected status";
     }

diff --git a/src/common/index.cc b/src/common/index.cc
@@ -126,9 +126,9 @@ Index<T>::RangeSearch(const DataSet& dataset, const Json& json, const BitsetView
     if (status != Status::success) {
         return expected<DataSetPtr>::Err(status, std::move(msg));
     }
-    status = cfg->CheckAndAdjustForRangeSearch();
+    status = cfg->CheckAndAdjustForRangeSearch(&msg);
     if (status != Status::success) {
-        return expected<DataSetPtr>::Err(status, "invalid params for range search");
+        return expected<DataSetPtr>::Err(status, std::move(msg));
     }
 
 #ifdef NOT_COMPILE_FOR_SWIG

diff --git a/src/index/hnsw/hnsw_config.h b/src/index/hnsw/hnsw_config.h
@@ -72,7 +72,7 @@ class HnswConfig : public BaseConfig {
     }
 
     inline Status
-    CheckAndAdjustForRangeSearch() override {
+    CheckAndAdjustForRangeSearch(std::string* err_msg) override {
         if (!ef.has_value()) {
             // if ef is not set by user, set it to default
             ef = kDefaultRangeSearchEf;

diff --git a/src/index/ivf/ivf_config.h b/src/index/ivf/ivf_config.h
@@ -13,6 +13,7 @@
 #define IVF_CONFIG_H
 
 #include "knowhere/config.h"
+#include "simd/hook.h"
 
 namespace knowhere {
 class IvfConfig : public BaseConfig {
@@ -75,6 +76,11 @@ class ScannConfig : public IvfFlatConfig {
 
     inline Status
     CheckAndAdjustForSearch(std::string* err_msg) override {
+        if (!faiss::support_pq_fast_scan) {
+            *err_msg = "SCANN index is not supported on the current CPU model, avx2 support is needed for x86 arch.";
+            LOG_KNOWHERE_ERROR_ << *err_msg;
+            return Status::invalid_instruction_set;
+        }
         if (!reorder_k.has_value()) {
             reorder_k = k.value();
         } else if (reorder_k.value() < k.value()) {
@@ -86,6 +92,26 @@ class ScannConfig : public IvfFlatConfig {
 
         return Status::success;
     }
+
+    inline Status
+    CheckAndAdjustForRangeSearch(std::string* err_msg) override {
+        if (!faiss::support_pq_fast_scan) {
+            *err_msg = "SCANN index is not supported on the current CPU model, avx2 support is needed for x86 arch.";
+            LOG_KNOWHERE_ERROR_ << *err_msg;
+            return Status::invalid_instruction_set;
+        }
+        return Status::success;
+    }
+
+    inline Status
+    CheckAndAdjustForBuild() override {
+        if (!faiss::support_pq_fast_scan) {
+            LOG_KNOWHERE_ERROR_
+                << "SCANN index is not supported on the current CPU model, avx2 support is needed for x86 arch.";
+            return Status::invalid_instruction_set;
+        }
+        return Status::success;
+    }
 };
 
 class IvfSqConfig : public IvfConfig {};

diff --git a/src/simd/hook.cc b/src/simd/hook.cc
@@ -39,6 +39,8 @@ bool use_avx2 = true;
 bool use_sse4_2 = true;
 #endif
 
+bool support_pq_fast_scan = true;
+
 decltype(fvec_inner_product) fvec_inner_product = fvec_inner_product_ref;
 decltype(fvec_L2sqr) fvec_L2sqr = fvec_L2sqr_ref;
 decltype(fvec_L1) fvec_L1 = fvec_L1_ref;
@@ -87,6 +89,7 @@ fvec_hook(std::string& simd_type) {
         fvec_madd_and_argmin = fvec_madd_and_argmin_sse;
 
         simd_type = "AVX512";
+        support_pq_fast_scan = true;
     } else if (use_avx2 && cpu_support_avx2()) {
         fvec_inner_product = fvec_inner_product_avx;
         fvec_L2sqr = fvec_L2sqr_avx;
@@ -100,6 +103,7 @@ fvec_hook(std::string& simd_type) {
         fvec_madd_and_argmin = fvec_madd_and_argmin_sse;
 
         simd_type = "AVX2";
+        support_pq_fast_scan = true;
     } else if (use_sse4_2 && cpu_support_sse4_2()) {
         fvec_inner_product = fvec_inner_product_sse;
         fvec_L2sqr = fvec_L2sqr_sse;
@@ -113,6 +117,7 @@ fvec_hook(std::string& simd_type) {
         fvec_madd_and_argmin = fvec_madd_and_argmin_sse;
 
         simd_type = "SSE4_2";
+        support_pq_fast_scan = false;
     } else {
         fvec_inner_product = fvec_inner_product_ref;
         fvec_L2sqr = fvec_L2sqr_ref;
@@ -126,6 +131,7 @@ fvec_hook(std::string& simd_type) {
         fvec_madd_and_argmin = fvec_madd_and_argmin_ref;
 
         simd_type = "GENERIC";
+        support_pq_fast_scan = false;
     }
 #endif
 
@@ -142,6 +148,7 @@ fvec_hook(std::string& simd_type) {
     fvec_madd_and_argmin = fvec_madd_and_argmin_neon;
 
     simd_type = "NEON";
+    support_pq_fast_scan = true;
 
 #endif
 }

diff --git a/src/simd/hook.h b/src/simd/hook.h
@@ -31,6 +31,8 @@ extern bool use_avx2;
 extern bool use_sse4_2;
 #endif
 
+extern bool support_pq_fast_scan;
+
 #if defined(__x86_64__)
 bool
 cpu_support_avx512();