refactor: turbo based refactor [WIP]

src/core/framework/index_factory.cc

@@ -257,5 +257,18 @@ std::vector<std::string> IndexFactory::AllRefiners(void) {
   return ailego::Factory<IndexRefiner>::Classes();
 }
 
+std::shared_ptr<turbo::Quantizer> IndexFactory::CreateQuantizer(
+    const std::string &name) {
+  return ailego::Factory<zvec::turbo::Quantizer>::MakeShared(name.c_str());
+}
+
+bool IndexFactory::HasQuantizer(const std::string &name) {
+  return ailego::Factory<turbo::Quantizer>::Has(name.c_str());
+}
+
+std::vector<std::string> IndexFactory::AllQuantizers(void) {
+  return ailego::Factory<turbo::Quantizer>::Classes();
+}
+
 }  // namespace core
 }  // namespace zvec

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

@@ -98,22 +98,23 @@ class QuantizedIntegerMetric : public IndexMetric {
         if (meta_.data_type() == IndexMeta::DataType::DT_INT8) {
           auto turbo_ret = turbo::get_distance_func(
               turbo::MetricType::kSquaredEuclidean, turbo::DataType::kInt8,
-              turbo::QuantizeType::kDefault);
+              static_cast<turbo::QuantizeType>(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) {
+
           return DistanceMatrixCompute<SquaredEuclidean, uint8_t>(m, n);
         }
         break;
 
       case MetricType::kInnerProduct:
         if (meta_.data_type() == IndexMeta::DataType::DT_INT8) {
-          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::kInt8,
+              static_cast<turbo::QuantizeType>(quantize_type_));
+          if (turbo_ret && m == 1 && n == 1) {
+            return turbo_ret;
+          }
           return DistanceMatrixCompute<MinusInnerProduct, uint8_t>(m, n);
         }
         break;
@@ -139,7 +140,7 @@ class QuantizedIntegerMetric : public IndexMetric {
         if (meta_.data_type() == IndexMeta::DataType::DT_INT8) {
           auto turbo_ret = turbo::get_distance_func(
               turbo::MetricType::kCosine, turbo::DataType::kInt8,
-              turbo::QuantizeType::kDefault);
+              static_cast<turbo::QuantizeType>(quantize_type_));
           if (turbo_ret) {
             return turbo_ret;
           }
@@ -160,7 +161,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);
+              static_cast<turbo::QuantizeType>(quantize_type_));
           if (turbo_ret) {
             return turbo_ret;
           }
@@ -215,7 +216,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);
+              static_cast<turbo::QuantizeType>(quantize_type_));
           if (turbo_ret) {
             return turbo_ret;
           }

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

@@ -14,6 +14,7 @@
 
 #pragma once
 
+#include <turbo/quantizer/quantizer.h>
 #include <zvec/ailego/pattern/factory.h>
 #include <zvec/core/framework/index_builder.h>
 #include <zvec/core/framework/index_cluster.h>
@@ -167,6 +168,16 @@ struct IndexFactory {
 
   //! Retrieve all refiner classes
   static std::vector<std::string> AllRefiners(void);
+
+  //! Create a quantizer by name
+  static std::shared_ptr<zvec::turbo::Quantizer> CreateQuantizer(
+      const std::string &name);
+
+  //! Test if the quantizer exists
+  static bool HasQuantizer(const std::string &name);
+
+  //! Retrieve all quantizer classes
+  static std::vector<std::string> AllQuantizers(void);
 };
 
 //! Register Index Metric
