Refactor/add scalar dist function and setup cmake with detailed cpu flag

CMakeLists.txt

@@ -21,8 +21,8 @@ include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake)
 include(${PROJECT_ROOT_DIR}/cmake/option.cmake)
 
 if(NOT ANDROID AND AUTO_DETECT_ARCH AND CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|i686|i386|x64")
-  setup_compiler_march_for_x86(MATH_MARCH_FLAG_SSE MATH_MARCH_FLAG_AVX2 MATH_MARCH_FLAG_AVX512)
-  message(STATUS "best compiler march, sse: " ${MATH_MARCH_FLAG_SSE} ", avx2: " ${MATH_MARCH_FLAG_AVX2} ", avx512: " ${MATH_MARCH_FLAG_AVX512})
+  setup_compiler_march_for_x86(MATH_MARCH_FLAG_SSE MATH_MARCH_FLAG_AVX2 MATH_MARCH_FLAG_AVX512 MATH_MARCH_FLAG_AVX512FP16)
+  message(STATUS "best compiler march, sse: " ${MATH_MARCH_FLAG_SSE} ", avx2: " ${MATH_MARCH_FLAG_AVX2} ", avx512: " ${MATH_MARCH_FLAG_AVX512} ", avx512fp16: " ${MATH_MARCH_FLAG_AVX512FP16})
 endif()
 
 include_directories(${PROJECT_ROOT_DIR}/src/include)

cmake/option.cmake

@@ -35,8 +35,8 @@ option(ENABLE_OPENMP "Enable OpenMP support" OFF)
 
 set(ARCH_OPTIONS
   ENABLE_NEHALEM ENABLE_SANDYBRIDGE ENABLE_HASWELL ENABLE_BROADWELL ENABLE_SKYLAKE
-  ENABLE_SKYLAKE_AVX512 ENABLE_ICELAKE ENABLE_SAPPHIRERAPIDS ENABLE_EMERALDRAPIDS
-  ENABLE_GRANITERAPIDS ENABLE_ZEN1 ENABLE_ZEN2 ENABLE_ZEN3
+  ENABLE_SKYLAKE_AVX512 ENABLE_ICELAKE ENABLE_SAPPHIRERAPIDS ENABLE_EMERALDRAPIDS ENABLE_GRANITERAPIDS
+  ENABLE_ZEN1 ENABLE_ZEN2 ENABLE_ZEN3
   ENABLE_ARMV8A ENABLE_ARMV8.1A ENABLE_ARMV8.2A ENABLE_ARMV8.3A ENABLE_ARMV8.4A
   ENABLE_ARMV8.5A ENABLE_ARMV8.6A
   ENABLE_NATIVE
@@ -103,30 +103,34 @@ function(_setup_x86_march)
   endif()
 endfunction()
 
-function(setup_compiler_march_for_x86 VAR_NAME_SSE VAR_NAME_AVX2 VAR_NAME_AVX512)
+function(setup_compiler_march_for_x86 VAR_NAME_SSE VAR_NAME_AVX2 VAR_NAME_AVX512 VAR_NAME_AVX512FP16)
   #sse
   set(${VAR_NAME_SSE} "-march=corei7" PARENT_SCOPE)
 
   #avx 2
   set(${VAR_NAME_AVX2} "-march=core-avx2" PARENT_SCOPE)
 
   #avx512
-  set(_x86_flags
-    "graniterapids" "emeraldrapids" "sapphirerapids"
-    "icelake-server" "skylake-avx512"
-  )
-  foreach(_arch IN LISTS _x86_flags)
-    check_c_compiler_flag("-march=${_arch}" _COMP_SUPP_${_arch})
-    if(_COMP_SUPP_${_arch})
-      set(${VAR_NAME_AVX512} "-march=${_arch}" PARENT_SCOPE)
-      return()
+  set(_x86_flags_avx512 "icelake-server" "skylake-avx512" "core-avx2" "x86-64")
+  foreach(_arch_avx512 IN LISTS _x86_flags_avx512)
+    check_c_compiler_flag("-march=${_arch_avx512}" _COMP_SUPP_${_arch_avx512})
+    if(_COMP_SUPP_${_arch_avx512})
+      set(${VAR_NAME_AVX512} "-march=${_arch_avx512}" PARENT_SCOPE)
+      break()
     endif()
   endforeach()
 
-
-  set(${VAR_NAME_AVX512} "-march=core-avx2" PARENT_SCOPE)
-  message(WARNING "No known avx512 microarchitecture flag found. Set up as core-avx2")
-
+  #avx512fp16
+  set(_x86_flags_avx512fp16
+    "sapphirerapids" "icelake-server" "skylake-avx512" "core-avx2" "x86-64"
+  )
+  foreach(_arch_avx512fp16 IN LISTS _x86_flags_avx512fp16)
+    check_c_compiler_flag("-march=${_arch_avx512fp16}" _COMP_SUPP_${_arch_avx512fp16})
+    if(_COMP_SUPP_${_arch_avx512fp16})
+      set(${VAR_NAME_AVX512FP16} "-march=${_arch_avx512fp16}" PARENT_SCOPE)
+      break()
+    endif()
+  endforeach()
 endfunction()
 
 if(MSVC)

src/ailego/CMakeLists.txt

@@ -20,8 +20,8 @@ endif()
 
 if(NOT ANDROID AND AUTO_DETECT_ARCH)
     if(CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|i686|i386|x64")
-        setup_compiler_march_for_x86(MATH_MARCH_FLAG_SSE MATH_MARCH_FLAG_AVX2 MATH_MARCH_FLAG_AVX512)
-        message(STATUS "best compiler march, sse: " ${MATH_MARCH_FLAG_SSE} ", avx2: " ${MATH_MARCH_FLAG_AVX2} ", avx512: " ${MATH_MARCH_FLAG_AVX512})
+        setup_compiler_march_for_x86(MATH_MARCH_FLAG_SSE MATH_MARCH_FLAG_AVX2 MATH_MARCH_FLAG_AVX512 MATH_MARCH_FLAG_AVX512FP16)
+        message(STATUS "best compiler march, sse: " ${MATH_MARCH_FLAG_SSE} ", avx2: " ${MATH_MARCH_FLAG_AVX2} ", avx512: " ${MATH_MARCH_FLAG_AVX512} ", avx512fp16: " ${MATH_MARCH_FLAG_AVX512FP16})
 
         file(GLOB_RECURSE MATH_FILES_SSE
             ${CMAKE_CURRENT_SOURCE_DIR}/math/*_sse.cc
@@ -42,16 +42,23 @@ if(NOT ANDROID AND AUTO_DETECT_ARCH)
         )
 
         file(GLOB_RECURSE MATH_FILES_AVX512
-            ${CMAKE_CURRENT_SOURCE_DIR}/math/*_dispatch.cc
-            ${CMAKE_CURRENT_SOURCE_DIR}/math/*_dispatch.c
             ${CMAKE_CURRENT_SOURCE_DIR}/math/*_avx512.cc
             ${CMAKE_CURRENT_SOURCE_DIR}/math/*_avx512.c
-            ${CMAKE_CURRENT_SOURCE_DIR}/math_batch/*_dispatch.cc
-            ${CMAKE_CURRENT_SOURCE_DIR}/math_batch/*_dispatch.c
             ${CMAKE_CURRENT_SOURCE_DIR}/math_batch/*_avx512.cc
             ${CMAKE_CURRENT_SOURCE_DIR}/math_batch/*_avx512.c
         )
 
+        file(GLOB_RECURSE MATH_FILES_AVX512FP16
+          ${CMAKE_CURRENT_SOURCE_DIR}/math/*_dispatch.cc
+          ${CMAKE_CURRENT_SOURCE_DIR}/math/*_dispatch.c
+          ${CMAKE_CURRENT_SOURCE_DIR}/math/*_avx512fp16.cc
+          ${CMAKE_CURRENT_SOURCE_DIR}/math/*_avx512fp16.c
+          ${CMAKE_CURRENT_SOURCE_DIR}/math_batch/*_dispatch.cc
+          ${CMAKE_CURRENT_SOURCE_DIR}/math_batch/*_dispatch.c
+          ${CMAKE_CURRENT_SOURCE_DIR}/math_batch/*_avx512fp16.cc
+          ${CMAKE_CURRENT_SOURCE_DIR}/math_batch/*_avx512fp16.c
+        )
+
         foreach(MATH_FILE ${MATH_FILES_SSE})
             set_source_files_properties(
                 ${MATH_FILE}
@@ -75,6 +82,14 @@ if(NOT ANDROID AND AUTO_DETECT_ARCH)
                 COMPILE_FLAGS "${MATH_MARCH_FLAG_AVX512}"
             )
         endforeach()
+
+        foreach(MATH_FILE ${MATH_FILES_AVX512FP16})
+        set_source_files_properties(
+            ${MATH_FILE}
+            PROPERTIES
+            COMPILE_FLAGS "${MATH_MARCH_FLAG_AVX512FP16}"
+        )
+    endforeach()
     elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64|arm64|ARM64")
       # set(CMAKE_CXX_FLAGS "-march=armv8-a")
       # set(CMAKE_C_FLAGS "-march=armv8-a")

src/ailego/math/euclidean_distance_matrix.h

@@ -22,6 +22,9 @@
 namespace zvec {
 namespace ailego {
 
+//--------------------------------------------------
+// Dense
+//--------------------------------------------------
 /*! Squared Euclidean Distance Matrix
  */
 template <typename T, size_t M, size_t N, typename = void>
@@ -48,6 +51,46 @@ struct SquaredEuclideanDistanceMatrix<
   }
 };
 
+template <>
+struct SquaredEuclideanDistanceMatrix<uint8_t, 1, 1> {
+  //! Type of value
+  using ValueType = uint8_t;
+
+  //! Compute the distance between matrix and query
+  static void Compute(const ValueType *m, const ValueType *q, size_t dim,
+                      float *out);
+};
+
+template <>
+struct SquaredEuclideanDistanceMatrix<int8_t, 1, 1> {
+  //! Type of value
+  using ValueType = int8_t;
+
+  //! Compute the distance between matrix and query
+  static void Compute(const ValueType *m, const ValueType *q, size_t dim,
+                      float *out);
+};
+
+template <>
+struct SquaredEuclideanDistanceMatrix<Float16, 1, 1> {
+  //! Type of value
+  using ValueType = Float16;
+
+  //! Compute the distance between matrix and query
+  static void Compute(const ValueType *m, const ValueType *q, size_t dim,
+                      float *out);
+};
+
+template <>
+struct SquaredEuclideanDistanceMatrix<float, 1, 1> {
+  //! Type of value
+  using ValueType = float;
+
+  //! Compute the distance between matrix and query
+  static void Compute(const ValueType *m, const ValueType *q, size_t dim,
+                      float *out);
+};
+
 /*! Squared Euclidean Distance Matrix
  */
 template <typename T, size_t M, size_t N>
@@ -353,32 +396,6 @@ struct SquaredEuclideanDistanceMatrix<uint8_t, M, 1,
   }
 };
 
