refactor: turbo quantizer (#546)

Author: richyreachy

examples/c++/CMakeLists.txt

@@ -12,8 +12,8 @@ set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
 if(NOT DEFINED HOST_BUILD_DIR)
     set(HOST_BUILD_DIR "build")
 endif()
-
 get_filename_component(ZVEC_ROOT_DIR "${CMAKE_CURRENT_LIST_DIR}/../.." ABSOLUTE)
+
 set(ZVEC_INCLUDE_DIR ${ZVEC_ROOT_DIR}/src/include)
 set(ZVEC_LIB_DIR ${ZVEC_ROOT_DIR}/${HOST_BUILD_DIR}/lib)
 

src/core/framework/index_factory.cc

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <turbo/quantizer/quantizer.h>
 #include <zvec/core/framework/index_factory.h>
 
 namespace zvec {
@@ -257,5 +258,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;
@@ -141,7 +143,9 @@ bool IndexMeta::deserialize(const void *data, size_t len) {
   data_type_ = static_cast<IndexMeta::DataType>(format->data_type);
   dimension_ = format->dimension;
   unit_size_ = format->unit_size;
-  element_size_ = IndexMeta::ElementSizeof(data_type_, unit_size_, dimension_);
+  extra_meta_size_ = format->extra_meta_size;
+  element_size_ = IndexMeta::ElementSizeof(data_type_, unit_size_, dimension_) +
+                  extra_meta_size_;
   space_id_ = format->space_id;
 
   // Read attachment

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

@@ -29,6 +29,12 @@
 #include <zvec/core/framework/index_streamer.h>
 #include <zvec/core/framework/index_trainer.h>
 
+namespace zvec {
+namespace turbo {
+class Quantizer;
+}  // namespace turbo
+}  // namespace zvec
+
 namespace zvec {
 namespace core {
 
@@ -167,6 +173,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 +299,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 {
@@ -77,6 +78,7 @@ class IndexMeta {
         dimension_(rhs.dimension_),
         unit_size_(rhs.unit_size_),
         element_size_(rhs.element_size_),
+        extra_meta_size_(rhs.extra_meta_size_),
         space_id_(rhs.space_id_),
         metric_revision_(rhs.metric_revision_),
         converter_revision_(rhs.converter_revision_),
@@ -112,6 +114,7 @@ class IndexMeta {
         dimension_(rhs.dimension_),
         unit_size_(rhs.unit_size_),
         element_size_(rhs.element_size_),
+        extra_meta_size_(rhs.extra_meta_size_),
         space_id_(rhs.space_id_),
         metric_revision_(rhs.metric_revision_),
         converter_revision_(rhs.converter_revision_),
@@ -173,6 +176,7 @@ class IndexMeta {
     searcher_params_ = std::move(rhs.searcher_params_);
     streamer_params_ = std::move(rhs.streamer_params_);
     attributes_ = std::move(rhs.attributes_);
+    extra_meta_size_ = rhs.extra_meta_size_;
 
     return *this;
   }
@@ -211,6 +215,7 @@ class IndexMeta {
     searcher_params_ = std::move(rhs.searcher_params_);
     streamer_params_ = std::move(rhs.streamer_params_);
     attributes_ = std::move(rhs.attributes_);
+    extra_meta_size_ = rhs.extra_meta_size_;
 
     return *this;
   }
@@ -249,6 +254,7 @@ class IndexMeta {
     searcher_params_.clear();
     streamer_params_.clear();
     attributes_.clear();
+    extra_meta_size_ = 0;
   }
 
   //! Retrieve major order information
@@ -281,6 +287,11 @@ class IndexMeta {
     return element_size_;
   }
 
+  //! Retrieve extra meta size in bytes
+  uint32_t extra_meta_size(void) const {
+    return extra_meta_size_;
+  }
+
   //! Retrieve space id
   uint64_t space_id(void) const {
     return space_id_;
@@ -429,7 +440,8 @@ class IndexMeta {
   //! 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 information of feature
@@ -443,14 +455,21 @@ class IndexMeta {
     data_type_ = data_type;
     dimension_ = dim;
     unit_size_ = unit;
-    element_size_ = ElementSizeof(data_type, unit, dim);
+    element_size_ = ElementSizeof(data_type, unit, dim) + extra_meta_size_;
   }
 
   //! Set meta information of feature
   void set_meta(DataType data_type, uint32_t dim) {
     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;
+    element_size_ =
+        ElementSizeof(data_type_, unit_size_, dimension_) + extra_meta_size_;
+  }
+
   //! Set information of metric
   template <typename TName, typename TParams>
   void set_metric(TName &&name, uint32_t rev, TParams &&params) {
@@ -586,6 +605,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 +652,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,
@@ -673,10 +706,21 @@ class IndexQueryMeta {
     return element_size_;
   }
 
+  //! Retrieve quantize type
+  uint32_t quantize_type(void) const {
+    return quantize_type_;
+  }
+
+  //! Retrieve extra meta size in bytes
+  uint32_t extra_meta_size(void) const {
+    return extra_meta_size_;
+  }
+
   //! 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
@@ -689,24 +733,47 @@ class IndexQueryMeta {
     data_type_ = data_type;
   }
 
+  //! Set extra meta size
+  void set_extra_meta_size(uint32_t size) {
+    extra_meta_size_ = size;
+    element_size_ =
+        IndexMeta::ElementSizeof(data_type_, unit_size_, dimension_) +
+        extra_meta_size_;
+  }
+
   //! Set meta information of feature
   void set_meta(IndexMeta::DataType data_type, uint32_t unit, uint32_t dim) {
     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/turbo/turbo.h

@@ -36,30 +36,56 @@ using UniformQuantizeFunc = void (*)(const float *in, size_t dim, float scale,
 enum class MetricType {
   kSquaredEuclidean,
   kCosine,
+  kInnerProduct,
   kMipsSquaredEuclidean,
   kUnknown,
 };
 
 enum class DataType {
+  kInt4,
   kInt8,
+  kFp16,
+  kFp32,
   kUnknown,
 };
 
 enum class QuantizeType {
   kDefault,
   kUniform,
+  kRecord,
+  kFp16,
+  kFp32,
+  kPQ,
+  kRabit
+};
+
+enum class CpuArchType {
+  kAuto,
+  kScalar,
+  // x86 SIMD
+  kSSE,
+  kAVX,
+  kAVX2,
+  kAVX512,
+  kAVX512VNNI,
+  kAVX512FP16,
+  // ARM SIMD
+  kNEON,
+  kSVE,
+  kSVE2
 };
 
 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);
 
 // Returns the SIMD kernel for the uniform quantizer on the current CPU for
 // the given output data_type, or nullptr if no SIMD implementation is

src/turbo/CMakeLists.txt

@@ -19,7 +19,7 @@ file(GLOB_RECURSE ALL_SRCS *.cc *.c *.h)
 # 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_VNNI_SRCS ${CMAKE_CURRENT_SOURCE_DIR}/distance/avx512_vnni/*.cc)
         set_source_files_properties(
             ${AVX512_VNNI_SRCS}
             PROPERTIES
@@ -29,8 +29,8 @@ if(NOT ANDROID AND AUTO_DETECT_ARCH)
 endif()
 
 cc_library(
-    NAME zvec_turbo STATIC STRICT PACKED
+    NAME zvec_turbo STATIC STRICT ALWAYS_LINK PACKED
     SRCS ${ALL_SRCS}
     LIBS zvec_ailego
-    INCS ${CMAKE_CURRENT_SOURCE_DIR} ${PROJECT_ROOT_DIR}/src/include
+    INCS ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/distance ${PROJECT_ROOT_DIR}/src/include
 )