@@ -283,5 +294,13 @@ struct IndexFactory {
 #define INDEX_FACTORY_REGISTER_REFINER(__IMPL__, ...) \
   INDEX_FACTORY_REGISTER_REFINER_ALIAS(__IMPL__, __IMPL__, ##__VA_ARGS__)
 
+//! Register Quantizer
+#define INDEX_FACTORY_REGISTER_QUANTIZER_ALIAS(__NAME__, __IMPL__, ...) \
+  AILEGO_FACTORY_REGISTER(__NAME__, turbo::Quantizer, __IMPL__, ##__VA_ARGS__)
+
+//! Register Quantizer
+#define INDEX_FACTORY_REGISTER_QUANTIZER(__IMPL__, ...) \
+  INDEX_FACTORY_REGISTER_QUANTIZER_ALIAS(__IMPL__, __IMPL__, ##__VA_ARGS__)
+
 }  // namespace core
 }  // namespace zvec

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

@@ -40,6 +40,7 @@ class IndexMeta {
     DT_BINARY64 = 8,
   };
 
+
   /*! Major Orders
    */
   enum MajorOrder {
@@ -451,6 +452,11 @@ class IndexMeta {
     this->set_meta(data_type, UnitSizeof(data_type), dim);
   }
 
+  //! Set extra meta size
+  void set_extra_meta_size(uint32_t size) {
+    extra_meta_size_ = size;
+  }
+
   //! Set information of metric
   template <typename TName, typename TParams>
   void set_metric(TName &&name, uint32_t rev, TParams &&params) {
@@ -586,6 +592,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};
@@ -632,6 +639,19 @@ class IndexQueryMeta {
         unit_size_(unit),
         element_size_(IndexMeta::ElementSizeof(data_type, unit, dim)) {}
 
+  //! Constructor
+  IndexQueryMeta(IndexMeta::MetaType meta_type, IndexMeta::DataType data_type,
+                 uint32_t unit, uint32_t dim, uint32_t quantize_type,
+                 uint32_t extra_meta_size)
+      : meta_type_(meta_type),
+        data_type_(data_type),
+        dimension_(dim),
+        unit_size_(unit),
+        quantize_type_(quantize_type),
+        extra_meta_size_(extra_meta_size),
+        element_size_(IndexMeta::ElementSizeof(data_type, unit, dim) +
+                      extra_meta_size_) {}
+
   //! Constructor
   IndexQueryMeta(IndexMeta::DataType data_type, uint32_t dim)
       : IndexQueryMeta{IndexMeta::MetaType::MT_DENSE, data_type,
@@ -676,7 +696,8 @@ class IndexQueryMeta {
   //! Set dimension of feature
   void set_dimension(uint32_t dim) {
     dimension_ = dim;
-    element_size_ = IndexMeta::ElementSizeof(data_type_, unit_size_, dim);
+    element_size_ = IndexMeta::ElementSizeof(data_type_, unit_size_, dim) +
+                    extra_meta_size_;
   }
 
   //! Set meta type
@@ -694,19 +715,34 @@ class IndexQueryMeta {
     data_type_ = data_type;
     dimension_ = dim;
     unit_size_ = unit;
-    element_size_ = IndexMeta::ElementSizeof(data_type, unit, dim);
+    element_size_ =
+        IndexMeta::ElementSizeof(data_type, unit, dim) + extra_meta_size_;
   }
 
   //! Set meta information of feature
   void set_meta(IndexMeta::DataType data_type, uint32_t dim) {
     this->set_meta(data_type, IndexMeta::UnitSizeof(data_type), dim);
   }
 
+  //! Set meta information of feature with quantize type and extra meta size
+  void set_meta(IndexMeta::DataType data_type, uint32_t dim,
+                uint32_t quantize_type, uint32_t extra_meta_size) {
+    data_type_ = data_type;
+    dimension_ = dim;
+    unit_size_ = IndexMeta::UnitSizeof(data_type);
+    quantize_type_ = quantize_type;
+    extra_meta_size_ = extra_meta_size;
+    element_size_ =
+        IndexMeta::ElementSizeof(data_type, unit_size_, dim) + extra_meta_size_;
+  }
+
  private:
   IndexMeta::MetaType meta_type_{IndexMeta::MetaType::MT_DENSE};
   IndexMeta::DataType data_type_{IndexMeta::DataType::DT_UNDEFINED};
   uint32_t dimension_{0};
   uint32_t unit_size_{0};
+  uint32_t quantize_type_{0};
+  uint32_t extra_meta_size_{0};
   uint32_t element_size_{0};
 };
 

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;
   }
+
+ protected:
+  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