-#if !defined(__SSE4_1__)
-/*! Squared Euclidean Distance Matrix (INT4, M=1, N=1)
- */
-template <>
-struct SquaredEuclideanDistanceMatrix<uint8_t, 1, 1> {
-  //! Type of value
-  using ValueType = uint8_t;
-
-  //! Compute the distance between matrix and query
-  static inline void Compute(const ValueType *m, const ValueType *q, size_t dim,
-                             float *out) {
-    ailego_assert(m && q && dim && !(dim & 1) && out);
-
-    float sum = 0.0;
-    for (size_t i = 0; i < (dim >> 1); ++i) {
-      uint8_t m_val = m[i];
-      uint8_t q_val = q[i];
-      sum +=
-          Int4SquaredDiffTable[((m_val << 4) & 0xf0) | ((q_val >> 0) & 0xf)] +
-          Int4SquaredDiffTable[((m_val >> 0) & 0xf0) | ((q_val >> 4) & 0xf)];
-    }
-    *out = sum;
-  }
-};
-#endif  // !__SSE4_1__
-
 /*! Euclidean Distance Matrix
  */
 template <typename T, size_t M, size_t N,
@@ -424,76 +441,26 @@ struct EuclideanDistanceMatrix<
   }
 };
 
-#if !defined(__SSE4_1__)
-/*! Euclidean Distance Matrix (INT4, M=1, N=1)
- */
 template <>
 struct EuclideanDistanceMatrix<uint8_t, 1, 1> {
   //! Type of value
   using ValueType = uint8_t;
 
-  //! Compute the distance between matrix and query
-  static inline void Compute(const ValueType *m, const ValueType *q, size_t dim,
-                             float *out) {
-    ailego_assert(m && q && dim && !(dim & 1) && out);
-
-    float sum = 0.0;
-    for (size_t i = 0; i < (dim >> 1); ++i) {
-      uint8_t m_val = m[i];
-      uint8_t q_val = q[i];
-      sum +=
-          Int4SquaredDiffTable[((m_val << 4) & 0xf0) | ((q_val >> 0) & 0xf)] +
-          Int4SquaredDiffTable[((m_val >> 0) & 0xf0) | ((q_val >> 4) & 0xf)];
-    }
-    *out = std::sqrt(sum);
-  }
-};
-#endif  // !__SSE4_1__
-
-#if defined(__SSE__) || defined(__ARM_NEON)
-/*! Squared Euclidean Distance Matrix (FP32, M=1, N=1)
- */
-template <>
-struct SquaredEuclideanDistanceMatrix<float, 1, 1> {
-  //! Type of value
-  using ValueType = float;
-
-  //! Compute the distance between matrix and query
-  static void Compute(const ValueType *m, const ValueType *q, size_t dim,
-                      float *out);
-};
-#endif  // __SSE__ || __ARM_NEON
-
-#if defined(__SSE__) || (defined(__ARM_NEON) && (defined(__aarch64__)))
-/*! Euclidean Distance Matrix (FP32, M=1, N=1)
- */
-template <>
-struct EuclideanDistanceMatrix<float, 1, 1> {
-  //! Type of value
-  using ValueType = float;
-
   //! Compute the distance between matrix and query
   static void Compute(const ValueType *m, const ValueType *q, size_t dim,
                       float *out);
 };
