refactor: turbo based refactor [WIP]

src/core/framework/index_meta.cc

@@ -30,7 +30,8 @@ struct IndexMetaFormatHeader {
   uint32_t space_id;
   uint32_t attachment_offset;
   uint32_t attachment_size;
-  uint8_t reserved_[4092];
+  uint32_t extra_meta_size;
+  uint8_t reserved_[4088];
 };
 
 static_assert(sizeof(IndexMetaFormatHeader) % 32 == 0,
@@ -47,6 +48,7 @@ void IndexMeta::serialize(std::string *out) const {
   format.dimension = dimension_;
   format.unit_size = unit_size_;
   format.space_id = space_id_;
+  format.extra_meta_size = extra_meta_size_;
 
   if (!metric_name_.empty()) {
     ailego::Params item;

src/core/metric/quantized_integer_metric.cc

@@ -96,24 +96,44 @@ class QuantizedIntegerMetric : public IndexMetric {
     switch (origin_metric_type_) {
       case MetricType::kSquaredEuclidean:
         if (meta_.data_type() == IndexMeta::DataType::DT_INT8) {
-          auto turbo_ret = turbo::get_distance_func(
-              turbo::MetricType::kSquaredEuclidean, turbo::DataType::kInt8,
-              turbo::QuantizeType::kDefault);
+          auto turbo_ret =
+              turbo::get_distance_func(turbo::MetricType::kSquaredEuclidean,
+                                       turbo::DataType::kInt8, quantize_type_);
           if (turbo_ret && m == 1 && n == 1) {
             return turbo_ret;
           }
           return DistanceMatrixCompute<SquaredEuclidean, int8_t>(m, n);
         }
         if (meta_.data_type() == IndexMeta::DataType::DT_INT4) {
+          auto turbo_ret =
+              turbo::get_distance_func(turbo::MetricType::kSquaredEuclidean,
+                                       turbo::DataType::kInt4, quantize_type_);
+          if (turbo_ret && m == 1 && n == 1) {
+            return turbo_ret;
+          }
+
           return DistanceMatrixCompute<SquaredEuclidean, uint8_t>(m, n);
         }
         break;
 
       case MetricType::kInnerProduct:
         if (meta_.data_type() == IndexMeta::DataType::DT_INT8) {
+          auto turbo_ret =
+              turbo::get_distance_func(turbo::MetricType::kInnerProduct,
+                                       turbo::DataType::kInt8, quantize_type_);
+          if (turbo_ret && m == 1 && n == 1) {
+            return turbo_ret;
+          }
           return DistanceMatrixCompute<MinusInnerProduct, int8_t>(m, n);
         }
+
         if (meta_.data_type() == IndexMeta::DataType::DT_INT4) {
+          auto turbo_ret =
+              turbo::get_distance_func(turbo::MetricType::kInnerProduct,
+                                       turbo::DataType::kInt4, quantize_type_);
+          if (turbo_ret && m == 1 && n == 1) {
+            return turbo_ret;
+          }
           return DistanceMatrixCompute<MinusInnerProduct, uint8_t>(m, n);
         }
         break;
@@ -137,9 +157,9 @@ class QuantizedIntegerMetric : public IndexMetric {
         break;
       case MetricType::kCosine:
         if (meta_.data_type() == IndexMeta::DataType::DT_INT8) {
-          auto turbo_ret = turbo::get_distance_func(
-              turbo::MetricType::kCosine, turbo::DataType::kInt8,
-              turbo::QuantizeType::kDefault);
+          auto turbo_ret =
+              turbo::get_distance_func(turbo::MetricType::kCosine,
+                                       turbo::DataType::kInt8, quantize_type_);
           if (turbo_ret) {
             return turbo_ret;
           }
@@ -160,7 +180,7 @@ class QuantizedIntegerMetric : public IndexMetric {
         if (meta_.data_type() == IndexMeta::DataType::DT_INT8) {
           auto turbo_ret = turbo::get_batch_distance_func(
               turbo::MetricType::kSquaredEuclidean, turbo::DataType::kInt8,
-              turbo::QuantizeType::kDefault);
+              quantize_type_);
           if (turbo_ret) {
             return turbo_ret;
           }
@@ -215,7 +235,7 @@ class QuantizedIntegerMetric : public IndexMetric {
         if (meta_.data_type() == IndexMeta::DataType::DT_INT8) {
           auto turbo_ret = turbo::get_batch_distance_func(
               turbo::MetricType::kCosine, turbo::DataType::kInt8,
-              turbo::QuantizeType::kDefault);
+              quantize_type_);
           if (turbo_ret) {
             return turbo_ret;
           }

src/include/zvec/core/framework/index_meta.h

@@ -40,6 +40,7 @@ class IndexMeta {
     DT_BINARY64 = 8,
   };
 
+
   /*! Major Orders
    */
   enum MajorOrder {
@@ -586,6 +587,7 @@ class IndexMeta {
   uint32_t dimension_{0};
   uint32_t unit_size_{0};
   uint32_t element_size_{0};
+  uint32_t extra_meta_size_{0};
   uint64_t space_id_{0};
   uint32_t metric_revision_{0};
   uint32_t converter_revision_{0};
@@ -708,6 +710,8 @@ class IndexQueryMeta {
   uint32_t dimension_{0};
   uint32_t unit_size_{0};
   uint32_t element_size_{0};
+  uint32_t extra_meta_size_{0};
+  uint32_t quantize_type_{0};
 };
 
 }  // namespace core

src/include/zvec/core/framework/index_metric.h

@@ -137,6 +137,9 @@ struct IndexMetric : public IndexModule {
   virtual DistanceBatchQueryPreprocessFunc get_query_preprocess_func() const {
     return nullptr;
   }
+
+ private:
+  int quantize_type_{0};
 };
 
 }  // namespace core

src/include/zvec/turbo/turbo.h

@@ -28,28 +28,51 @@ using QueryPreprocessFunc =
 enum class MetricType {
   kSquaredEuclidean,
   kCosine,
+  kInnerProduct,
   kMipsSquaredEuclidean,
   kUnknown,
 };
 
 enum class DataType {
+  kInt4,
   kInt8,
+  kFp16,
+  kFp32,
   kUnknown,
 };
 
 enum class QuantizeType {
   kDefault,
+  kRecordInt8,
+  kRecordInt4,
+  kInt8,
+  kInt4,
+  kFp16,
+  kPQ,
+  kRabit
+};
+
+enum class CpuArchType {
+  kAuto,
+  kScalar,
+  kSSE,
+  kAVX,
+  kAVX2,
+  kAVX512,
+  kAVX512VNNI,
+  kAVX512FP16
 };
 
 DistanceFunc get_distance_func(MetricType metric_type, DataType data_type,
-                               QuantizeType quantize_type);
+                               QuantizeType quantize_type,
+                               CpuArchType cpu_arch_type = CpuArchType::kAuto);
 
-BatchDistanceFunc get_batch_distance_func(MetricType metric_type,
-                                          DataType data_type,
-                                          QuantizeType quantize_type);
+BatchDistanceFunc get_batch_distance_func(
+    MetricType metric_type, DataType data_type, QuantizeType quantize_type,
+    CpuArchType cpu_arch_type = CpuArchType::kAuto);
 
-QueryPreprocessFunc get_query_preprocess_func(MetricType metric_type,
-                                              DataType data_type,
-                                              QuantizeType quantize_type);
+QueryPreprocessFunc get_query_preprocess_func(
+    MetricType metric_type, DataType data_type, QuantizeType quantize_type,
+    CpuArchType cpu_arch_type = CpuArchType::kAuto);
 
 }  // namespace zvec::turbo

src/turbo/CMakeLists.txt

@@ -13,18 +13,71 @@ endif()
 
 file(GLOB_RECURSE ALL_SRCS *.cc *.c *.h)
 
-# Set per-file compile flags for AVX512-VNNI sources.
-# set_source_files_properties is directory-scoped, so it must be called in the
-# same directory that adds the sources to a target (i.e. here, not in a
-# subdirectory).
 if(NOT ANDROID AND AUTO_DETECT_ARCH)
     if (HOST_ARCH MATCHES "^(x86|x64)$")
-        file(GLOB_RECURSE AVX512_VNNI_SRCS ${CMAKE_CURRENT_SOURCE_DIR}/avx512_vnni/*.cc)
+        file(GLOB_RECURSE AVX512_AVX512FP16_SRCS 
+          ${CMAKE_CURRENT_SOURCE_DIR}/avx512_fp16/*.cc
+          ${CMAKE_CURRENT_SOURCE_DIR}/avx512_fp16/*.c)
+        set_source_files_properties(
+            ${AVX512_AVX512FP16_SRCS}
+            PROPERTIES
+            COMPILE_FLAGS "${TURBO_MARCH_FLAG_AVX512FP16}"
+        )
+
+        # Set per-file compile flags for AVX512-VNNI sources.
+        # set_source_files_properties is directory-scoped, so it must be called in the
+        # same directory that adds the sources to a target (i.e. here, not in a
+        # subdirectory).
+        file(GLOB_RECURSE AVX512_VNNI_SRCS 
+          ${CMAKE_CURRENT_SOURCE_DIR}/avx512_vnni/*.cc
+          ${CMAKE_CURRENT_SOURCE_DIR}/avx512_vnni/*.c)
         set_source_files_properties(
             ${AVX512_VNNI_SRCS}
             PROPERTIES
             COMPILE_FLAGS "${TURBO_MARCH_FLAG_AVX512}"
         )
+
+        file(GLOB_RECURSE AVX512_SRCS 
+          ${CMAKE_CURRENT_SOURCE_DIR}/avx512/*.cc
+          ${CMAKE_CURRENT_SOURCE_DIR}/avx512/*.c)
+        set_source_files_properties(
+            ${AVX512_SRCS}
+            PROPERTIES
+            COMPILE_FLAGS "${TURBO_MARCH_FLAG_AVX512}"
+        )
+
+        file(GLOB_RECURSE AVX2_SRCS 
+          ${CMAKE_CURRENT_SOURCE_DIR}/avx2/*.cc 
+          ${CMAKE_CURRENT_SOURCE_DIR}/avx2/*.c
+          ${CMAKE_CURRENT_SOURCE_DIR}/avx/*.cc
+          ${CMAKE_CURRENT_SOURCE_DIR}/avx/*.c)
+        set_source_files_properties(
+            ${AVX2_SRCS}
+            PROPERTIES
+            COMPILE_FLAGS "${TURBO_MARCH_FLAG_AVX2}"
+        )
+
+        file(GLOB_RECURSE SSE_SRCS 
+          ${CMAKE_CURRENT_SOURCE_DIR}/sse/*.cc
+          ${CMAKE_CURRENT_SOURCE_DIR}/sse/*.c)
+        set_source_files_properties(
+            ${SSE_SRCS}
+            PROPERTIES
+            COMPILE_FLAGS "${TURBO_MARCH_FLAG_SSE}"
+        )
+    elseif (HOST_ARCH MATCHES "^(arm|arm64)$")
+        set(TURBO_MARCH_FLAG_NEON "-march=armv8-a")
+
+        file(GLOB_RECURSE NEON_SRCS
+          ${CMAKE_CURRENT_SOURCE_DIR}/armv8/*.cc
+          ${CMAKE_CURRENT_SOURCE_DIR}/armv8/*.c
+        )
+
+        set_source_files_properties(
+            ${NEON_SRCS}
+            PROPERTIES
+            COMPILE_FLAGS "${TURBO_MARCH_FLAG_NEON}"
+        )
     endif()
 endif()
 

src/turbo/armv8/float32/cosine.cc

@@ -0,0 +1,62 @@
+// 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 "armv8/float32/cosine.h"
+#include "armv8/float32/inner_product.h"
+#include "armv8/float32/inner_product_common.h"
+
+namespace zvec::turbo::armv8 {
+
+void cosine_fp32_distance(const void *a, const void *b, size_t dim,
+                          size_t extra_size, float *distance) {
+#if defined(__ARM_NEON)
+  constexpr size_t extra_dim = 2;
+  size_t original_dim = dim - extra_dim;
+
+  float ip;
+  internal::inner_product_fp32_armv8(a, b, original_dim, &ip);
+
+  *distance = 1 - ip;
+#else
+  (void)a;
+  (void)b;
+  (void)dim;
+  (void)distance;
+#endif  // __ARM_NEON
+}
+
+void cosine_fp32_batch_distance(const void *const *vectors, const void *query,
+                                size_t n, size_t dim, float *distances) {
+#if defined(__ARM_NEON)
+  const int original_dim = dim - 1;
+  if (original_dim <= 0) {
+    return;
+  }
+
+  internal::inner_product_fp32_batch_armv8(vectors, query, n, original_dim,
+                                           distances);
+
+  for (int i = 0; i < n; ++i) {
+    distances[i] = 1 - distances[i];
+  }
+#else
+  (void)vectors;
+  (void)query;
+  (void)n;
+  (void)dim;
+  (void)distances;
+#endif  //__ARM_NEON
+}
+
+}  // namespace zvec::turbo::armv8

src/turbo/armv8/float32/cosine.h

@@ -0,0 +1,30 @@
+// 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.
+
+#pragma once
+
+#include <cstddef>
+
+namespace zvec::turbo::armv8 {
+
+// Compute cosine distance (negative inner product after normalization) between
+// a single quantized FP32 vector pair.
+void cosine_fp32_distance(const void *a, const void *b, size_t dim,
+                          float *distance);
+
+// Batch version of cosine_fp32_distance.
+void cosine_fp32_batch_distance(const void *const *vectors, const void *query,
+                                size_t n, size_t dim, float *distances);
+
+}  // namespace zvec::turbo::armv8