-#endif  // __SSE__ || __ARM_NEON  && __aarch64__
 
-#if (defined(__F16C__) && defined(__AVX__)) || \
-    (defined(__ARM_NEON) && defined(__aarch64__))
-/*! Squared Euclidean Distance Matrix (FP16, M=1, N=1)
- */
 template <>
-struct SquaredEuclideanDistanceMatrix<Float16, 1, 1> {
+struct EuclideanDistanceMatrix<int8_t, 1, 1> {
   //! Type of value
-  using ValueType = Float16;
+  using ValueType = int8_t;
 
   //! Compute the distance between matrix and query
   static void Compute(const ValueType *m, const ValueType *q, size_t dim,
                       float *out);
 };
 
-/*! Euclidean Distance Matrix (FP16, M=1, N=1)
- */
 template <>
 struct EuclideanDistanceMatrix<Float16, 1, 1> {
   //! Type of value
@@ -503,58 +470,21 @@ struct EuclideanDistanceMatrix<Float16, 1, 1> {
   static void Compute(const ValueType *m, const ValueType *q, size_t dim,
                       float *out);
 };
-#endif  // (__F16C__ && __AVX__) || (__ARM_NEON && __aarch64__)
 
-#if defined(__SSE4_1__)
-/*! Squared Euclidean Distance Matrix (INT8, M=1, N=1)
- */
 template <>
-struct SquaredEuclideanDistanceMatrix<int8_t, 1, 1> {
-  //! Type of value
-  using ValueType = int8_t;
-
-  //! Compute the distance between matrix and query
-  static void Compute(const ValueType *m, const ValueType *q, size_t dim,
-                      float *out);
-};
-
-/*! Euclidean Distance Matrix (INT8, M=1, N=1)
- */
-template <>
-struct EuclideanDistanceMatrix<int8_t, 1, 1> {
-  //! Type of value
-  using ValueType = int8_t;
-
-  //! Compute the distance between matrix and query
-  static void Compute(const ValueType *m, const ValueType *q, size_t dim,
-                      float *out);
-};
-
-/*! Squared Euclidean Distance Matrix (INT4, M=1, N=1)
- */
-template <>
-struct SquaredEuclideanDistanceMatrix<uint8_t, 1, 1> {
+struct EuclideanDistanceMatrix<float, 1, 1> {
   //! Type of value
-  using ValueType = uint8_t;
+  using ValueType = float;
 
   //! Compute the distance between matrix and query
   static void Compute(const ValueType *m, const ValueType *q, size_t dim,
                       float *out);
 };
 
-/*! Euclidean Distance Matrix (INT4, M=1, N=1)
- */
-template <>
-struct EuclideanDistanceMatrix<uint8_t, 1, 1> {
-  //! Type of value
-  using ValueType = uint8_t;
-
-  //! Compute the distance between matrix and query
-  static void Compute(const ValueType *m, const ValueType *q, size_t dim,
-                      float *out);
-};
-#endif  // __SSE4_1__
 
+//--------------------------------------------------
+// Sparse
+//--------------------------------------------------
 /*! Squared Euclidean Distance Sparse Matrix
  */
 template <typename T>

src/ailego/math/euclidean_distance_matrix_fp16_avx.cc

@@ -21,15 +21,13 @@ namespace ailego {
 
 #if defined(__AVX__)
 
-void SquaredEuclideanDistanceAVX(const Float16 *lhs, const Float16 *rhs,
-                                 size_t size, float *out) {
-  ACCUM_FP16_1X1_AVX(lhs, rhs, size, out, 0ull, )
-}
+float SquaredEuclideanDistanceFp16AVX(const Float16 *lhs, const Float16 *rhs,
+                                      size_t size) {
+  float score{0.0f};
+
+  ACCUM_FP16_1X1_AVX(lhs, rhs, size, &score, 0ull, )
 
-//! EuclideanDistance
-void EuclideanDistanceAVX(const Float16 *lhs, const Float16 *rhs, size_t size,
-                          float *out) {
-  ACCUM_FP16_1X1_AVX(lhs, rhs, size, out, 0ull, std::sqrt)
+  return score;
 }
 
 #endif  // __AVX__