diff --git a/.gitignore b/.gitignore index 65859ca54..99623a413 100644 --- a/.gitignore +++ b/.gitignore @@ -60,9 +60,9 @@ allure-* !build_android.sh !build_ios.sh -# congfig +# config doc/ config/ examples/python/ examples/c_api/ -logs/ \ No newline at end of file +logs/ diff --git a/examples/c++/CMakeLists.txt b/examples/c++/CMakeLists.txt index d58814d52..4de64b45c 100644 --- a/examples/c++/CMakeLists.txt +++ b/examples/c++/CMakeLists.txt @@ -14,7 +14,10 @@ if(NOT DEFINED HOST_BUILD_DIR) endif() get_filename_component(ZVEC_ROOT_DIR "${CMAKE_CURRENT_LIST_DIR}/../.." ABSOLUTE) -set(ZVEC_INCLUDE_DIR ${ZVEC_ROOT_DIR}/src/include) +# Include both src/include (public headers) and src (internal headers that +# are transitively reachable from public headers, e.g. +# turbo/quantizer/quantizer.h via index_factory.h). +set(ZVEC_INCLUDE_DIR ${ZVEC_ROOT_DIR}/src/include ${ZVEC_ROOT_DIR}/src) set(ZVEC_LIB_DIR ${ZVEC_ROOT_DIR}/${HOST_BUILD_DIR}/lib) # Add include and library search paths diff --git a/src/core/algorithm/hnsw_sparse/hnsw_sparse_entity.h b/src/core/algorithm/hnsw_sparse/hnsw_sparse_entity.h index 7e6f814b1..ee4730a7e 100644 --- a/src/core/algorithm/hnsw_sparse/hnsw_sparse_entity.h +++ b/src/core/algorithm/hnsw_sparse/hnsw_sparse_entity.h @@ -73,7 +73,7 @@ struct HnswSparseHeader { struct SparseData { public: - SparseData() {}; + SparseData(){}; SparseData(uint32_t sparse_count, const uint32_t *sparse_indices, const void *sparse_vec) diff --git a/src/core/framework/index_factory.cc b/src/core/framework/index_factory.cc index 69fe0e98d..e93f57bc7 100644 --- a/src/core/framework/index_factory.cc +++ b/src/core/framework/index_factory.cc @@ -257,5 +257,18 @@ std::vector IndexFactory::AllRefiners(void) { return ailego::Factory::Classes(); } +std::shared_ptr IndexFactory::CreateQuantizer( + const std::string &name) { + return ailego::Factory::MakeShared(name.c_str()); +} + +bool IndexFactory::HasQuantizer(const std::string &name) { + return ailego::Factory::Has(name.c_str()); +} + +std::vector IndexFactory::AllQuantizers(void) { + return ailego::Factory::Classes(); +} + } // namespace core } // namespace zvec diff --git a/src/core/framework/index_meta.cc b/src/core/framework/index_meta.cc index 11d54cb63..d0eadb02d 100644 --- a/src/core/framework/index_meta.cc +++ b/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; diff --git a/src/core/utility/visit_filter.h b/src/core/utility/visit_filter.h index 959a0b450..853110ac1 100644 --- a/src/core/utility/visit_filter.h +++ b/src/core/utility/visit_filter.h @@ -44,7 +44,7 @@ class VisitBloomFilter { static constexpr int N = 5; struct Context { Context() - : mt(std::chrono::system_clock::now().time_since_epoch().count()) {}; + : mt(std::chrono::system_clock::now().time_since_epoch().count()){}; VisitFilterHeader h; std::mt19937 mt; ailego::BloomFilter *filter{nullptr}; @@ -367,7 +367,7 @@ class VisitFilter { ByteMap = VisitByteMap::mode }; - VisitFilter() : mode_(0), ctx_(nullptr) {}; + VisitFilter() : mode_(0), ctx_(nullptr){}; inline bool visited(id_t idx) { PROXIMA_HNSW_VISITFILTER_CALL_IMPL(visited, idx); diff --git a/src/db/index/column/inverted_column/inverted_column_indexer.h b/src/db/index/column/inverted_column/inverted_column_indexer.h index 94e3c05b1..56b8b6f57 100644 --- a/src/db/index/column/inverted_column/inverted_column_indexer.h +++ b/src/db/index/column/inverted_column/inverted_column_indexer.h @@ -68,7 +68,7 @@ class InvertedColumnIndexer { field_(field), path_(context.db_path_), ctx_(context), - read_only_(read_only) {}; + read_only_(read_only){}; InvertedColumnIndexer(const InvertedColumnIndexer &) = delete; InvertedColumnIndexer(InvertedColumnIndexer &&) = delete; diff --git a/src/db/index/column/inverted_column/inverted_indexer.h b/src/db/index/column/inverted_column/inverted_indexer.h index 19c7b392d..13884ab43 100644 --- a/src/db/index/column/inverted_column/inverted_indexer.h +++ b/src/db/index/column/inverted_column/inverted_indexer.h @@ -33,7 +33,7 @@ class InvertedIndexer { const std::vector &fields) : collection_name_(collection_name), working_dir_(working_dir), - fields_(fields) {}; + fields_(fields){}; virtual ~InvertedIndexer() { diff --git a/src/db/index/column/inverted_column/inverted_search_result.h b/src/db/index/column/inverted_column/inverted_search_result.h index 3c9720a6e..3bf391f9f 100644 --- a/src/db/index/column/inverted_column/inverted_search_result.h +++ b/src/db/index/column/inverted_column/inverted_search_result.h @@ -34,7 +34,7 @@ class InvertedSearchResult class Filter : public IndexFilter { public: explicit Filter(std::shared_ptr result) - : result_(std::move(result)) {}; + : result_(std::move(result)){}; bool is_filtered(uint64_t id) const override { return !result_->contains(id); diff --git a/src/db/index/common/delete_store.h b/src/db/index/common/delete_store.h index e6b742758..56ac8733c 100644 --- a/src/db/index/common/delete_store.h +++ b/src/db/index/common/delete_store.h @@ -29,7 +29,7 @@ class DeleteStore : public std::enable_shared_from_this { using Ptr = std::shared_ptr; explicit DeleteStore(std::string collection_name) - : collection_name_(std::move(collection_name)) {}; + : collection_name_(std::move(collection_name)){}; ~DeleteStore() { LOG_INFO("Closed delete store"); diff --git a/src/db/index/common/id_map.h b/src/db/index/common/id_map.h index 59e6b3ecc..67b0e8ce0 100644 --- a/src/db/index/common/id_map.h +++ b/src/db/index/common/id_map.h @@ -30,7 +30,7 @@ class IDMap { using Ptr = std::shared_ptr; explicit IDMap(std::string collection_name) - : collection_name_(std::move(collection_name)) {}; + : collection_name_(std::move(collection_name)){}; ~IDMap() { if (opened_) { diff --git a/src/db/index/common/index_params.cc b/src/db/index/common/index_params.cc index 0b696956a..8f08bcdbc 100644 --- a/src/db/index/common/index_params.cc +++ b/src/db/index/common/index_params.cc @@ -26,8 +26,7 @@ namespace zvec { std::string InvertIndexParams::to_string() const { std::ostringstream oss; - oss << "InvertIndexParams{" - << "enable_range_optimization:" + oss << "InvertIndexParams{" << "enable_range_optimization:" << (enable_range_optimization_ ? "true" : "false") << ", enable_extended_wildcard:" << (enable_extended_wildcard_ ? "true" : "false") << "}"; diff --git a/src/db/index/common/meta.h b/src/db/index/common/meta.h index 272e831ef..81a5f0fc9 100644 --- a/src/db/index/common/meta.h +++ b/src/db/index/common/meta.h @@ -124,8 +124,8 @@ class BlockMeta { std::string to_string() const { std::ostringstream oss; - oss << "BlockMeta{" - << "id:" << id_ << ",type:" << BlockTypeCodeBook::AsString(type_) + oss << "BlockMeta{" << "id:" << id_ + << ",type:" << BlockTypeCodeBook::AsString(type_) << ",min_doc_id:" << min_doc_id_ << ",max_doc_id:" << max_doc_id_ << ",doc_count:" << doc_count_ << ",columns:["; @@ -179,7 +179,7 @@ class SegmentMeta { using Ptr = std::shared_ptr; public: - SegmentMeta() {}; + SegmentMeta(){}; explicit SegmentMeta(SegmentID id) : id_(id) {} @@ -368,8 +368,7 @@ class SegmentMeta { std::string to_string() const { std::ostringstream oss; - oss << "SegmentMeta{" - << "id:" << id_ << ",persisted_blocks:["; + oss << "SegmentMeta{" << "id:" << id_ << ",persisted_blocks:["; for (size_t i = 0; i < persisted_blocks_.size(); ++i) { if (i > 0) oss << ","; diff --git a/src/db/index/common/schema.cc b/src/db/index/common/schema.cc index 06958ccc3..12be291cb 100644 --- a/src/db/index/common/schema.cc +++ b/src/db/index/common/schema.cc @@ -296,8 +296,7 @@ Status FieldSchema::validate() const { std::string FieldSchema::to_string() const { std::ostringstream oss; - oss << "FieldSchema{" - << "name:'" << name_ << "'" + oss << "FieldSchema{" << "name:'" << name_ << "'" << ",data_type:" << DataTypeCodeBook::AsString(data_type_) << ",nullable:" << (nullable_ ? "true" : "false") << ",dimension:" << dimension_; @@ -382,8 +381,7 @@ Status CollectionSchema::validate() const { std::string CollectionSchema::to_string() const { std::ostringstream oss; - oss << "CollectionSchema{" - << "name:'" << name_ << "'" + oss << "CollectionSchema{" << "name:'" << name_ << "'" << ",max_doc_count_per_segment:" << max_doc_count_per_segment_ << ",fields:["; diff --git a/src/db/index/common/stats.cc b/src/db/index/common/stats.cc index fdb26f25a..ff1223ec1 100644 --- a/src/db/index/common/stats.cc +++ b/src/db/index/common/stats.cc @@ -19,8 +19,8 @@ namespace zvec { std::string CollectionStats::to_string() const { std::ostringstream oss; - oss << "CollectionStats{" - << "doc_count:" << doc_count << ",index_completeness:{"; + oss << "CollectionStats{" << "doc_count:" << doc_count + << ",index_completeness:{"; size_t i = 0; for (const auto &pair : index_completeness) { diff --git a/src/db/sqlengine/parser/zvec_parser.h b/src/db/sqlengine/parser/zvec_parser.h index 44102d0aa..2e1152b76 100644 --- a/src/db/sqlengine/parser/zvec_parser.h +++ b/src/db/sqlengine/parser/zvec_parser.h @@ -26,7 +26,7 @@ class ZVecParser { using Ptr = std::shared_ptr; ZVecParser() = default; - virtual ~ZVecParser() {}; + virtual ~ZVecParser(){}; virtual SQLInfo::Ptr parse(const std::string &query, bool formatted_tree = false) = 0; diff --git a/src/include/zvec/ailego/pattern/closure.h b/src/include/zvec/ailego/pattern/closure.h index df624ac92..554e8a77e 100644 --- a/src/include/zvec/ailego/pattern/closure.h +++ b/src/include/zvec/ailego/pattern/closure.h @@ -309,7 +309,7 @@ class Callback : public Callback { protected: //! Constructor - Callback(void) {}; + Callback(void){}; }; /*! Callback Implementation diff --git a/src/include/zvec/core/framework/index_factory.h b/src/include/zvec/core/framework/index_factory.h index d891eaa5a..00e77894c 100644 --- a/src/include/zvec/core/framework/index_factory.h +++ b/src/include/zvec/core/framework/index_factory.h @@ -14,6 +14,7 @@ #pragma once +#include #include #include #include @@ -167,6 +168,16 @@ struct IndexFactory { //! Retrieve all refiner classes static std::vector AllRefiners(void); + + //! Create a quantizer by name + static std::shared_ptr CreateQuantizer( + const std::string &name); + + //! Test if the quantizer exists + static bool HasQuantizer(const std::string &name); + + //! Retrieve all quantizer classes + static std::vector 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 diff --git a/src/include/zvec/core/framework/index_meta.h b/src/include/zvec/core/framework/index_meta.h index 3a09aaefb..c6c5137db 100644 --- a/src/include/zvec/core/framework/index_meta.h +++ b/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_; @@ -451,6 +462,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 void set_metric(TName &&name, uint32_t rev, TParams &¶ms) { @@ -586,6 +602,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 +649,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 +706,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,7 +725,8 @@ 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 @@ -702,11 +734,25 @@ class IndexQueryMeta { 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}; }; diff --git a/src/include/zvec/turbo/turbo.h b/src/include/zvec/turbo/turbo.h index 2fbf6d680..36dd8aae1 100644 --- a/src/include/zvec/turbo/turbo.h +++ b/src/include/zvec/turbo/turbo.h @@ -13,6 +13,8 @@ // limitations under the License. #pragma once +#include +#include #include #include @@ -33,33 +35,107 @@ using QueryPreprocessFunc = using UniformQuantizeFunc = void (*)(const float *in, size_t dim, float scale, float bias, int8_t *out); +// FHT primitive function pointer types. +using FhtFlipSignFunc = void (*)(const uint8_t *flip, float *data, size_t dim); +using FhtKacsWalkFunc = void (*)(float *data, size_t len); +using FhtInplaceFunc = void (*)(float *data, size_t n); +using FhtVecRescaleFunc = void (*)(float *data, size_t n, float factor); + +// PQ kernel function pointer types. +// +// ADC: LUT look-up distance between a PQ code and a query (via LUT). +// pq_code: [num_subquantizers] uint8_t +// lut: [num_subquantizers * 256] float +// Uses void* to match DistanceFunc signature for direct assignment. +using PqAdcDistanceFunc = void (*)(const void *pq_code, const void *lut, + size_t num_subquantizers, float *out); + +// SDC kernel: centroid-to-centroid distance between two PQ codes. +// a, b: [num_subquantizers] uint8_t +// dist_table: [num_subquantizers * 256 * 256] float +// Uses void* for consistency with DistanceFunc / PqAdcDistanceFunc. +using PqSdcKernelFunc = void (*)(const void *a, const void *b, + const void *dist_table, + size_t num_subquantizers, float *out); + +// Batch ADC: compute distances for multiple PQ codes against a shared LUT. +// Signature matches BatchDistanceFunc for direct assignment (no lambda). +using PqBatchAdcFunc = void (*)(const void **candidates, const void *lut, + size_t num, size_t num_subquantizers, + float *out); + +// Aggregate of all FHT kernels needed by FhtRotator, dispatched by ISA. +struct FhtKernels { + FhtFlipSignFunc flip_sign; + FhtKacsWalkFunc kacs_walk; + FhtKacsWalkFunc inv_kacs_walk; + FhtInplaceFunc inplace; + FhtVecRescaleFunc rescale; +}; + +// Aggregate of all PQ-specific kernels needed by PqInt8Quantizer / +// PqInt4Quantizer, dispatched by ISA and data_type. LUT computation +// (compute_distance_table) is NOT included here — it reuses the existing +// fp32 BatchDistanceFunc from get_batch_distance_func() which is metric- +// aware and already SIMD-optimized. +// +// data_type selects the code packing layout: +// kInt8: one uint8 per sub-quantizer (256 centroids, stride=256) +// kInt4: two sub-quantizers packed into one uint8 (16 centroids, stride=16) +struct PqKernels { + PqAdcDistanceFunc adc_distance; + PqSdcKernelFunc sdc_distance; + PqBatchAdcFunc batch_adc_distance; +}; + 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, + 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); // 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 @@ -69,4 +145,13 @@ QueryPreprocessFunc get_query_preprocess_func(MetricType metric_type, // interface can grow to cover other output types (e.g. fp16) in the future. UniformQuantizeFunc get_uniform_quantize_func(DataType data_type); +// Returns all FHT kernels dispatched for the current CPU. +FhtKernels get_fht_kernels(); + +// Returns all PQ kernels dispatched for the given data_type, quantize_type +// and CPU arch. data_type selects the code packing layout (kInt8 vs kInt4). +PqKernels get_pq_kernels(DataType data_type, + QuantizeType quantize_type = QuantizeType::kPQ, + CpuArchType cpu_arch_type = CpuArchType::kAuto); + } // namespace zvec::turbo diff --git a/src/turbo/CMakeLists.txt b/src/turbo/CMakeLists.txt index 9cbb2fac7..f34f24e66 100644 --- a/src/turbo/CMakeLists.txt +++ b/src/turbo/CMakeLists.txt @@ -13,15 +13,32 @@ endif() file(GLOB_RECURSE ALL_SRCS *.cc *.c *.h) -# Set per-file compile flags for AVX512-VNNI sources. +# Set per-file compile flags for SIMD 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) + # SSE + file(GLOB_RECURSE SSE_SRCS ${CMAKE_CURRENT_SOURCE_DIR}/distance/sse/*.cc) set_source_files_properties( - ${AVX512_VNNI_SRCS} + ${SSE_SRCS} + PROPERTIES + COMPILE_FLAGS "${TURBO_MARCH_FLAG_SSE}" + ) + + # AVX2 + file(GLOB_RECURSE AVX2_SRCS ${CMAKE_CURRENT_SOURCE_DIR}/distance/avx2/*.cc) + set_source_files_properties( + ${AVX2_SRCS} + PROPERTIES + COMPILE_FLAGS "${TURBO_MARCH_FLAG_AVX2}" + ) + + # AVX512 (avx512_vnni and avx512/fht both use AVX512 march) + file(GLOB_RECURSE AVX512_SRCS ${CMAKE_CURRENT_SOURCE_DIR}/distance/avx512_vnni/*.cc ${CMAKE_CURRENT_SOURCE_DIR}/distance/avx512/*.cc) + set_source_files_properties( + ${AVX512_SRCS} PROPERTIES COMPILE_FLAGS "${TURBO_MARCH_FLAG_AVX512}" ) @@ -29,8 +46,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 ) diff --git a/src/turbo/distance/avx2/fht/fht.cc b/src/turbo/distance/avx2/fht/fht.cc new file mode 100644 index 000000000..a6d3edc36 --- /dev/null +++ b/src/turbo/distance/avx2/fht/fht.cc @@ -0,0 +1,137 @@ +// 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. + +#if defined(__AVX2__) + +#include "fht.h" +#include +#include +#include +#include +#include + +namespace zvec::turbo::avx2 { + +void fht_flip_sign_avx2(const uint8_t *flip, float *data, size_t dim) { + size_t simd_end = dim & ~31u; + constexpr size_t kChunk = 32; + const __m256i bit_select = + _mm256_setr_epi32(0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80); + const __m256 sign_flip = _mm256_castsi256_ps(_mm256_set1_epi32(0x80000000)); + for (size_t i = 0; i < simd_end; i += kChunk) { + uint32_t mask_bits; + std::memcpy(&mask_bits, &flip[i / 8], sizeof(mask_bits)); + for (int b = 0; b < 4; ++b) { + __m256i mb = _mm256_set1_epi32((mask_bits >> (b * 8)) & 0xFF); + __m256i test = _mm256_and_si256(mb, bit_select); + __m256i cmp = _mm256_cmpeq_epi32(test, bit_select); + __m256 xor_mask = _mm256_and_ps(_mm256_castsi256_ps(cmp), sign_flip); + __m256 v = _mm256_loadu_ps(&data[i + b * 8]); + v = _mm256_xor_ps(v, xor_mask); + _mm256_storeu_ps(&data[i + b * 8], v); + } + } + // Scalar tail + for (size_t i = simd_end; i < dim; ++i) { + if (flip[i / 8] & (1u << (i % 8))) { + data[i] = -data[i]; + } + } +} + +void fht_kacs_walk_avx2(float *data, size_t len) { + size_t half = len / 2; + size_t base = len % 2; + size_t offset = base + half; + size_t half_end = half & ~7u; + for (size_t i = 0; i < half_end; i += 8) { + __m256 x = _mm256_loadu_ps(&data[i]); + __m256 y = _mm256_loadu_ps(&data[i + offset]); + _mm256_storeu_ps(&data[i], _mm256_add_ps(x, y)); + _mm256_storeu_ps(&data[i + offset], _mm256_sub_ps(x, y)); + } + // Scalar tail + for (size_t i = half_end; i < half; ++i) { + float x = data[i]; + float y = data[i + offset]; + data[i] = x + y; + data[i + offset] = x - y; + } + if (base != 0) { + data[half] *= std::sqrt(2.0f); + } +} + +void fht_inv_kacs_walk_avx2(float *data, size_t len) { + size_t half = len / 2; + size_t base = len % 2; + size_t offset = base + half; + if (base != 0) { + data[half] *= std::sqrt(0.5f); + } + size_t half_end = half & ~7u; + const __m256 half_fac = _mm256_set1_ps(0.5f); + for (size_t i = 0; i < half_end; i += 8) { + __m256 a = _mm256_loadu_ps(&data[i]); + __m256 b = _mm256_loadu_ps(&data[i + offset]); + _mm256_storeu_ps(&data[i], _mm256_mul_ps(_mm256_add_ps(a, b), half_fac)); + _mm256_storeu_ps(&data[i + offset], + _mm256_mul_ps(_mm256_sub_ps(a, b), half_fac)); + } + // Scalar tail + for (size_t i = half_end; i < half; ++i) { + float a = data[i]; + float b = data[i + offset]; + data[i] = (a + b) * 0.5f; + data[i + offset] = (a - b) * 0.5f; + } +} + +void fht_inplace_avx2(float *data, size_t n) { + for (size_t len = 1; len < n; len <<= 1) { + size_t step = len << 1; + size_t simd_end = len & ~7u; + for (size_t i = 0; i < n; i += step) { + for (size_t j = 0; j < simd_end; j += 8) { + __m256 u = _mm256_loadu_ps(&data[i + j]); + __m256 v = _mm256_loadu_ps(&data[i + j + len]); + _mm256_storeu_ps(&data[i + j], _mm256_add_ps(u, v)); + _mm256_storeu_ps(&data[i + j + len], _mm256_sub_ps(u, v)); + } + for (size_t j = simd_end; j < len; ++j) { + float u = data[i + j]; + float v = data[i + j + len]; + data[i + j] = u + v; + data[i + j + len] = u - v; + } + } + } +} + +void fht_vec_rescale_avx2(float *data, size_t n, float factor) { + const __m256 fac = _mm256_set1_ps(factor); + size_t simd_end = n & ~7u; + for (size_t i = 0; i < simd_end; i += 8) { + __m256 v = _mm256_loadu_ps(&data[i]); + _mm256_storeu_ps(&data[i], _mm256_mul_ps(v, fac)); + } + // Scalar tail + for (size_t i = simd_end; i < n; ++i) { + data[i] *= factor; + } +} + +} // namespace zvec::turbo::avx2 + +#endif // __AVX2__ diff --git a/src/turbo/distance/avx2/fht/fht.h b/src/turbo/distance/avx2/fht/fht.h new file mode 100644 index 000000000..5e636a54f --- /dev/null +++ b/src/turbo/distance/avx2/fht/fht.h @@ -0,0 +1,37 @@ +// 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 +#include + +namespace zvec::turbo::avx2 { + +//! Apply bitwise sign-flip mask to a float vector (AVX2). +void fht_flip_sign_avx2(const uint8_t *flip, float *data, size_t dim); + +//! Apply KacsWalk butterfly operation (AVX2). +void fht_kacs_walk_avx2(float *data, size_t len); + +//! Inverse KacsWalk butterfly operation (AVX2). +void fht_inv_kacs_walk_avx2(float *data, size_t len); + +//! In-place Fast Hadamard Transform (AVX2, n must be power-of-2). +void fht_inplace_avx2(float *data, size_t n); + +//! Element-wise rescale: data[i] *= factor (AVX2). +void fht_vec_rescale_avx2(float *data, size_t n, float factor); + +} // namespace zvec::turbo::avx2 diff --git a/src/turbo/distance/avx2/pq_quantizer_int4/pq_distance.cc b/src/turbo/distance/avx2/pq_quantizer_int4/pq_distance.cc new file mode 100644 index 000000000..657c335d3 --- /dev/null +++ b/src/turbo/distance/avx2/pq_quantizer_int4/pq_distance.cc @@ -0,0 +1,224 @@ +// 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 "avx2/pq_quantizer_int4/pq_distance.h" +#include + +namespace zvec::turbo::avx2 { + +namespace { + +constexpr int kNumCentroids = 16; +constexpr int kTablePerSub = kNumCentroids * kNumCentroids; // 256 +constexpr int kChunkSize = 16; // process 16 subs per iteration (8 bytes) + +// Horizontal sum of 8 floats in a __m256 register. +inline float horizontal_sum_avx2(__m256 v) { + __m128 hi = _mm256_extractf128_ps(v, 1); + __m128 lo = _mm256_castps256_ps128(v); + __m128 sum128 = _mm_add_ps(lo, hi); + __m128 shuf = _mm_movehdup_ps(sum128); + __m128 sum64 = _mm_add_ps(sum128, shuf); + __m128 shuf32 = _mm_movehl_ps(sum64, sum64); + __m128 sum32 = _mm_add_ss(sum64, shuf32); + return _mm_cvtss_f32(sum32); +} + +// Unpack 8 bytes of packed int4 codes into 16 nibbles (0-15). +// Returns: +// lo_nib: __m128i with bytes [nib0, nib1, ..., nib7, 0, ..., 0] +// hi_nib: __m128i with bytes [nib8, nib9, ..., nib15, 0, ..., 0] +// +// Uses vpshufb to extract high nibbles efficiently: +// low nibbles = packed & 0x0F (subs 0-7 in bytes 0-7) +// high nibbles = (packed >> 4) & 0x0F (subs 8-15 in bytes 0-7) +inline void unpack_nibbles(const uint8_t *packed, __m128i &lo_nib, + __m128i &hi_nib) { + __m128i packed128 = + _mm_loadl_epi64(reinterpret_cast(packed)); + __m256i packed256 = _mm256_castsi128_si256(packed128); + __m256i mask_0f = _mm256_set1_epi8(0x0F); + __m256i lo256 = _mm256_and_si256(packed256, mask_0f); + + // Extract high nibbles via vpshufb: + // mask[1,1,2,2,3,3,4,4, 5,5,6,6,7,7,8,8] selects bytes 1-8 + // (byte 8 is zero since we only loaded 8 bytes into the 128-bit reg). + // After >> 4, positions 0-7 hold high nibbles of bytes 1-8 = subs 8-15. + __m256i hi_mask = + _mm256_setr_epi8(1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8, 1, 1, 2, + 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8); + __m256i hi256 = _mm256_and_si256( + _mm256_srli_epi16(_mm256_shuffle_epi8(packed256, hi_mask), 4), mask_0f); + + lo_nib = _mm256_castsi256_si128(lo256); // subs 0-7 in bytes 0-7 + // For hi_nib: bytes 8-11 of the lower 128-bit lane contain subs 8-15. + // Shift right by 8 bytes to move them to positions 0-3. + __m128i hi128 = _mm256_extracti128_si256(hi256, 0); + hi_nib = _mm_srli_si128(hi128, 8); +} + +// Accumulate 16 sub-quantizer distances into acc using the precomputed +// lo_nib / hi_nib nibble vectors. +inline void accumulate_adc(__m256 &acc, const float *lut, __m128i lo_nib, + __m128i hi_nib) { + // base_offsets = [0, 16, 32, ..., 7*16] — sub index offsets in float units + const __m256i base_offsets = _mm256_setr_epi32( + 0, kNumCentroids, 2 * kNumCentroids, 3 * kNumCentroids, 4 * kNumCentroids, + 5 * kNumCentroids, 6 * kNumCentroids, 7 * kNumCentroids); + + // Low half: subs m .. m+7 + __m256i lo32 = _mm256_cvtepu8_epi32(lo_nib); + __m256i lo_idx = _mm256_add_epi32(lo32, base_offsets); + acc = _mm256_add_ps(acc, _mm256_i32gather_ps(lut, lo_idx, 4)); + + // High half: subs m+8 .. m+15 + __m256i hi32 = _mm256_cvtepu8_epi32(hi_nib); + __m256i hi_idx = _mm256_add_epi32(hi32, base_offsets); + acc = _mm256_add_ps(acc, + _mm256_i32gather_ps(lut + 8 * kNumCentroids, hi_idx, 4)); +} + +} // namespace + +void pq_adc_int4_distance_avx2(const void *pq_code_v, const void *lut_v, + size_t num_subquantizers, float *out) { + const auto *pq_code = reinterpret_cast(pq_code_v); + const auto *lut = reinterpret_cast(lut_v); + __m256 acc = _mm256_setzero_ps(); + + size_t m = 0; + for (; m + kChunkSize <= num_subquantizers; m += kChunkSize) { + __m128i lo_nib, hi_nib; + unpack_nibbles(pq_code + (m >> 1), lo_nib, hi_nib); + accumulate_adc(acc, lut + m * kNumCentroids, lo_nib, hi_nib); + } + + float sum = horizontal_sum_avx2(acc); + // Scalar leftover + for (; m < num_subquantizers; ++m) { + uint8_t byte = pq_code[m >> 1]; + uint8_t idx = (m & 1u) ? (byte >> 4) : (byte & 0x0Fu); + sum += lut[m * kNumCentroids + idx]; + } + *out = sum; +} + +void pq_sdc_int4_distance_avx2(const void *a_v, const void *b_v, + const void *dist_table_v, + size_t num_subquantizers, float *out) { + const auto *a = reinterpret_cast(a_v); + const auto *b = reinterpret_cast(b_v); + const auto *dist_table = reinterpret_cast(dist_table_v); + + const __m256i base_offsets = _mm256_setr_epi32( + 0, kTablePerSub, 2 * kTablePerSub, 3 * kTablePerSub, 4 * kTablePerSub, + 5 * kTablePerSub, 6 * kTablePerSub, 7 * kTablePerSub); + const __m256i mul16 = _mm256_set1_epi32(kNumCentroids); + + __m256 acc = _mm256_setzero_ps(); + size_t m = 0; + for (; m + kChunkSize <= num_subquantizers; m += kChunkSize) { + const float *dt_base = dist_table + m * kTablePerSub; + + __m128i a_lo, a_hi, b_lo, b_hi; + unpack_nibbles(a + (m >> 1), a_lo, a_hi); + unpack_nibbles(b + (m >> 1), b_lo, b_hi); + + // Low half: index = m_local*256 + a_nib*16 + b_nib + __m256i a_lo32 = _mm256_cvtepu8_epi32(a_lo); + __m256i b_lo32 = _mm256_cvtepu8_epi32(b_lo); + __m256i lo_idx = _mm256_add_epi32( + _mm256_add_epi32(_mm256_mullo_epi32(a_lo32, mul16), b_lo32), + base_offsets); + acc = _mm256_add_ps(acc, _mm256_i32gather_ps(dt_base, lo_idx, 4)); + + // High half + __m256i a_hi32 = _mm256_cvtepu8_epi32(a_hi); + __m256i b_hi32 = _mm256_cvtepu8_epi32(b_hi); + __m256i hi_idx = _mm256_add_epi32( + _mm256_add_epi32(_mm256_mullo_epi32(a_hi32, mul16), b_hi32), + base_offsets); + acc = _mm256_add_ps( + acc, _mm256_i32gather_ps(dt_base + 8 * kTablePerSub, hi_idx, 4)); + } + + float sum = horizontal_sum_avx2(acc); + for (; m < num_subquantizers; ++m) { + uint8_t ab = a[m >> 1], bb = b[m >> 1]; + uint8_t ai = (m & 1u) ? (ab >> 4) : (ab & 0x0Fu); + uint8_t bi = (m & 1u) ? (bb >> 4) : (bb & 0x0Fu); + sum += dist_table[m * kTablePerSub + ai * kNumCentroids + bi]; + } + *out = sum; +} + +void pq_adc_int4_batch_distance_avx2(const void **candidates_v, + const void *lut_v, size_t num, + size_t num_subquantizers, float *out) { + const auto *lut = reinterpret_cast(lut_v); + const auto *candidates = + reinterpret_cast(candidates_v); + + size_t i = 0; + for (; i + 4 <= num; i += 4) { + const uint8_t *c0 = candidates[i]; + const uint8_t *c1 = candidates[i + 1]; + const uint8_t *c2 = candidates[i + 2]; + const uint8_t *c3 = candidates[i + 3]; + __m256 acc0 = _mm256_setzero_ps(); + __m256 acc1 = _mm256_setzero_ps(); + __m256 acc2 = _mm256_setzero_ps(); + __m256 acc3 = _mm256_setzero_ps(); + + size_t m = 0; + for (; m + kChunkSize <= num_subquantizers; m += kChunkSize) { + const float *lut_base = lut + m * kNumCentroids; + __m128i lo0, hi0, lo1, hi1, lo2, hi2, lo3, hi3; + unpack_nibbles(c0 + (m >> 1), lo0, hi0); + unpack_nibbles(c1 + (m >> 1), lo1, hi1); + unpack_nibbles(c2 + (m >> 1), lo2, hi2); + unpack_nibbles(c3 + (m >> 1), lo3, hi3); + accumulate_adc(acc0, lut_base, lo0, hi0); + accumulate_adc(acc1, lut_base, lo1, hi1); + accumulate_adc(acc2, lut_base, lo2, hi2); + accumulate_adc(acc3, lut_base, lo3, hi3); + } + + float s0 = horizontal_sum_avx2(acc0); + float s1 = horizontal_sum_avx2(acc1); + float s2 = horizontal_sum_avx2(acc2); + float s3 = horizontal_sum_avx2(acc3); + + // Scalar leftover for remaining sub-quantizers. + for (; m < num_subquantizers; ++m) { + const float *tab = lut + m * kNumCentroids; + uint8_t b0 = c0[m >> 1], b1 = c1[m >> 1], b2 = c2[m >> 1], + b3 = c3[m >> 1]; + s0 += tab[(m & 1u) ? (b0 >> 4) : (b0 & 0x0Fu)]; + s1 += tab[(m & 1u) ? (b1 >> 4) : (b1 & 0x0Fu)]; + s2 += tab[(m & 1u) ? (b2 >> 4) : (b2 & 0x0Fu)]; + s3 += tab[(m & 1u) ? (b3 >> 4) : (b3 & 0x0Fu)]; + } + out[i] = s0; + out[i + 1] = s1; + out[i + 2] = s2; + out[i + 3] = s3; + } + // Remaining candidates: use single ADC kernel. + for (; i < num; ++i) { + pq_adc_int4_distance_avx2(candidates[i], lut, num_subquantizers, out + i); + } +} + +} // namespace zvec::turbo::avx2 diff --git a/src/turbo/distance/avx2/pq_quantizer_int4/pq_distance.h b/src/turbo/distance/avx2/pq_quantizer_int4/pq_distance.h new file mode 100644 index 000000000..1d4164685 --- /dev/null +++ b/src/turbo/distance/avx2/pq_quantizer_int4/pq_distance.h @@ -0,0 +1,41 @@ +// 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 +#include + +namespace zvec::turbo::avx2 { + +// ADC (Asymmetric Distance Computation) for int4 PQ codes via AVX2. +// Uses vpshufb to efficiently unpack nibbles from packed bytes, then +// processes 16 sub-quantizers per iteration with two _mm256_i32gather_ps. +void pq_adc_int4_distance_avx2(const void *pq_code, const void *lut, + size_t num_subquantizers, float *out); + +// SDC (Symmetric Distance Computation) for int4 PQ codes via AVX2. +// Unpacks nibbles from both codes, computes index = m*256 + a*16 + b, +// gathers from the precomputed dist_table. +void pq_sdc_int4_distance_avx2(const void *a, const void *b, + const void *dist_table, size_t num_subquantizers, + float *out); + +// Batch ADC via AVX2: process 4 candidates per iteration, +// each using the 16-sub vpshufb + gather kernel. +void pq_adc_int4_batch_distance_avx2(const void **candidates, const void *lut, + size_t num, size_t num_subquantizers, + float *out); + +} // namespace zvec::turbo::avx2 diff --git a/src/turbo/distance/avx2/pq_quantizer_int8/pq_distance.cc b/src/turbo/distance/avx2/pq_quantizer_int8/pq_distance.cc new file mode 100644 index 000000000..18c177069 --- /dev/null +++ b/src/turbo/distance/avx2/pq_quantizer_int8/pq_distance.cc @@ -0,0 +1,215 @@ +// 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 "avx2/pq_quantizer_int8/pq_distance.h" +#include + +namespace zvec::turbo::avx2 { + +namespace { + +// Horizontal sum of 8 floats in a __m256 register. +inline float horizontal_sum_avx2(__m256 v) { + // High 128 bits + low 128 bits + __m128 hi = _mm256_extractf128_ps(v, 1); + __m128 lo = _mm256_castps256_ps128(v); + __m128 sum128 = _mm_add_ps(lo, hi); + // Shuffle and add: [a+b, c+d, a+b, c+d] + __m128 shuf = _mm_movehdup_ps(sum128); // [b, b, d, d] + __m128 sum64 = _mm_add_ps(sum128, shuf); + // Final: [a+b+c+d, ..., ...] + __m128 shuf32 = _mm_movehl_ps(sum64, sum64); + __m128 sum32 = _mm_add_ss(sum64, shuf32); + return _mm_cvtss_f32(sum32); +} + +} // namespace + +void pq_adc_int8_distance_avx2(const void *pq_code_v, const void *lut_v, + size_t num_subquantizers, float *out) { + constexpr int kNumCentroids = 256; + constexpr int kChunkSize = 8; // AVX2 processes 8 floats at once + const auto *pq_code = reinterpret_cast(pq_code_v); + const auto *lut = reinterpret_cast(lut_v); + + __m256 acc = _mm256_setzero_ps(); + + // Base offsets: [0, 256, 512, 768, 1024, 1280, 1536, 1792] + // These represent m * 256 for m = 0..7 within each chunk. + const __m256i base_offsets = _mm256_setr_epi32( + 0, kNumCentroids, 2 * kNumCentroids, 3 * kNumCentroids, 4 * kNumCentroids, + 5 * kNumCentroids, 6 * kNumCentroids, 7 * kNumCentroids); + + size_t m = 0; + + // Main loop: process 8 subquantizers per iteration + for (; m + kChunkSize <= num_subquantizers; m += kChunkSize) { + // Load 8 uint8 codes and zero-extend to int32 + // pq_code[m..m+7] -> 8 int32 indices + __m128i codes_8x8 = + _mm_loadl_epi64(reinterpret_cast(pq_code + m)); + __m256i codes_8x32 = _mm256_cvtepu8_epi32(codes_8x8); + + // Add base offsets: indices[m] = m * 256 + code[m] + __m256i indices = _mm256_add_epi32(codes_8x32, base_offsets); + + // Gather 8 floats from lut using computed indices + // lut_ptr + indices[i] * scale(4 bytes per float) + __m256 gathered = _mm256_i32gather_ps(lut + m * kNumCentroids, indices, 4); + + acc = _mm256_add_ps(acc, gathered); + } + + float sum = horizontal_sum_avx2(acc); + + // Scalar leftover: process remaining subquantizers + for (; m < num_subquantizers; ++m) { + sum += lut[m * kNumCentroids + pq_code[m]]; + } + + *out = sum; +} + +void pq_sdc_int8_distance_avx2(const void *a_v, const void *b_v, + const void *dist_table_v, + size_t num_subquantizers, float *out) { + constexpr int kNumCentroids = 256; + constexpr int kTablePerSub = kNumCentroids * kNumCentroids; // 65536 + constexpr int kChunkSize = 8; + const auto *a = reinterpret_cast(a_v); + const auto *b = reinterpret_cast(b_v); + const auto *dist_table = reinterpret_cast(dist_table_v); + + __m256 acc = _mm256_setzero_ps(); + + // Base offsets for SDC: m * 65536 (float indices into dist_table) + const __m256i base_offsets = _mm256_setr_epi32( + 0, kTablePerSub, 2 * kTablePerSub, 3 * kTablePerSub, 4 * kTablePerSub, + 5 * kTablePerSub, 6 * kTablePerSub, 7 * kTablePerSub); + + // Multiplier for a[m] * 256 + const __m256i a_multiplier = _mm256_set1_epi32(kNumCentroids); + + size_t m = 0; + + // Main loop: process 8 subquantizers per iteration + for (; m + kChunkSize <= num_subquantizers; m += kChunkSize) { + // Load a[m..m+7] and b[m..m+7], zero-extend to int32 + __m128i a_8x8 = _mm_loadl_epi64(reinterpret_cast(a + m)); + __m128i b_8x8 = _mm_loadl_epi64(reinterpret_cast(b + m)); + __m256i a_8x32 = _mm256_cvtepu8_epi32(a_8x8); + __m256i b_8x32 = _mm256_cvtepu8_epi32(b_8x8); + + // Compute index: a[m] * 256 + b[m] + m * 65536 + __m256i a_shifted = _mm256_mullo_epi32(a_8x32, a_multiplier); + __m256i indices = _mm256_add_epi32(a_shifted, b_8x32); + indices = _mm256_add_epi32(indices, base_offsets); + + // Gather 8 floats from dist_table + __m256 gathered = _mm256_i32gather_ps(dist_table, indices, 4); + + acc = _mm256_add_ps(acc, gathered); + } + + float sum = horizontal_sum_avx2(acc); + + // Scalar leftover + for (; m < num_subquantizers; ++m) { + size_t idx = m * kTablePerSub + static_cast(a[m]) * kNumCentroids + + static_cast(b[m]); + sum += dist_table[idx]; + } + + *out = sum; +} + +void pq_adc_int8_batch_distance_avx2(const void **candidates_v, + const void *lut_v, size_t num, + size_t num_subquantizers, float *out) { + constexpr int kNumCentroids = 256; + constexpr int kChunkSize = 8; + constexpr int kBatch = 4; + const auto *lut = reinterpret_cast(lut_v); + const auto *candidates = + reinterpret_cast(candidates_v); + + // Base offsets: [0, 256, 512, ..., 7*256] — reused for all candidates. + const __m256i base_offsets = _mm256_setr_epi32( + 0, kNumCentroids, 2 * kNumCentroids, 3 * kNumCentroids, 4 * kNumCentroids, + 5 * kNumCentroids, 6 * kNumCentroids, 7 * kNumCentroids); + + size_t i = 0; + for (; i + kBatch <= num; i += kBatch) { + const uint8_t *c0 = candidates[i]; + const uint8_t *c1 = candidates[i + 1]; + const uint8_t *c2 = candidates[i + 2]; + const uint8_t *c3 = candidates[i + 3]; + __m256 acc0 = _mm256_setzero_ps(); + __m256 acc1 = _mm256_setzero_ps(); + __m256 acc2 = _mm256_setzero_ps(); + __m256 acc3 = _mm256_setzero_ps(); + + size_t m = 0; + for (; m + kChunkSize <= num_subquantizers; m += kChunkSize) { + const float *lut_base = lut + m * kNumCentroids; + + __m128i codes0 = + _mm_loadl_epi64(reinterpret_cast(c0 + m)); + __m128i codes1 = + _mm_loadl_epi64(reinterpret_cast(c1 + m)); + __m128i codes2 = + _mm_loadl_epi64(reinterpret_cast(c2 + m)); + __m128i codes3 = + _mm_loadl_epi64(reinterpret_cast(c3 + m)); + + __m256i idx0 = + _mm256_add_epi32(_mm256_cvtepu8_epi32(codes0), base_offsets); + __m256i idx1 = + _mm256_add_epi32(_mm256_cvtepu8_epi32(codes1), base_offsets); + __m256i idx2 = + _mm256_add_epi32(_mm256_cvtepu8_epi32(codes2), base_offsets); + __m256i idx3 = + _mm256_add_epi32(_mm256_cvtepu8_epi32(codes3), base_offsets); + + acc0 = _mm256_add_ps(acc0, _mm256_i32gather_ps(lut_base, idx0, 4)); + acc1 = _mm256_add_ps(acc1, _mm256_i32gather_ps(lut_base, idx1, 4)); + acc2 = _mm256_add_ps(acc2, _mm256_i32gather_ps(lut_base, idx2, 4)); + acc3 = _mm256_add_ps(acc3, _mm256_i32gather_ps(lut_base, idx3, 4)); + } + + float s0 = horizontal_sum_avx2(acc0); + float s1 = horizontal_sum_avx2(acc1); + float s2 = horizontal_sum_avx2(acc2); + float s3 = horizontal_sum_avx2(acc3); + + // Scalar leftover for remaining subquantizers. + for (; m < num_subquantizers; ++m) { + const float *tab = lut + m * kNumCentroids; + s0 += tab[c0[m]]; + s1 += tab[c1[m]]; + s2 += tab[c2[m]]; + s3 += tab[c3[m]]; + } + out[i] = s0; + out[i + 1] = s1; + out[i + 2] = s2; + out[i + 3] = s3; + } + // Remaining candidates: use single ADC kernel. + for (; i < num; ++i) { + pq_adc_int8_distance_avx2(candidates[i], lut, num_subquantizers, out + i); + } +} + +} // namespace zvec::turbo::avx2 diff --git a/src/turbo/distance/avx2/pq_quantizer_int8/pq_distance.h b/src/turbo/distance/avx2/pq_quantizer_int8/pq_distance.h new file mode 100644 index 000000000..1d1b412cd --- /dev/null +++ b/src/turbo/distance/avx2/pq_quantizer_int8/pq_distance.h @@ -0,0 +1,42 @@ +// 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 +#include + +namespace zvec::turbo::avx2 { + +// ADC (Asymmetric Distance Computation) via AVX2 gather. +// Processes 8 subquantizers per _mm256_i32gather_ps iteration. +// For general M: loop in chunks of 8, scalar leftover. +void pq_adc_int8_distance_avx2(const void *pq_code, const void *lut, + size_t num_subquantizers, float *out); + +// SDC (Symmetric Distance Computation) via AVX2 gather. +// Computes indices (a[m]*256 + b[m]) as int32, adds per-subquantizer +// base offsets, gathers 8 floats per iteration. +void pq_sdc_int8_distance_avx2(const void *a, const void *b, + const void *dist_table, size_t num_subquantizers, + float *out); + +// Batch ADC via AVX2 gather: process 4 candidates per iteration, +// each using 8-wide _mm256_i32gather_ps. 4 independent __m256 +// accumulators maximize ILP. +void pq_adc_int8_batch_distance_avx2(const void **candidates, const void *lut, + size_t num, size_t num_subquantizers, + float *out); + +} // namespace zvec::turbo::avx2 diff --git a/src/turbo/distance/avx512/fht/fht.cc b/src/turbo/distance/avx512/fht/fht.cc new file mode 100644 index 000000000..2ba9e4b3b --- /dev/null +++ b/src/turbo/distance/avx512/fht/fht.cc @@ -0,0 +1,142 @@ +// 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. + +#if defined(__AVX512F__) + +#include "fht.h" +#include +#include +#include +#include +#include + +namespace zvec::turbo::avx512 { + +void fht_flip_sign_avx512(const uint8_t *flip, float *data, size_t dim) { + size_t simd_end = dim & ~63u; + constexpr size_t kChunk = 64; + const __m512 sign_flip = _mm512_castsi512_ps(_mm512_set1_epi32(0x80000000)); + for (size_t i = 0; i < simd_end; i += kChunk) { + uint64_t mask_bits; + std::memcpy(&mask_bits, &flip[i / 8], sizeof(mask_bits)); + const __mmask16 m0 = _cvtu32_mask16(mask_bits & 0xFFFF); + const __mmask16 m1 = _cvtu32_mask16((mask_bits >> 16) & 0xFFFF); + const __mmask16 m2 = _cvtu32_mask16((mask_bits >> 32) & 0xFFFF); + const __mmask16 m3 = _cvtu32_mask16((mask_bits >> 48) & 0xFFFF); + __m512 v0 = _mm512_loadu_ps(&data[i]); + v0 = _mm512_mask_xor_ps(v0, m0, v0, sign_flip); + _mm512_storeu_ps(&data[i], v0); + __m512 v1 = _mm512_loadu_ps(&data[i + 16]); + v1 = _mm512_mask_xor_ps(v1, m1, v1, sign_flip); + _mm512_storeu_ps(&data[i + 16], v1); + __m512 v2 = _mm512_loadu_ps(&data[i + 32]); + v2 = _mm512_mask_xor_ps(v2, m2, v2, sign_flip); + _mm512_storeu_ps(&data[i + 32], v2); + __m512 v3 = _mm512_loadu_ps(&data[i + 48]); + v3 = _mm512_mask_xor_ps(v3, m3, v3, sign_flip); + _mm512_storeu_ps(&data[i + 48], v3); + } + // Scalar tail + for (size_t i = simd_end; i < dim; ++i) { + if (flip[i / 8] & (1u << (i % 8))) { + data[i] = -data[i]; + } + } +} + +void fht_kacs_walk_avx512(float *data, size_t len) { + size_t half = len / 2; + size_t base = len % 2; + size_t offset = base + half; + size_t half_end = half & ~15u; + for (size_t i = 0; i < half_end; i += 16) { + __m512 x = _mm512_loadu_ps(&data[i]); + __m512 y = _mm512_loadu_ps(&data[i + offset]); + _mm512_storeu_ps(&data[i], _mm512_add_ps(x, y)); + _mm512_storeu_ps(&data[i + offset], _mm512_sub_ps(x, y)); + } + // Scalar tail + for (size_t i = half_end; i < half; ++i) { + float x = data[i]; + float y = data[i + offset]; + data[i] = x + y; + data[i + offset] = x - y; + } + if (base != 0) { + data[half] *= std::sqrt(2.0f); + } +} + +void fht_inv_kacs_walk_avx512(float *data, size_t len) { + size_t half = len / 2; + size_t base = len % 2; + size_t offset = base + half; + if (base != 0) { + data[half] *= std::sqrt(0.5f); + } + size_t half_end = half & ~15u; + const __m512 half_fac = _mm512_set1_ps(0.5f); + for (size_t i = 0; i < half_end; i += 16) { + __m512 a = _mm512_loadu_ps(&data[i]); + __m512 b = _mm512_loadu_ps(&data[i + offset]); + _mm512_storeu_ps(&data[i], _mm512_mul_ps(_mm512_add_ps(a, b), half_fac)); + _mm512_storeu_ps(&data[i + offset], + _mm512_mul_ps(_mm512_sub_ps(a, b), half_fac)); + } + // Scalar tail + for (size_t i = half_end; i < half; ++i) { + float a = data[i]; + float b = data[i + offset]; + data[i] = (a + b) * 0.5f; + data[i + offset] = (a - b) * 0.5f; + } +} + +void fht_inplace_avx512(float *data, size_t n) { + for (size_t len = 1; len < n; len <<= 1) { + size_t step = len << 1; + size_t simd_end = len & ~15u; + for (size_t i = 0; i < n; i += step) { + for (size_t j = 0; j < simd_end; j += 16) { + __m512 u = _mm512_loadu_ps(&data[i + j]); + __m512 v = _mm512_loadu_ps(&data[i + j + len]); + _mm512_storeu_ps(&data[i + j], _mm512_add_ps(u, v)); + _mm512_storeu_ps(&data[i + j + len], _mm512_sub_ps(u, v)); + } + for (size_t j = simd_end; j < len; ++j) { + float u = data[i + j]; + float v = data[i + j + len]; + data[i + j] = u + v; + data[i + j + len] = u - v; + } + } + } +} + +void fht_vec_rescale_avx512(float *data, size_t n, float factor) { + const __m512 fac = _mm512_set1_ps(factor); + size_t simd_end = n & ~15u; + for (size_t i = 0; i < simd_end; i += 16) { + __m512 v = _mm512_loadu_ps(&data[i]); + _mm512_storeu_ps(&data[i], _mm512_mul_ps(v, fac)); + } + // Scalar tail + for (size_t i = simd_end; i < n; ++i) { + data[i] *= factor; + } +} + +} // namespace zvec::turbo::avx512 + +#endif // __AVX512F__ diff --git a/src/turbo/distance/avx512/fht/fht.h b/src/turbo/distance/avx512/fht/fht.h new file mode 100644 index 000000000..53241a0c5 --- /dev/null +++ b/src/turbo/distance/avx512/fht/fht.h @@ -0,0 +1,37 @@ +// 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 +#include + +namespace zvec::turbo::avx512 { + +//! Apply bitwise sign-flip mask to a float vector (AVX512). +void fht_flip_sign_avx512(const uint8_t *flip, float *data, size_t dim); + +//! Apply KacsWalk butterfly operation (AVX512). +void fht_kacs_walk_avx512(float *data, size_t len); + +//! Inverse KacsWalk butterfly operation (AVX512). +void fht_inv_kacs_walk_avx512(float *data, size_t len); + +//! In-place Fast Hadamard Transform (AVX512, n must be power-of-2). +void fht_inplace_avx512(float *data, size_t n); + +//! Element-wise rescale: data[i] *= factor (AVX512). +void fht_vec_rescale_avx512(float *data, size_t n, float factor); + +} // namespace zvec::turbo::avx512 diff --git a/src/turbo/distance/avx512/pq_quantizer_int4/pq_distance.cc b/src/turbo/distance/avx512/pq_quantizer_int4/pq_distance.cc new file mode 100644 index 000000000..e93c4f92c --- /dev/null +++ b/src/turbo/distance/avx512/pq_quantizer_int4/pq_distance.cc @@ -0,0 +1,208 @@ +// 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 "avx512/pq_quantizer_int4/pq_distance.h" +#include + +namespace zvec::turbo::avx512 { + +namespace { + +constexpr int kNumCentroids = 16; +constexpr int kTablePerSub = kNumCentroids * kNumCentroids; // 256 +constexpr int kChunkSize = 16; // process 16 subs per iteration (8 bytes) + +// Unpack 8 bytes of packed int4 codes into 16 nibbles (0-15), +// zero-extended to 16 x int32 in a single __m512i. +// +// Uses vpshufb to extract high nibbles: +// low nibbles = packed & 0x0F (subs 0-7 in bytes 0-7) +// high nibbles = (packed >> 4) & 0x0F (subs 8-15 in bytes 8-15) +inline __m512i unpack_nibbles_512(const uint8_t *packed) { + __m128i packed128 = + _mm_loadl_epi64(reinterpret_cast(packed)); + __m512i packed512 = _mm512_castsi128_si512(packed128); + + __m512i mask_0f = _mm512_set1_epi8(0x0F); + __m512i lo512 = _mm512_and_si512(packed512, mask_0f); + // Low nibbles occupy bytes 0-7 (= subs 0-7), bytes 8-15 are zero. + + // Extract high nibbles into bytes 8-15 via vpshufb: + // mask[1,1,2,2,...,8,8] selects bytes 1-8; after >>4 and &0x0F + // positions 8-15 hold the high nibbles of bytes 1-8 = subs 8-15. + __m512i hi_mask = _mm512_set_epi8( + 8, 8, 7, 7, 6, 6, 5, 5, 4, 4, 3, 3, 2, 2, 1, 1, 8, 8, 7, 7, 6, 6, 5, 5, 4, + 4, 3, 3, 2, 2, 1, 1, 8, 8, 7, 7, 6, 6, 5, 5, 4, 4, 3, 3, 2, 2, 1, 1, 8, 8, + 7, 7, 6, 6, 5, 5, 4, 4, 3, 3, 2, 2, 1, 1); + __m512i hi512 = _mm512_and_si512( + _mm512_srli_epi16(_mm512_shuffle_epi8(packed512, hi_mask), 4), mask_0f); + + // Blend: low nibbles (bytes 0-7) + high nibbles (bytes 8-15) + // 0x00FF = bits set in bytes 0-7, clear in bytes 8-15 + __m512i blend_mask = _mm512_set1_epi16(0x00FF); + __m512i nibbles = _mm512_or_si512(_mm512_and_si512(lo512, blend_mask), + _mm512_andnot_si512(blend_mask, hi512)); + + // Zero-extend 16 bytes (lower 128 bits) to 16 x int32. + // _mm512_cvtepu8_epi32 takes __m128i (16 bytes) -> __m512i (16 x int32). + return _mm512_cvtepu8_epi32(_mm512_castsi512_si128(nibbles)); +} + +} // namespace + +void pq_adc_int4_distance_avx512(const void *pq_code_v, const void *lut_v, + size_t num_subquantizers, float *out) { + const auto *pq_code = reinterpret_cast(pq_code_v); + const auto *lut = reinterpret_cast(lut_v); + + // base_offsets: [0, 16, 32, ..., 15*16] = sub index offsets in float units + const __m512i base_offsets = _mm512_setr_epi32( + 0, kNumCentroids, 2 * kNumCentroids, 3 * kNumCentroids, 4 * kNumCentroids, + 5 * kNumCentroids, 6 * kNumCentroids, 7 * kNumCentroids, + 8 * kNumCentroids, 9 * kNumCentroids, 10 * kNumCentroids, + 11 * kNumCentroids, 12 * kNumCentroids, 13 * kNumCentroids, + 14 * kNumCentroids, 15 * kNumCentroids); + + __m512 acc = _mm512_setzero_ps(); + size_t m = 0; + + for (; m + kChunkSize <= num_subquantizers; m += kChunkSize) { + __m512i nibbles = unpack_nibbles_512(pq_code + (m >> 1)); + __m512i indices = _mm512_add_epi32(nibbles, base_offsets); + __m512 gathered = _mm512_i32gather_ps(indices, lut + m * kNumCentroids, 4); + acc = _mm512_add_ps(acc, gathered); + } + + float sum = _mm512_reduce_add_ps(acc); + // Scalar leftover + for (; m < num_subquantizers; ++m) { + uint8_t byte = pq_code[m >> 1]; + uint8_t idx = (m & 1u) ? (byte >> 4) : (byte & 0x0Fu); + sum += lut[m * kNumCentroids + idx]; + } + *out = sum; +} + +void pq_sdc_int4_distance_avx512(const void *a_v, const void *b_v, + const void *dist_table_v, + size_t num_subquantizers, float *out) { + const auto *a = reinterpret_cast(a_v); + const auto *b = reinterpret_cast(b_v); + const auto *dist_table = reinterpret_cast(dist_table_v); + + const __m512i base_offsets = _mm512_setr_epi32( + 0, kTablePerSub, 2 * kTablePerSub, 3 * kTablePerSub, 4 * kTablePerSub, + 5 * kTablePerSub, 6 * kTablePerSub, 7 * kTablePerSub, 8 * kTablePerSub, + 9 * kTablePerSub, 10 * kTablePerSub, 11 * kTablePerSub, 12 * kTablePerSub, + 13 * kTablePerSub, 14 * kTablePerSub, 15 * kTablePerSub); + const __m512i mul16 = _mm512_set1_epi32(kNumCentroids); + + __m512 acc = _mm512_setzero_ps(); + size_t m = 0; + + for (; m + kChunkSize <= num_subquantizers; m += kChunkSize) { + __m512i a_nibs = unpack_nibbles_512(a + (m >> 1)); + __m512i b_nibs = unpack_nibbles_512(b + (m >> 1)); + // index = m_local * 256 + a_nib * 16 + b_nib + __m512i indices = _mm512_add_epi32( + _mm512_add_epi32(_mm512_mullo_epi32(a_nibs, mul16), b_nibs), + base_offsets); + __m512 gathered = + _mm512_i32gather_ps(indices, dist_table + m * kTablePerSub, 4); + acc = _mm512_add_ps(acc, gathered); + } + + float sum = _mm512_reduce_add_ps(acc); + for (; m < num_subquantizers; ++m) { + uint8_t ab = a[m >> 1], bb = b[m >> 1]; + uint8_t ai = (m & 1u) ? (ab >> 4) : (ab & 0x0Fu); + uint8_t bi = (m & 1u) ? (bb >> 4) : (bb & 0x0Fu); + sum += dist_table[m * kTablePerSub + ai * kNumCentroids + bi]; + } + *out = sum; +} + +void pq_adc_int4_batch_distance_avx512(const void **candidates_v, + const void *lut_v, size_t num, + size_t num_subquantizers, float *out) { + const auto *lut = reinterpret_cast(lut_v); + const auto *candidates = + reinterpret_cast(candidates_v); + + const __m512i base_offsets = _mm512_setr_epi32( + 0, kNumCentroids, 2 * kNumCentroids, 3 * kNumCentroids, 4 * kNumCentroids, + 5 * kNumCentroids, 6 * kNumCentroids, 7 * kNumCentroids, + 8 * kNumCentroids, 9 * kNumCentroids, 10 * kNumCentroids, + 11 * kNumCentroids, 12 * kNumCentroids, 13 * kNumCentroids, + 14 * kNumCentroids, 15 * kNumCentroids); + + size_t i = 0; + for (; i + 4 <= num; i += 4) { + const uint8_t *c0 = candidates[i]; + const uint8_t *c1 = candidates[i + 1]; + const uint8_t *c2 = candidates[i + 2]; + const uint8_t *c3 = candidates[i + 3]; + __m512 acc0 = _mm512_setzero_ps(); + __m512 acc1 = _mm512_setzero_ps(); + __m512 acc2 = _mm512_setzero_ps(); + __m512 acc3 = _mm512_setzero_ps(); + + size_t m = 0; + for (; m + kChunkSize <= num_subquantizers; m += kChunkSize) { + const float *lut_base = lut + m * kNumCentroids; + __m512i nib0 = unpack_nibbles_512(c0 + (m >> 1)); + __m512i nib1 = unpack_nibbles_512(c1 + (m >> 1)); + __m512i nib2 = unpack_nibbles_512(c2 + (m >> 1)); + __m512i nib3 = unpack_nibbles_512(c3 + (m >> 1)); + acc0 = _mm512_add_ps( + acc0, _mm512_i32gather_ps(_mm512_add_epi32(nib0, base_offsets), + lut_base, 4)); + acc1 = _mm512_add_ps( + acc1, _mm512_i32gather_ps(_mm512_add_epi32(nib1, base_offsets), + lut_base, 4)); + acc2 = _mm512_add_ps( + acc2, _mm512_i32gather_ps(_mm512_add_epi32(nib2, base_offsets), + lut_base, 4)); + acc3 = _mm512_add_ps( + acc3, _mm512_i32gather_ps(_mm512_add_epi32(nib3, base_offsets), + lut_base, 4)); + } + + float s0 = _mm512_reduce_add_ps(acc0); + float s1 = _mm512_reduce_add_ps(acc1); + float s2 = _mm512_reduce_add_ps(acc2); + float s3 = _mm512_reduce_add_ps(acc3); + + // Scalar leftover for remaining sub-quantizers. + for (; m < num_subquantizers; ++m) { + const float *tab = lut + m * kNumCentroids; + uint8_t b0 = c0[m >> 1], b1 = c1[m >> 1], b2 = c2[m >> 1], + b3 = c3[m >> 1]; + s0 += tab[(m & 1u) ? (b0 >> 4) : (b0 & 0x0Fu)]; + s1 += tab[(m & 1u) ? (b1 >> 4) : (b1 & 0x0Fu)]; + s2 += tab[(m & 1u) ? (b2 >> 4) : (b2 & 0x0Fu)]; + s3 += tab[(m & 1u) ? (b3 >> 4) : (b3 & 0x0Fu)]; + } + out[i] = s0; + out[i + 1] = s1; + out[i + 2] = s2; + out[i + 3] = s3; + } + // Remaining candidates: use single ADC kernel. + for (; i < num; ++i) { + pq_adc_int4_distance_avx512(candidates[i], lut, num_subquantizers, out + i); + } +} + +} // namespace zvec::turbo::avx512 diff --git a/src/turbo/distance/avx512/pq_quantizer_int4/pq_distance.h b/src/turbo/distance/avx512/pq_quantizer_int4/pq_distance.h new file mode 100644 index 000000000..6a84bfe5b --- /dev/null +++ b/src/turbo/distance/avx512/pq_quantizer_int4/pq_distance.h @@ -0,0 +1,40 @@ +// 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 +#include + +namespace zvec::turbo::avx512 { + +// ADC (Asymmetric Distance Computation) for int4 PQ codes via AVX512. +// Uses vpshufb to unpack 16 nibbles from 8 bytes, then a single +// _mm512_i32gather_ps to process all 16 sub-quantizers at once. +void pq_adc_int4_distance_avx512(const void *pq_code, const void *lut, + size_t num_subquantizers, float *out); + +// SDC (Symmetric Distance Computation) for int4 PQ codes via AVX512. +// 16-wide index computation (a*16 + b + m*256) + single gather. +void pq_sdc_int4_distance_avx512(const void *a, const void *b, + const void *dist_table, + size_t num_subquantizers, float *out); + +// Batch ADC via AVX512: process 4 candidates per iteration, +// each using the 16-sub vpshufb + gather kernel. +void pq_adc_int4_batch_distance_avx512(const void **candidates, const void *lut, + size_t num, size_t num_subquantizers, + float *out); + +} // namespace zvec::turbo::avx512 diff --git a/src/turbo/distance/avx512/pq_quantizer_int8/pq_distance.cc b/src/turbo/distance/avx512/pq_quantizer_int8/pq_distance.cc new file mode 100644 index 000000000..8ad40e1e4 --- /dev/null +++ b/src/turbo/distance/avx512/pq_quantizer_int8/pq_distance.cc @@ -0,0 +1,215 @@ +// 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 "avx512/pq_quantizer_int8/pq_distance.h" +#include + +namespace zvec::turbo::avx512 { + +namespace { + +// Horizontal sum of 16 floats in a __m512 register. +inline float horizontal_sum_avx512(__m512 v) { + // Use _mm512_reduce_add_ps which is available in AVX512F + return _mm512_reduce_add_ps(v); +} + +} // namespace + +void pq_adc_int8_distance_avx512(const void *pq_code_v, const void *lut_v, + size_t num_subquantizers, float *out) { + constexpr int kNumCentroids = 256; + constexpr int kChunkSize = 16; // AVX512 processes 16 floats at once + const auto *pq_code = reinterpret_cast(pq_code_v); + const auto *lut = reinterpret_cast(lut_v); + + __m512 acc = _mm512_setzero_ps(); + + // Base offsets: [0, 256, 512, ..., 15*256] = m * 256 for m = 0..15 + const __m512i base_offsets = _mm512_setr_epi32( + 0 * kNumCentroids, 1 * kNumCentroids, 2 * kNumCentroids, + 3 * kNumCentroids, 4 * kNumCentroids, 5 * kNumCentroids, + 6 * kNumCentroids, 7 * kNumCentroids, 8 * kNumCentroids, + 9 * kNumCentroids, 10 * kNumCentroids, 11 * kNumCentroids, + 12 * kNumCentroids, 13 * kNumCentroids, 14 * kNumCentroids, + 15 * kNumCentroids); + + size_t m = 0; + + // Main loop: process 16 subquantizers per iteration + for (; m + kChunkSize <= num_subquantizers; m += kChunkSize) { + // Load 16 uint8 codes and zero-extend to int32 + // Use unaligned load of 16 bytes + __m128i codes_16x8 = + _mm_loadu_si128(reinterpret_cast(pq_code + m)); + __m512i codes_16x32 = _mm512_cvtepu8_epi32(codes_16x8); + + // Add base offsets: indices[m] = m * 256 + code[m] + __m512i indices = _mm512_add_epi32(codes_16x32, base_offsets); + + // Gather 16 floats from lut using computed indices + __m512 gathered = _mm512_i32gather_ps(indices, lut + m * kNumCentroids, 4); + + acc = _mm512_add_ps(acc, gathered); + } + + float sum = horizontal_sum_avx512(acc); + + // Scalar leftover: process remaining subquantizers + for (; m < num_subquantizers; ++m) { + sum += lut[m * kNumCentroids + pq_code[m]]; + } + + *out = sum; +} + +void pq_sdc_int8_distance_avx512(const void *a_v, const void *b_v, + const void *dist_table_v, + size_t num_subquantizers, float *out) { + constexpr int kNumCentroids = 256; + constexpr int kTablePerSub = kNumCentroids * kNumCentroids; // 65536 + constexpr int kChunkSize = 16; + const auto *a = reinterpret_cast(a_v); + const auto *b = reinterpret_cast(b_v); + const auto *dist_table = reinterpret_cast(dist_table_v); + + __m512 acc = _mm512_setzero_ps(); + + // Base offsets for SDC: m * 65536 + const __m512i base_offsets = _mm512_setr_epi32( + 0 * kTablePerSub, 1 * kTablePerSub, 2 * kTablePerSub, 3 * kTablePerSub, + 4 * kTablePerSub, 5 * kTablePerSub, 6 * kTablePerSub, 7 * kTablePerSub, + 8 * kTablePerSub, 9 * kTablePerSub, 10 * kTablePerSub, 11 * kTablePerSub, + 12 * kTablePerSub, 13 * kTablePerSub, 14 * kTablePerSub, + 15 * kTablePerSub); + + // Multiplier for a[m] * 256 + const __m512i a_multiplier = _mm512_set1_epi32(kNumCentroids); + + size_t m = 0; + + // Main loop: process 16 subquantizers per iteration + for (; m + kChunkSize <= num_subquantizers; m += kChunkSize) { + // Load a[m..m+15] and b[m..m+15], zero-extend to int32 + __m128i a_16x8 = _mm_loadu_si128(reinterpret_cast(a + m)); + __m128i b_16x8 = _mm_loadu_si128(reinterpret_cast(b + m)); + __m512i a_16x32 = _mm512_cvtepu8_epi32(a_16x8); + __m512i b_16x32 = _mm512_cvtepu8_epi32(b_16x8); + + // Compute index: a[m] * 256 + b[m] + m * 65536 + __m512i a_shifted = _mm512_mullo_epi32(a_16x32, a_multiplier); + __m512i indices = _mm512_add_epi32(a_shifted, b_16x32); + indices = _mm512_add_epi32(indices, base_offsets); + + // Gather 16 floats from dist_table + __m512 gathered = _mm512_i32gather_ps(indices, dist_table, 4); + + acc = _mm512_add_ps(acc, gathered); + } + + float sum = horizontal_sum_avx512(acc); + + // Scalar leftover + for (; m < num_subquantizers; ++m) { + size_t idx = m * kTablePerSub + static_cast(a[m]) * kNumCentroids + + static_cast(b[m]); + sum += dist_table[idx]; + } + + *out = sum; +} + +void pq_adc_int8_batch_distance_avx512(const void **candidates_v, + const void *lut_v, size_t num, + size_t num_subquantizers, float *out) { + constexpr int kNumCentroids = 256; + constexpr int kChunkSize = 16; + constexpr int kBatch = 4; + const auto *lut = reinterpret_cast(lut_v); + const auto *candidates = + reinterpret_cast(candidates_v); + + // Base offsets: [0, 256, 512, ..., 15*256] — reused for all candidates. + const __m512i base_offsets = _mm512_setr_epi32( + 0 * kNumCentroids, 1 * kNumCentroids, 2 * kNumCentroids, + 3 * kNumCentroids, 4 * kNumCentroids, 5 * kNumCentroids, + 6 * kNumCentroids, 7 * kNumCentroids, 8 * kNumCentroids, + 9 * kNumCentroids, 10 * kNumCentroids, 11 * kNumCentroids, + 12 * kNumCentroids, 13 * kNumCentroids, 14 * kNumCentroids, + 15 * kNumCentroids); + + size_t i = 0; + for (; i + kBatch <= num; i += kBatch) { + const uint8_t *c0 = candidates[i]; + const uint8_t *c1 = candidates[i + 1]; + const uint8_t *c2 = candidates[i + 2]; + const uint8_t *c3 = candidates[i + 3]; + __m512 acc0 = _mm512_setzero_ps(); + __m512 acc1 = _mm512_setzero_ps(); + __m512 acc2 = _mm512_setzero_ps(); + __m512 acc3 = _mm512_setzero_ps(); + + size_t m = 0; + for (; m + kChunkSize <= num_subquantizers; m += kChunkSize) { + const float *lut_base = lut + m * kNumCentroids; + + __m128i codes0 = + _mm_loadu_si128(reinterpret_cast(c0 + m)); + __m128i codes1 = + _mm_loadu_si128(reinterpret_cast(c1 + m)); + __m128i codes2 = + _mm_loadu_si128(reinterpret_cast(c2 + m)); + __m128i codes3 = + _mm_loadu_si128(reinterpret_cast(c3 + m)); + + __m512i idx0 = + _mm512_add_epi32(_mm512_cvtepu8_epi32(codes0), base_offsets); + __m512i idx1 = + _mm512_add_epi32(_mm512_cvtepu8_epi32(codes1), base_offsets); + __m512i idx2 = + _mm512_add_epi32(_mm512_cvtepu8_epi32(codes2), base_offsets); + __m512i idx3 = + _mm512_add_epi32(_mm512_cvtepu8_epi32(codes3), base_offsets); + + acc0 = _mm512_add_ps(acc0, _mm512_i32gather_ps(idx0, lut_base, 4)); + acc1 = _mm512_add_ps(acc1, _mm512_i32gather_ps(idx1, lut_base, 4)); + acc2 = _mm512_add_ps(acc2, _mm512_i32gather_ps(idx2, lut_base, 4)); + acc3 = _mm512_add_ps(acc3, _mm512_i32gather_ps(idx3, lut_base, 4)); + } + + float s0 = _mm512_reduce_add_ps(acc0); + float s1 = _mm512_reduce_add_ps(acc1); + float s2 = _mm512_reduce_add_ps(acc2); + float s3 = _mm512_reduce_add_ps(acc3); + + // Scalar leftover for remaining subquantizers. + for (; m < num_subquantizers; ++m) { + const float *tab = lut + m * kNumCentroids; + s0 += tab[c0[m]]; + s1 += tab[c1[m]]; + s2 += tab[c2[m]]; + s3 += tab[c3[m]]; + } + out[i] = s0; + out[i + 1] = s1; + out[i + 2] = s2; + out[i + 3] = s3; + } + // Remaining candidates: use single ADC kernel. + for (; i < num; ++i) { + pq_adc_int8_distance_avx512(candidates[i], lut, num_subquantizers, out + i); + } +} + +} // namespace zvec::turbo::avx512 diff --git a/src/turbo/distance/avx512/pq_quantizer_int8/pq_distance.h b/src/turbo/distance/avx512/pq_quantizer_int8/pq_distance.h new file mode 100644 index 000000000..2c2ea92d5 --- /dev/null +++ b/src/turbo/distance/avx512/pq_quantizer_int8/pq_distance.h @@ -0,0 +1,40 @@ +// 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 +#include + +namespace zvec::turbo::avx512 { + +// ADC (Asymmetric Distance Computation) via AVX512 gather. +// Processes 16 subquantizers per _mm512_i32gather_ps iteration. +void pq_adc_int8_distance_avx512(const void *pq_code, const void *lut, + size_t num_subquantizers, float *out); + +// SDC (Symmetric Distance Computation) via AVX512 gather. +// 16-wide index computation + gather. +void pq_sdc_int8_distance_avx512(const void *a, const void *b, + const void *dist_table, + size_t num_subquantizers, float *out); + +// Batch ADC via AVX512 gather: process 4 candidates per iteration, +// each using 16-wide _mm512_i32gather_ps. 4 independent __m512 +// accumulators maximize ILP. +void pq_adc_int8_batch_distance_avx512(const void **candidates, const void *lut, + size_t num, size_t num_subquantizers, + float *out); + +} // namespace zvec::turbo::avx512 diff --git a/src/turbo/avx512_vnni/record_quantized_int8/common.h b/src/turbo/distance/avx512_vnni/record_quantized_int8/common.h similarity index 100% rename from src/turbo/avx512_vnni/record_quantized_int8/common.h rename to src/turbo/distance/avx512_vnni/record_quantized_int8/common.h diff --git a/src/turbo/avx512_vnni/record_quantized_int8/cosine.cc b/src/turbo/distance/avx512_vnni/record_quantized_int8/cosine.cc similarity index 100% rename from src/turbo/avx512_vnni/record_quantized_int8/cosine.cc rename to src/turbo/distance/avx512_vnni/record_quantized_int8/cosine.cc diff --git a/src/turbo/avx512_vnni/record_quantized_int8/cosine.h b/src/turbo/distance/avx512_vnni/record_quantized_int8/cosine.h similarity index 100% rename from src/turbo/avx512_vnni/record_quantized_int8/cosine.h rename to src/turbo/distance/avx512_vnni/record_quantized_int8/cosine.h diff --git a/src/turbo/avx512_vnni/record_quantized_int8/squared_euclidean.cc b/src/turbo/distance/avx512_vnni/record_quantized_int8/squared_euclidean.cc similarity index 100% rename from src/turbo/avx512_vnni/record_quantized_int8/squared_euclidean.cc rename to src/turbo/distance/avx512_vnni/record_quantized_int8/squared_euclidean.cc diff --git a/src/turbo/avx512_vnni/record_quantized_int8/squared_euclidean.h b/src/turbo/distance/avx512_vnni/record_quantized_int8/squared_euclidean.h similarity index 100% rename from src/turbo/avx512_vnni/record_quantized_int8/squared_euclidean.h rename to src/turbo/distance/avx512_vnni/record_quantized_int8/squared_euclidean.h diff --git a/src/turbo/avx512_vnni/uniform_int8/quantize.cc b/src/turbo/distance/avx512_vnni/uniform_int8/quantize.cc similarity index 100% rename from src/turbo/avx512_vnni/uniform_int8/quantize.cc rename to src/turbo/distance/avx512_vnni/uniform_int8/quantize.cc diff --git a/src/turbo/avx512_vnni/uniform_int8/quantize.h b/src/turbo/distance/avx512_vnni/uniform_int8/quantize.h similarity index 100% rename from src/turbo/avx512_vnni/uniform_int8/quantize.h rename to src/turbo/distance/avx512_vnni/uniform_int8/quantize.h diff --git a/src/turbo/avx512_vnni/uniform_int8/squared_euclidean.cc b/src/turbo/distance/avx512_vnni/uniform_int8/squared_euclidean.cc similarity index 100% rename from src/turbo/avx512_vnni/uniform_int8/squared_euclidean.cc rename to src/turbo/distance/avx512_vnni/uniform_int8/squared_euclidean.cc diff --git a/src/turbo/avx512_vnni/uniform_int8/squared_euclidean.h b/src/turbo/distance/avx512_vnni/uniform_int8/squared_euclidean.h similarity index 100% rename from src/turbo/avx512_vnni/uniform_int8/squared_euclidean.h rename to src/turbo/distance/avx512_vnni/uniform_int8/squared_euclidean.h diff --git a/src/turbo/distance/scalar/fht/fht.cc b/src/turbo/distance/scalar/fht/fht.cc new file mode 100644 index 000000000..f1cdffd78 --- /dev/null +++ b/src/turbo/distance/scalar/fht/fht.cc @@ -0,0 +1,79 @@ +// 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 "fht.h" +#include +#include +#include + +namespace zvec::turbo::scalar { + +void fht_flip_sign(const uint8_t *flip, float *data, size_t dim) { + for (size_t i = 0; i < dim; ++i) { + if (flip[i / 8] & (1u << (i % 8))) { + data[i] = -data[i]; + } + } +} + +void fht_kacs_walk(float *data, size_t len) { + size_t half = len / 2; + size_t base = len % 2; + size_t offset = base + half; + for (size_t i = 0; i < half; ++i) { + float x = data[i]; + float y = data[i + offset]; + data[i] = x + y; + data[i + offset] = x - y; + } + if (base != 0) { + data[half] *= std::sqrt(2.0f); + } +} + +void fht_inv_kacs_walk(float *data, size_t len) { + size_t half = len / 2; + size_t base = len % 2; + size_t offset = base + half; + if (base != 0) { + data[half] *= std::sqrt(0.5f); + } + for (size_t i = 0; i < half; ++i) { + float a = data[i]; + float b = data[i + offset]; + data[i] = (a + b) * 0.5f; + data[i + offset] = (a - b) * 0.5f; + } +} + +void fht_inplace(float *data, size_t n) { + for (size_t len = 1; len < n; len <<= 1) { + for (size_t i = 0; i < n; i += len << 1) { + for (size_t j = i; j < i + len; ++j) { + float u = data[j]; + float v = data[j + len]; + data[j] = u + v; + data[j + len] = u - v; + } + } + } +} + +void fht_vec_rescale(float *data, size_t n, float factor) { + for (size_t i = 0; i < n; ++i) { + data[i] *= factor; + } +} + +} // namespace zvec::turbo::scalar diff --git a/src/turbo/distance/scalar/fht/fht.h b/src/turbo/distance/scalar/fht/fht.h new file mode 100644 index 000000000..7ba98be5d --- /dev/null +++ b/src/turbo/distance/scalar/fht/fht.h @@ -0,0 +1,38 @@ +// 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 +#include + +namespace zvec::turbo::scalar { + +//! Apply bitwise sign-flip mask to a float vector. +//! Each bit in \p flip controls one element: bit set means negate. +void fht_flip_sign(const uint8_t *flip, float *data, size_t dim); + +//! Apply KacsWalk butterfly operation to non-power-of-2 FHT. +void fht_kacs_walk(float *data, size_t len); + +//! Inverse KacsWalk butterfly operation. +void fht_inv_kacs_walk(float *data, size_t len); + +//! In-place Fast Hadamard Transform on \p n elements (must be power-of-2). +void fht_inplace(float *data, size_t n); + +//! Element-wise rescale: data[i] *= factor. +void fht_vec_rescale(float *data, size_t n, float factor); + +} // namespace zvec::turbo::scalar diff --git a/src/turbo/distance/scalar/fp32/cosine.cc b/src/turbo/distance/scalar/fp32/cosine.cc new file mode 100644 index 000000000..6dd9fc532 --- /dev/null +++ b/src/turbo/distance/scalar/fp32/cosine.cc @@ -0,0 +1,40 @@ +// 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 "scalar/fp32/cosine.h" +#include "scalar/fp32/inner_product.h" + +namespace zvec::turbo::scalar { + +void cosine_fp32_distance(const void *a, const void *b, size_t dim, + float *distance) { + // inner_product_fp32_distance returns -real_IP; cosine = 1 - real_IP = 1 + + // ip. + float ip; + inner_product_fp32_distance(a, b, dim, &ip); + + *distance = 1 + ip; +} + +void cosine_fp32_batch_distance(const void *const *vectors, const void *query, + size_t n, size_t dim, float *distances) { + inner_product_fp32_batch_distance(vectors, query, n, dim, distances); + // inner_product batch returns -real_IP per element; cosine = 1 - real_IP = 1 + // + d. + for (size_t i = 0; i < n; i++) { + distances[i] = 1 + distances[i]; + } +} + +} // namespace zvec::turbo::scalar \ No newline at end of file diff --git a/src/turbo/distance/scalar/fp32/cosine.h b/src/turbo/distance/scalar/fp32/cosine.h new file mode 100644 index 000000000..b5e4f4eee --- /dev/null +++ b/src/turbo/distance/scalar/fp32/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 + +namespace zvec::turbo::scalar { + +// 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::scalar \ No newline at end of file diff --git a/src/turbo/distance/scalar/fp32/inner_product.cc b/src/turbo/distance/scalar/fp32/inner_product.cc new file mode 100644 index 000000000..1f966d788 --- /dev/null +++ b/src/turbo/distance/scalar/fp32/inner_product.cc @@ -0,0 +1,43 @@ +// 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 "scalar/fp32/inner_product.h" + +namespace zvec::turbo::scalar { + +// Compute negated inner product between a single FP32 vector pair. +// Returns -dot(a, b) so that callers can derive cosine distance as 1 + ip. +void inner_product_fp32_distance(const void *a, const void *b, size_t dim, + float *distance) { + const float *m = reinterpret_cast(a); + const float *q = reinterpret_cast(b); + + float sum = 0.0; + for (size_t i = 0; i < dim; ++i) { + sum += static_cast(m[i] * q[i]); + } + + *distance = -sum; +} + +// Batch version of inner_product_fp32_distance. +void inner_product_fp32_batch_distance(const void *const *vectors, + const void *query, size_t n, size_t dim, + float *distances) { + for (size_t i = 0; i < n; ++i) { + inner_product_fp32_distance(vectors[i], query, dim, &distances[i]); + } +} + +} // namespace zvec::turbo::scalar \ No newline at end of file diff --git a/src/turbo/distance/scalar/fp32/inner_product.h b/src/turbo/distance/scalar/fp32/inner_product.h new file mode 100644 index 000000000..d4e03418e --- /dev/null +++ b/src/turbo/distance/scalar/fp32/inner_product.h @@ -0,0 +1,31 @@ +// 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 + +namespace zvec::turbo::scalar { + +// Compute inner product distance between a single quantized FP32 +// vector pair. +void inner_product_fp32_distance(const void *a, const void *b, size_t dim, + float *distance); + +// Batch version of inner_product_fp32_distance. +void inner_product_fp32_batch_distance(const void *const *vectors, + const void *query, size_t n, size_t dim, + float *distances); + +} // namespace zvec::turbo::scalar diff --git a/src/turbo/distance/scalar/fp32/squared_euclidean.cc b/src/turbo/distance/scalar/fp32/squared_euclidean.cc new file mode 100644 index 000000000..602d8bd7f --- /dev/null +++ b/src/turbo/distance/scalar/fp32/squared_euclidean.cc @@ -0,0 +1,41 @@ +// 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 "scalar/fp32/squared_euclidean.h" +#include + +namespace zvec::turbo::scalar { + +void squared_euclidean_fp32_distance(const void *a, const void *b, size_t dim, + float *distance) { + const float *m = reinterpret_cast(a); + const float *q = reinterpret_cast(b); + + float sum = 0.0; + for (size_t i = 0; i < dim; ++i) { + sum += zvec::ailego::MathHelper::SquaredDifference(m[i], q[i]); + } + + *distance = sum; +} + +void squared_euclidean_fp32_batch_distance(const void *const *vectors, + const void *query, size_t n, + size_t dim, float *distances) { + for (size_t i = 0; i < n; ++i) { + squared_euclidean_fp32_distance(vectors[i], query, dim, &distances[i]); + } +} + +} // namespace zvec::turbo::scalar \ No newline at end of file diff --git a/src/turbo/distance/scalar/fp32/squared_euclidean.h b/src/turbo/distance/scalar/fp32/squared_euclidean.h new file mode 100644 index 000000000..bf319c1d2 --- /dev/null +++ b/src/turbo/distance/scalar/fp32/squared_euclidean.h @@ -0,0 +1,31 @@ +// 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 + +namespace zvec::turbo::scalar { + +// Compute squared euclidean distance between a single quantized FP32 +// vector pair. +void squared_euclidean_fp32_distance(const void *a, const void *b, size_t dim, + float *distance); + +// Batch version of squared euclidean FP32. +void squared_euclidean_fp32_batch_distance(const void *const *vectors, + const void *query, size_t n, + size_t dim, float *distances); + +} // namespace zvec::turbo::scalar diff --git a/src/turbo/distance/scalar/pq_quantizer_int4/pq_distance.cc b/src/turbo/distance/scalar/pq_quantizer_int4/pq_distance.cc new file mode 100644 index 000000000..4684fe67c --- /dev/null +++ b/src/turbo/distance/scalar/pq_quantizer_int4/pq_distance.cc @@ -0,0 +1,108 @@ +// 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 "scalar/pq_quantizer_int4/pq_distance.h" + +namespace zvec::turbo::scalar { + +namespace { + +// Decode the 4-bit index for sub-quantizer m from a packed int4 code. +// Layout: byte[m/2] = (code[2*(m/2)+1] << 4) | code[2*(m/2)] +// even m -> low nibble; odd m -> high nibble. +inline uint8_t decode_nibble(const uint8_t *code, size_t m) { + uint8_t byte = code[m >> 1]; + return (m & 1u) ? static_cast(byte >> 4) + : static_cast(byte & 0x0Fu); +} + +} // namespace + +void pq_adc_int4_distance(const void *pq_code_v, const void *lut_v, + size_t num_subquantizers, float *out) { + constexpr size_t kNumCentroids = 16; + const auto *pq_code = reinterpret_cast(pq_code_v); + const auto *lut = reinterpret_cast(lut_v); + float sum = 0.0f; + for (size_t m = 0; m < num_subquantizers; ++m) { + uint8_t idx = decode_nibble(pq_code, m); + sum += lut[m * kNumCentroids + idx]; + } + *out = sum; +} + +void pq_sdc_int4_distance(const void *a_v, const void *b_v, + const void *dist_table_v, size_t num_subquantizers, + float *out) { + constexpr size_t kNumCentroids = 16; + constexpr size_t kTablePerSub = kNumCentroids * kNumCentroids; // 256 + const auto *a = reinterpret_cast(a_v); + const auto *b = reinterpret_cast(b_v); + const auto *dist_table = reinterpret_cast(dist_table_v); + float sum = 0.0f; + for (size_t m = 0; m < num_subquantizers; ++m) { + uint8_t ai = decode_nibble(a, m); + uint8_t bi = decode_nibble(b, m); + size_t idx = m * kTablePerSub + static_cast(ai) * kNumCentroids + + static_cast(bi); + sum += dist_table[idx]; + } + *out = sum; +} + +void pq_adc_int4_batch_distance(const void **candidates_v, const void *lut_v, + size_t num, size_t num_subquantizers, + float *out) { + constexpr size_t kNumCentroids = 16; + const auto *lut = reinterpret_cast(lut_v); + const auto *candidates = + reinterpret_cast(candidates_v); + + size_t i = 0; + // Main loop: process 4 candidates per iteration (batch4 ILP). + for (; i + 4 <= num; i += 4) { + const uint8_t *c0 = candidates[i]; + const uint8_t *c1 = candidates[i + 1]; + const uint8_t *c2 = candidates[i + 2]; + const uint8_t *c3 = candidates[i + 3]; + float d0 = 0.0f, d1 = 0.0f, d2 = 0.0f, d3 = 0.0f; + for (size_t m = 0; m < num_subquantizers; ++m) { + const float *tab = lut + m * kNumCentroids; + uint8_t n0 = decode_nibble(c0, m); + uint8_t n1 = decode_nibble(c1, m); + uint8_t n2 = decode_nibble(c2, m); + uint8_t n3 = decode_nibble(c3, m); + d0 += tab[n0]; + d1 += tab[n1]; + d2 += tab[n2]; + d3 += tab[n3]; + } + out[i] = d0; + out[i + 1] = d1; + out[i + 2] = d2; + out[i + 3] = d3; + } + // Scalar leftover: remaining candidates processed one at a time. + for (; i < num; ++i) { + const uint8_t *code = candidates[i]; + float d = 0.0f; + for (size_t m = 0; m < num_subquantizers; ++m) { + uint8_t idx = decode_nibble(code, m); + d += lut[m * kNumCentroids + idx]; + } + out[i] = d; + } +} + +} // namespace zvec::turbo::scalar diff --git a/src/turbo/distance/scalar/pq_quantizer_int4/pq_distance.h b/src/turbo/distance/scalar/pq_quantizer_int4/pq_distance.h new file mode 100644 index 000000000..561df9b6b --- /dev/null +++ b/src/turbo/distance/scalar/pq_quantizer_int4/pq_distance.h @@ -0,0 +1,51 @@ +// 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 +#include + +namespace zvec::turbo::scalar { + +// ADC (Asymmetric Distance Computation) for int4 PQ codes. +// +// pq_code layout: num_subquantizers 4-bit indices packed into +// num_subquantizers/2 bytes. byte[i] = (code[2*i+1] << 4) | code[2*i]. +// num_subquantizers MUST be even. +// +// lut layout: [num_subquantizers * 16] floats (kNumCentroids=16 for int4). +// lut[m * 16 + c] = distance from query sub-vector m to centroid c. +// +// distance = sum_{m=0}^{nsq-1} lut[m * 16 + nibble(m)] +void pq_adc_int4_distance(const void *pq_code, const void *lut, + size_t num_subquantizers, float *out); + +// SDC (Symmetric Distance Computation) for int4 PQ codes. +// +// dist_table layout: [num_subquantizers * 16 * 16] floats. +// dist_table[m * 256 + i * 16 + j] = ||centroid[m][i] - centroid[m][j]||^2 +// +// distance = sum_{m=0}^{nsq-1} dist_table[m*256 + nibble_a(m)*16 + nibble_b(m)] +void pq_sdc_int4_distance(const void *a, const void *b, const void *dist_table, + size_t num_subquantizers, float *out); + +// Batch ADC: compute distances for multiple int4 PQ codes against a shared +// LUT. Processes 4 candidates per iteration (batch4) with 4 independent +// accumulators for ILP. Scalar leftover for remaining candidates. +void pq_adc_int4_batch_distance(const void **candidates, const void *lut, + size_t num, size_t num_subquantizers, + float *out); + +} // namespace zvec::turbo::scalar diff --git a/src/turbo/distance/scalar/pq_quantizer_int8/pq_distance.cc b/src/turbo/distance/scalar/pq_quantizer_int8/pq_distance.cc new file mode 100644 index 000000000..3fb46f1a8 --- /dev/null +++ b/src/turbo/distance/scalar/pq_quantizer_int8/pq_distance.cc @@ -0,0 +1,90 @@ +// 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 "scalar/pq_quantizer_int8/pq_distance.h" + +namespace zvec::turbo::scalar { + +void pq_adc_int8_distance(const void *pq_code_v, const void *lut_v, + size_t num_subquantizers, float *out) { + constexpr size_t kNumCentroids = 256; + const auto *pq_code = reinterpret_cast(pq_code_v); + const auto *lut = reinterpret_cast(lut_v); + float sum = 0.0f; + for (size_t m = 0; m < num_subquantizers; ++m) { + sum += lut[m * kNumCentroids + pq_code[m]]; + } + *out = sum; +} + +void pq_sdc_int8_distance(const void *a_v, const void *b_v, + const void *dist_table_v, size_t num_subquantizers, + float *out) { + constexpr size_t kNumCentroids = 256; + constexpr size_t kTablePerSub = kNumCentroids * kNumCentroids; // 65536 + const auto *a = reinterpret_cast(a_v); + const auto *b = reinterpret_cast(b_v); + const auto *dist_table = reinterpret_cast(dist_table_v); + float sum = 0.0f; + for (size_t m = 0; m < num_subquantizers; ++m) { + size_t idx = m * kTablePerSub + static_cast(a[m]) * kNumCentroids + + static_cast(b[m]); + sum += dist_table[idx]; + } + *out = sum; +} + +void pq_adc_int8_batch_distance(const void **candidates_v, const void *lut_v, + size_t num, size_t num_subquantizers, + float *out) { + constexpr size_t kNumCentroids = 256; + const auto *lut = reinterpret_cast(lut_v); + // candidates_v is const void**, but we need const uint8_t** + // Use an intermediate cast through const char** to avoid aliasing issues. + auto candidates = reinterpret_cast(candidates_v); + + size_t i = 0; + // Main loop: process 4 candidates per iteration. + // Shared LUT base pointer (tab) is computed once per subquantizer, + // reducing redundant pointer arithmetic across the 4 candidates. + for (; i + 4 <= num; i += 4) { + const uint8_t *c0 = candidates[i]; + const uint8_t *c1 = candidates[i + 1]; + const uint8_t *c2 = candidates[i + 2]; + const uint8_t *c3 = candidates[i + 3]; + float d0 = 0.0f, d1 = 0.0f, d2 = 0.0f, d3 = 0.0f; + for (size_t m = 0; m < num_subquantizers; ++m) { + const float *tab = lut + m * kNumCentroids; + d0 += tab[c0[m]]; + d1 += tab[c1[m]]; + d2 += tab[c2[m]]; + d3 += tab[c3[m]]; + } + out[i] = d0; + out[i + 1] = d1; + out[i + 2] = d2; + out[i + 3] = d3; + } + // Scalar leftover: remaining candidates processed one at a time. + for (; i < num; ++i) { + const uint8_t *code = candidates[i]; + float d = 0.0f; + for (size_t m = 0; m < num_subquantizers; ++m) { + d += lut[m * kNumCentroids + code[m]]; + } + out[i] = d; + } +} + +} // namespace zvec::turbo::scalar diff --git a/src/turbo/distance/scalar/pq_quantizer_int8/pq_distance.h b/src/turbo/distance/scalar/pq_quantizer_int8/pq_distance.h new file mode 100644 index 000000000..d968aad7d --- /dev/null +++ b/src/turbo/distance/scalar/pq_quantizer_int8/pq_distance.h @@ -0,0 +1,50 @@ +// 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 +#include + +namespace zvec::turbo::scalar { + +// ADC (Asymmetric Distance Computation): compute the distance between a +// PQ-encoded datapoint and a query using a precomputed LUT. +// +// distance = sum_{m=0}^{num_subquantizers-1} lut[m * 256 + pq_code[m]] +void pq_adc_int8_distance(const void *pq_code, const void *lut, + size_t num_subquantizers, float *out); + +// SDC (Symmetric Distance Computation): compute the distance between two +// PQ-encoded datapoints using a precomputed centroid-to-centroid distance +// table. +// +// dist_table layout: [num_subquantizers * 256 * 256] +// dist_table[m * 256 * 256 + i * 256 + j] = +// ||centroid[m][i] - centroid[m][j]||^2 +// +// distance = sum_{m=0}^{num_subquantizers-1} +// dist_table[m * 65536 + a[m] * 256 + b[m]] +void pq_sdc_int8_distance(const void *a, const void *b, const void *dist_table, + size_t num_subquantizers, float *out); + +// Batch ADC: compute distances for multiple PQ codes against a shared LUT. +// Processes 4 candidates per iteration (batch4) with shared LUT pointer +// offsets and 4 independent accumulators for ILP. +// Falls back to scalar per-code loop for the remaining candidates. +void pq_adc_int8_batch_distance(const void **candidates, const void *lut, + size_t num, size_t num_subquantizers, + float *out); + +} // namespace zvec::turbo::scalar diff --git a/src/turbo/distance/sse/fht/fht.cc b/src/turbo/distance/sse/fht/fht.cc new file mode 100644 index 000000000..18b12c545 --- /dev/null +++ b/src/turbo/distance/sse/fht/fht.cc @@ -0,0 +1,118 @@ +// 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. + +#if defined(__SSE2__) + +#include "fht.h" +#include +#include +#include +#include +#include + +namespace zvec::turbo::sse { + +void fht_flip_sign_sse(const uint8_t *flip, float *data, size_t dim) { + size_t simd_end = dim & ~3u; + size_t flip_bytes = (dim + 7) / 8; + for (size_t i = 0; i < simd_end; i += 4) { + uint16_t bits16; + size_t byte_pos = i / 8; + if (byte_pos + 1 < flip_bytes) { + std::memcpy(&bits16, &flip[byte_pos], sizeof(bits16)); + } else { + bits16 = flip[byte_pos]; + } + bits16 >>= (i % 8); + uint32_t b0 = bits16 & 1u; + uint32_t b1 = (bits16 >> 1) & 1u; + uint32_t b2 = (bits16 >> 2) & 1u; + uint32_t b3 = (bits16 >> 3) & 1u; + __m128i bit_mask = _mm_set_epi32(b3, b2, b1, b0); + __m128i sign_mask = _mm_slli_epi32(bit_mask, 31); + __m128 v = _mm_loadu_ps(&data[i]); + v = _mm_xor_ps(v, _mm_castsi128_ps(sign_mask)); + _mm_storeu_ps(&data[i], v); + } + // Scalar tail + for (size_t i = simd_end; i < dim; ++i) { + if (flip[i / 8] & (1u << (i % 8))) { + data[i] = -data[i]; + } + } +} + +void fht_kacs_walk_sse(float *data, size_t len) { + size_t half = len / 2; + size_t base = len % 2; + size_t offset = base + half; + size_t half_end = half & ~3u; + for (size_t i = 0; i < half_end; i += 4) { + __m128 x = _mm_loadu_ps(&data[i]); + __m128 y = _mm_loadu_ps(&data[i + offset]); + _mm_storeu_ps(&data[i], _mm_add_ps(x, y)); + _mm_storeu_ps(&data[i + offset], _mm_sub_ps(x, y)); + } + // Scalar tail + for (size_t i = half_end; i < half; ++i) { + float x = data[i]; + float y = data[i + offset]; + data[i] = x + y; + data[i + offset] = x - y; + } + if (base != 0) { + data[half] *= std::sqrt(2.0f); + } +} + +void fht_inv_kacs_walk_sse(float *data, size_t len) { + size_t half = len / 2; + size_t base = len % 2; + size_t offset = base + half; + if (base != 0) { + data[half] *= std::sqrt(0.5f); + } + size_t half_end = half & ~3u; + const __m128 half_fac = _mm_set1_ps(0.5f); + for (size_t i = 0; i < half_end; i += 4) { + __m128 a = _mm_loadu_ps(&data[i]); + __m128 b = _mm_loadu_ps(&data[i + offset]); + _mm_storeu_ps(&data[i], _mm_mul_ps(_mm_add_ps(a, b), half_fac)); + _mm_storeu_ps(&data[i + offset], _mm_mul_ps(_mm_sub_ps(a, b), half_fac)); + } + // Scalar tail + for (size_t i = half_end; i < half; ++i) { + float a = data[i]; + float b = data[i + offset]; + data[i] = (a + b) * 0.5f; + data[i + offset] = (a - b) * 0.5f; + } +} + +void fht_vec_rescale_sse(float *data, size_t n, float factor) { + const __m128 fac = _mm_set1_ps(factor); + size_t simd_end = n & ~3u; + for (size_t i = 0; i < simd_end; i += 4) { + __m128 v = _mm_loadu_ps(&data[i]); + _mm_storeu_ps(&data[i], _mm_mul_ps(v, fac)); + } + // Scalar tail + for (size_t i = simd_end; i < n; ++i) { + data[i] *= factor; + } +} + +} // namespace zvec::turbo::sse + +#endif // __SSE2__ diff --git a/src/turbo/distance/sse/fht/fht.h b/src/turbo/distance/sse/fht/fht.h new file mode 100644 index 000000000..9e1cd20a5 --- /dev/null +++ b/src/turbo/distance/sse/fht/fht.h @@ -0,0 +1,34 @@ +// 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 +#include + +namespace zvec::turbo::sse { + +//! Apply bitwise sign-flip mask to a float vector (SSE). +void fht_flip_sign_sse(const uint8_t *flip, float *data, size_t dim); + +//! Apply KacsWalk butterfly operation (SSE). +void fht_kacs_walk_sse(float *data, size_t len); + +//! Inverse KacsWalk butterfly operation (SSE). +void fht_inv_kacs_walk_sse(float *data, size_t len); + +//! Element-wise rescale: data[i] *= factor (SSE). +void fht_vec_rescale_sse(float *data, size_t n, float factor); + +} // namespace zvec::turbo::sse diff --git a/src/turbo/quantizer/distance.h b/src/turbo/quantizer/distance.h new file mode 100644 index 000000000..ecc43f21a --- /dev/null +++ b/src/turbo/quantizer/distance.h @@ -0,0 +1,128 @@ +// 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 +#include +#include +#include + +namespace zvec { +namespace turbo { + +//! A callable distance handle bound to a quantized query vector. +//! +//! DistanceImpl owns the quantized query bytes and a dispatched +//! DistanceFunc. Invoking `operator()(candidate)` computes the distance +//! between the stored query and the given candidate vector, which is +//! expected to already be in the same quantized layout. +class DistanceImpl { + public: + DistanceImpl() = default; + + DistanceImpl(DistanceFunc func, std::string quantized_query, size_t dim) + : func_(std::move(func)), + sym_func_(func_), + query_storage_(std::move(quantized_query)), + dim_(dim) {} + + DistanceImpl(DistanceFunc func, BatchDistanceFunc batch_func, + std::string quantized_query, size_t dim) + : func_(std::move(func)), + sym_func_(func_), + batch_func_(std::move(batch_func)), + query_storage_(std::move(quantized_query)), + dim_(dim) {} + + //! PQ-specific constructor: sym_func_ is the SDC kernel (distinct + //! from func_ which is the ADC kernel). Non-PQ quantizers should use + //! the constructors above where sym_func_ defaults to func_. + DistanceImpl(DistanceFunc func, DistanceFunc sym_func, + BatchDistanceFunc batch_func, std::string quantized_query, + size_t dim) + : func_(std::move(func)), + sym_func_(std::move(sym_func)), + batch_func_(std::move(batch_func)), + query_storage_(std::move(quantized_query)), + dim_(dim) {} + + //! Whether the handle is ready to compute distances. + bool valid() const { + return static_cast(func_); + } + + //! Whether a batch distance function is available. + bool batch_valid() const { + return static_cast(batch_func_); + } + + //! Compute the distance between the stored query and `candidate`. + float operator()(const void *candidate) const { + float d = 0.0f; + func_(candidate, query_storage_.data(), dim_, &d); + return d; + } + + //! Compute distances for a batch of `num` candidates against the + //! stored query. Falls back to the scalar path when no batch function + //! is bound. + void batch(const void **candidates, size_t num, float *out) const { + if (batch_func_) { + batch_func_(candidates, query_storage_.data(), num, dim_, out); + return; + } + for (size_t i = 0; i < num; ++i) { + out[i] = 0.0f; + func_(candidates[i], query_storage_.data(), dim_, out + i); + } + } + + //! Access the quantized query bytes (for pairwise helpers). + const std::string &query_storage() const { + return query_storage_; + } + + size_t dim() const { + return dim_; + } + + //! Raw scalar distance function (operates on already-quantized + //! candidates). This is the ADC path for PQ. + const DistanceFunc &func() const { + return func_; + } + + //! Symmetric (SDC) distance function. For non-PQ quantizers this + //! defaults to func(); for PQ quantizers it points to the SDC kernel + //! that operates on two PQ codes via a centroid-to-centroid table. + const DistanceFunc &sym_func() const { + return sym_func_; + } + + //! Raw batch distance function. + const BatchDistanceFunc &batch_func() const { + return batch_func_; + } + + private: + DistanceFunc func_{}; + DistanceFunc sym_func_{}; // defaults to func_; PQ sets SDC kernel + BatchDistanceFunc batch_func_{}; + std::string query_storage_{}; + size_t dim_{0}; +}; + +} // namespace turbo +} // namespace zvec diff --git a/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.cc b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.cc new file mode 100644 index 000000000..839137e6d --- /dev/null +++ b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.cc @@ -0,0 +1,192 @@ +// 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 "quantizer/fp32_quantizer/fp32_quantizer.h" +#include +#include +#include +#include +#include +#include +#include "core/quantizer/record_quantizer.h" + +namespace zvec { +namespace turbo { + +int Fp32Quantizer::init(const IndexMeta &meta, + const ailego::Params & /*params*/) { + meta_ = meta; + + meta_.set_meta(IndexMeta::DataType::DT_FP32, meta.dimension()); + + + original_dim_ = meta.dimension(); + auto metric_name = meta.metric_name(); + if (metric_name == "Cosine") { + extra_meta_size_ = EXTRA_META_SIZE_COSINE; + meta_.set_extra_meta_size(extra_meta_size_); + } + + // Cache the distance dispatch for the new Quantizer interface. + dp_query_func_ = + get_distance_func(metric_from_name(metric_name), DataType::kFp32, + QuantizeType::kDefault, CpuArchType::kAuto); + dp_query_batch_func_ = + get_batch_distance_func(metric_from_name(metric_name), DataType::kFp32, + QuantizeType::kDefault, CpuArchType::kAuto); + + return 0; +} + +int Fp32Quantizer::quantize(const void *query, const IndexQueryMeta &qmeta, + std::string *out, IndexQueryMeta *ometa) const { + if (qmeta.unit_size() != sizeof(float)) { + return kErrUnsupported; + } + + // qmeta.dimension() may be the inflated (data + extras) dimension when the + // caller uses meta_.dimension() directly (e.g. HnswDistCalculator). Use the + // raw original dim we recorded at init() to avoid over-reading the query. + size_t raw_dim = (original_dim_ != 0 && qmeta.dimension() >= original_dim_) + ? original_dim_ + : qmeta.dimension(); + size_t byte_size = raw_dim * sizeof(float) + extra_meta_size_; + out->resize(byte_size); + + if (meta_.metric_name() == "Cosine") { + // L2-normalize the vector in-place and store the norm at the end so the + // original vector can be reconstructed during dequantize. + float *buf = reinterpret_cast(&(*out)[0]); + std::memcpy(buf, query, raw_dim * sizeof(float)); + float norm = 0.0f; + ailego::Normalizer::L2(buf, raw_dim, &norm); + std::memcpy( + reinterpret_cast(&(*out)[0]) + raw_dim * sizeof(float), + &norm, extra_meta_size_); + } else { + std::memcpy(&(*out)[0], query, byte_size); + } + + *ometa = qmeta; + ometa->set_meta(IndexMeta::DataType::DT_FP32, raw_dim, + static_cast(type_), extra_meta_size_); + + return 0; +} + +int Fp32Quantizer::dequantize(const void *in, const IndexQueryMeta &qmeta, + std::string *out) const { + size_t raw_dim = (original_dim_ != 0 && qmeta.dimension() >= original_dim_) + ? original_dim_ + : qmeta.dimension(); + size_t byte_size = raw_dim * sizeof(float); + + if (meta_.metric_name() == "Cosine") { + // Denormalize the vector using the stored norm. + out->resize(byte_size); + const float *in_buf = reinterpret_cast(in); + float *out_buf = reinterpret_cast(&(*out)[0]); + float norm = 0.0f; + std::memcpy(&norm, reinterpret_cast(in) + byte_size, + extra_meta_size_); + for (size_t i = 0; i < raw_dim; ++i) { + out_buf[i] = in_buf[i] * norm; + } + } else { + out->resize(byte_size); + std::memcpy(out->data(), in, byte_size); + } + return 0; +} + +DistanceImpl Fp32Quantizer::distance(const void *query, + const IndexQueryMeta &qmeta) const { + auto metric = metric_from_name(meta_.metric_name()); + auto func = get_distance_func(metric, DataType::kFp32, QuantizeType::kDefault, + CpuArchType::kAuto); + if (!func) { + return DistanceImpl{}; + } + auto batch_func = get_batch_distance_func( + metric, DataType::kFp32, QuantizeType::kDefault, CpuArchType::kAuto); + + // The query is assumed to be already quantized — copy it directly. + std::string quantized_query(static_cast(query), + qmeta.element_size()); + return DistanceImpl(std::move(func), std::move(batch_func), + std::move(quantized_query), original_dim_); +} + +void Fp32Quantizer::quantize_one(const void *input, void *output) const { + size_t byte_size = static_cast(original_dim_) * sizeof(float); + + if (meta_.metric_name() == "Cosine") { + // L2-normalize and store the norm at the end. + std::memcpy(output, input, byte_size); + float *buf = reinterpret_cast(output); + float norm = 0.0f; + ailego::Normalizer::L2(buf, original_dim_, &norm); + std::memcpy(reinterpret_cast(output) + byte_size, &norm, + extra_meta_size_); + } else { + std::memcpy(output, input, byte_size); + } +} + +float Fp32Quantizer::calc_distance_dp_query(const void *dp, + const void *query) const { + float d = 0.0f; + if (dp_query_func_) { + dp_query_func_(dp, query, original_dim_, &d); + } + return d; +} + +void Fp32Quantizer::calc_distance_dp_query_batch(const void *const *dp_list, + int dp_num, const void *query, + float *dist_list) const { + if (dp_query_batch_func_) { + dp_query_batch_func_(const_cast(dp_list), query, + static_cast(dp_num), original_dim_, dist_list); + return; + } + for (int i = 0; i < dp_num; ++i) { + dist_list[i] = calc_distance_dp_query(dp_list[i], query); + } +} + +float Fp32Quantizer::calc_distance_dp_query_unquantized( + const void *dp, const void *query) const { + std::string buf(quantized_length(), '\0'); + quantize_one(query, &buf[0]); + return calc_distance_dp_query(dp, buf.data()); +} + +void Fp32Quantizer::calc_distance_dp_query_batch_unquantized( + const void *const *dp_list, int dp_num, const void *query, + float *dist_list) const { + std::string buf(quantized_length(), '\0'); + quantize_one(query, &buf[0]); + calc_distance_dp_query_batch(dp_list, dp_num, buf.data(), dist_list); +} + +float Fp32Quantizer::calc_distance_dp_dp(const void *dp1, + const void *dp2) const { + return calc_distance_dp_query(dp1, dp2); +} + +INDEX_FACTORY_REGISTER_QUANTIZER(Fp32Quantizer); + +} // namespace turbo +} // namespace zvec diff --git a/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h new file mode 100644 index 000000000..885390639 --- /dev/null +++ b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h @@ -0,0 +1,127 @@ +// 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 +#include +#include +#include +#include "quantizer/quantizer.h" + +namespace zvec { +namespace turbo { + +using namespace zvec::core; + +class Fp32Quantizer : public Quantizer { + public: + Fp32Quantizer() { + type_ = QuantizeType::kFp32; + } + + virtual ~Fp32Quantizer() {} + + public: + int init(const core::IndexMeta &meta, const ailego::Params ¶ms) override; + + const core::IndexMeta &meta(void) const override { + return meta_; + } + + DataType input_data_type() const override { + return DataType::kFp32; + } + + QuantizeType type() const override { + return type_; + } + + int dim() const override { + return static_cast(original_dim_); + } + + bool require_train() const override { + return false; + } + + int train(core::IndexHolder::Pointer /*holder*/) override { + return 0; + } + + size_t quantized_datapoint_vector_length() const override { + return quantized_length(); + } + + size_t quantized_query_vector_length() const override { + return quantized_length(); + } + + void quantize_data(const void *input, void *output) const override { + quantize_one(input, output); + } + + void quantize_query(const void *input, void *output) const override { + quantize_one(input, output); + } + + float calc_distance_dp_query(const void *dp, + const void *query) const override; + + void calc_distance_dp_query_batch(const void *const *dp_list, int dp_num, + const void *query, + float *dist_list) const override; + + float calc_distance_dp_query_unquantized(const void *dp, + const void *query) const override; + + void calc_distance_dp_query_batch_unquantized( + const void *const *dp_list, int dp_num, const void *query, + float *dist_list) const override; + + float calc_distance_dp_dp(const void *dp1, const void *dp2) const override; + + int quantize(const void *query, const core::IndexQueryMeta &qmeta, + std::string *out, core::IndexQueryMeta *ometa) const override; + + int dequantize(const void *in, const core::IndexQueryMeta &qmeta, + std::string *out) const override; + + DistanceImpl distance(const void *query, + const core::IndexQueryMeta &qmeta) const override; + + private: + //! Byte length of a quantized vector (raw fp32 data + extra meta). + size_t quantized_length() const { + return static_cast(original_dim_) * sizeof(float) + + extra_meta_size_; + } + + //! Quantize a single fp32 vector into a caller-provided buffer of + //! quantized_length() bytes. + void quantize_one(const void *input, void *output) const; + + static constexpr uint32_t EXTRA_META_SIZE_COSINE = 4; + + IndexMeta meta_{}; + uint32_t original_dim_{0}; + + //! Cached distance dispatch (bound in init()). + DistanceFunc dp_query_func_{}; + BatchDistanceFunc dp_query_batch_func_{}; +}; + + +} // namespace turbo +} // namespace zvec diff --git a/src/turbo/quantizer/pq_int4_quantizer/pq_int4_quantizer.cc b/src/turbo/quantizer/pq_int4_quantizer/pq_int4_quantizer.cc new file mode 100644 index 000000000..c86492a79 --- /dev/null +++ b/src/turbo/quantizer/pq_int4_quantizer/pq_int4_quantizer.cc @@ -0,0 +1,642 @@ +// 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 "quantizer/pq_int4_quantizer/pq_int4_quantizer.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace zvec { +namespace turbo { + +// --------------------------------------------------------------------------- +// PQ int4 serialization payload (follows the QuantizerSerHeader). +// --------------------------------------------------------------------------- +struct PqInt4SerPayload { + uint32_t original_dim; + uint32_t num_subquantizers; + uint32_t sub_dim; + uint32_t num_centroids; // always 16 for int4 +}; + +int PqInt4Quantizer::init(const IndexMeta &meta, const ailego::Params ¶ms) { + meta_ = meta; + + uint32_t d = meta.dimension(); + original_dim_ = d; + + // Read num_subquantizers from params (required). + uint32_t nsq = 0; + if (!params.get("num_subquantizers", &nsq) || nsq == 0) { + LOG_ERROR("PqInt4Quantizer: num_subquantizers not set or zero"); + return kErrUnsupported; + } + if (d % nsq != 0) { + LOG_ERROR( + "PqInt4Quantizer: dim (%u) is not divisible by num_subquantizers (%u)", + d, nsq); + return kErrUnsupported; + } + // int4 codes are packed nibbles: 2 sub-quantizers per byte. + // num_subquantizers MUST be even for clean byte packing. + if (nsq % 2 != 0) { + LOG_ERROR( + "PqInt4Quantizer: num_subquantizers (%u) must be even for int4 " + "packed nibble encoding", + nsq); + return kErrUnsupported; + } + + num_subquantizers_ = nsq; + sub_dim_ = d / nsq; + + // Pre-allocate centroids (filled by train()). + centroids_.resize( + static_cast(num_subquantizers_) * kNumCentroids * sub_dim_, 0.0f); + + // Dispatch ISA kernels: int4 packed nibble path. + auto pq_k = get_pq_kernels(DataType::kInt4, QuantizeType::kPQ); + adc_fn_ = pq_k.adc_distance; + sdc_fn_ = pq_k.sdc_distance; + batch_adc_fn_ = pq_k.batch_adc_distance; + + // Resolve the configured metric type. + auto mt = metric_from_name(meta_.metric_name()); + + // L2-only batch distance: always used for encoding (quantize_data) and + // KMeans training (train_subquantizer). + fp32_l2_batch_fn_ = + get_batch_distance_func(MetricType::kSquaredEuclidean, DataType::kFp32, + QuantizeType::kDefault, CpuArchType::kAuto); + + // Metric-aware batch distance for search-side LUT and SDC table. + fp32_batch_fn_ = get_batch_distance_func( + mt, DataType::kFp32, QuantizeType::kDefault, CpuArchType::kAuto); + + // For Cosine: fall back to IP (normalization applied explicitly). + if (meta_.metric_name() == "Cosine") { + fp32_batch_fn_ = + get_batch_distance_func(MetricType::kInnerProduct, DataType::kFp32, + QuantizeType::kDefault, CpuArchType::kAuto); + } + + // For Cosine: reserve extra space to store the L2 norm (one float). + if (meta_.metric_name() == "Cosine") { + extra_meta_size_ = kExtraMetaSizeCosine; + meta_.set_extra_meta_size(extra_meta_size_); + } + + meta_.set_meta(IndexMeta::DataType::DT_FP32, d); + return 0; +} + +// --------------------------------------------------------------------------- +// Simple Lloyd's KMeans for one sub-quantizer (k=16). +// --------------------------------------------------------------------------- + +void PqInt4Quantizer::build_centroid_ptrs_cache() { + const size_t k = kNumCentroids; + const size_t d = sub_dim_; + centroid_ptrs_cache_.resize(num_subquantizers_); + for (uint32_t m = 0; m < num_subquantizers_; ++m) { + const float *centroids_m = + centroids_.data() + static_cast(m) * k * d; + auto &ptrs = centroid_ptrs_cache_[m]; + ptrs.resize(k); + for (size_t c = 0; c < k; ++c) { + ptrs[c] = centroids_m + c * d; + } + } +} + +void PqInt4Quantizer::train_subquantizer(const float *data, size_t num, + size_t stride, size_t sub_idx) { + const size_t k = kNumCentroids; + const size_t d = sub_dim_; + float *centroids_m = centroids_.data() + static_cast(sub_idx) * k * d; + + // -- Initialization: pick k distinct random data points ----------------- + std::mt19937 rng(42 + static_cast(sub_idx)); + std::vector perm(num); + for (size_t i = 0; i < num; ++i) perm[i] = i; + std::shuffle(perm.begin(), perm.end(), rng); + + size_t n_init = std::min(static_cast(k), num); + for (size_t i = 0; i < n_init; ++i) { + const float *src = + reinterpret_cast( + reinterpret_cast(data) + perm[i] * stride) + + sub_idx * d; + std::memcpy(centroids_m + i * d, src, d * sizeof(float)); + } + for (size_t i = n_init; i < k; ++i) { + std::memcpy(centroids_m + i * d, centroids_m + (i % n_init) * d, + d * sizeof(float)); + } + + // -- Lloyd iterations --------------------------------------------------- + std::vector assignments(num); + std::vector new_centroids(k * d); + std::vector counts(k); + + std::vector centroid_ptrs(k); + for (size_t c = 0; c < k; ++c) { + centroid_ptrs[c] = centroids_m + c * d; + } + std::vector dists(k); + + for (uint32_t iter = 0; iter < kMaxKmeansIters; ++iter) { + bool changed = false; + + for (size_t i = 0; i < num; ++i) { + const float *sub_vec = + reinterpret_cast( + reinterpret_cast(data) + i * stride) + + sub_idx * d; + + fp32_l2_batch_fn_(centroid_ptrs.data(), + reinterpret_cast(sub_vec), k, d, + dists.data()); + uint32_t best_idx = static_cast( + std::min_element(dists.begin(), dists.end()) - dists.begin()); + + if (assignments[i] != best_idx) { + changed = true; + assignments[i] = best_idx; + } + } + + if (!changed) { + LOG_INFO(" sub[%zu] converged at iter %u", sub_idx, iter + 1); + break; + } + + std::fill(new_centroids.begin(), new_centroids.end(), 0.0f); + std::fill(counts.begin(), counts.end(), 0); + + for (size_t i = 0; i < num; ++i) { + const float *sub_vec = + reinterpret_cast( + reinterpret_cast(data) + i * stride) + + sub_idx * d; + uint32_t c = assignments[i]; + counts[c]++; + float *cent = new_centroids.data() + c * d; + for (size_t j = 0; j < d; ++j) { + cent[j] += sub_vec[j]; + } + } + + for (size_t c = 0; c < k; ++c) { + if (counts[c] == 0) continue; + float inv = 1.0f / static_cast(counts[c]); + float *cent = new_centroids.data() + c * d; + for (size_t j = 0; j < d; ++j) { + cent[j] *= inv; + } + } + + for (size_t c = 0; c < k; ++c) { + if (counts[c] > 0) { + std::memcpy(centroids_m + c * d, new_centroids.data() + c * d, + d * sizeof(float)); + } + } + } +} + +int PqInt4Quantizer::train(IndexHolder::Pointer holder) { + return train(holder, 1); +} + +int PqInt4Quantizer::train(IndexHolder::Pointer holder, int thread_count) { + if (!holder) { + return kErrUnsupported; + } + + size_t num = holder->count(); + + // Collect all data into a contiguous buffer. + auto iter = holder->create_iterator(); + std::vector all_data(num * original_dim_); + size_t row = 0; + for (; iter->is_valid(); iter->next(), ++row) { + const float *src = reinterpret_cast(iter->data()); + std::memcpy(all_data.data() + row * original_dim_, src, + original_dim_ * sizeof(float)); + } + + // Subsample if the dataset exceeds the training limit. + if (num > kMaxTrainVectors) { + LOG_INFO("PQ int4 training: subsampling %zu -> %zu vectors", num, + kMaxTrainVectors); + std::mt19937 rng(42); + for (size_t i = 0; i < kMaxTrainVectors; ++i) { + std::uniform_int_distribution dist(i, num - 1); + size_t j = dist(rng); + if (i != j) { + for (uint32_t d = 0; d < original_dim_; ++d) { + std::swap(all_data[i * original_dim_ + d], + all_data[j * original_dim_ + d]); + } + } + } + num = kMaxTrainVectors; + all_data.resize(num * original_dim_); + all_data.shrink_to_fit(); + } + + size_t data_stride = original_dim_ * sizeof(float); + + // For Cosine: normalize training data to unit length. + if (meta_.metric_name() == "Cosine") { + for (size_t i = 0; i < num; ++i) { + float *v = all_data.data() + i * original_dim_; + float norm = 0.0f; + ailego::Normalizer::L2(v, original_dim_, &norm); + } + } + + thread_count = + std::max(1, std::min(thread_count, static_cast(num_subquantizers_))); + + LOG_INFO( + "PQ int4 training: %zu vectors, dim=%u, nsq=%u, sub_dim=%u, " + "max_iters=%u, threads=%d", + num, original_dim_, num_subquantizers_, sub_dim_, kMaxKmeansIters, + thread_count); + + if (thread_count == 1) { + for (uint32_t m = 0; m < num_subquantizers_; ++m) { + train_subquantizer(all_data.data(), num, data_stride, m); + LOG_INFO(" sub-quantizer [%u/%u] done", m + 1, num_subquantizers_); + } + } else { + std::vector threads; + threads.reserve(thread_count); + + for (int t = 0; t < thread_count; ++t) { + uint32_t m_begin = static_cast( + (static_cast(t) * num_subquantizers_) / thread_count); + uint32_t m_end = static_cast( + (static_cast(t + 1) * num_subquantizers_) / thread_count); + + threads.emplace_back( + [this, &all_data, num, data_stride, m_begin, m_end]() { + for (uint32_t m = m_begin; m < m_end; ++m) { + train_subquantizer(all_data.data(), num, data_stride, m); + } + }); + } + + for (auto &th : threads) { + th.join(); + } + + LOG_INFO(" all %u sub-quantizers trained (%d threads)", num_subquantizers_, + thread_count); + } + + build_centroid_ptrs_cache(); + + LOG_INFO("Computing SDC dist_table ..."); + compute_dist_table(); + + LOG_INFO("PQ int4 training complete."); + return 0; +} + +void PqInt4Quantizer::compute_dist_table() { + const size_t k = kNumCentroids; + const size_t d = sub_dim_; + dist_table_.resize(static_cast(num_subquantizers_) * k * k, 0.0f); + + for (uint32_t m = 0; m < num_subquantizers_; ++m) { + const float *centroids_m = centroids_.data() + m * k * d; + float *table_m = dist_table_.data() + m * k * k; + + const auto ¢roid_ptrs = centroid_ptrs_cache_[m]; + for (uint32_t i = 0; i < k; ++i) { + fp32_batch_fn_(const_cast(centroid_ptrs.data()), + reinterpret_cast(centroids_m + i * d), k, d, + table_m + i * k); + } + } +} + +void PqInt4Quantizer::quantize_data(const void *input, void *output) const { + const float *vec = reinterpret_cast(input); + uint8_t *code = reinterpret_cast(output); + + // For Cosine: normalize the vector before encoding. + std::vector norm_vec_storage; + float vec_norm = 0.0f; + if (meta_.metric_name() == "Cosine") { + norm_vec_storage.assign(vec, vec + original_dim_); + ailego::Normalizer::L2(norm_vec_storage.data(), original_dim_, + &vec_norm); + vec = norm_vec_storage.data(); + } + + // Clear packed nibble code buffer (num_subquantizers / 2 bytes). + size_t code_bytes = static_cast(num_subquantizers_) / 2; + std::memset(code, 0, code_bytes); + + // Encode each sub-quantizer and pack into nibbles. + float dists[kNumCentroids]; + + for (uint32_t m = 0; m < num_subquantizers_; ++m) { + const float *sub_vec = vec + m * sub_dim_; + const auto ¢roid_ptrs = centroid_ptrs_cache_[m]; + + fp32_l2_batch_fn_(const_cast(centroid_ptrs.data()), + reinterpret_cast(sub_vec), kNumCentroids, + sub_dim_, dists); + + // Argmin: find nearest centroid (0..15). + float best_dist = dists[0]; + uint32_t best_idx = 0; + for (uint32_t j = 1; j < kNumCentroids; ++j) { + if (dists[j] < best_dist) { + best_dist = dists[j]; + best_idx = j; + } + } + + // Pack nibble: even m -> low 4 bits; odd m -> high 4 bits. + if (m & 1u) { + code[m >> 1] |= static_cast(best_idx << 4); + } else { + code[m >> 1] |= static_cast(best_idx & 0x0Fu); + } + } + + // Store norm after PQ code for Cosine dequantize support. + if (meta_.metric_name() == "Cosine") { + float *norm_out = reinterpret_cast(code + code_bytes); + *norm_out = vec_norm; + } +} + +void PqInt4Quantizer::quantize_query(const void *input, void *output) const { + const float *query = reinterpret_cast(input); + float *lut = reinterpret_cast(output); + + // For Cosine: normalize query, then use IP as LUT metric. + std::vector norm_query_storage; + if (meta_.metric_name() == "Cosine") { + norm_query_storage.assign(query, query + original_dim_); + float norm = 0.0f; + ailego::Normalizer::L2(norm_query_storage.data(), original_dim_, + &norm); + query = norm_query_storage.data(); + } + + // For each sub-quantizer, compute distance from query sub-vector to all + // 16 centroids. LUT stride = kNumCentroids = 16. + for (uint32_t m = 0; m < num_subquantizers_; ++m) { + const auto ¢roid_ptrs = centroid_ptrs_cache_[m]; + fp32_batch_fn_(const_cast(centroid_ptrs.data()), + reinterpret_cast(query + m * sub_dim_), + kNumCentroids, sub_dim_, lut + m * kNumCentroids); + } +} + +float PqInt4Quantizer::calc_distance_dp_query(const void *dp, + const void *query) const { + float d = 0.0f; + adc_fn_(dp, query, num_subquantizers_, &d); + if (meta_.metric_name() == "Cosine") { + d = 1.0f + d; + } + return d; +} + +void PqInt4Quantizer::calc_distance_dp_query_batch(const void *const *dp_list, + int dp_num, + const void *query, + float *dist_list) const { + batch_adc_fn_(const_cast(dp_list), query, + static_cast(dp_num), num_subquantizers_, dist_list); + if (meta_.metric_name() == "Cosine") { + for (int i = 0; i < dp_num; ++i) { + dist_list[i] = 1.0f + dist_list[i]; + } + } +} + +float PqInt4Quantizer::calc_distance_dp_query_unquantized( + const void *dp, const void *query) const { + std::vector lut(static_cast(num_subquantizers_) * + kNumCentroids); + quantize_query(query, lut.data()); + float d = 0.0f; + adc_fn_(dp, lut.data(), num_subquantizers_, &d); + if (meta_.metric_name() == "Cosine") { + d = 1.0f + d; + } + return d; +} + +void PqInt4Quantizer::calc_distance_dp_query_batch_unquantized( + const void *const *dp_list, int dp_num, const void *query, + float *dist_list) const { + std::vector lut(static_cast(num_subquantizers_) * + kNumCentroids); + quantize_query(query, lut.data()); + batch_adc_fn_(const_cast(dp_list), lut.data(), + static_cast(dp_num), num_subquantizers_, dist_list); + if (meta_.metric_name() == "Cosine") { + for (int i = 0; i < dp_num; ++i) { + dist_list[i] = 1.0f + dist_list[i]; + } + } +} + +float PqInt4Quantizer::calc_distance_dp_dp(const void *dp1, + const void *dp2) const { + float d = 0.0f; + sdc_fn_(dp1, dp2, dist_table_.data(), num_subquantizers_, &d); + return d; +} + +int PqInt4Quantizer::quantize(const void *query, const IndexQueryMeta &qmeta, + std::string *out, IndexQueryMeta *ometa) const { + if (qmeta.unit_size() != sizeof(float)) { + return kErrUnsupported; + } + + size_t lut_bytes = quantized_query_vector_length(); + out->resize(lut_bytes); + quantize_query(query, &(*out)[0]); + + *ometa = qmeta; + ometa->set_meta(IndexMeta::DataType::DT_FP32, original_dim_, + static_cast(type_), 0); + return 0; +} + +int PqInt4Quantizer::dequantize(const void *in, const IndexQueryMeta &qmeta, + std::string *out) const { + (void)qmeta; + const uint8_t *code = reinterpret_cast(in); + size_t byte_size = static_cast(original_dim_) * sizeof(float); + out->resize(byte_size); + float *result = reinterpret_cast(&(*out)[0]); + + const size_t k = kNumCentroids; + const size_t d = sub_dim_; + for (uint32_t m = 0; m < num_subquantizers_; ++m) { + const float *centroids_m = + centroids_.data() + static_cast(m) * k * d; + + // Decode nibble index from packed code. + uint8_t byte = code[m >> 1]; + uint8_t idx = (m & 1u) ? static_cast(byte >> 4) + : static_cast(byte & 0x0Fu); + + const float *centroid = centroids_m + static_cast(idx) * d; + std::memcpy(result + m * d, centroid, d * sizeof(float)); + } + + // For Cosine: rescale by stored norm. + if (meta_.metric_name() == "Cosine") { + size_t code_bytes = static_cast(num_subquantizers_) / 2; + float norm = 0.0f; + std::memcpy(&norm, code + code_bytes, sizeof(float)); + for (uint32_t j = 0; j < original_dim_; ++j) { + result[j] *= norm; + } + } + return 0; +} + +DistanceImpl PqInt4Quantizer::distance(const void *query, + const IndexQueryMeta &qmeta) const { + (void)qmeta; + + DistanceFunc adc_func = adc_fn_; + + auto sdc = sdc_fn_; + const void *dt = dist_table_.data(); + DistanceFunc sdc_func = [sdc, dt](const void *a, const void *b, size_t dim, + float *out) { sdc(a, b, dt, dim, out); }; + + BatchDistanceFunc batch_func = batch_adc_fn_; + + size_t lut_bytes = quantized_query_vector_length(); + std::string lut_storage(static_cast(query), lut_bytes); + + return DistanceImpl(std::move(adc_func), std::move(sdc_func), + std::move(batch_func), std::move(lut_storage), + static_cast(num_subquantizers_)); +} + +// --------------------------------------------------------------------------- +// Serialization +// --------------------------------------------------------------------------- +int PqInt4Quantizer::serialize(std::string *out) const { + if (!out) return kErrUnsupported; + + QuantizerSerHeader hdr{}; + hdr.magic = kQuantizerMagic; + hdr.version = kQuantizerSerVersion; + hdr.quant_type = static_cast(QuantizeType::kPQ); + hdr.dim = original_dim_; + hdr.metric = static_cast(metric_from_name(meta_.metric_name())); + + PqInt4SerPayload payload{}; + payload.original_dim = original_dim_; + payload.num_subquantizers = num_subquantizers_; + payload.sub_dim = sub_dim_; + payload.num_centroids = kNumCentroids; + + size_t centroids_bytes = centroids_.size() * sizeof(float); + hdr.payload_size = static_cast(sizeof(payload) + centroids_bytes); + + out->clear(); + out->append(reinterpret_cast(&hdr), sizeof(hdr)); + out->append(reinterpret_cast(&payload), sizeof(payload)); + out->append(reinterpret_cast(centroids_.data()), + centroids_bytes); + return 0; +} + +int PqInt4Quantizer::deserialize(std::string &in) { + return deserialize(in.data(), in.size()); +} + +int PqInt4Quantizer::deserialize(const void *data, size_t len) { + if (len < sizeof(QuantizerSerHeader) + sizeof(PqInt4SerPayload)) { + return kErrUnsupported; + } + + const char *ptr = reinterpret_cast(data); + QuantizerSerHeader hdr; + std::memcpy(&hdr, ptr, sizeof(hdr)); + ptr += sizeof(hdr); + + if (hdr.magic != kQuantizerMagic) return kErrUnsupported; + + PqInt4SerPayload payload; + std::memcpy(&payload, ptr, sizeof(payload)); + ptr += sizeof(payload); + + original_dim_ = payload.original_dim; + num_subquantizers_ = payload.num_subquantizers; + sub_dim_ = payload.sub_dim; + + meta_.set_meta(IndexMeta::DataType::DT_FP32, original_dim_); + + size_t centroids_bytes = static_cast(num_subquantizers_) * + kNumCentroids * sub_dim_ * sizeof(float); + + centroids_.resize(centroids_bytes / sizeof(float)); + std::memcpy(centroids_.data(), ptr, centroids_bytes); + + // Re-dispatch int4 kernels. + auto pq_k = get_pq_kernels(DataType::kInt4, QuantizeType::kPQ); + adc_fn_ = pq_k.adc_distance; + sdc_fn_ = pq_k.sdc_distance; + batch_adc_fn_ = pq_k.batch_adc_distance; + + fp32_l2_batch_fn_ = + get_batch_distance_func(MetricType::kSquaredEuclidean, DataType::kFp32, + QuantizeType::kDefault, CpuArchType::kAuto); + + fp32_batch_fn_ = get_batch_distance_func( + metric_from_name(meta_.metric_name()), DataType::kFp32, + QuantizeType::kDefault, CpuArchType::kAuto); + if (meta_.metric_name() == "Cosine") { + fp32_batch_fn_ = + get_batch_distance_func(MetricType::kInnerProduct, DataType::kFp32, + QuantizeType::kDefault, CpuArchType::kAuto); + } + + build_centroid_ptrs_cache(); + + return 0; +} + +INDEX_FACTORY_REGISTER_QUANTIZER(PqInt4Quantizer); + +} // namespace turbo +} // namespace zvec diff --git a/src/turbo/quantizer/pq_int4_quantizer/pq_int4_quantizer.h b/src/turbo/quantizer/pq_int4_quantizer/pq_int4_quantizer.h new file mode 100644 index 000000000..2a1649cfc --- /dev/null +++ b/src/turbo/quantizer/pq_int4_quantizer/pq_int4_quantizer.h @@ -0,0 +1,163 @@ +// 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 +#include +#include +#include +#include "quantizer/quantizer.h" + +namespace zvec { +namespace turbo { + +using namespace zvec::core; + +//! Product Quantizer with 4-bit sub-codes (num_bits=4, 16 centroids). +//! +//! Datapoints are encoded as packed nibbles: two 4-bit indices per byte. +//! byte[i] = (code[2*i+1] << 4) | code[2*i] +//! num_subquantizers MUST be even for clean byte packing. +//! +//! Queries are encoded as a float LUT of size [num_subquantizers * 16] +//! via quantize_query(). Distance between a PQ code and a query uses +//! ADC (LUT look-up); distance between two PQ codes uses SDC +//! (centroid-to-centroid distance table of size [nsq * 16 * 16]). +class PqInt4Quantizer : public Quantizer { + public: + PqInt4Quantizer() { + type_ = QuantizeType::kPQ; // shares kPQ with int8; DataType distinguishes + } + + ~PqInt4Quantizer() override = default; + + int init(const IndexMeta &meta, const ailego::Params ¶ms) override; + + const IndexMeta &meta() const override { + return meta_; + } + + DataType input_data_type() const override { + return DataType::kFp32; + } + + QuantizeType type() const override { + return type_; + } + + int dim() const override { + return static_cast(original_dim_); + } + + bool require_train() const override { + return true; + } + + int train(IndexHolder::Pointer holder) override; + + int train(IndexHolder::Pointer holder, int thread_count) override; + + // Packed nibbles: num_subquantizers / 2 bytes + optional Cosine norm. + size_t quantized_datapoint_vector_length() const override { + return static_cast(num_subquantizers_) / 2 + extra_meta_size_; + } + + // LUT: [num_subquantizers * 16] floats. + size_t quantized_query_vector_length() const override { + return static_cast(num_subquantizers_) * kNumCentroids * + sizeof(float); + } + + void quantize_data(const void *input, void *output) const override; + + void quantize_query(const void *input, void *output) const override; + + float calc_distance_dp_query(const void *dp, + const void *query) const override; + + void calc_distance_dp_query_batch(const void *const *dp_list, int dp_num, + const void *query, + float *dist_list) const override; + + float calc_distance_dp_query_unquantized(const void *dp, + const void *query) const override; + + void calc_distance_dp_query_batch_unquantized( + const void *const *dp_list, int dp_num, const void *query, + float *dist_list) const override; + + float calc_distance_dp_dp(const void *dp1, const void *dp2) const override; + + int quantize(const void *query, const IndexQueryMeta &qmeta, std::string *out, + IndexQueryMeta *ometa) const override; + + int dequantize(const void *in, const IndexQueryMeta &qmeta, + std::string *out) const override; + + DistanceImpl distance(const void *query, + const IndexQueryMeta &qmeta) const override; + + int serialize(std::string *out) const override; + + int deserialize(std::string &in) override; + + int deserialize(const void *data, size_t len) override; + + private: + //! Train a single sub-quantizer (KMeans, k=16) on the sub-vectors. + void train_subquantizer(const float *data, size_t num, size_t stride, + size_t sub_idx); + + //! Compute the centroid-to-centroid distance table for SDC. + void compute_dist_table(); + + //! Build centroid_ptrs_cache_ from current centroids_. + void build_centroid_ptrs_cache(); + + static constexpr uint32_t kNumCentroids = 16; + static constexpr uint32_t kMaxKmeansIters = 25; + static constexpr size_t kMaxTrainVectors = 65536; + static constexpr uint32_t kExtraMetaSizeCosine = sizeof(float); + + IndexMeta meta_{}; + uint32_t original_dim_{0}; + uint32_t num_subquantizers_{0}; + uint32_t sub_dim_{0}; + + //! Centroids: [num_subquantizers * kNumCentroids * sub_dim] + std::vector centroids_; + + //! Centroid-to-centroid distance table for SDC: + //! [num_subquantizers * kNumCentroids * kNumCentroids] + std::vector dist_table_; + + //! Pre-built centroid pointer arrays for each sub-quantizer. + std::vector> centroid_ptrs_cache_; + + //! ISA-dispatched kernel function pointers (ADC / SDC / Batch ADC). + //! These point to int4-specific kernels (packed nibble codes). + PqAdcDistanceFunc adc_fn_{nullptr}; + PqSdcKernelFunc sdc_fn_{nullptr}; + PqBatchAdcFunc batch_adc_fn_{nullptr}; + + //! Metric-aware fp32 batch distance for search-side LUT and SDC table. + BatchDistanceFunc fp32_batch_fn_{}; + + //! L2-only fp32 batch distance for encoding and KMeans training. + BatchDistanceFunc fp32_l2_batch_fn_{}; +}; + +} // namespace turbo +} // namespace zvec diff --git a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc new file mode 100644 index 000000000..2dd724ccf --- /dev/null +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc @@ -0,0 +1,701 @@ +// 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 "quantizer/pq_int8_quantizer/pq_int8_quantizer.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace zvec { +namespace turbo { + +// --------------------------------------------------------------------------- +// PQ serialization payload (follows the QuantizerSerHeader). +// --------------------------------------------------------------------------- +struct PqInt8SerPayload { + uint32_t original_dim; + uint32_t num_subquantizers; + uint32_t sub_dim; + uint32_t num_centroids; // always 256 for int8 +}; + +int PqInt8Quantizer::init(const IndexMeta &meta, const ailego::Params ¶ms) { + meta_ = meta; + + uint32_t d = meta.dimension(); + original_dim_ = d; + + // Read num_subquantizers from params (required). + uint32_t nsq = 0; + if (!params.get("num_subquantizers", &nsq) || nsq == 0) { + LOG_ERROR("PqInt8Quantizer: num_subquantizers not set or zero"); + return kErrUnsupported; + } + if (d % nsq != 0) { + LOG_ERROR( + "PqInt8Quantizer: dim (%u) is not divisible by num_subquantizers (%u)", + d, nsq); + return kErrUnsupported; + } + + num_subquantizers_ = nsq; + sub_dim_ = d / nsq; + + // Pre-allocate centroids (filled by train()). + centroids_.resize( + static_cast(num_subquantizers_) * kNumCentroids * sub_dim_, 0.0f); + + // Dispatch ISA kernels (scalar only for now). + auto pq_k = get_pq_kernels(DataType::kInt8); + adc_fn_ = pq_k.adc_distance; + sdc_fn_ = pq_k.sdc_distance; + batch_adc_fn_ = pq_k.batch_adc_distance; + + // Resolve the configured metric type (local only — not stored). + auto mt = metric_from_name(meta_.metric_name()); + + // L2-only batch distance: always used for encoding (quantize_data) and + // KMeans training (train_subquantizer). PQ codebook is trained in L2 + // space, so encoding must minimize L2 quantization error regardless + // of the search metric. + fp32_l2_batch_fn_ = + get_batch_distance_func(MetricType::kSquaredEuclidean, DataType::kFp32, + QuantizeType::kDefault, CpuArchType::kAuto); + + // Metric-aware batch distance: used for search-side LUT computation + // (quantize_query) and SDC dist_table. For IP/Cosine this computes + // inner-product-based distances instead of L2. + fp32_batch_fn_ = get_batch_distance_func( + mt, DataType::kFp32, QuantizeType::kDefault, CpuArchType::kAuto); + + // For Cosine: fall back to IP for the metric-aware batch function, + // because cosine = normalize + IP. The normalization is applied + // explicitly in quantize_query. + if (meta_.metric_name() == "Cosine") { + fp32_batch_fn_ = + get_batch_distance_func(MetricType::kInnerProduct, DataType::kFp32, + QuantizeType::kDefault, CpuArchType::kAuto); + } + + // For Cosine: reserve extra space to store the L2 norm (one float) + // alongside each PQ code, enabling dequantize() to rescale. + if (meta_.metric_name() == "Cosine") { + extra_meta_size_ = kExtraMetaSizeCosine; + meta_.set_extra_meta_size(extra_meta_size_); + } + + meta_.set_meta(IndexMeta::DataType::DT_FP32, d); + return 0; +} + +// --------------------------------------------------------------------------- +// Simple Lloyd's KMeans for one sub-quantizer. +// --------------------------------------------------------------------------- + +void PqInt8Quantizer::build_centroid_ptrs_cache() { + const size_t k = kNumCentroids; + const size_t d = sub_dim_; + centroid_ptrs_cache_.resize(num_subquantizers_); + for (uint32_t m = 0; m < num_subquantizers_; ++m) { + const float *centroids_m = + centroids_.data() + static_cast(m) * k * d; + auto &ptrs = centroid_ptrs_cache_[m]; + ptrs.resize(k); + for (size_t c = 0; c < k; ++c) { + ptrs[c] = centroids_m + c * d; + } + } +} + +void PqInt8Quantizer::train_subquantizer(const float *data, size_t num, + size_t stride, size_t sub_idx) { + const size_t k = kNumCentroids; + const size_t d = sub_dim_; + float *centroids_m = centroids_.data() + static_cast(sub_idx) * k * d; + + // -- Initialization: pick k distinct random data points ----------------- + std::mt19937 rng(42 + static_cast(sub_idx)); + std::vector perm(num); + for (size_t i = 0; i < num; ++i) perm[i] = i; + std::shuffle(perm.begin(), perm.end(), rng); + + size_t n_init = std::min(static_cast(k), num); + for (size_t i = 0; i < n_init; ++i) { + const float *src = + reinterpret_cast( + reinterpret_cast(data) + perm[i] * stride) + + sub_idx * d; + std::memcpy(centroids_m + i * d, src, d * sizeof(float)); + } + // Duplicate if num < k (rare edge case). + for (size_t i = n_init; i < k; ++i) { + std::memcpy(centroids_m + i * d, centroids_m + (i % n_init) * d, + d * sizeof(float)); + } + + // -- Lloyd iterations --------------------------------------------------- + std::vector assignments(num); + std::vector new_centroids(k * d); + std::vector counts(k); + + // Pre-build centroid pointer array for SIMD batch distance. + std::vector centroid_ptrs(k); + for (size_t c = 0; c < k; ++c) { + centroid_ptrs[c] = centroids_m + c * d; + } + std::vector dists(k); + + for (uint32_t iter = 0; iter < kMaxKmeansIters; ++iter) { + bool changed = false; + + // Assignment step — always use L2 distance (fp32_l2_batch_fn_) to + // minimize L2 quantization error, regardless of the search metric. + // KMeans operates in L2 space; the search metric only affects LUT + // construction in quantize_query. + for (size_t i = 0; i < num; ++i) { + const float *sub_vec = + reinterpret_cast( + reinterpret_cast(data) + i * stride) + + sub_idx * d; + + fp32_l2_batch_fn_(centroid_ptrs.data(), + reinterpret_cast(sub_vec), k, d, + dists.data()); + uint32_t best_idx = static_cast( + std::min_element(dists.begin(), dists.end()) - dists.begin()); + + if (assignments[i] != best_idx) { + changed = true; + assignments[i] = best_idx; + } + } + + if (!changed) { + LOG_INFO(" sub[%zu] converged at iter %u", sub_idx, iter + 1); + break; + } + + // Update step. + std::fill(new_centroids.begin(), new_centroids.end(), 0.0f); + std::fill(counts.begin(), counts.end(), 0); + + for (size_t i = 0; i < num; ++i) { + const float *sub_vec = + reinterpret_cast( + reinterpret_cast(data) + i * stride) + + sub_idx * d; + uint32_t c = assignments[i]; + counts[c]++; + float *cent = new_centroids.data() + c * d; + for (size_t j = 0; j < d; ++j) { + cent[j] += sub_vec[j]; + } + } + + for (size_t c = 0; c < k; ++c) { + if (counts[c] == 0) continue; + float inv = 1.0f / static_cast(counts[c]); + float *cent = new_centroids.data() + c * d; + for (size_t j = 0; j < d; ++j) { + cent[j] *= inv; + } + } + + // Copy back; keep old centroid for empty clusters. + for (size_t c = 0; c < k; ++c) { + if (counts[c] > 0) { + std::memcpy(centroids_m + c * d, new_centroids.data() + c * d, + d * sizeof(float)); + } + } + } +} + +int PqInt8Quantizer::train(IndexHolder::Pointer holder) { + return train(holder, 1); +} + +int PqInt8Quantizer::train(IndexHolder::Pointer holder, int thread_count) { + if (!holder) { + return kErrUnsupported; + } + + size_t num = holder->count(); + + // Collect all data into a contiguous buffer. + auto iter = holder->create_iterator(); + std::vector all_data(num * original_dim_); + size_t row = 0; + for (; iter->is_valid(); iter->next(), ++row) { + const float *src = reinterpret_cast(iter->data()); + std::memcpy(all_data.data() + row * original_dim_, src, + original_dim_ * sizeof(float)); + } + + // Subsample if the dataset exceeds the training limit (aligned with + // faiss/vsag: 256 centroids * 256 max_points_per_centroid ≈ 65535). + if (num > kMaxTrainVectors) { + LOG_INFO("PQ training: subsampling %zu -> %zu vectors", num, + kMaxTrainVectors); + std::mt19937 rng(42); + // Fisher-Yates partial shuffle: randomly place kMaxTrainVectors vectors + // at the front of the buffer. + for (size_t i = 0; i < kMaxTrainVectors; ++i) { + std::uniform_int_distribution dist(i, num - 1); + size_t j = dist(rng); + if (i != j) { + // Swap full vectors (dim-sized chunks). + for (uint32_t d = 0; d < original_dim_; ++d) { + std::swap(all_data[i * original_dim_ + d], + all_data[j * original_dim_ + d]); + } + } + } + num = kMaxTrainVectors; + all_data.resize(num * original_dim_); + all_data.shrink_to_fit(); + } + + size_t data_stride = original_dim_ * sizeof(float); + + // For Cosine: normalize training data to unit length so that KMeans + // centroids are learned in normalized space. After normalization, + // L2 minimization is equivalent to maximizing cosine similarity. + if (meta_.metric_name() == "Cosine") { + for (size_t i = 0; i < num; ++i) { + float *v = all_data.data() + i * original_dim_; + float norm = 0.0f; + ailego::Normalizer::L2(v, original_dim_, &norm); + } + } + + // Clamp thread count to [1, num_subquantizers_]. + thread_count = + std::max(1, std::min(thread_count, static_cast(num_subquantizers_))); + + LOG_INFO( + "PQ training: %zu vectors, dim=%u, nsq=%u, sub_dim=%u, " + "max_iters=%u, threads=%d", + num, original_dim_, num_subquantizers_, sub_dim_, kMaxKmeansIters, + thread_count); + + if (thread_count == 1) { + // Single-threaded path. + for (uint32_t m = 0; m < num_subquantizers_; ++m) { + train_subquantizer(all_data.data(), num, data_stride, m); + LOG_INFO(" sub-quantizer [%u/%u] done", m + 1, num_subquantizers_); + } + } else { + // Multi-threaded path: each thread handles a contiguous range of + // sub-quantizers. Sub-quantizers are fully independent — no + // synchronization needed. + std::vector threads; + threads.reserve(thread_count); + + for (int t = 0; t < thread_count; ++t) { + uint32_t m_begin = static_cast( + (static_cast(t) * num_subquantizers_) / thread_count); + uint32_t m_end = static_cast( + (static_cast(t + 1) * num_subquantizers_) / thread_count); + + threads.emplace_back( + [this, &all_data, num, data_stride, m_begin, m_end]() { + for (uint32_t m = m_begin; m < m_end; ++m) { + train_subquantizer(all_data.data(), num, data_stride, m); + } + }); + } + + for (auto &th : threads) { + th.join(); + } + + LOG_INFO(" all %u sub-quantizers trained (%d threads)", num_subquantizers_, + thread_count); + } + + // Pre-build centroid pointer cache (needed by compute_dist_table). + build_centroid_ptrs_cache(); + + // Pre-compute SDC dist_table. + LOG_INFO("Computing SDC dist_table ..."); + compute_dist_table(); + + LOG_INFO("PQ training complete."); + return 0; +} + +void PqInt8Quantizer::compute_dist_table() { + const size_t k = kNumCentroids; + const size_t d = sub_dim_; + dist_table_.resize(static_cast(num_subquantizers_) * k * k, 0.0f); + + // For each sub-quantizer, compute centroid-to-centroid distances via + // the metric-aware fp32 batch distance function (fp32_batch_fn_). + // L2: dist_table[m][i][j] = ||c_m[i] - c_m[j]||^2 + // IP: dist_table[m][i][j] = -dot(c_m[i], c_m[j]) + // Cosine: centroids trained on normalized data, same as IP. + for (uint32_t m = 0; m < num_subquantizers_; ++m) { + const float *centroids_m = centroids_.data() + m * k * d; + float *table_m = dist_table_.data() + m * k * k; + + // Use pre-built centroid pointer cache. + // const_cast: .data() returns const void* const* but fp32_batch_fn_ + // expects const void**. The kernel never modifies the pointer array. + const auto ¢roid_ptrs = centroid_ptrs_cache_[m]; + for (uint32_t i = 0; i < k; ++i) { + fp32_batch_fn_(const_cast(centroid_ptrs.data()), + reinterpret_cast(centroids_m + i * d), k, d, + table_m + i * k); + } + } +} + +void PqInt8Quantizer::quantize_data(const void *input, void *output) const { + const float *vec = reinterpret_cast(input); + uint8_t *code = reinterpret_cast(output); + + // For Cosine: normalize the vector to unit length before encoding, + // because the codebook was trained on normalized data. The original + // norm is stored after the PQ code for later dequantize(). + std::vector norm_vec_storage; + float vec_norm = 0.0f; + if (meta_.metric_name() == "Cosine") { + norm_vec_storage.assign(vec, vec + original_dim_); + ailego::Normalizer::L2(norm_vec_storage.data(), original_dim_, + &vec_norm); + vec = norm_vec_storage.data(); + } + + // Encode with L2-only batch distance (search-metric independent). + // Process one sub-quantizer at a time, reusing a single 256-float + // scratch buffer and fusing argmin into the distance loop to avoid + // a second pass over the data. + float dists[kNumCentroids]; + + for (uint32_t m = 0; m < num_subquantizers_; ++m) { + const float *sub_vec = vec + m * sub_dim_; + const auto ¢roid_ptrs = centroid_ptrs_cache_[m]; + + // Compute L2 distances from this sub-vector to all 256 centroids. + fp32_l2_batch_fn_(const_cast(centroid_ptrs.data()), + reinterpret_cast(sub_vec), kNumCentroids, + sub_dim_, dists); + + // Argmin: find nearest centroid. + float best_dist = dists[0]; + uint32_t best_idx = 0; + for (uint32_t j = 1; j < kNumCentroids; ++j) { + if (dists[j] < best_dist) { + best_dist = dists[j]; + best_idx = j; + } + } + code[m] = static_cast(best_idx); + } + + // Store norm after PQ code for Cosine dequantize support. + if (meta_.metric_name() == "Cosine") { + float *norm_out = reinterpret_cast(code + num_subquantizers_); + *norm_out = vec_norm; + } +} + +void PqInt8Quantizer::quantize_query(const void *input, void *output) const { + const float *query = reinterpret_cast(input); + float *lut = reinterpret_cast(output); + + // For Cosine: normalize the query to unit length first, then use IP + // as the LUT metric. After normalization, dot(q_norm, c) equals the + // cosine similarity when centroids were trained on normalized data. + // IP is additive across subspaces, so sum of sub-IPs = full IP. + std::vector norm_query_storage; + if (meta_.metric_name() == "Cosine") { + norm_query_storage.assign(query, query + original_dim_); + float norm = 0.0f; + ailego::Normalizer::L2(norm_query_storage.data(), original_dim_, + &norm); + query = norm_query_storage.data(); + } + + // For each sub-quantizer, compute distance from query sub-vector to all + // 256 centroids via the metric-aware fp32 batch distance function. + // For L2: LUT[m][j] = ||q_m - c_m[j]||^2 + // For IP: LUT[m][j] = -dot(q_m, c_m[j]) + // For Cosine: (query is normalized) LUT[m][j] = -dot(q_norm_m, c_m[j]) + // const_cast: see compute_dist_table for rationale. + for (uint32_t m = 0; m < num_subquantizers_; ++m) { + const auto ¢roid_ptrs = centroid_ptrs_cache_[m]; + fp32_batch_fn_(const_cast(centroid_ptrs.data()), + reinterpret_cast(query + m * sub_dim_), + kNumCentroids, sub_dim_, lut + m * kNumCentroids); + } +} + +float PqInt8Quantizer::calc_distance_dp_query(const void *dp, + const void *query) const { + // dp = pq_code (uint8_t[num_subquantizers]) + // query = LUT (float[num_subquantizers * 256]) + float d = 0.0f; + adc_fn_(reinterpret_cast(dp), + reinterpret_cast(query), num_subquantizers_, &d); + // For Cosine: ADC sum = -cos_sim; convert to cosine distance = 1 - cos_sim. + if (meta_.metric_name() == "Cosine") { + d = 1.0f + d; + } + return d; +} + +void PqInt8Quantizer::calc_distance_dp_query_batch(const void *const *dp_list, + int dp_num, + const void *query, + float *dist_list) const { + // Use ISA-dispatched batch4 ADC kernel (4-way ILP + SIMD gather). + // const_cast: dp_list is const void* const* (outer const from vtable + // signature), but batch_adc_fn_ expects const void**. Kernel is read-only + // on the pointer array. + batch_adc_fn_(const_cast(dp_list), query, + static_cast(dp_num), num_subquantizers_, dist_list); + // For Cosine: convert -cos_sim to cosine distance. + if (meta_.metric_name() == "Cosine") { + for (int i = 0; i < dp_num; ++i) { + dist_list[i] = 1.0f + dist_list[i]; + } + } +} + +float PqInt8Quantizer::calc_distance_dp_query_unquantized( + const void *dp, const void *query) const { + // Build LUT on the fly, then use ADC. + std::vector lut(static_cast(num_subquantizers_) * + kNumCentroids); + quantize_query(query, lut.data()); + float d = 0.0f; + adc_fn_(reinterpret_cast(dp), lut.data(), num_subquantizers_, + &d); + if (meta_.metric_name() == "Cosine") { + d = 1.0f + d; + } + return d; +} + +void PqInt8Quantizer::calc_distance_dp_query_batch_unquantized( + const void *const *dp_list, int dp_num, const void *query, + float *dist_list) const { + std::vector lut(static_cast(num_subquantizers_) * + kNumCentroids); + quantize_query(query, lut.data()); + // Use ISA-dispatched batch4 ADC kernel (4-way ILP + SIMD gather). + // const_cast: see calc_distance_dp_query_batch for rationale. + batch_adc_fn_(const_cast(dp_list), lut.data(), + static_cast(dp_num), num_subquantizers_, dist_list); + if (meta_.metric_name() == "Cosine") { + for (int i = 0; i < dp_num; ++i) { + dist_list[i] = 1.0f + dist_list[i]; + } + } +} + +float PqInt8Quantizer::calc_distance_dp_dp(const void *dp1, + const void *dp2) const { + float d = 0.0f; + sdc_fn_(reinterpret_cast(dp1), + reinterpret_cast(dp2), dist_table_.data(), + num_subquantizers_, &d); + return d; +} + +int PqInt8Quantizer::quantize(const void *query, const IndexQueryMeta &qmeta, + std::string *out, IndexQueryMeta *ometa) const { + if (qmeta.unit_size() != sizeof(float)) { + return kErrUnsupported; + } + + size_t lut_bytes = quantized_query_vector_length(); + out->resize(lut_bytes); + quantize_query(query, &(*out)[0]); + + *ometa = qmeta; + ometa->set_meta(IndexMeta::DataType::DT_FP32, original_dim_, + static_cast(type_), 0); + return 0; +} + +int PqInt8Quantizer::dequantize(const void *in, const IndexQueryMeta &qmeta, + std::string *out) const { + (void)qmeta; + const uint8_t *code = reinterpret_cast(in); + size_t byte_size = static_cast(original_dim_) * sizeof(float); + out->resize(byte_size); + float *result = reinterpret_cast(&(*out)[0]); + + // Reconstruct by concatenating the selected centroids from each + // sub-quantizer. This yields the PQ approximation of the vector + // in the space the codebook was trained in (normalized for Cosine). + const size_t k = kNumCentroids; + const size_t d = sub_dim_; + for (uint32_t m = 0; m < num_subquantizers_; ++m) { + const float *centroids_m = + centroids_.data() + static_cast(m) * k * d; + const float *centroid = centroids_m + static_cast(code[m]) * d; + std::memcpy(result + m * d, centroid, d * sizeof(float)); + } + + // For Cosine: the codebook was trained on normalized data and the + // stored norm lets us rescale back to the original vector's magnitude. + if (meta_.metric_name() == "Cosine") { + float norm = 0.0f; + std::memcpy(&norm, code + num_subquantizers_, sizeof(float)); + for (uint32_t j = 0; j < original_dim_; ++j) { + result[j] *= norm; + } + } + return 0; +} + +DistanceImpl PqInt8Quantizer::distance(const void *query, + const IndexQueryMeta &qmeta) const { + (void)qmeta; + + // ADC: PqAdcDistanceFunc signature now matches DistanceFunc directly + // (both use void*), so we can assign without a lambda wrapper. + DistanceFunc adc_func = adc_fn_; + + // SDC: still needs a lambda because the kernel has 5 parameters + // (extra dist_table pointer), vs DistanceFunc's 4. + auto sdc = sdc_fn_; + const void *dt = dist_table_.data(); + DistanceFunc sdc_func = [sdc, dt](const void *a, const void *b, size_t dim, + float *out) { sdc(a, b, dt, dim, out); }; + + // Batch ADC: ISA-dispatched batch4 kernel, no lambda needed. + BatchDistanceFunc batch_func = batch_adc_fn_; + + // The query is already quantized (LUT) by the caller (reset_query) + // — copy it directly into DistanceImpl storage. + size_t lut_bytes = quantized_query_vector_length(); + std::string lut_storage(static_cast(query), lut_bytes); + + return DistanceImpl(std::move(adc_func), std::move(sdc_func), + std::move(batch_func), std::move(lut_storage), + static_cast(num_subquantizers_)); +} + +// --------------------------------------------------------------------------- +// Serialization +// --------------------------------------------------------------------------- +int PqInt8Quantizer::serialize(std::string *out) const { + if (!out) return kErrUnsupported; + + QuantizerSerHeader hdr{}; + hdr.magic = kQuantizerMagic; + hdr.version = kQuantizerSerVersion; + hdr.quant_type = static_cast(QuantizeType::kPQ); + hdr.dim = original_dim_; + hdr.metric = static_cast(metric_from_name(meta_.metric_name())); + + PqInt8SerPayload payload{}; + payload.original_dim = original_dim_; + payload.num_subquantizers = num_subquantizers_; + payload.sub_dim = sub_dim_; + payload.num_centroids = kNumCentroids; + + size_t centroids_bytes = centroids_.size() * sizeof(float); + hdr.payload_size = static_cast(sizeof(payload) + centroids_bytes); + + out->clear(); + out->append(reinterpret_cast(&hdr), sizeof(hdr)); + out->append(reinterpret_cast(&payload), sizeof(payload)); + out->append(reinterpret_cast(centroids_.data()), + centroids_bytes); + // dist_table_ is NOT serialized: it is a build-phase-only derivative + // of the codebook (used by SDC during graph construction). Search + // uses ADC exclusively, so the table can be recomputed on demand if + // ever needed, but is unnecessary after deserialization. + return 0; +} + +int PqInt8Quantizer::deserialize(std::string &in) { + return deserialize(in.data(), in.size()); +} + +int PqInt8Quantizer::deserialize(const void *data, size_t len) { + if (len < sizeof(QuantizerSerHeader) + sizeof(PqInt8SerPayload)) { + return kErrUnsupported; + } + + const char *ptr = reinterpret_cast(data); + QuantizerSerHeader hdr; + std::memcpy(&hdr, ptr, sizeof(hdr)); + ptr += sizeof(hdr); + + if (hdr.magic != kQuantizerMagic) return kErrUnsupported; + + PqInt8SerPayload payload; + std::memcpy(&payload, ptr, sizeof(payload)); + ptr += sizeof(payload); + + original_dim_ = payload.original_dim; + num_subquantizers_ = payload.num_subquantizers; + sub_dim_ = payload.sub_dim; + + meta_.set_meta(IndexMeta::DataType::DT_FP32, original_dim_); + + size_t centroids_bytes = static_cast(num_subquantizers_) * + kNumCentroids * sub_dim_ * sizeof(float); + + centroids_.resize(centroids_bytes / sizeof(float)); + std::memcpy(centroids_.data(), ptr, centroids_bytes); + // dist_table_ is intentionally not restored: SDC is only needed during + // offline build, not after deserialization (search uses ADC). + + // Re-dispatch kernels. + auto pq_k = get_pq_kernels(DataType::kInt8); + adc_fn_ = pq_k.adc_distance; + sdc_fn_ = pq_k.sdc_distance; + batch_adc_fn_ = pq_k.batch_adc_distance; + + // L2-only batch distance for encoding (always L2 regardless of metric). + fp32_l2_batch_fn_ = + get_batch_distance_func(MetricType::kSquaredEuclidean, DataType::kFp32, + QuantizeType::kDefault, CpuArchType::kAuto); + + // Metric-aware batch distance for search LUT. For Cosine, use IP + // (normalization is applied explicitly in quantize_query). + fp32_batch_fn_ = get_batch_distance_func( + metric_from_name(meta_.metric_name()), DataType::kFp32, + QuantizeType::kDefault, CpuArchType::kAuto); + if (meta_.metric_name() == "Cosine") { + fp32_batch_fn_ = + get_batch_distance_func(MetricType::kInnerProduct, DataType::kFp32, + QuantizeType::kDefault, CpuArchType::kAuto); + } + + // Pre-build centroid pointer cache for fast encode/search. + build_centroid_ptrs_cache(); + + return 0; +} + +INDEX_FACTORY_REGISTER_QUANTIZER(PqInt8Quantizer); + +} // namespace turbo +} // namespace zvec diff --git a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h new file mode 100644 index 000000000..deb661941 --- /dev/null +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h @@ -0,0 +1,168 @@ +// 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 +#include +#include +#include +#include "quantizer/quantizer.h" + +namespace zvec { +namespace turbo { + +using namespace zvec::core; + +//! Product Quantizer with 8-bit sub-codes (num_bits=8, 256 centroids). +//! +//! Datapoints are encoded as uint8_t[num_subquantizers] codes. +//! Queries are encoded as a float LUT of size [num_subquantizers * 256] +//! via compute_distance_table(). Distance between a PQ code and a +//! query uses ADC (LUT look-up); distance between two PQ codes uses +//! SDC (centroid-to-centroid distance table). +class PqInt8Quantizer : public Quantizer { + public: + PqInt8Quantizer() { + type_ = QuantizeType::kPQ; + } + + ~PqInt8Quantizer() override = default; + + int init(const IndexMeta &meta, const ailego::Params ¶ms) override; + + const IndexMeta &meta() const override { + return meta_; + } + + DataType input_data_type() const override { + return DataType::kFp32; + } + + QuantizeType type() const override { + return type_; + } + + int dim() const override { + return static_cast(original_dim_); + } + + bool require_train() const override { + return true; + } + + int train(IndexHolder::Pointer holder) override; + + int train(IndexHolder::Pointer holder, int thread_count) override; + + size_t quantized_datapoint_vector_length() const override { + return num_subquantizers_ + extra_meta_size_; + } + + size_t quantized_query_vector_length() const override { + return static_cast(num_subquantizers_) * kNumCentroids * + sizeof(float); + } + + void quantize_data(const void *input, void *output) const override; + + void quantize_query(const void *input, void *output) const override; + + float calc_distance_dp_query(const void *dp, + const void *query) const override; + + void calc_distance_dp_query_batch(const void *const *dp_list, int dp_num, + const void *query, + float *dist_list) const override; + + float calc_distance_dp_query_unquantized(const void *dp, + const void *query) const override; + + void calc_distance_dp_query_batch_unquantized( + const void *const *dp_list, int dp_num, const void *query, + float *dist_list) const override; + + float calc_distance_dp_dp(const void *dp1, const void *dp2) const override; + + int quantize(const void *query, const IndexQueryMeta &qmeta, std::string *out, + IndexQueryMeta *ometa) const override; + + int dequantize(const void *in, const IndexQueryMeta &qmeta, + std::string *out) const override; + + DistanceImpl distance(const void *query, + const IndexQueryMeta &qmeta) const override; + + int serialize(std::string *out) const override; + + int deserialize(std::string &in) override; + + int deserialize(const void *data, size_t len) override; + + private: + //! Train a single sub-quantizer (KMeans, k=256) on the sub-vectors + //! extracted from holder. sub_idx selects which sub-quantizer to train. + void train_subquantizer(const float *data, size_t num, size_t stride, + size_t sub_idx); + + //! Compute the centroid-to-centroid distance table for SDC. + void compute_dist_table(); + + //! Build centroid_ptrs_cache_ from current centroids_. + //! Called after train() and deserialize() when centroids are available. + void build_centroid_ptrs_cache(); + + static constexpr uint32_t kNumCentroids = 256; + static constexpr uint32_t kMaxKmeansIters = 25; + static constexpr size_t kMaxTrainVectors = 65536; + static constexpr uint32_t kExtraMetaSizeCosine = sizeof(float); + + IndexMeta meta_{}; + uint32_t original_dim_{0}; + uint32_t num_subquantizers_{0}; + uint32_t sub_dim_{0}; + + //! Centroids: [num_subquantizers * kNumCentroids * sub_dim] + std::vector centroids_; + + //! Centroid-to-centroid distance table for SDC: + //! [num_subquantizers * kNumCentroids * kNumCentroids] + std::vector dist_table_; + + //! Pre-built centroid pointer arrays for each sub-quantizer. + //! Layout: centroid_ptrs_cache_[sub_idx][centroid_idx] = pointer to centroid. + //! Built once during init/deserialize, reused by compute_dist_table + //! and quantize_query to avoid repeated allocations. + std::vector> centroid_ptrs_cache_; + + //! ISA-dispatched kernel function pointers (ADC / SDC / Batch ADC). + PqAdcDistanceFunc adc_fn_{nullptr}; + PqSdcKernelFunc sdc_fn_{nullptr}; + PqBatchAdcFunc batch_adc_fn_{nullptr}; + + //! Metric-aware fp32 batch distance function for search-side LUT + //! computation and SDC dist_table. Obtained from + //! get_batch_distance_func() with the configured metric. + BatchDistanceFunc fp32_batch_fn_{}; + + //! L2-only fp32 batch distance function for encoding (quantize_data) + //! and KMeans training (train_subquantizer). PQ encoding must always + //! minimize L2 quantization error regardless of the search metric. + //! This separation ensures encoding and search LUT use the correct + //! distance semantics independently. + BatchDistanceFunc fp32_l2_batch_fn_{}; +}; + +} // namespace turbo +} // namespace zvec diff --git a/src/turbo/quantizer/preprocessor/fht_rotator.cc b/src/turbo/quantizer/preprocessor/fht_rotator.cc new file mode 100644 index 000000000..4d31f6949 --- /dev/null +++ b/src/turbo/quantizer/preprocessor/fht_rotator.cc @@ -0,0 +1,260 @@ +// 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 "fht_rotator.h" +#include +#include +#include +#include "../quantizer.h" +namespace zvec { +namespace turbo { + +// ============================================================================ +// FhtRotator method implementations +// ============================================================================ + +size_t FhtRotator::floor_pow2(size_t n) { + if (n == 0) return 0; + size_t p = 1; + while (p * 2 <= n) p *= 2; + return p; +} + +FhtRotator::Pointer FhtRotator::create(int dim) { + if (dim <= 0) return nullptr; + + Pointer r(new FhtRotator()); + r->in_dim_ = dim; + r->out_dim_ = dim; + r->trunc_dim_ = floor_pow2(static_cast(dim)); + r->fac_ = 1.0f / std::sqrt(static_cast(r->trunc_dim_)); + r->flip_offset_ = (static_cast(dim) + kByteLen - 1) / kByteLen; + auto k = get_fht_kernels(); + r->flip_sign_fn_ = k.flip_sign; + r->kacs_walk_fn_ = k.kacs_walk; + r->inv_kacs_walk_fn_ = k.inv_kacs_walk; + r->inplace_fn_ = k.inplace; + r->rescale_fn_ = k.rescale; + + // Generate 4 rounds of random flip-sign arrays so the rotator is + // immediately usable after create(). + r->flip_.resize(4 * r->flip_offset_); + std::random_device rd; + std::mt19937 gen(rd()); + std::uniform_int_distribution dist(0, 255); + for (auto &b : r->flip_) b = static_cast(dist(gen)); + + return r; +} + +FhtRotator::Pointer FhtRotator::from_blob(const void *data, size_t len) { + if (!data || len < sizeof(RotatorSerHeader)) return nullptr; + + const auto *hdr = reinterpret_cast(data); + if (hdr->magic != kRotatorMagic) return nullptr; + if (hdr->version != kRotatorSerVersion) return nullptr; + if (static_cast(hdr->rotator_type) != RotatorType::kFht) { + return nullptr; + } + + Pointer r(new FhtRotator()); + const size_t expected_total = + sizeof(RotatorSerHeader) + static_cast(hdr->payload_size); + if (len < expected_total) return nullptr; + + int rc = r->deserialize(data, len); + if (rc != 0) return nullptr; + return r; +} + +void FhtRotator::train(const void * /*data*/, size_t /*num*/, + size_t /*stride*/) { + // No-op: flip-sign arrays are generated in create(). +} + +// --------------------------------------------------------------------------- +// apply (forward rotation) +// --------------------------------------------------------------------------- + +void FhtRotator::apply(const float *in, float *out) const { + const size_t dim = static_cast(in_dim_); + std::memcpy(out, in, sizeof(float) * dim); + + if (trunc_dim_ == dim) { + // Exact power-of-2: 4 rounds of (flip -> FHT -> rescale) + flip_sign_fn_(flip_.data(), out, dim); + inplace_fn_(out, trunc_dim_); + rescale_fn_(out, trunc_dim_, fac_); + + flip_sign_fn_(flip_.data() + flip_offset_, out, dim); + inplace_fn_(out, trunc_dim_); + rescale_fn_(out, trunc_dim_, fac_); + + flip_sign_fn_(flip_.data() + 2 * flip_offset_, out, dim); + inplace_fn_(out, trunc_dim_); + rescale_fn_(out, trunc_dim_, fac_); + + flip_sign_fn_(flip_.data() + 3 * flip_offset_, out, dim); + inplace_fn_(out, trunc_dim_); + rescale_fn_(out, trunc_dim_, fac_); + + return; + } + + // Non-power-of-2: 4 rounds with kacs_walk + size_t start = dim - trunc_dim_; + float *trunc_ptr = out + start; + + // Round 1: FHT on [0, trunc_dim) + flip_sign_fn_(flip_.data(), out, dim); + inplace_fn_(out, trunc_dim_); + rescale_fn_(out, trunc_dim_, fac_); + kacs_walk_fn_(out, dim); + + // Round 2: FHT on [start, start + trunc_dim) + flip_sign_fn_(flip_.data() + flip_offset_, out, dim); + inplace_fn_(trunc_ptr, trunc_dim_); + rescale_fn_(trunc_ptr, trunc_dim_, fac_); + kacs_walk_fn_(out, dim); + + // Round 3: FHT on [0, trunc_dim) + flip_sign_fn_(flip_.data() + 2 * flip_offset_, out, dim); + inplace_fn_(out, trunc_dim_); + rescale_fn_(out, trunc_dim_, fac_); + kacs_walk_fn_(out, dim); + + // Round 4: FHT on [start, start + trunc_dim) + flip_sign_fn_(flip_.data() + 3 * flip_offset_, out, dim); + inplace_fn_(trunc_ptr, trunc_dim_); + rescale_fn_(trunc_ptr, trunc_dim_, fac_); + kacs_walk_fn_(out, dim); + + // Final rescale: combine the 4 kacs_walk reductions + rescale_fn_(out, dim, 0.25f); +} + +// --------------------------------------------------------------------------- +// apply_inverse (inverse rotation) +// --------------------------------------------------------------------------- + +void FhtRotator::apply_inverse(const float *in, float *out) const { + const size_t dim = static_cast(in_dim_); + // Copy input into working buffer + std::vector data(in, in + dim); + + if (trunc_dim_ == dim) { + // Exact power-of-2: reverse 4 rounds in reverse order. + const float inv_fac = 1.0f / std::sqrt(static_cast(trunc_dim_)); + for (int round = 3; round >= 0; --round) { + inplace_fn_(data.data(), trunc_dim_); + rescale_fn_(data.data(), trunc_dim_, inv_fac); + flip_sign_fn_(flip_.data() + static_cast(round) * flip_offset_, + data.data(), dim); + } + std::memcpy(out, data.data(), dim * sizeof(float)); + return; + } + + // Non-power-of-2: undo final rescale(0.25) first + rescale_fn_(data.data(), dim, 4.0f); + + const float inv_fac = 1.0f / std::sqrt(static_cast(trunc_dim_)); + size_t start = dim - trunc_dim_; + float *trunc_ptr = data.data() + start; + + // Undo Round 4 (FHT on [start, start+trunc_dim)) + inv_kacs_walk_fn_(data.data(), dim); + inplace_fn_(trunc_ptr, trunc_dim_); + rescale_fn_(trunc_ptr, trunc_dim_, inv_fac); + flip_sign_fn_(flip_.data() + 3 * flip_offset_, data.data(), dim); + + // Undo Round 3 (FHT on [0, trunc_dim)) + inv_kacs_walk_fn_(data.data(), dim); + inplace_fn_(data.data(), trunc_dim_); + rescale_fn_(data.data(), trunc_dim_, inv_fac); + flip_sign_fn_(flip_.data() + 2 * flip_offset_, data.data(), dim); + + // Undo Round 2 (FHT on [start, start+trunc_dim)) + inv_kacs_walk_fn_(data.data(), dim); + inplace_fn_(trunc_ptr, trunc_dim_); + rescale_fn_(trunc_ptr, trunc_dim_, inv_fac); + flip_sign_fn_(flip_.data() + flip_offset_, data.data(), dim); + + // Undo Round 1 (FHT on [0, trunc_dim)) + inv_kacs_walk_fn_(data.data(), dim); + inplace_fn_(data.data(), trunc_dim_); + rescale_fn_(data.data(), trunc_dim_, inv_fac); + flip_sign_fn_(flip_.data(), data.data(), dim); + + std::memcpy(out, data.data(), dim * sizeof(float)); +} + +// --------------------------------------------------------------------------- +// serialize / deserialize +// --------------------------------------------------------------------------- + +int FhtRotator::serialize(std::string *out) const { + if (!out) return kErrInvalidArgument; + if (flip_.empty()) return kErrRuntime; + + RotatorSerHeader hdr{}; + hdr.magic = kRotatorMagic; + hdr.version = kRotatorSerVersion; + hdr.rotator_type = static_cast(RotatorType::kFht); + hdr.in_dim = static_cast(in_dim_); + hdr.out_dim = static_cast(out_dim_); + hdr.payload_size = static_cast(flip_.size()); + hdr.reserved = 0; + + out->resize(sizeof(hdr) + flip_.size()); + std::memcpy(&(*out)[0], &hdr, sizeof(hdr)); + std::memcpy(&(*out)[sizeof(hdr)], flip_.data(), flip_.size()); + return 0; +} + +int FhtRotator::deserialize(const void *data, size_t len) { + if (!data || len < sizeof(RotatorSerHeader)) return kErrInvalidArgument; + + const auto *hdr = reinterpret_cast(data); + if (hdr->magic != kRotatorMagic) return kErrUnsupported; + if (hdr->version != kRotatorSerVersion) return kErrUnsupported; + if (static_cast(hdr->rotator_type) != RotatorType::kFht) { + return kErrUnsupported; + } + + const size_t total = sizeof(RotatorSerHeader) + hdr->payload_size; + if (len < total) return kErrInvalidArgument; + + in_dim_ = static_cast(hdr->in_dim); + out_dim_ = static_cast(hdr->out_dim); + trunc_dim_ = floor_pow2(static_cast(in_dim_)); + fac_ = 1.0f / std::sqrt(static_cast(trunc_dim_)); + flip_offset_ = (static_cast(in_dim_) + kByteLen - 1) / kByteLen; + auto k = get_fht_kernels(); + flip_sign_fn_ = k.flip_sign; + kacs_walk_fn_ = k.kacs_walk; + inv_kacs_walk_fn_ = k.inv_kacs_walk; + inplace_fn_ = k.inplace; + rescale_fn_ = k.rescale; + + flip_.resize(hdr->payload_size); + std::memcpy(flip_.data(), + reinterpret_cast(data) + sizeof(RotatorSerHeader), + hdr->payload_size); + + return 0; +} + +} // namespace turbo +} // namespace zvec diff --git a/src/turbo/quantizer/preprocessor/fht_rotator.h b/src/turbo/quantizer/preprocessor/fht_rotator.h new file mode 100644 index 000000000..1bad6e924 --- /dev/null +++ b/src/turbo/quantizer/preprocessor/fht_rotator.h @@ -0,0 +1,104 @@ +// 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 +#include +#include +#include +#include +#include "preprocessor.h" + +namespace zvec { +namespace turbo { + +// ============================================================================ +// FhtRotator - O(d log d) FHT-based Kac random rotation +// +// Works with any dimension (non-power-of-2 uses trunc_dim + KacsWalk). +// When dimension is a power of 2, uses 4 rounds of (flip -> FHT -> rescale). +// When dimension is NOT a power of 2, uses kacs_walk reduction. +// ============================================================================ + +class FhtRotator : public Preprocessor { + public: + using Pointer = std::shared_ptr; + + //! Create a fully-initialized rotator for \p in_dim dimensions. + //! Random flip-sign arrays are generated during creation; the returned + //! object is immediately usable for apply() / apply_inverse(). + static Pointer create(int in_dim); + + //! Create and restore a rotator from a serialized blob (reads the type from + //! the embedded RotatorSerHeader). Returns nullptr on malformed input. + static Pointer from_blob(const void *data, size_t len); + + // -- Preprocessor interface ------------------------------------------------ + + int in_dim() const override { + return in_dim_; + } + int out_dim() const override { + return out_dim_; + } + + void apply(const float *in, float *out) const override; + void apply_inverse(const float *in, float *out) const override; + + //! No-op for FhtRotator. Flip-sign arrays are generated in create(). + //! Provided for interface compatibility with the Preprocessor contract. + void train(const void *data, size_t num, size_t stride) override; + + int serialize(std::string *out) const override; + int deserialize(const void *data, size_t len) override; + + //! Rotator type tag (kFht = 1). + RotatorType rotator_type() const { + return RotatorType::kFht; + } + + private: + FhtRotator() = default; + + //! Largest power of 2 <= dim. + static size_t floor_pow2(size_t n); + + int in_dim_{0}; + int out_dim_{0}; + + //! Packed flip-sign bits: 4 rounds, each ceil(in_dim / 8) bytes. + std::vector flip_; + + //! Bytes per round: ceil(in_dim / 8). + size_t flip_offset_{0}; + + //! Largest power of 2 <= in_dim (used for FHT length). + size_t trunc_dim_{0}; + + //! Rescale factor: 1 / sqrt(trunc_dim). + float fac_{0}; + + //! ISA-dispatched FHT kernels (cached at construction time). + FhtFlipSignFunc flip_sign_fn_{nullptr}; + FhtKacsWalkFunc kacs_walk_fn_{nullptr}; + FhtKacsWalkFunc inv_kacs_walk_fn_{nullptr}; + FhtInplaceFunc inplace_fn_{nullptr}; + FhtVecRescaleFunc rescale_fn_{nullptr}; + + static constexpr size_t kByteLen = 8; +}; + +} // namespace turbo +} // namespace zvec diff --git a/src/turbo/quantizer/preprocessor/preprocessor.h b/src/turbo/quantizer/preprocessor/preprocessor.h new file mode 100644 index 000000000..07ed313ea --- /dev/null +++ b/src/turbo/quantizer/preprocessor/preprocessor.h @@ -0,0 +1,93 @@ +// 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 +#include +#include +#include + +namespace zvec { +namespace turbo { + +//! Magic number ('ROTR') stamped at the start of a serialized rotator blob. +constexpr uint32_t kRotatorMagic = 0x52544F52u; +//! Current rotator serialization format version. +constexpr uint16_t kRotatorSerVersion = 1; + +//! Kind of rotator (stored in RotatorSerHeader::rotator_type). +enum class RotatorType : uint16_t { + kFht = 0, //!< O(d log d) FHT-based Kac random rotation +}; + +//! Self-describing, fixed-size header that prefixes every serialized rotator. +//! The type-specific payload (flip signs, rotation matrix, ...) follows +//! immediately after this header. +struct RotatorSerHeader { + uint32_t magic; // kRotatorMagic + uint16_t version; // kRotatorSerVersion + uint16_t rotator_type; // RotatorType + uint32_t in_dim; // input dimensionality + uint32_t out_dim; // output dimensionality + uint32_t payload_size; // bytes following the header + uint32_t reserved; // 0, for future use / alignment +}; +static_assert(sizeof(RotatorSerHeader) == 24, + "RotatorSerHeader must be 24 bytes"); + +//! Abstract preprocessor interface. +//! +//! A Preprocessor applies a deterministic, invertible transform to each +//! vector (e.g. random rotation). Concrete subclasses (FhtRotator, ...) +//! implement the actual algorithm. +class Preprocessor { + public: + using Pointer = std::shared_ptr; + + virtual ~Preprocessor() = default; + + //! Input dimensionality accepted by apply(). + virtual int in_dim() const = 0; + + //! Output dimensionality produced by apply(). May differ from in_dim() + //! (e.g. FHT pads to the next power of two), but for FhtRotator they are + //! equal. + virtual int out_dim() const = 0; + + //! Forward transform: map an input vector to the preprocessed space. + //! \p out must hold at least out_dim() elements. + virtual void apply(const float *in, float *out) const = 0; + + //! Inverse transform: recover the original-space vector from a preprocessed + //! one. \p out must hold at least in_dim() elements. + virtual void apply_inverse(const float *in, float *out) const = 0; + + //! Fit / initialize the preprocessor from a contiguous batch of training + //! data. For FhtRotator this generates the random flip-sign arrays. + //! \p data pointer to the first element of the batch. + //! \p num number of vectors in the batch. + //! \p stride byte offset between consecutive vectors (0 => packed). + virtual void train(const void *data, size_t num, size_t stride) = 0; + + //! Serialize the preprocessor into a self-contained blob + //! (RotatorSerHeader + payload). + virtual int serialize(std::string *out) const = 0; + + //! Deserialize the preprocessor from a raw, possibly mmap-backed buffer. + virtual int deserialize(const void *data, size_t len) = 0; +}; + +} // namespace turbo +} // namespace zvec diff --git a/src/turbo/quantizer/quantizer.h b/src/turbo/quantizer/quantizer.h new file mode 100644 index 000000000..4ef239bd5 --- /dev/null +++ b/src/turbo/quantizer/quantizer.h @@ -0,0 +1,202 @@ +// 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 +#include +#include +#include +#include +#include +#include +#include "distance.h" + +namespace zvec { +namespace turbo { + +using namespace zvec::core; + +//! Error code literals mirroring core::IndexError::Code integer values. +//! +//! Turbo quantizer sources use these directly instead of the +//! `IndexError_NotImplemented` / `IndexError_Unsupported` const objects +//! because MSVC's WINDOWS_EXPORT_ALL_SYMBOLS does not export const data +//! with constructors from zvec_shared.dll. zvec_turbo is a static library +//! linked with /WHOLEARCHIVE, so referencing those unexported symbols across +//! the DLL boundary triggers LNK2019 on Windows. +//! +//! IndexError::Code stores -val in its constructor, so NotImplemented(11) +//! yields -11 and Unsupported(12) yields -12. +constexpr int kErrRuntime = -1; +constexpr int kErrNotImplemented = -11; +constexpr int kErrUnsupported = -12; +constexpr int kErrInvalidArgument = -31; + +//! Magic number ('QTZR') stamped at the start of a serialized quantizer blob. +constexpr uint32_t kQuantizerMagic = 0x52545A51u; + +//! Current quantizer serialization format version. +constexpr uint16_t kQuantizerSerVersion = 1; + +//! Self-describing, fixed-size header that prefixes every serialized quantizer. +//! The type-specific payload (scalar params, codebook, rotation matrix, ...) +//! follows immediately after this header. +struct QuantizerSerHeader { + uint32_t magic; // kQuantizerMagic + uint16_t version; // kQuantizerSerVersion + uint16_t quant_type; // QuantizeType + uint32_t dim; // original dim (sanity check) + uint32_t metric; // MetricType (sanity check) + uint32_t payload_size; // bytes following the header + uint32_t reserved; // 0, for future use / alignment +}; +static_assert(sizeof(QuantizerSerHeader) == 24, + "QuantizerSerHeader must be 24 bytes"); + +class Quantizer { + public: + typedef std::shared_ptr Pointer; + + Quantizer() {} + virtual ~Quantizer() {} + + //! Initialize quantizer with index metadata and parameters + virtual int init(const IndexMeta &meta, const ailego::Params ¶ms) = 0; + + //! Get the output metadata after initialization + virtual const IndexMeta &meta() const = 0; + + //! Input data type accepted by the quantizer + virtual DataType input_data_type() const = 0; + + //! Data type + virtual QuantizeType type() const { + return type_; + } + + //! Dimensionality of the input vectors + virtual int dim() const = 0; + + //! Train the quantizer with a contiguous batch of data + virtual int train(const void * /*data*/, size_t /*num*/, size_t /*stride*/) { + return 0; + } + + //! Whether the quantizer requires training before use + virtual bool require_train() const = 0; + + //! Train the quantizer with data from an IndexHolder + virtual int train(IndexHolder::Pointer /*holder*/) { + return 0; + } + + //! Train the quantizer with data from an IndexHolder using multiple threads. + //! Default implementation ignores thread_count and delegates to train(). + virtual int train(IndexHolder::Pointer holder, int /*thread_count*/) { + return train(holder); + } + + //! Byte length of a quantized datapoint vector + virtual size_t quantized_datapoint_vector_length() const = 0; + + //! Byte length of a quantized query vector + virtual size_t quantized_query_vector_length() const = 0; + + //! Quantize a datapoint vector + virtual void quantize_data(const void *input, void *output) const = 0; + + //! Quantize a query vector + virtual void quantize_query(const void *input, void *output) const = 0; + + //! Distance between a quantized datapoint and a quantized query + virtual float calc_distance_dp_query(const void *dp, + const void *query) const = 0; + + //! Batched distance between quantized datapoints and a quantized query + virtual void calc_distance_dp_query_batch(const void *const *dp_list, + int dp_num, const void *query, + float *dist_list) const = 0; + + //! Distance between a quantized datapoint and an unquantized query + virtual float calc_distance_dp_query_unquantized(const void *dp, + const void *query) const = 0; + + //! Batched distance between quantized datapoints and an unquantized query + virtual void calc_distance_dp_query_batch_unquantized( + const void *const *dp_list, int dp_num, const void *query, + float *dist_list) const = 0; + + //! Distance between two quantized datapoints + virtual float calc_distance_dp_dp(const void *dp1, const void *dp2) const = 0; + + //! Quantize a query vector for search + virtual int quantize(const void * /*query*/, const IndexQueryMeta & /*qmeta*/, + std::string * /*out*/, + IndexQueryMeta * /*ometa*/) const { + return 0; + } + + //! Dequantize a result vector back to original format + virtual int dequantize(const void * /*in*/, const IndexQueryMeta & /*qmeta*/, + std::string * /*out*/) const { + return 0; + } + + virtual DistanceImpl distance(const void * /*query*/, + const IndexQueryMeta & /*qmeta*/) const { + return DistanceImpl{}; + } + + //! Serialize quantizer parameters + virtual int serialize(std::string * /*out*/) const { + return 0; + } + + //! Deserialize quantizer parameters + virtual int deserialize(std::string & /*in*/) { + return 0; + } + + //! Deserialize quantizer parameters from a raw, possibly mmap-backed buffer + //! (zero-copy entry point for large payloads such as codebooks/matrices). + virtual int deserialize(const void * /*data*/, size_t /*len*/) { + return 0; + } + + protected: + //! Map a metric name (e.g. "SquaredEuclidean", "Cosine", + //! "InnerProduct", "MipsSquaredEuclidean") to its MetricType. + static MetricType metric_from_name(const std::string &name) { + if (name == "SquaredEuclidean") { + return MetricType::kSquaredEuclidean; + } + if (name == "Cosine") { + return MetricType::kCosine; + } + if (name == "InnerProduct") { + return MetricType::kInnerProduct; + } + if (name == "MipsSquaredEuclidean") { + return MetricType::kMipsSquaredEuclidean; + } + return MetricType::kUnknown; + } + + QuantizeType type_{QuantizeType::kDefault}; + uint32_t extra_meta_size_{0}; +}; + +} // namespace turbo +} // namespace zvec diff --git a/src/turbo/turbo.cc b/src/turbo/turbo.cc index adc9b785e..c7655a4dc 100644 --- a/src/turbo/turbo.cc +++ b/src/turbo/turbo.cc @@ -18,14 +18,52 @@ #include "avx512_vnni/record_quantized_int8/squared_euclidean.h" #include "avx512_vnni/uniform_int8/quantize.h" #include "avx512_vnni/uniform_int8/squared_euclidean.h" +#include "scalar/fht/fht.h" +#include "scalar/fp32/cosine.h" +#include "scalar/fp32/inner_product.h" +#include "scalar/fp32/squared_euclidean.h" +#include "scalar/pq_quantizer_int4/pq_distance.h" +#include "scalar/pq_quantizer_int8/pq_distance.h" + +#if defined(__SSE2__) +#include "sse/fht/fht.h" +#endif +#if defined(__AVX2__) +#include "avx2/fht/fht.h" +#include "avx2/pq_quantizer_int4/pq_distance.h" +#include "avx2/pq_quantizer_int8/pq_distance.h" +#endif +#if defined(__AVX512F__) +#include "avx512/fht/fht.h" +#include "avx512/pq_quantizer_int4/pq_distance.h" +#include "avx512/pq_quantizer_int8/pq_distance.h" +#endif namespace zvec::turbo { DistanceFunc get_distance_func(MetricType metric_type, DataType data_type, - QuantizeType quantize_type) { + QuantizeType quantize_type, + CpuArchType cpu_arch_type) { + if (data_type == DataType::kFp32) { + if (quantize_type == QuantizeType::kDefault || + quantize_type == QuantizeType::kFp32) { + if (metric_type == MetricType::kCosine) { + return scalar::cosine_fp32_distance; + } + if (metric_type == MetricType::kSquaredEuclidean) { + return scalar::squared_euclidean_fp32_distance; + } + if (metric_type == MetricType::kInnerProduct) { + return scalar::inner_product_fp32_distance; + } + } + return nullptr; + } if (data_type == DataType::kInt8) { if (quantize_type == QuantizeType::kDefault) { - if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI) { + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI && + (cpu_arch_type == CpuArchType::kAuto || + cpu_arch_type == CpuArchType::kAVX512VNNI)) { if (metric_type == MetricType::kSquaredEuclidean) { return avx512_vnni::squared_euclidean_int8_distance; } @@ -47,10 +85,28 @@ DistanceFunc get_distance_func(MetricType metric_type, DataType data_type, BatchDistanceFunc get_batch_distance_func(MetricType metric_type, DataType data_type, - QuantizeType quantize_type) { + QuantizeType quantize_type, + CpuArchType cpu_arch_type) { + if (data_type == DataType::kFp32) { + if (quantize_type == QuantizeType::kDefault || + quantize_type == QuantizeType::kFp32) { + if (metric_type == MetricType::kCosine) { + return scalar::cosine_fp32_batch_distance; + } + if (metric_type == MetricType::kSquaredEuclidean) { + return scalar::squared_euclidean_fp32_batch_distance; + } + if (metric_type == MetricType::kInnerProduct) { + return scalar::inner_product_fp32_batch_distance; + } + } + return nullptr; + } if (data_type == DataType::kInt8) { if (quantize_type == QuantizeType::kDefault) { - if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI) { + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI && + (cpu_arch_type == CpuArchType::kAuto || + cpu_arch_type == CpuArchType::kAVX512VNNI)) { if (metric_type == MetricType::kSquaredEuclidean) { return avx512_vnni::squared_euclidean_int8_batch_distance; } @@ -67,15 +123,19 @@ BatchDistanceFunc get_batch_distance_func(MetricType metric_type, } } } + return nullptr; } QueryPreprocessFunc get_query_preprocess_func(MetricType metric_type, DataType data_type, - QuantizeType quantize_type) { + QuantizeType quantize_type, + CpuArchType cpu_arch_type) { if (data_type == DataType::kInt8) { if (quantize_type == QuantizeType::kDefault) { - if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI) { + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI && + (cpu_arch_type == CpuArchType::kAuto || + cpu_arch_type == CpuArchType::kAVX512VNNI)) { if (metric_type == MetricType::kSquaredEuclidean) { return avx512_vnni::squared_euclidean_int8_query_preprocess; } @@ -100,4 +160,99 @@ UniformQuantizeFunc get_uniform_quantize_func(DataType data_type) { return nullptr; } -} // namespace zvec::turbo +FhtKernels get_fht_kernels() { + FhtKernels k; + // Default: scalar fallback for all + k.flip_sign = scalar::fht_flip_sign; + k.kacs_walk = scalar::fht_kacs_walk; + k.inv_kacs_walk = scalar::fht_inv_kacs_walk; + k.inplace = scalar::fht_inplace; + k.rescale = scalar::fht_vec_rescale; + +#if defined(__AVX512F__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F && + zvec::ailego::internal::CpuFeatures::static_flags_.AVX512DQ) { + k.flip_sign = avx512::fht_flip_sign_avx512; + k.kacs_walk = avx512::fht_kacs_walk_avx512; + k.inv_kacs_walk = avx512::fht_inv_kacs_walk_avx512; + k.inplace = avx512::fht_inplace_avx512; + k.rescale = avx512::fht_vec_rescale_avx512; + return k; + } +#endif +#if defined(__AVX2__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) { + k.flip_sign = avx2::fht_flip_sign_avx2; + k.kacs_walk = avx2::fht_kacs_walk_avx2; + k.inv_kacs_walk = avx2::fht_inv_kacs_walk_avx2; + k.inplace = avx2::fht_inplace_avx2; + k.rescale = avx2::fht_vec_rescale_avx2; + return k; + } +#endif +#if defined(__SSE2__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.SSE2) { + k.flip_sign = sse::fht_flip_sign_sse; + k.kacs_walk = sse::fht_kacs_walk_sse; + k.inv_kacs_walk = sse::fht_inv_kacs_walk_sse; + k.rescale = sse::fht_vec_rescale_sse; + // inplace fallback to scalar (SSE has no fht_inplace) + return k; + } +#endif + return k; // scalar +} + +PqKernels get_pq_kernels(DataType data_type, QuantizeType quantize_type, + CpuArchType cpu_arch_type) { + (void)cpu_arch_type; // currently unused, reserved for future use + PqKernels k{}; + if (quantize_type == QuantizeType::kPQ) { + if (data_type == DataType::kInt4) { + // int4 packed nibble path — scalar by default. + k.adc_distance = scalar::pq_adc_int4_distance; + k.sdc_distance = scalar::pq_sdc_int4_distance; + k.batch_adc_distance = scalar::pq_adc_int4_batch_distance; + +#if defined(__AVX512F__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) { + k.adc_distance = avx512::pq_adc_int4_distance_avx512; + k.sdc_distance = avx512::pq_sdc_int4_distance_avx512; + k.batch_adc_distance = avx512::pq_adc_int4_batch_distance_avx512; + return k; + } +#endif +#if defined(__AVX2__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) { + k.adc_distance = avx2::pq_adc_int4_distance_avx2; + k.sdc_distance = avx2::pq_sdc_int4_distance_avx2; + k.batch_adc_distance = avx2::pq_adc_int4_batch_distance_avx2; + } +#endif + return k; + } + // Default (kInt8 / fallback): scalar int8 kernels. + k.adc_distance = scalar::pq_adc_int8_distance; + k.sdc_distance = scalar::pq_sdc_int8_distance; + k.batch_adc_distance = scalar::pq_adc_int8_batch_distance; + +#if defined(__AVX512F__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) { + k.adc_distance = avx512::pq_adc_int8_distance_avx512; + k.sdc_distance = avx512::pq_sdc_int8_distance_avx512; + k.batch_adc_distance = avx512::pq_adc_int8_batch_distance_avx512; + return k; + } +#endif +#if defined(__AVX2__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) { + k.adc_distance = avx2::pq_adc_int8_distance_avx2; + k.sdc_distance = avx2::pq_sdc_int8_distance_avx2; + k.batch_adc_distance = avx2::pq_adc_int8_batch_distance_avx2; + } +#endif + } + return k; +} + +} // namespace zvec::turbo \ No newline at end of file diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 7be2294dd..e3b54ee24 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -7,6 +7,7 @@ include_directories(${PROJECT_ROOT_DIR}) cc_directories(ailego) cc_directories(db) cc_directories(core) +cc_directories(turbo) if(BUILD_C_BINDINGS) cc_directories(c) endif() diff --git a/tests/ailego/parallel/multi_thread_list_test.cc b/tests/ailego/parallel/multi_thread_list_test.cc index 17a896f1b..3a94afebf 100644 --- a/tests/ailego/parallel/multi_thread_list_test.cc +++ b/tests/ailego/parallel/multi_thread_list_test.cc @@ -30,7 +30,7 @@ using namespace std; struct Item { uint32_t a_; std::string b_; - Item() {}; + Item(){}; Item(uint32_t a, std::string b) : a_(a), b_(b) {} }; @@ -182,7 +182,7 @@ TEST(MultiThreadListTest, ConsumeStopResume) { struct MoveableItem { uint32_t a_; std::string b_; - MoveableItem() {}; + MoveableItem(){}; MoveableItem(uint32_t a, std::string b) : a_(a), b_(b) {} MoveableItem(const MoveableItem &) = delete; diff --git a/tests/ailego/parallel/thread_queue_test.cc b/tests/ailego/parallel/thread_queue_test.cc index 9dc8fc7e7..439fdc3ae 100644 --- a/tests/ailego/parallel/thread_queue_test.cc +++ b/tests/ailego/parallel/thread_queue_test.cc @@ -66,7 +66,7 @@ TEST(ThreadQueue, MutliThread) { } TEST(ThreadQueue, MultiThreadWithHighPriority) { -// TODO(windows): add it back + // TODO(windows): add it back GTEST_SKIP(); ThreadQueue queue; diff --git a/tests/android_gmock_main.cc b/tests/android_gmock_main.cc index 20cd45da7..1538da6b9 100644 --- a/tests/android_gmock_main.cc +++ b/tests/android_gmock_main.cc @@ -5,10 +5,10 @@ // (exit code 139) or aborts (134/135) during teardown even when all tests // passed. Using _exit() skips the static destructor phase entirely. +#include #include #include #include -#include #include "gmock/gmock.h" GTEST_API_ int main(int argc, char **argv) { diff --git a/tests/android_gtest_main.cc b/tests/android_gtest_main.cc index 372e2b463..81ee6e36e 100644 --- a/tests/android_gtest_main.cc +++ b/tests/android_gtest_main.cc @@ -5,10 +5,10 @@ // (exit code 139) or aborts (134/135) during teardown even when all tests // passed. Using _exit() skips the static destructor phase entirely. +#include #include #include #include -#include #include "gtest/gtest.h" GTEST_API_ int main(int argc, char **argv) { diff --git a/tests/core/algorithm/hnsw/hnsw_streamer_buffer_test.cc b/tests/core/algorithm/hnsw/hnsw_streamer_buffer_test.cc index 17c3c8704..46365cbaf 100644 --- a/tests/core/algorithm/hnsw/hnsw_streamer_buffer_test.cc +++ b/tests/core/algorithm/hnsw/hnsw_streamer_buffer_test.cc @@ -334,7 +334,8 @@ TEST_F(HnswStreamerTest, TestHnswSearchBufferMMap) { auto read_storage = IndexFactory::CreateStorage("MMapFileStorage"); ASSERT_NE(nullptr, read_storage); ASSERT_EQ(0, read_storage->init(stg_params)); - ASSERT_EQ(0, read_storage->open(dir_ + "Test/TestHnswSearchBufferMMap", false)); + ASSERT_EQ(0, + read_storage->open(dir_ + "Test/TestHnswSearchBufferMMap", false)); ASSERT_EQ(0, read_streamer->open(read_storage)); size_t topk = 3; auto provider = read_streamer->create_provider(); diff --git a/tests/core/utility/buffer_storage_write_test.cc b/tests/core/utility/buffer_storage_write_test.cc index 894c68f1e..1079e01ff 100644 --- a/tests/core/utility/buffer_storage_write_test.cc +++ b/tests/core/utility/buffer_storage_write_test.cc @@ -42,7 +42,9 @@ class BufferStorageWriteTest : public ::testing::Test { ailego::File::MakePath("buffer_storage_write_test_dir"); } - void TearDown() override { ailego::File::Delete(file_path_); } + void TearDown() override { + ailego::File::Delete(file_path_); + } // Open BufferStorage in writable mode (create_if_missing=true) IndexStorage::Pointer OpenWritable() { @@ -69,7 +71,8 @@ class BufferStorageWriteTest : public ::testing::Test { // ===== Basic Write Tests ===== -// Test: Create new index via BufferStorage, append segment, write data, read back +// Test: Create new index via BufferStorage, append segment, write data, read +// back TEST_F(BufferStorageWriteTest, WriteBasicCreateAndWrite) { auto storage = OpenWritable(); ASSERT_TRUE(storage); @@ -275,8 +278,7 @@ TEST_F(BufferStorageWriteTest, WriteMultipleFlushCycles) { EXPECT_EQ(0, storage->flush()); // Second write at a different offset + flush - EXPECT_EQ(data2.size(), - seg->write(200, data2.data(), data2.size())); + EXPECT_EQ(data2.size(), seg->write(200, data2.data(), data2.size())); EXPECT_EQ(0, storage->flush()); EXPECT_EQ(0, storage->close()); } @@ -353,8 +355,7 @@ TEST_F(BufferStorageWriteTest, WriteReadOnlyNoOp) { std::string new_data = "overwrite_attempt"; // Should return len (silent no-op) - EXPECT_EQ(new_data.size(), - seg->write(0, new_data.data(), new_data.size())); + EXPECT_EQ(new_data.size(), seg->write(0, new_data.data(), new_data.size())); // Data should remain unchanged (still "initial") std::vector buf(7); @@ -881,7 +882,8 @@ TEST_F(BufferStorageWriteTest, CR_ConcurrentWriteAndResize) { // chain split. After reopen, ALL segments must be findable. // (Tests fix for reserve()-induced dangling pointer in append_segment.) TEST_F(BufferStorageWriteTest, CR_ChainSplitAllSegmentsAccessible) { - const int kNumSegments = 50; // Enough to trigger chain split with default 4096 meta capacity + const int kNumSegments = + 50; // Enough to trigger chain split with default 4096 meta capacity { auto storage = OpenWritable(); @@ -892,7 +894,8 @@ TEST_F(BufferStorageWriteTest, CR_ChainSplitAllSegmentsAccessible) { ASSERT_EQ(0, storage->append(name, 4096)) << "Failed to append segment " << i; auto seg = storage->get(name); - ASSERT_TRUE(seg) << "Failed to get segment " << name << " right after append"; + ASSERT_TRUE(seg) << "Failed to get segment " << name + << " right after append"; // Write a marker so we can verify on reopen std::string marker = "marker_" + std::to_string(i); EXPECT_EQ(marker.size(), seg->write(0, marker.data(), marker.size())); @@ -908,7 +911,8 @@ TEST_F(BufferStorageWriteTest, CR_ChainSplitAllSegmentsAccessible) { for (int i = 0; i < kNumSegments; ++i) { std::string name = "chain_seg_" + std::to_string(i); auto seg = storage->get(name); - ASSERT_TRUE(seg) << "Segment " << name << " missing after reopen (chain-split bug?)"; + ASSERT_TRUE(seg) << "Segment " << name + << " missing after reopen (chain-split bug?)"; std::string expected = "marker_" + std::to_string(i); std::vector buf(expected.size()); EXPECT_EQ(expected.size(), seg->fetch(0, buf.data(), buf.size())); @@ -1047,8 +1051,9 @@ TEST_F(BufferStorageWriteTest, CR_DirtyFlagNotLostAfterFlush) { } } -// Stress test: Concurrent flush + write interleaving to expose dirty flag races. -// All writes that return successfully MUST be visible after final close+reopen. +// Stress test: Concurrent flush + write interleaving to expose dirty flag +// races. All writes that return successfully MUST be visible after final +// close+reopen. TEST_F(BufferStorageWriteTest, CR_ConcurrentFlushWriteDirtyFlagStress) { auto storage = OpenWritable(); ASSERT_TRUE(storage); @@ -1112,7 +1117,8 @@ TEST_F(BufferStorageWriteTest, CR_PointerStabilityAcrossAppend) { // Write initial data std::string initial = "before_append"; - EXPECT_EQ(initial.size(), seg_first->write(0, initial.data(), initial.size())); + EXPECT_EQ(initial.size(), + seg_first->write(0, initial.data(), initial.size())); // Append many more segments (may trigger internal rehash/resize) for (int i = 0; i < 20; ++i) { diff --git a/tests/db/collection_test.cc b/tests/db/collection_test.cc index 74407af6e..68cbdc8b2 100644 --- a/tests/db/collection_test.cc +++ b/tests/db/collection_test.cc @@ -211,10 +211,10 @@ TEST_F(CollectionTest, Feature_OpenReadOnly_WithReadOnlyLockFile) { // Use std::filesystem to set read-only permissions (cross-platform) std::error_code ec; - fs::permissions(lock_path, - fs::perms::owner_read | fs::perms::group_read | - fs::perms::others_read, - fs::perm_options::replace, ec); + fs::permissions( + lock_path, + fs::perms::owner_read | fs::perms::group_read | fs::perms::others_read, + fs::perm_options::replace, ec); ASSERT_FALSE(ec) << "Failed to set read-only permissions: " << ec.message(); // Open with read_only=true should succeed even with read-only LOCK file diff --git a/tests/db/crash_recovery/utility.h b/tests/db/crash_recovery/utility.h index 826111ca2..ceb4b7364 100644 --- a/tests/db/crash_recovery/utility.h +++ b/tests/db/crash_recovery/utility.h @@ -22,7 +22,6 @@ #include #include #include - #include #include @@ -156,7 +155,6 @@ inline Doc CreateTestDoc(uint64_t doc_id, int version) { } - /** * @brief Locate a binary by name, searching common paths and TEST_BINARY_DIR. * diff --git a/tests/db/index/column/inverted_column/inverted_column_indexer_array_numbers_test.cc b/tests/db/index/column/inverted_column/inverted_column_indexer_array_numbers_test.cc index 620b4ef6f..e58f41793 100644 --- a/tests/db/index/column/inverted_column/inverted_column_indexer_array_numbers_test.cc +++ b/tests/db/index/column/inverted_column/inverted_column_indexer_array_numbers_test.cc @@ -43,7 +43,7 @@ class TestHelper { public: TestHelper(uint32_t num_docs, uint32_t num_write_threads = 10) : num_docs_(num_docs / 100 * 100), - num_write_threads_(num_write_threads) {}; + num_write_threads_(num_write_threads){}; template diff --git a/tests/db/index/column/inverted_column/inverted_column_indexer_bool_test.cc b/tests/db/index/column/inverted_column/inverted_column_indexer_bool_test.cc index b26fe7123..c51e6f733 100644 --- a/tests/db/index/column/inverted_column/inverted_column_indexer_bool_test.cc +++ b/tests/db/index/column/inverted_column/inverted_column_indexer_bool_test.cc @@ -43,7 +43,7 @@ class TestHelper { public: TestHelper(uint32_t num_docs, uint32_t num_write_threads = 10) : num_docs_(num_docs / 100 * 100), - num_write_threads_(num_write_threads) {}; + num_write_threads_(num_write_threads){}; void insert_bools(InvertedColumnIndexer::Ptr indexer) { @@ -337,7 +337,8 @@ TEST_F(InvertedIndexTest, SEALED) { TEST_F(InvertedIndexTest, SNAPSHOT) { #ifdef __ANDROID__ - GTEST_SKIP() << "Skipped on Android: emulator filesystem lacks hardlink support (needed by RocksDB checkpoint)"; + GTEST_SKIP() << "Skipped on Android: emulator filesystem lacks hardlink " + "support (needed by RocksDB checkpoint)"; #endif ASSERT_TRUE(indexer_); diff --git a/tests/db/index/column/inverted_column/inverted_column_indexer_cyclic_numbers_test.cc b/tests/db/index/column/inverted_column/inverted_column_indexer_cyclic_numbers_test.cc index 07bc1e7da..a64d662c3 100644 --- a/tests/db/index/column/inverted_column/inverted_column_indexer_cyclic_numbers_test.cc +++ b/tests/db/index/column/inverted_column/inverted_column_indexer_cyclic_numbers_test.cc @@ -44,7 +44,7 @@ class TestHelper { public: TestHelper(uint32_t num_docs, uint32_t num_write_threads = 10) : num_docs_(num_docs / 100 * 100), - num_write_threads_(num_write_threads) {}; + num_write_threads_(num_write_threads){}; template diff --git a/tests/db/index/column/inverted_column/inverted_column_indexer_sequential_numbers_test.cc b/tests/db/index/column/inverted_column/inverted_column_indexer_sequential_numbers_test.cc index 0d2bda43a..ee634cff2 100644 --- a/tests/db/index/column/inverted_column/inverted_column_indexer_sequential_numbers_test.cc +++ b/tests/db/index/column/inverted_column/inverted_column_indexer_sequential_numbers_test.cc @@ -45,7 +45,7 @@ class TestHelper { public: TestHelper(uint32_t num_docs, uint32_t num_write_threads = 10) : num_docs_(num_docs / 100 * 100), - num_write_threads_(num_write_threads) {}; + num_write_threads_(num_write_threads){}; template @@ -698,7 +698,8 @@ TEST_F(InvertedIndexTest, SEALED) { TEST_F(InvertedIndexTest, CREATE_SNAPSHOT) { #ifdef __ANDROID__ - GTEST_SKIP() << "Skipped on Android: emulator filesystem lacks hardlink support (needed by RocksDB checkpoint)"; + GTEST_SKIP() << "Skipped on Android: emulator filesystem lacks hardlink " + "support (needed by RocksDB checkpoint)"; #endif ASSERT_TRUE(indexer_); diff --git a/tests/db/index/column/inverted_column/inverted_column_indexer_string_test.cc b/tests/db/index/column/inverted_column/inverted_column_indexer_string_test.cc index b1da86595..9d1fbfd21 100644 --- a/tests/db/index/column/inverted_column/inverted_column_indexer_string_test.cc +++ b/tests/db/index/column/inverted_column/inverted_column_indexer_string_test.cc @@ -44,7 +44,7 @@ class TestHelper { public: TestHelper(uint32_t num_docs, uint32_t num_write_threads = 10) : num_docs_(num_docs / 100 * 100), - num_write_threads_(num_write_threads) {}; + num_write_threads_(num_write_threads){}; void insert_strings(InvertedColumnIndexer::Ptr indexer) { diff --git a/tests/db/index/common/query_params_test.cc b/tests/db/index/common/query_params_test.cc index ceb510746..13e8bd00b 100644 --- a/tests/db/index/common/query_params_test.cc +++ b/tests/db/index/common/query_params_test.cc @@ -21,7 +21,6 @@ TEST(QueryParamsTest, QueryParamsBaseClass) { // Test constructor QueryParams params(IndexType::HNSW); EXPECT_EQ(params.type(), IndexType::HNSW); - } TEST(QueryParamsTest, HnswQueryParams) { diff --git a/tests/db/index/segment/segment_test.cc b/tests/db/index/segment/segment_test.cc index 9582d2bf3..c5f4c7b1b 100644 --- a/tests/db/index/segment/segment_test.cc +++ b/tests/db/index/segment/segment_test.cc @@ -35,16 +35,16 @@ #include "db/index/common/id_map.h" #include "db/index/common/version_manager.h" #include "db/index/storage/wal/wal_file.h" -#include "segment_test_fixture.h" #include "utils/utils.h" #include "zvec/db/options.h" +#include "segment_test_fixture.h" using namespace zvec; TEST_P(SegmentTest, EmptySchema) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 0); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 0); ASSERT_TRUE(segment != nullptr); EXPECT_EQ(segment->id(), 0); @@ -56,8 +56,8 @@ TEST_P(SegmentTest, General) { options_.max_buffer_size_ = 1 * 1024; auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 25); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 25); ASSERT_TRUE(segment != nullptr); auto combined_reader = segment->scan({LOCAL_ROW_ID, "id", "name", "age"}); @@ -106,8 +106,8 @@ TEST_P(SegmentTest, General) { TEST_P(SegmentTest, InsertMoreData) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 0); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 0); ASSERT_TRUE(segment != nullptr); uint64_t MAX_DOC = 1000; @@ -137,14 +137,14 @@ TEST_P(SegmentTest, InsertScalarTypes) { auto invert_params = std::make_shared(false); schema_->add_field(std::make_shared("binary", DataType::BINARY, - false, invert_params)); + false, invert_params)); schema_->add_field(std::make_shared( "array_binary", DataType::ARRAY_BINARY, false, invert_params)); auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 10); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 10); ASSERT_TRUE(segment != nullptr); } @@ -167,8 +167,8 @@ TEST_P(SegmentTest, InsertVectorTypes) { { Version v = version_manager_->get_current_version(); auto result = - Segment::Open(col_path_, *tmp_schema, *v.writing_segment_meta(), id_map_, - delete_store_, version_manager_, options_); + Segment::Open(col_path_, *tmp_schema, *v.writing_segment_meta(), + id_map_, delete_store_, version_manager_, options_); ASSERT_TRUE(result.has_value()); auto segment = result.value(); @@ -179,8 +179,8 @@ TEST_P(SegmentTest, InsertVectorTypes) { TEST_P(SegmentTest, FetchByGlobalDocID) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 1); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 1); ASSERT_TRUE(segment != nullptr); auto ret_doc = segment->Fetch(0); @@ -192,8 +192,8 @@ TEST_P(SegmentTest, FetchByGlobalDocID) { TEST_P(SegmentTest, FetchSingleRow) { int doc_count = 10; auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, doc_count); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, doc_count); ASSERT_TRUE(segment != nullptr); auto func = [&](int index) -> void { @@ -219,8 +219,8 @@ TEST_P(SegmentTest, FetchSingleRowWithPersistStore) { int doc_count = 1000; { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, doc_count); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, doc_count); ASSERT_TRUE(segment != nullptr); } @@ -229,9 +229,9 @@ TEST_P(SegmentTest, FetchSingleRowWithPersistStore) { Version v = version_manager_->get_current_version(); SegmentOptions open_options; open_options.read_only_ = false; - auto result = Segment::Open(col_path_, *schema_, *v.writing_segment_meta(), - id_map_, delete_store_, version_manager_, - open_options); + auto result = + Segment::Open(col_path_, *schema_, *v.writing_segment_meta(), id_map_, + delete_store_, version_manager_, open_options); ASSERT_TRUE(result.has_value()); auto segment = result.value(); @@ -259,8 +259,8 @@ TEST_P(SegmentTest, FetchSingleRowWithPersistStore) { TEST_P(SegmentTest, FetchSingleRowWithUserID) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 10); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 10); ASSERT_TRUE(segment != nullptr); ExecBatchPtr batch = segment->fetch({USER_ID, "id", "name"}, 2); @@ -276,8 +276,8 @@ TEST_P(SegmentTest, FetchSingleRowWithUserID) { TEST_P(SegmentTest, FetchSingleRowWithGlobalDocID) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 10); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 10); ASSERT_TRUE(segment != nullptr); ExecBatchPtr batch = segment->fetch({GLOBAL_DOC_ID, "id", "name"}, 4); @@ -293,8 +293,8 @@ TEST_P(SegmentTest, FetchSingleRowWithGlobalDocID) { TEST_P(SegmentTest, FetchSingleRowWithNegativeIndex) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 10); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 10); ASSERT_TRUE(segment != nullptr); ExecBatchPtr batch = segment->fetch({"id", "name"}, -1); @@ -303,8 +303,8 @@ TEST_P(SegmentTest, FetchSingleRowWithNegativeIndex) { TEST_P(SegmentTest, FetchSingleRowWithOutOfRangeIndex) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 10); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 10); ASSERT_TRUE(segment != nullptr); ExecBatchPtr batch = segment->fetch({"id", "name"}, 15); @@ -313,8 +313,8 @@ TEST_P(SegmentTest, FetchSingleRowWithOutOfRangeIndex) { TEST_P(SegmentTest, FetchSingleRowWithInvalidColumn) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 10); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 10); ASSERT_TRUE(segment != nullptr); ExecBatchPtr batch = segment->fetch({"id", "invalid_column"}, 0); @@ -323,8 +323,8 @@ TEST_P(SegmentTest, FetchSingleRowWithInvalidColumn) { TEST_P(SegmentTest, FetchSingleRowWithEmptyColumns) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 10); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 10); ASSERT_TRUE(segment != nullptr); ExecBatchPtr batch = segment->fetch({}, 0); @@ -333,8 +333,8 @@ TEST_P(SegmentTest, FetchSingleRowWithEmptyColumns) { TEST_P(SegmentTest, FetchSingleRowFromEmptySegment) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 0); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 0); ASSERT_TRUE(segment != nullptr); ExecBatchPtr batch = segment->fetch({"id", "name"}, 0); @@ -343,8 +343,8 @@ TEST_P(SegmentTest, FetchSingleRowFromEmptySegment) { TEST_P(SegmentTest, FetchSingleRowWithBinaryFields) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 10); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 10); ASSERT_TRUE(segment != nullptr); ExecBatchPtr batch = segment->fetch({"binary", "array_binary"}, 1); @@ -368,8 +368,8 @@ TEST_P(SegmentTest, Recover) { int doc_count = 100; { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, doc_count); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, doc_count); ASSERT_TRUE(segment != nullptr); } @@ -424,9 +424,9 @@ TEST_P(SegmentTest, Recover) { Version v = version_manager_->get_current_version(); SegmentOptions open_options; open_options.read_only_ = false; - auto result = Segment::Open(col_path_, *schema_, *v.writing_segment_meta(), - id_map_, delete_store_, version_manager_, - open_options); + auto result = + Segment::Open(col_path_, *schema_, *v.writing_segment_meta(), id_map_, + delete_store_, version_manager_, open_options); ASSERT_TRUE(result.has_value()); auto segment = result.value(); @@ -452,8 +452,8 @@ TEST_P(SegmentTest, Recover) { TEST_P(SegmentTest, UpdateDoc) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 10); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 10); ASSERT_TRUE(segment != nullptr); // before update @@ -484,8 +484,8 @@ TEST_P(SegmentTest, UpdateDoc) { TEST_P(SegmentTest, UpdateDocBatch) { int doc_count = 10; auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, doc_count); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, doc_count); ASSERT_TRUE(segment != nullptr); // before update uint64_t count = segment->doc_count(delete_store_->make_filter()); @@ -513,8 +513,8 @@ TEST_P(SegmentTest, UpdateDocBatch) { TEST_P(SegmentTest, DeleteDoc) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 10); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 10); ASSERT_TRUE(segment != nullptr); // before update @@ -541,8 +541,8 @@ TEST_P(SegmentTest, DeleteDoc) { TEST_P(SegmentTest, DeleteBatch) { int doc_count = 10; auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, doc_count); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, doc_count); ASSERT_TRUE(segment != nullptr); // before update @@ -562,8 +562,8 @@ TEST_P(SegmentTest, DeleteBatch) { TEST_P(SegmentTest, UpsertDoc) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 5); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 5); ASSERT_TRUE(segment != nullptr); // before update @@ -603,8 +603,8 @@ TEST_P(SegmentTest, UpsertDoc) { TEST_P(SegmentTest, UpsertDocBatch) { int doc_count = 10; auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, doc_count); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, doc_count); ASSERT_TRUE(segment != nullptr); // before update @@ -648,8 +648,8 @@ TEST_P(SegmentTest, UpsertDocBatch) { TEST_P(SegmentTest, Flush) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 100); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 100); ASSERT_TRUE(segment != nullptr); // Flush the segment @@ -659,8 +659,8 @@ TEST_P(SegmentTest, Flush) { TEST_P(SegmentTest, FlushAfterInsert) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 100); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 100); ASSERT_TRUE(segment != nullptr); // Flush the segment @@ -686,8 +686,8 @@ TEST_P(SegmentTest, FlushAfterInsert) { TEST_P(SegmentTest, Dump) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 100); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 100); ASSERT_TRUE(segment != nullptr); // Dump the segment @@ -701,8 +701,8 @@ TEST_P(SegmentTest, Dump) { TEST_P(SegmentTest, DocCount) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 50); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 50); ASSERT_TRUE(segment != nullptr); // Get document count @@ -760,8 +760,8 @@ TEST_P(SegmentTest, CombinedVectorColumnIndexer) { options_.max_buffer_size_ = 10 * 1024; auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 0); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 0); ASSERT_TRUE(segment != nullptr); @@ -947,8 +947,8 @@ TEST_P(SegmentTest, CombinedVectorColumnIndexerQueryWithPks) { TEST_P(SegmentTest, ConcurrentInsertOperations) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 0); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 0); ASSERT_TRUE(segment != nullptr); const int num_threads = 4; @@ -980,8 +980,8 @@ TEST_P(SegmentTest, ConcurrentInsertOperations) { TEST_P(SegmentTest, ConcurrentMixedOperations) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 100); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 100); ASSERT_TRUE(segment != nullptr); std::vector threads; @@ -1022,8 +1022,8 @@ TEST_P(SegmentTest, ConcurrentMixedOperations) { // corner cases TEST_P(SegmentTest, DuplicateInsert) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 0); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 0); ASSERT_TRUE(segment != nullptr); Doc doc1 = test::TestHelper::CreateDoc(0, *schema_); @@ -1051,8 +1051,8 @@ TEST_P(SegmentTest, DuplicateInsert) { TEST_P(SegmentTest, DuplicateDelete) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 5); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 5); ASSERT_TRUE(segment != nullptr); auto status1 = segment->Delete("pk_2"); @@ -1068,8 +1068,8 @@ TEST_P(SegmentTest, DuplicateDelete) { TEST_P(SegmentTest, DeleteNonExistentDoc) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 5); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 5); ASSERT_TRUE(segment != nullptr); auto status1 = segment->Delete("pk_999"); @@ -1078,8 +1078,8 @@ TEST_P(SegmentTest, DeleteNonExistentDoc) { TEST_P(SegmentTest, UpdateNonExistentDoc) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 5); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 5); ASSERT_TRUE(segment != nullptr); Doc doc = test::TestHelper::CreateDoc(999, *schema_); @@ -1091,8 +1091,8 @@ TEST_P(SegmentTest, UpdateNonExistentDoc) { TEST_P(SegmentTest, UpsertNonExistentDoc) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 5); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 5); ASSERT_TRUE(segment != nullptr); Doc doc = test::TestHelper::CreateDoc(999, *schema_); @@ -1109,8 +1109,8 @@ TEST_P(SegmentTest, UpsertNonExistentDoc) { TEST_P(SegmentTest, ScanWithEmptyColumns) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 5); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 5); ASSERT_TRUE(segment != nullptr); auto reader = segment->scan({}); @@ -1119,8 +1119,8 @@ TEST_P(SegmentTest, ScanWithEmptyColumns) { TEST_P(SegmentTest, ScanWithInvalidColumns) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 10); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 10); ASSERT_TRUE(segment != nullptr); // Try to scan with invalid column name @@ -1130,8 +1130,8 @@ TEST_P(SegmentTest, ScanWithInvalidColumns) { TEST_P(SegmentTest, FetchNonExistentDoc) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 5); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 5); ASSERT_TRUE(segment != nullptr); auto doc = segment->Fetch(999); @@ -1140,8 +1140,8 @@ TEST_P(SegmentTest, FetchNonExistentDoc) { TEST_P(SegmentTest, FetchWithInvalidSegmentDocIDs) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 5); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 5); ASSERT_TRUE(segment != nullptr); std::vector invalid_segment_doc_ids = {999, 1000}; @@ -1152,8 +1152,8 @@ TEST_P(SegmentTest, FetchWithInvalidSegmentDocIDs) { TEST_P(SegmentTest, FetchWithInvalidColumns) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 10); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 10); ASSERT_TRUE(segment != nullptr); // Try to fetch with invalid column name @@ -1166,8 +1166,8 @@ TEST_P(SegmentTest, InsertEmptyDocWithNullableSchema) { auto nullable_schema = test::TestHelper::CreateNormalSchema(true, col_name_); auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *nullable_schema, 0, 0, id_map_, delete_store_, version_manager_, - options_, 0, 0); + col_path_, *nullable_schema, 0, 0, id_map_, delete_store_, + version_manager_, options_, 0, 0); ASSERT_TRUE(segment != nullptr); Doc empty_doc; @@ -1178,8 +1178,8 @@ TEST_P(SegmentTest, InsertEmptyDocWithNullableSchema) { TEST_P(SegmentTest, MultipleDuplicateDeletes) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 5); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 5); ASSERT_TRUE(segment != nullptr); auto status1 = segment->Delete("pk_1"); @@ -1202,8 +1202,8 @@ TEST_P(SegmentTest, FetchWithTwoVectorFields) { int doc_count = 1000; auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, doc_count); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, doc_count); ASSERT_TRUE(segment != nullptr); segment.reset(); version_manager_.reset(); @@ -1262,8 +1262,8 @@ TEST_P(SegmentTest, FetchPerf) { int doc_count = 1000; options_.max_buffer_size_ = 100 * 1024; auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, doc_count); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, doc_count); ASSERT_TRUE(segment != nullptr); segment->dump(); @@ -1367,8 +1367,8 @@ TEST_P(SegmentTest, AddColumn) { options_.max_buffer_size_ = 10 * 1024 * 1024; int doc_count = 1000; auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, doc_count); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, doc_count); ASSERT_TRUE(segment != nullptr); auto s = segment->add_column( @@ -1600,8 +1600,8 @@ TEST_P(SegmentTest, AddNullableColumnWithoutExpressionMultiBlock) { options_.max_buffer_size_ = 1 * 1024; int doc_count = 100; auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, doc_count); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, doc_count); ASSERT_TRUE(segment != nullptr); segment->dump(); @@ -1674,9 +1674,8 @@ TEST_P(SegmentTest, AddNullableColumnWithoutExpressionMultiBlock) { auto field_schema = std::make_shared(nullable_col_name, data_type, true); s = segment->add_column(field_schema, "", AddColumnOptions()); - ASSERT_TRUE(s.ok()) - << "Failed to add nullable column " << nullable_col_name << ": " - << s.message(); + ASSERT_TRUE(s.ok()) << "Failed to add nullable column " << nullable_col_name + << ": " << s.message(); auto combined_reader = segment->scan({"id", "name", "age", nullable_col_name}); @@ -1699,8 +1698,8 @@ TEST_P(SegmentTest, AddColumnWithExpressionMultiBlock) { options_.max_buffer_size_ = 1 * 1024; int doc_count = 100; auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, doc_count); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, doc_count); ASSERT_TRUE(segment != nullptr); segment->dump(); @@ -1807,8 +1806,8 @@ TEST_P(SegmentTest, AlterColumnMultiBlock) { options_.max_buffer_size_ = 1 * 1024; int doc_count = 100; auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, doc_count); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, doc_count); ASSERT_TRUE(segment != nullptr); segment->dump(); @@ -1891,8 +1890,8 @@ TEST_P(SegmentTest, DropColumnMultiBlock) { options_.max_buffer_size_ = 1 * 1024; int doc_count = 100; auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, doc_count); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, doc_count); ASSERT_TRUE(segment != nullptr); segment->dump(); @@ -1978,8 +1977,8 @@ TEST_P(SegmentTest, AddNullableThenAlterDropMultiBlock) { options_.max_buffer_size_ = 1 * 1024; int doc_count = 100; auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, doc_count); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, doc_count); ASSERT_TRUE(segment != nullptr); segment->dump(); @@ -2087,8 +2086,8 @@ TEST_P(SegmentTest, AlterColumn) { // create segment int doc_count = 1000; auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, doc_count); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, doc_count); ASSERT_TRUE(segment != nullptr); auto s = segment->alter_column( @@ -2203,8 +2202,8 @@ TEST_P(SegmentTest, DropColumn) { // create segment int doc_count = 1000; auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, doc_count); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, doc_count); ASSERT_TRUE(segment != nullptr); auto s = segment->drop_column("int32"); diff --git a/tests/ios_test_sandbox.cc b/tests/ios_test_sandbox.cc index ca8bed2c9..7356438f5 100644 --- a/tests/ios_test_sandbox.cc +++ b/tests/ios_test_sandbox.cc @@ -7,39 +7,38 @@ #include #if TARGET_OS_IOS || TARGET_OS_SIMULATOR +#include #include -#include #include -#include +#include -__attribute__((constructor)) -static void ios_test_sandbox_setup() { - // 1. Raise file descriptor limit (iOS default is very low) - struct rlimit rl; - if (getrlimit(RLIMIT_NOFILE, &rl) == 0) { - rlim_t target = 65536; - if (rl.rlim_cur < target) { - rl.rlim_cur = (rl.rlim_max >= target) ? target : rl.rlim_max; - if (setrlimit(RLIMIT_NOFILE, &rl) == 0) { - fprintf(stderr, "[iOS] File descriptor limit raised to: %llu\n", - (unsigned long long)rl.rlim_cur); - } - } +__attribute__((constructor)) static void ios_test_sandbox_setup() { + // 1. Raise file descriptor limit (iOS default is very low) + struct rlimit rl; + if (getrlimit(RLIMIT_NOFILE, &rl) == 0) { + rlim_t target = 65536; + if (rl.rlim_cur < target) { + rl.rlim_cur = (rl.rlim_max >= target) ? target : rl.rlim_max; + if (setrlimit(RLIMIT_NOFILE, &rl) == 0) { + fprintf(stderr, "[iOS] File descriptor limit raised to: %llu\n", + (unsigned long long)rl.rlim_cur); + } } + } - // 2. Change CWD to writable sandbox directory - // TMPDIR is set by iOS to the app's writable sandbox tmp directory - const char* tmpdir = getenv("TMPDIR"); - if (tmpdir && chdir(tmpdir) == 0) { - fprintf(stderr, "[iOS] Working directory set to: %s\n", tmpdir); - } else { - // Fallback: try HOME directory - const char* home = getenv("HOME"); - if (home && chdir(home) == 0) { - fprintf(stderr, "[iOS] Working directory set to: %s\n", home); - } + // 2. Change CWD to writable sandbox directory + // TMPDIR is set by iOS to the app's writable sandbox tmp directory + const char *tmpdir = getenv("TMPDIR"); + if (tmpdir && chdir(tmpdir) == 0) { + fprintf(stderr, "[iOS] Working directory set to: %s\n", tmpdir); + } else { + // Fallback: try HOME directory + const char *home = getenv("HOME"); + if (home && chdir(home) == 0) { + fprintf(stderr, "[iOS] Working directory set to: %s\n", home); } + } } -#endif // TARGET_OS_IOS || TARGET_OS_SIMULATOR -#endif // __APPLE__ +#endif // TARGET_OS_IOS || TARGET_OS_SIMULATOR +#endif // __APPLE__ diff --git a/tests/turbo/CMakeLists.txt b/tests/turbo/CMakeLists.txt new file mode 100644 index 000000000..8a3527d41 --- /dev/null +++ b/tests/turbo/CMakeLists.txt @@ -0,0 +1,14 @@ +include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) + +file(GLOB_RECURSE ALL_TEST_SRCS *_test.cc) + +foreach(CC_SRCS ${ALL_TEST_SRCS}) + get_filename_component(CC_TARGET ${CC_SRCS} NAME_WE) + cc_gtest( + NAME ${CC_TARGET} + STRICT + LIBS zvec_ailego core_framework core_metric core_quantizer zvec_turbo + SRCS ${CC_SRCS} + INCS . ${PROJECT_ROOT_DIR}/src/core/ ${PROJECT_ROOT_DIR}/src/turbo/ + ) +endforeach() \ No newline at end of file diff --git a/tests/turbo/turbo_fht_rotator_test.cc b/tests/turbo/turbo_fht_rotator_test.cc new file mode 100644 index 000000000..4d5ca82a6 --- /dev/null +++ b/tests/turbo/turbo_fht_rotator_test.cc @@ -0,0 +1,447 @@ +// 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 +#include +#include +#include +#include +#include "quantizer/preprocessor/fht_rotator.h" + +using namespace zvec::turbo; + +namespace { + +// Helper: fill a vector with random floats. +void fill_random(float *data, size_t dim, std::mt19937 &gen) { + std::uniform_real_distribution dist(-1.0f, 1.0f); + for (size_t i = 0; i < dim; ++i) data[i] = dist(gen); +} + +// Helper: check round-trip (apply_inverse(apply(x)) == x) within tolerance. +void check_round_trip(const FhtRotator &rot, const std::vector &input, + float tol = 1e-3f) { + const int dim = rot.in_dim(); + std::vector rotated(dim); + std::vector recovered(dim); + + rot.apply(input.data(), rotated.data()); + rot.apply_inverse(rotated.data(), recovered.data()); + + for (int i = 0; i < dim; ++i) { + EXPECT_NEAR(input[i], recovered[i], tol) + << "mismatch at i=" << i << " dim=" << dim; + } +} + +} // anonymous namespace + +// --------------------------------------------------------------------------- +// Power-of-2 dimensions +// --------------------------------------------------------------------------- + +TEST(FhtRotator, PowerOf2RoundTrip) { + std::mt19937 gen(42); + for (int dim : {1, 2, 4, 8, 16, 32, 64, 128, 256}) { + auto rot = FhtRotator::create(dim); + ASSERT_TRUE(rot) << "create failed for dim=" << dim; + + std::vector input(dim); + fill_random(input.data(), dim, gen); + check_round_trip(*rot, input); + } +} + +// --------------------------------------------------------------------------- +// Non-power-of-2 dimensions +// --------------------------------------------------------------------------- + +TEST(FhtRotator, NonPowerOf2RoundTrip) { + std::mt19937 gen(123); + for (int dim : {3, 5, 7, 10, 13, 31, 50, 97, 100, 127, 192, 320}) { + auto rot = FhtRotator::create(dim); + ASSERT_TRUE(rot) << "create failed for dim=" << dim; + + std::vector input(dim); + fill_random(input.data(), dim, gen); + check_round_trip(*rot, input); + } +} + +// --------------------------------------------------------------------------- +// Serialize / Deserialize round-trip +// --------------------------------------------------------------------------- + +TEST(FhtRotator, SerializeDeserialize) { + std::mt19937 gen(999); + for (int dim : {32, 97, 128}) { + // Build original rotator. + auto rot = FhtRotator::create(dim); + ASSERT_TRUE(rot); + + // Serialize. + std::string blob; + ASSERT_EQ(0, rot->serialize(&blob)); + ASSERT_GT(blob.size(), sizeof(RotatorSerHeader)); + + // Restore from blob. + auto rot2 = FhtRotator::from_blob(blob.data(), blob.size()); + ASSERT_TRUE(rot2) << "from_blob failed for dim=" << dim; + + // Dimensions must match. + EXPECT_EQ(rot2->in_dim(), dim); + EXPECT_EQ(rot2->out_dim(), dim); + + // Round-trip via the restored rotator must produce the same result + // as the original (same flip signs). + std::vector input(dim); + fill_random(input.data(), dim, gen); + + std::vector r1(dim), r2(dim); + rot->apply(input.data(), r1.data()); + rot2->apply(input.data(), r2.data()); + for (int i = 0; i < dim; ++i) { + EXPECT_FLOAT_EQ(r1[i], r2[i]) << "apply mismatch at i=" << i; + } + + // Inverse via restored rotator must recover the input. + check_round_trip(*rot2, input); + } +} + +// --------------------------------------------------------------------------- +// Dimension preserved +// --------------------------------------------------------------------------- + +TEST(FhtRotator, DimensionPreserved) { + for (int dim : {1, 7, 64, 97, 128}) { + auto rot = FhtRotator::create(dim); + ASSERT_TRUE(rot); + EXPECT_EQ(rot->in_dim(), dim); + EXPECT_EQ(rot->out_dim(), dim); + } +} + +// --------------------------------------------------------------------------- +// Train generates non-zero flip signs +// --------------------------------------------------------------------------- + +TEST(FhtRotator, CreateGeneratesFlip) { + for (int dim : {8, 64, 97}) { + auto rot = FhtRotator::create(dim); + ASSERT_TRUE(rot); + + // After create, flip is already populated, so serialize must succeed. + + std::string blob; + EXPECT_EQ(0, rot->serialize(&blob)) + << "serialize failed after create for dim=" << dim; + + // Verify the payload is not all zeros (extremely unlikely for random bits). + const auto *hdr = reinterpret_cast(blob.data()); + EXPECT_EQ(hdr->magic, kRotatorMagic); + EXPECT_EQ(hdr->version, kRotatorSerVersion); + EXPECT_EQ(static_cast(hdr->rotator_type), RotatorType::kFht); + EXPECT_EQ(static_cast(hdr->in_dim), dim); + EXPECT_EQ(static_cast(hdr->out_dim), dim); + EXPECT_GT(hdr->payload_size, 0u); + + // Check that the flip payload is not all-zero. + const uint8_t *payload = reinterpret_cast(blob.data()) + + sizeof(RotatorSerHeader); + bool any_nonzero = false; + for (uint32_t i = 0; i < hdr->payload_size; ++i) { + if (payload[i] != 0) { + any_nonzero = true; + break; + } + } + EXPECT_TRUE(any_nonzero) << "flip payload is all-zero for dim=" << dim; + } +} + +// --------------------------------------------------------------------------- +// Create with invalid dimension returns nullptr +// --------------------------------------------------------------------------- + +TEST(FhtRotator, InvalidDimension) { + EXPECT_EQ(FhtRotator::create(0), nullptr); + EXPECT_EQ(FhtRotator::create(-1), nullptr); +} + +// --------------------------------------------------------------------------- +// from_blob with malformed input returns nullptr +// --------------------------------------------------------------------------- + +TEST(FhtRotator, FromBlobMalformed) { + EXPECT_EQ(FhtRotator::from_blob(nullptr, 0), nullptr); + + // Too short. + char buf[4] = {}; + EXPECT_EQ(FhtRotator::from_blob(buf, sizeof(buf)), nullptr); + + // Wrong magic. + RotatorSerHeader hdr{}; + hdr.magic = 0xDEADBEEF; + hdr.version = kRotatorSerVersion; + hdr.rotator_type = static_cast(RotatorType::kFht); + hdr.payload_size = 0; + EXPECT_EQ(FhtRotator::from_blob(&hdr, sizeof(hdr)), nullptr); +} + +// --------------------------------------------------------------------------- +// L2 distance preserved (orthogonal transform) +// --------------------------------------------------------------------------- + +TEST(FhtRotator, L2DistancePreserved) { + std::mt19937 gen(2024); + std::uniform_real_distribution dist(-1.0f, 1.0f); + + for (int dim : {32, 64, 97, 128}) { + auto rot = FhtRotator::create(dim); + ASSERT_TRUE(rot); + + const int N = 50; + std::vector> raw(N, std::vector(dim)); + std::vector> rotated(N, std::vector(dim)); + for (int i = 0; i < N; ++i) { + fill_random(raw[i].data(), dim, gen); + rot->apply(raw[i].data(), rotated[i].data()); + } + + // Check that ||rotated[i] - rotated[j]|| ≈ ||raw[i] - raw[j]||. + for (int i = 1; i < N; ++i) { + float d_raw = 0.0f, d_rot = 0.0f; + for (int j = 0; j < dim; ++j) { + float dr = raw[i][j] - raw[0][j]; + float dt = rotated[i][j] - rotated[0][j]; + d_raw += dr * dr; + d_rot += dt * dt; + } + EXPECT_NEAR(d_raw, d_rot, 1e-2f) + << "L2 mismatch for dim=" << dim << " i=" << i; + } + } +} + +// --------------------------------------------------------------------------- +// Cosine distance preserved (orthogonal transform) +// --------------------------------------------------------------------------- + +TEST(FhtRotator, CosineDistancePreserved) { + std::mt19937 gen(777); + std::uniform_real_distribution dist(0.1f, 1.0f); + + auto cosine_dist = [](const float *a, const float *b, int dim) { + float dot = 0, na = 0, nb = 0; + for (int i = 0; i < dim; ++i) { + dot += a[i] * b[i]; + na += a[i] * a[i]; + nb += b[i] * b[i]; + } + float denom = std::sqrt(na) * std::sqrt(nb); + return (denom < 1e-12f) ? 1.0f : 1.0f - dot / denom; + }; + + for (int dim : {32, 97, 128}) { + auto rot = FhtRotator::create(dim); + ASSERT_TRUE(rot); + + const int N = 50; + std::vector> raw(N, std::vector(dim)); + std::vector> rotated(N, std::vector(dim)); + for (int i = 0; i < N; ++i) { + fill_random(raw[i].data(), dim, gen); + rot->apply(raw[i].data(), rotated[i].data()); + } + + for (int i = 1; i < N; ++i) { + float d_raw = cosine_dist(raw[i].data(), raw[0].data(), dim); + float d_rot = cosine_dist(rotated[i].data(), rotated[0].data(), dim); + EXPECT_NEAR(d_raw, d_rot, 1e-3f) + << "Cosine mismatch for dim=" << dim << " i=" << i; + } + } +} + +// --------------------------------------------------------------------------- +// Apply is non-trivial (not identity) +// --------------------------------------------------------------------------- + +TEST(FhtRotator, ApplyIsNonTrivial) { + std::mt19937 gen(42); + for (int dim : {32, 97, 128}) { + auto rot = FhtRotator::create(dim); + ASSERT_TRUE(rot); + + std::vector input(dim); + fill_random(input.data(), dim, gen); + + std::vector output(dim); + rot->apply(input.data(), output.data()); + + // At least some elements should differ from the input. + bool any_diff = false; + for (int i = 0; i < dim; ++i) { + if (std::abs(input[i] - output[i]) > 1e-6f) { + any_diff = true; + break; + } + } + EXPECT_TRUE(any_diff) << "apply is identity for dim=" << dim; + } +} + +// --------------------------------------------------------------------------- +// Apply is deterministic (same input → same output) +// --------------------------------------------------------------------------- + +TEST(FhtRotator, ApplyDeterministic) { + std::mt19937 gen(55); + for (int dim : {32, 97, 128}) { + auto rot = FhtRotator::create(dim); + ASSERT_TRUE(rot); + + std::vector input(dim); + fill_random(input.data(), dim, gen); + + std::vector r1(dim), r2(dim); + rot->apply(input.data(), r1.data()); + rot->apply(input.data(), r2.data()); + + for (int i = 0; i < dim; ++i) { + EXPECT_FLOAT_EQ(r1[i], r2[i]) + << "non-deterministic apply at i=" << i << " dim=" << dim; + } + } +} + +// --------------------------------------------------------------------------- +// Deserialize on existing object (init → serialize → deserialize on new object) +// --------------------------------------------------------------------------- + +TEST(FhtRotator, DeserializeOnExistingObject) { + std::mt19937 gen(314); + for (int dim : {32, 97, 128}) { + // Build original. + auto rot1 = FhtRotator::create(dim); + ASSERT_TRUE(rot1); + + std::string blob; + ASSERT_EQ(0, rot1->serialize(&blob)); + + // Create a fresh rotator, then call deserialize() on it. + auto rot2 = FhtRotator::create(dim); + ASSERT_TRUE(rot2); + ASSERT_EQ(0, rot2->deserialize(blob.data(), blob.size())); + + EXPECT_EQ(rot2->in_dim(), dim); + EXPECT_EQ(rot2->out_dim(), dim); + + // Both rotators should produce identical results. + std::vector input(dim); + fill_random(input.data(), dim, gen); + + std::vector r1(dim), r2(dim); + rot1->apply(input.data(), r1.data()); + rot2->apply(input.data(), r2.data()); + for (int i = 0; i < dim; ++i) { + EXPECT_FLOAT_EQ(r1[i], r2[i]) + << "apply mismatch at i=" << i << " dim=" << dim; + } + + // Inverse via rot2 should recover input. + check_round_trip(*rot2, input); + } +} + +// --------------------------------------------------------------------------- +// Deserialize with truncated payload fails +// --------------------------------------------------------------------------- + +TEST(FhtRotator, DeserializeTruncatedPayload) { + std::mt19937 gen(42); + auto rot = FhtRotator::create(64); + ASSERT_TRUE(rot); + + std::string blob; + ASSERT_EQ(0, rot->serialize(&blob)); + + // Truncate the blob: keep header but cut half the payload. + auto *hdr = reinterpret_cast(blob.data()); + size_t truncated_len = sizeof(RotatorSerHeader) + hdr->payload_size / 2; + + auto rot2 = FhtRotator::create(64); + ASSERT_TRUE(rot2); + EXPECT_NE(0, rot2->deserialize(blob.data(), truncated_len)); + + // Also test from_blob with truncated data. + EXPECT_EQ(FhtRotator::from_blob(blob.data(), truncated_len), nullptr); +} + +// --------------------------------------------------------------------------- +// Large dimension stress test +// --------------------------------------------------------------------------- + +TEST(FhtRotator, LargeDimension) { + std::mt19937 gen(2025); + for (int dim : {1024, 2048, 4096}) { + auto rot = FhtRotator::create(dim); + ASSERT_TRUE(rot) << "create failed for dim=" << dim; + + std::vector input(dim); + fill_random(input.data(), dim, gen); + check_round_trip(*rot, input, 1e-2f); + + // Verify serialize/deserialize round-trip. + std::string blob; + ASSERT_EQ(0, rot->serialize(&blob)); + auto rot2 = FhtRotator::from_blob(blob.data(), blob.size()); + ASSERT_TRUE(rot2); + + std::vector r1(dim), r2(dim); + rot->apply(input.data(), r1.data()); + rot2->apply(input.data(), r2.data()); + for (int i = 0; i < dim; ++i) { + EXPECT_FLOAT_EQ(r1[i], r2[i]) << "mismatch at i=" << i << " dim=" << dim; + } + } +} + +// --------------------------------------------------------------------------- +// Norm preserved (orthogonal transform preserves vector norm) +// --------------------------------------------------------------------------- + +TEST(FhtRotator, NormPreserved) { + std::mt19937 gen(99); + for (int dim : {32, 97, 128}) { + auto rot = FhtRotator::create(dim); + ASSERT_TRUE(rot); + + std::vector input(dim); + fill_random(input.data(), dim, gen); + + float norm_in = 0.0f; + for (int i = 0; i < dim; ++i) norm_in += input[i] * input[i]; + + std::vector output(dim); + rot->apply(input.data(), output.data()); + + float norm_out = 0.0f; + for (int i = 0; i < dim; ++i) norm_out += output[i] * output[i]; + + EXPECT_NEAR(norm_in, norm_out, 1e-2f) + << "norm not preserved for dim=" << dim; + } +} diff --git a/tests/turbo/turbo_fp32_quantizer_test.cc b/tests/turbo/turbo_fp32_quantizer_test.cc new file mode 100644 index 000000000..c180de4a4 --- /dev/null +++ b/tests/turbo/turbo_fp32_quantizer_test.cc @@ -0,0 +1,202 @@ +// 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 +#include +#include +#include +#include +#include +#include "zvec/core/framework/index_factory.h" + +using namespace zvec; +using namespace zvec::core; +using namespace zvec::ailego; + +// Helper: reference cosine distance between two raw fp32 vectors. +static float reference_cosine(const float *a, const float *b, size_t dim) { + float dot = 0.0f, na = 0.0f, nb = 0.0f; + for (size_t i = 0; i < dim; ++i) { + dot += a[i] * b[i]; + na += a[i] * a[i]; + nb += b[i] * b[i]; + } + float denom = std::sqrt(na) * std::sqrt(nb); + return (denom < 1e-12f) ? 1.0f : 1.0f - dot / denom; +} + +TEST(Fp32Quantizer, General) { + std::mt19937 gen(15583); + std::uniform_real_distribution dist(0.0, 1.0); + + const size_t COUNT = 10000; + const size_t DIMENSION = 12; + + IndexMeta meta; + meta.set_meta(IndexMeta::DataType::DT_FP32, DIMENSION); + meta.set_metric("Cosine", 0, Params()); + + auto quantizer = IndexFactory::CreateQuantizer("Fp32Quantizer"); + ASSERT_TRUE(quantizer); + zvec::ailego::Params params; + ASSERT_EQ(0u, quantizer->init(meta, params)); + + auto holder = + std::make_shared>( + DIMENSION); + for (size_t i = 0; i < COUNT; ++i) { + zvec::ailego::NumericalVector vec(DIMENSION); + for (size_t j = 0; j < DIMENSION; ++j) { + vec[j] = dist(gen); + } + holder->emplace(i + 1, vec); + } + EXPECT_EQ(COUNT, holder->count()); + EXPECT_EQ(IndexMeta::DataType::DT_FP32, holder->data_type()); + + ASSERT_EQ(0u, quantizer->train(holder)); + + auto iter = holder->create_iterator(); + std::string quant_buffer; + std::string dequant_buffer; + + for (; iter->is_valid(); iter->next()) { + EXPECT_TRUE(iter->data()); + + IndexQueryMeta qmeta; + quant_buffer.clear(); + EXPECT_EQ(0, quantizer->quantize( + iter->data(), + IndexQueryMeta(holder->data_type(), holder->dimension()), + &quant_buffer, &qmeta)); + EXPECT_EQ(IndexMeta::DataType::DT_FP32, qmeta.data_type()); + EXPECT_EQ(holder->dimension(), qmeta.dimension()); + + dequant_buffer.clear(); + EXPECT_EQ( + 0, quantizer->dequantize(quant_buffer.data(), qmeta, &dequant_buffer)); + + const float *original_data = reinterpret_cast(iter->data()); + const float *dequantize_data = + reinterpret_cast(dequant_buffer.data()); + for (size_t i = 0; i < holder->dimension(); ++i) { + EXPECT_NEAR(original_data[i], dequantize_data[i], 1e-3); + } + } +} + +TEST(Fp32Quantizer, Score) { + std::mt19937 gen(42); + std::uniform_real_distribution dist(0.0, 1.0); + + const size_t DIMENSION = 12; + const size_t COUNT = 100; + + IndexMeta meta; + meta.set_meta(IndexMeta::DataType::DT_FP32, DIMENSION); + meta.set_metric("Cosine", 0, Params()); + + auto quantizer = IndexFactory::CreateQuantizer("Fp32Quantizer"); + ASSERT_TRUE(quantizer); + zvec::ailego::Params params; + ASSERT_EQ(0u, quantizer->init(meta, params)); + + // Generate raw vectors and quantize them. + std::vector> raw_vecs(COUNT); + std::vector quant_vecs(COUNT); + for (size_t i = 0; i < COUNT; ++i) { + raw_vecs[i].resize(DIMENSION); + for (size_t j = 0; j < DIMENSION; ++j) { + raw_vecs[i][j] = dist(gen); + } + IndexQueryMeta ometa; + EXPECT_EQ(0, quantizer->quantize( + raw_vecs[i].data(), + IndexQueryMeta(IndexMeta::DataType::DT_FP32, DIMENSION), + &quant_vecs[i], &ometa)); + } + + // --- calc_distance_dp_query (single) --- + for (size_t i = 1; i < COUNT; ++i) { + float d = quantizer->calc_distance_dp_query(quant_vecs[i].data(), + quant_vecs[0].data()); + float expected = + reference_cosine(raw_vecs[i].data(), raw_vecs[0].data(), DIMENSION); + EXPECT_NEAR(d, expected, 1e-4) << "i=" << i; + } + + // --- calc_distance_dp_query_batch --- + { + std::vector dp_list(COUNT - 1); + for (size_t i = 1; i < COUNT; ++i) { + dp_list[i - 1] = quant_vecs[i].data(); + } + std::vector results(COUNT - 1); + quantizer->calc_distance_dp_query_batch( + dp_list.data(), static_cast(dp_list.size()), quant_vecs[0].data(), + results.data()); + + for (size_t i = 0; i < dp_list.size(); ++i) { + float expected = reference_cosine(raw_vecs[i + 1].data(), + raw_vecs[0].data(), DIMENSION); + EXPECT_NEAR(results[i], expected, 1e-4) << "i=" << i; + } + } + + // --- distance() + DistanceImpl (single + batch) --- + { + IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, DIMENSION); + auto dist_impl = quantizer->distance(quant_vecs[0].data(), qmeta); + ASSERT_TRUE(dist_impl.valid()); + + for (size_t i = 1; i < COUNT; ++i) { + float d = dist_impl(quant_vecs[i].data()); + float expected = + reference_cosine(raw_vecs[0].data(), raw_vecs[i].data(), DIMENSION); + EXPECT_NEAR(d, expected, 1e-4) << "i=" << i; + } + + // Batch via DistanceImpl. + ASSERT_TRUE(dist_impl.batch_valid()); + std::vector dp_list(COUNT - 1); + for (size_t i = 1; i < COUNT; ++i) { + dp_list[i - 1] = quant_vecs[i].data(); + } + std::vector batch_results(COUNT - 1); + dist_impl.batch(dp_list.data(), dp_list.size(), batch_results.data()); + for (size_t i = 0; i < dp_list.size(); ++i) { + float expected = reference_cosine(raw_vecs[0].data(), + raw_vecs[i + 1].data(), DIMENSION); + EXPECT_NEAR(batch_results[i], expected, 1e-4) << "i=" << i; + } + } + + // --- calc_distance_dp_dp (pairwise) --- + for (size_t i = 1; i < 10; ++i) { + float d = quantizer->calc_distance_dp_dp(quant_vecs[i].data(), + quant_vecs[0].data()); + float expected = + reference_cosine(raw_vecs[i].data(), raw_vecs[0].data(), DIMENSION); + EXPECT_NEAR(d, expected, 1e-4) << "i=" << i; + } + + // --- calc_distance_dp_query_unquantized --- + for (size_t i = 1; i < 10; ++i) { + float d = quantizer->calc_distance_dp_query_unquantized( + quant_vecs[i].data(), raw_vecs[0].data()); + float expected = + reference_cosine(raw_vecs[i].data(), raw_vecs[0].data(), DIMENSION); + EXPECT_NEAR(d, expected, 1e-4) << "i=" << i; + } +} \ No newline at end of file diff --git a/tests/turbo/turbo_pq_int4_quantizer_test.cc b/tests/turbo/turbo_pq_int4_quantizer_test.cc new file mode 100644 index 000000000..86249f156 --- /dev/null +++ b/tests/turbo/turbo_pq_int4_quantizer_test.cc @@ -0,0 +1,587 @@ +// 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 +#include +#include +#include +#include +#include +#include +#include +#include "distance/scalar/pq_quantizer_int4/pq_distance.h" +#include "quantizer/pq_int4_quantizer/pq_int4_quantizer.h" +#include "zvec/core/framework/index_factory.h" + +#if defined(__AVX2__) +#include "distance/avx2/pq_quantizer_int4/pq_distance.h" +#endif +#if defined(__AVX512F__) +#include "distance/avx512/pq_quantizer_int4/pq_distance.h" +#endif + +using namespace zvec; +using namespace zvec::core; +using namespace zvec::ailego; + +// Reference squared Euclidean distance between two raw fp32 vectors. +static float reference_sq_euclidean(const float *a, const float *b, + size_t dim) { + float sum = 0.0f; + for (size_t i = 0; i < dim; ++i) { + float diff = a[i] - b[i]; + sum += diff * diff; + } + return sum; +} + +// Helper to create a PqInt4Quantizer via the factory. +static std::shared_ptr make_pq4_quantizer( + size_t dim, size_t num_subquantizers) { + auto q = IndexFactory::CreateQuantizer("PqInt4Quantizer"); + if (!q) return nullptr; + + IndexMeta meta; + meta.set_meta(IndexMeta::DataType::DT_FP32, dim); + meta.set_metric("SquaredEuclidean", 0, Params()); + + Params params; + params.set("num_subquantizers", static_cast(num_subquantizers)); + if (q->init(meta, params) != 0) return nullptr; + return q; +} + +// Helper: build a holder with random fp32 vectors. +static std::shared_ptr> +make_random_holder(size_t count, size_t dim, uint32_t seed = 42) { + auto holder = + std::make_shared>(dim); + std::mt19937 gen(seed); + std::uniform_real_distribution dist(-1.0f, 1.0f); + for (size_t i = 0; i < count; ++i) { + NumericalVector vec(dim); + for (size_t j = 0; j < dim; ++j) vec[j] = dist(gen); + holder->emplace(i + 1, vec); + } + return holder; +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +TEST(PqInt4Quantizer, InitInvalidParams) { + // dim not divisible by num_subquantizers + auto q = make_pq4_quantizer(10, 4); + EXPECT_EQ(q, nullptr); + + // num_subquantizers = 0 + auto q2 = IndexFactory::CreateQuantizer("PqInt4Quantizer"); + ASSERT_TRUE(q2); + IndexMeta meta; + meta.set_meta(IndexMeta::DataType::DT_FP32, 16); + meta.set_metric("SquaredEuclidean", 0, Params()); + Params params; + params.set("num_subquantizers", static_cast(0)); + EXPECT_NE(0, q2->init(meta, params)); + + // Odd num_subquantizers: must be rejected for int4 packed nibble. + auto q3 = IndexFactory::CreateQuantizer("PqInt4Quantizer"); + ASSERT_TRUE(q3); + IndexMeta meta3; + meta3.set_meta(IndexMeta::DataType::DT_FP32, 12); + meta3.set_metric("SquaredEuclidean", 0, Params()); + Params params3; + params3.set("num_subquantizers", static_cast(3)); + EXPECT_NE(0, q3->init(meta3, params3)); +} + +TEST(PqInt4Quantizer, TrainAndEncode) { + const size_t DIM = 16; + const size_t NSQ = 4; // must be even + const size_t COUNT = 1000; + + auto quantizer = make_pq4_quantizer(DIM, NSQ); + ASSERT_TRUE(quantizer); + EXPECT_TRUE(quantizer->require_train()); + + // Packed nibble: code length = NSQ/2 bytes (no extra meta for L2). + EXPECT_EQ(quantizer->quantized_datapoint_vector_length(), NSQ / 2); + + auto holder = make_random_holder(COUNT, DIM); + ASSERT_EQ(0, quantizer->train(holder)); + + // Quantize a few vectors and check code length + nibble range. + auto iter = holder->create_iterator(); + size_t checked = 0; + std::vector code(quantizer->quantized_datapoint_vector_length()); + for (; iter->is_valid() && checked < 10; iter->next(), ++checked) { + quantizer->quantize_data(iter->data(), code.data()); + // Each nibble (low and high) should be in [0, 15]. + for (size_t b = 0; b < NSQ / 2; ++b) { + uint8_t lo = code[b] & 0x0Fu; + uint8_t hi = (code[b] >> 4) & 0x0Fu; + EXPECT_LE(lo, 15u); + EXPECT_LE(hi, 15u); + } + } + EXPECT_EQ(10u, checked); +} + +TEST(PqInt4Quantizer, AdcDistance) { + const size_t DIM = 32; + const size_t NSQ = 8; // even + const size_t COUNT = 2000; + + auto quantizer = make_pq4_quantizer(DIM, NSQ); + ASSERT_TRUE(quantizer); + + auto holder = make_random_holder(COUNT, DIM); + ASSERT_EQ(0, quantizer->train(holder)); + + std::vector> raw_vecs(COUNT); + std::vector> pq_codes(COUNT); + size_t code_len = quantizer->quantized_datapoint_vector_length(); + size_t lut_len = quantizer->quantized_query_vector_length(); + + auto iter = holder->create_iterator(); + for (size_t i = 0; iter->is_valid(); iter->next(), ++i) { + const float *v = reinterpret_cast(iter->data()); + raw_vecs[i].assign(v, v + DIM); + pq_codes[i].resize(code_len); + quantizer->quantize_data(iter->data(), pq_codes[i].data()); + } + + // Build LUT for query = raw_vecs[0]. + std::vector lut(lut_len / sizeof(float)); + quantizer->quantize_query(raw_vecs[0].data(), lut.data()); + + float max_rel_error = 0.0f; + for (size_t i = 1; i < COUNT; ++i) { + float adc_dist = + quantizer->calc_distance_dp_query(pq_codes[i].data(), lut.data()); + float true_dist = + reference_sq_euclidean(raw_vecs[i].data(), raw_vecs[0].data(), DIM); + if (true_dist > 1e-6f) { + float rel = std::fabs(adc_dist - true_dist) / true_dist; + max_rel_error = std::max(max_rel_error, rel); + } + EXPECT_GE(adc_dist, 0.0f) << "i=" << i; + } + // int4 PQ has higher quantization error than int8 (only 16 centroids). + // Generous bound: relative error < 2.0 (200%). + EXPECT_LT(max_rel_error, 2.0f) << "max_rel_error=" << max_rel_error; +} + +TEST(PqInt4Quantizer, SdcDistance) { + const size_t DIM = 16; + const size_t NSQ = 4; + const size_t COUNT = 2000; + + auto quantizer = make_pq4_quantizer(DIM, NSQ); + ASSERT_TRUE(quantizer); + + auto holder = make_random_holder(COUNT, DIM); + ASSERT_EQ(0, quantizer->train(holder)); + + auto iter = holder->create_iterator(); + std::vector code1(quantizer->quantized_datapoint_vector_length()); + std::vector code2(quantizer->quantized_datapoint_vector_length()); + + iter->is_valid(); + quantizer->quantize_data(iter->data(), code1.data()); + iter->next(); + iter->is_valid(); + quantizer->quantize_data(iter->data(), code2.data()); + + float sdc_dist = quantizer->calc_distance_dp_dp(code1.data(), code2.data()); + EXPECT_GE(sdc_dist, 0.0f); +} + +TEST(PqInt4Quantizer, DistanceImplAdcAndSdc) { + const size_t DIM = 16; + const size_t NSQ = 4; + const size_t COUNT = 1000; + + auto quantizer = make_pq4_quantizer(DIM, NSQ); + ASSERT_TRUE(quantizer); + + auto holder = make_random_holder(COUNT, DIM); + ASSERT_EQ(0, quantizer->train(holder)); + + auto iter = holder->create_iterator(); + iter->is_valid(); + const float *query_raw = reinterpret_cast(iter->data()); + + size_t lut_bytes = quantizer->quantized_query_vector_length(); + std::string lut_storage(lut_bytes, '\0'); + quantizer->quantize_query(query_raw, &lut_storage[0]); + + IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, DIM); + auto dist_impl = quantizer->distance(lut_storage.data(), qmeta); + ASSERT_TRUE(dist_impl.valid()); + EXPECT_TRUE(static_cast(dist_impl.func())); + EXPECT_TRUE(static_cast(dist_impl.sym_func())); + + // Encode a candidate and compute distance via ADC. + iter->next(); + iter->is_valid(); + std::vector code(quantizer->quantized_datapoint_vector_length()); + quantizer->quantize_data(iter->data(), code.data()); + + float d = dist_impl(code.data()); + EXPECT_GE(d, 0.0f); + + // SDC (symmetric) via sym_func(). + std::vector code2(quantizer->quantized_datapoint_vector_length()); + iter->next(); + iter->is_valid(); + quantizer->quantize_data(iter->data(), code2.data()); + + float sdc_d = 0.0f; + dist_impl.sym_func()(code.data(), code2.data(), NSQ, &sdc_d); + EXPECT_GE(sdc_d, 0.0f); +} + +TEST(PqInt4Quantizer, SerializeDeserialize) { + const size_t DIM = 16; + const size_t NSQ = 4; + const size_t COUNT = 500; + + auto quantizer = make_pq4_quantizer(DIM, NSQ); + ASSERT_TRUE(quantizer); + + auto holder = make_random_holder(COUNT, DIM); + ASSERT_EQ(0, quantizer->train(holder)); + + std::string blob; + ASSERT_EQ(0, quantizer->serialize(&blob)); + EXPECT_GT(blob.size(), sizeof(zvec::turbo::QuantizerSerHeader)); + + auto q2 = IndexFactory::CreateQuantizer("PqInt4Quantizer"); + ASSERT_TRUE(q2); + ASSERT_EQ(0, q2->deserialize(blob)); + + auto iter = holder->create_iterator(); + iter->is_valid(); + std::vector code1(quantizer->quantized_datapoint_vector_length()); + std::vector code2(q2->quantized_datapoint_vector_length()); + quantizer->quantize_data(iter->data(), code1.data()); + q2->quantize_data(iter->data(), code2.data()); + + size_t code_bytes = NSQ / 2; + for (size_t b = 0; b < code_bytes; ++b) { + EXPECT_EQ(code1[b], code2[b]) << "b=" << b; + } + + // ADC distances should also match. + size_t lut_len = quantizer->quantized_query_vector_length(); + std::vector lut1(lut_len / sizeof(float)); + std::vector lut2(lut_len / sizeof(float)); + quantizer->quantize_query(iter->data(), lut1.data()); + q2->quantize_query(iter->data(), lut2.data()); + + float adc1 = quantizer->calc_distance_dp_query(code1.data(), lut1.data()); + float adc2 = q2->calc_distance_dp_query(code2.data(), lut2.data()); + EXPECT_NEAR(adc1, adc2, 1e-6f); +} + +// --------------------------------------------------------------------------- +// Scalar int4 kernel direct tests +// --------------------------------------------------------------------------- + +namespace { + +// Decode nibble helper for test verification. +inline uint8_t test_decode_nibble(const uint8_t *code, size_t m) { + uint8_t byte = code[m >> 1]; + return (m & 1u) ? static_cast(byte >> 4) + : static_cast(byte & 0x0Fu); +} + +void fill_random_int4_codes(uint8_t *codes, size_t num_subquantizers, + std::mt19937 &gen) { + std::uniform_int_distribution dist(0, 15); + size_t code_bytes = num_subquantizers / 2; + std::memset(codes, 0, code_bytes); + for (size_t m = 0; m < num_subquantizers; ++m) { + uint8_t idx = static_cast(dist(gen)); + if (m & 1u) { + codes[m >> 1] |= static_cast(idx << 4); + } else { + codes[m >> 1] |= static_cast(idx & 0x0Fu); + } + } +} + +void fill_random_int4_lut(float *lut, size_t num_subquantizers, + std::mt19937 &gen) { + constexpr size_t kNumCentroids = 16; + std::uniform_real_distribution dist(0.0f, 1.0f); + for (size_t m = 0; m < num_subquantizers; ++m) { + for (size_t c = 0; c < kNumCentroids; ++c) { + lut[m * kNumCentroids + c] = dist(gen); + } + } +} + +} // anonymous namespace + +// Test ADC int4 kernel directly. +TEST(PqInt4Kernel, AdcDistance) { + std::mt19937 gen(2024); + + for (size_t nsq : {4, 8, 12, 16}) { + constexpr size_t kNumCentroids = 16; + std::vector codes(nsq / 2); + std::vector lut(nsq * kNumCentroids); + + fill_random_int4_codes(codes.data(), nsq, gen); + fill_random_int4_lut(lut.data(), nsq, gen); + + // Reference: manual sum. + float ref = 0.0f; + for (size_t m = 0; m < nsq; ++m) { + uint8_t idx = test_decode_nibble(codes.data(), m); + ref += lut[m * kNumCentroids + idx]; + } + + float result = 0.0f; + zvec::turbo::scalar::pq_adc_int4_distance(codes.data(), lut.data(), nsq, + &result); + EXPECT_NEAR(ref, result, 1e-5f) << "ADC mismatch for nsq=" << nsq; + } +} + +// Test SDC int4 kernel directly. +TEST(PqInt4Kernel, SdcDistance) { + std::mt19937 gen(2025); + constexpr size_t kNumCentroids = 16; + constexpr size_t kTablePerSub = kNumCentroids * kNumCentroids; // 256 + + for (size_t nsq : {4, 8}) { + std::vector codes_a(nsq / 2); + std::vector codes_b(nsq / 2); + std::vector dist_table(nsq * kTablePerSub); + + fill_random_int4_codes(codes_a.data(), nsq, gen); + fill_random_int4_codes(codes_b.data(), nsq, gen); + + std::uniform_real_distribution dist(0.0f, 1.0f); + for (auto &v : dist_table) v = dist(gen); + + // Reference. + float ref = 0.0f; + for (size_t m = 0; m < nsq; ++m) { + uint8_t ai = test_decode_nibble(codes_a.data(), m); + uint8_t bi = test_decode_nibble(codes_b.data(), m); + size_t idx = m * kTablePerSub + static_cast(ai) * kNumCentroids + + static_cast(bi); + ref += dist_table[idx]; + } + + float result = 0.0f; + zvec::turbo::scalar::pq_sdc_int4_distance(codes_a.data(), codes_b.data(), + dist_table.data(), nsq, &result); + EXPECT_NEAR(ref, result, 1e-5f) << "SDC mismatch for nsq=" << nsq; + } +} + +// Test Batch ADC int4 kernel. +TEST(PqInt4Kernel, BatchAdcDistance) { + std::mt19937 gen(2026); + constexpr size_t kNumCentroids = 16; + const size_t nsq = 8; + const size_t num_candidates = 7; // not multiple of 4 to test leftover + + std::vector> codes(num_candidates, + std::vector(nsq / 2)); + std::vector lut(nsq * kNumCentroids); + + for (auto &c : codes) fill_random_int4_codes(c.data(), nsq, gen); + fill_random_int4_lut(lut.data(), nsq, gen); + + // Build pointer array. + std::vector ptrs(num_candidates); + for (size_t i = 0; i < num_candidates; ++i) { + ptrs[i] = codes[i].data(); + } + + // Reference per-candidate distances. + std::vector ref(num_candidates, 0.0f); + for (size_t i = 0; i < num_candidates; ++i) { + for (size_t m = 0; m < nsq; ++m) { + uint8_t idx = test_decode_nibble(codes[i].data(), m); + ref[i] += lut[m * kNumCentroids + idx]; + } + } + + std::vector result(num_candidates, 0.0f); + zvec::turbo::scalar::pq_adc_int4_batch_distance( + ptrs.data(), lut.data(), num_candidates, nsq, result.data()); + + for (size_t i = 0; i < num_candidates; ++i) { + EXPECT_NEAR(ref[i], result[i], 1e-5f) << "batch ADC mismatch at i=" << i; + } +} + +// --------------------------------------------------------------------------- +// SIMD Consistency Tests +// --------------------------------------------------------------------------- + +// Helper to generate random SDC table for int4 (kTablePerSub = 256) +namespace { +void fill_random_int4_sdc_table(float *table, size_t num_subquantizers, + std::mt19937 &gen) { + constexpr size_t kTablePerSub = 16 * 16; + std::uniform_real_distribution dist(0.0f, 1.0f); + for (size_t m = 0; m < num_subquantizers; ++m) { + for (size_t i = 0; i < kTablePerSub; ++i) { + table[m * kTablePerSub + i] = dist(gen); + } + } +} +} // anonymous namespace + +// Test ADC SIMD consistency across multiple M values (must be even for int4). +TEST(PqInt4SimdConsistency, AdcDistance) { + std::mt19937 gen(2024); + constexpr size_t kNumCentroids = 16; + // M=4: less than AVX2/AVX512 chunk (16) — scalar only + // M=8: less than chunk — scalar only + // M=16: exact chunk boundary + // M=20: chunk + scalar leftover + for (size_t num_sq : {4, 8, 16, 20}) { + std::vector codes(num_sq / 2); + std::vector lut(num_sq * kNumCentroids); + + fill_random_int4_codes(codes.data(), num_sq, gen); + fill_random_int4_lut(lut.data(), num_sq, gen); + + float scalar_result = 0.0f; + zvec::turbo::scalar::pq_adc_int4_distance(codes.data(), lut.data(), num_sq, + &scalar_result); + +#if defined(__AVX2__) + { + float avx2_result = 0.0f; + zvec::turbo::avx2::pq_adc_int4_distance_avx2(codes.data(), lut.data(), + num_sq, &avx2_result); + EXPECT_NEAR(scalar_result, avx2_result, 1e-5f) + << "AVX2 ADC mismatch for M=" << num_sq; + } +#endif + +#if defined(__AVX512F__) + { + float avx512_result = 0.0f; + zvec::turbo::avx512::pq_adc_int4_distance_avx512(codes.data(), lut.data(), + num_sq, &avx512_result); + EXPECT_NEAR(scalar_result, avx512_result, 1e-5f) + << "AVX512 ADC mismatch for M=" << num_sq; + } +#endif + } +} + +// Test SDC SIMD consistency across multiple M values. +TEST(PqInt4SimdConsistency, SdcDistance) { + std::mt19937 gen(2025); + constexpr size_t kTablePerSub = 16 * 16; + + for (size_t num_sq : {4, 8, 16, 20}) { + std::vector codes_a(num_sq / 2); + std::vector codes_b(num_sq / 2); + std::vector dist_table(num_sq * kTablePerSub); + + fill_random_int4_codes(codes_a.data(), num_sq, gen); + fill_random_int4_codes(codes_b.data(), num_sq, gen); + fill_random_int4_sdc_table(dist_table.data(), num_sq, gen); + + float scalar_result = 0.0f; + zvec::turbo::scalar::pq_sdc_int4_distance(codes_a.data(), codes_b.data(), + dist_table.data(), num_sq, + &scalar_result); + +#if defined(__AVX2__) + { + float avx2_result = 0.0f; + zvec::turbo::avx2::pq_sdc_int4_distance_avx2( + codes_a.data(), codes_b.data(), dist_table.data(), num_sq, + &avx2_result); + EXPECT_NEAR(scalar_result, avx2_result, 1e-5f) + << "AVX2 SDC mismatch for M=" << num_sq; + } +#endif + +#if defined(__AVX512F__) + { + float avx512_result = 0.0f; + zvec::turbo::avx512::pq_sdc_int4_distance_avx512( + codes_a.data(), codes_b.data(), dist_table.data(), num_sq, + &avx512_result); + EXPECT_NEAR(scalar_result, avx512_result, 1e-5f) + << "AVX512 SDC mismatch for M=" << num_sq; + } +#endif + } +} + +// Test Batch ADC SIMD consistency. +TEST(PqInt4SimdConsistency, BatchAdcDistance) { + std::mt19937 gen(2026); + constexpr size_t kNumCentroids = 16; + const size_t nsq = 16; // chunk-aligned + const size_t num_candidates = 7; // not multiple of 4 to test leftover + + std::vector> codes(num_candidates, + std::vector(nsq / 2)); + std::vector lut(nsq * kNumCentroids); + + for (auto &c : codes) fill_random_int4_codes(c.data(), nsq, gen); + fill_random_int4_lut(lut.data(), nsq, gen); + + // Reference via scalar. + std::vector ptrs(num_candidates); + std::vector scalar_result(num_candidates, 0.0f); + for (size_t i = 0; i < num_candidates; ++i) { + ptrs[i] = codes[i].data(); + } + zvec::turbo::scalar::pq_adc_int4_batch_distance( + ptrs.data(), lut.data(), num_candidates, nsq, scalar_result.data()); + +#if defined(__AVX2__) + { + std::vector avx2_result(num_candidates, 0.0f); + zvec::turbo::avx2::pq_adc_int4_batch_distance_avx2( + ptrs.data(), lut.data(), num_candidates, nsq, avx2_result.data()); + for (size_t i = 0; i < num_candidates; ++i) { + EXPECT_NEAR(scalar_result[i], avx2_result[i], 1e-5f) + << "AVX2 batch ADC mismatch at i=" << i; + } + } +#endif + +#if defined(__AVX512F__) + { + std::vector avx512_result(num_candidates, 0.0f); + zvec::turbo::avx512::pq_adc_int4_batch_distance_avx512( + ptrs.data(), lut.data(), num_candidates, nsq, avx512_result.data()); + for (size_t i = 0; i < num_candidates; ++i) { + EXPECT_NEAR(scalar_result[i], avx512_result[i], 1e-5f) + << "AVX512 batch ADC mismatch at i=" << i; + } + } +#endif +} diff --git a/tests/turbo/turbo_pq_int8_quantizer_test.cc b/tests/turbo/turbo_pq_int8_quantizer_test.cc new file mode 100644 index 000000000..438b19c96 --- /dev/null +++ b/tests/turbo/turbo_pq_int8_quantizer_test.cc @@ -0,0 +1,817 @@ +// 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 +#include +#include +#include +#include +#include +#include +#include +#include "distance/scalar/pq_quantizer_int8/pq_distance.h" +#include "quantizer/pq_int8_quantizer/pq_int8_quantizer.h" +#include "zvec/core/framework/index_factory.h" + +#if defined(__AVX2__) +#include "distance/avx2/pq_quantizer_int8/pq_distance.h" +#endif +#if defined(__AVX512F__) +#include "distance/avx512/pq_quantizer_int8/pq_distance.h" +#endif + +using namespace zvec; +using namespace zvec::core; +using namespace zvec::ailego; + +// Reference squared Euclidean distance between two raw fp32 vectors. +static float reference_sq_euclidean(const float *a, const float *b, + size_t dim) { + float sum = 0.0f; + for (size_t i = 0; i < dim; ++i) { + float diff = a[i] - b[i]; + sum += diff * diff; + } + return sum; +} + +// Helper to create a PqInt8Quantizer via the factory. +static std::shared_ptr make_pq_quantizer( + size_t dim, size_t num_subquantizers) { + auto q = IndexFactory::CreateQuantizer("PqInt8Quantizer"); + if (!q) return nullptr; + + IndexMeta meta; + meta.set_meta(IndexMeta::DataType::DT_FP32, dim); + meta.set_metric("SquaredEuclidean", 0, Params()); + + Params params; + params.set("num_subquantizers", static_cast(num_subquantizers)); + if (q->init(meta, params) != 0) return nullptr; + return q; +} + +// Helper: build a holder with random fp32 vectors. +static std::shared_ptr> +make_random_holder(size_t count, size_t dim, uint32_t seed = 42) { + auto holder = + std::make_shared>(dim); + std::mt19937 gen(seed); + std::uniform_real_distribution dist(-1.0f, 1.0f); + for (size_t i = 0; i < count; ++i) { + NumericalVector vec(dim); + for (size_t j = 0; j < dim; ++j) vec[j] = dist(gen); + holder->emplace(i + 1, vec); + } + return holder; +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +TEST(PqInt8Quantizer, InitInvalidParams) { + // dim not divisible by num_subquantizers + auto q = make_pq_quantizer(10, 3); + EXPECT_EQ(q, nullptr); + + // num_subquantizers = 0 + auto q2 = IndexFactory::CreateQuantizer("PqInt8Quantizer"); + ASSERT_TRUE(q2); + IndexMeta meta; + meta.set_meta(IndexMeta::DataType::DT_FP32, 16); + meta.set_metric("SquaredEuclidean", 0, Params()); + Params params; + params.set("num_subquantizers", static_cast(0)); + EXPECT_NE(0, q2->init(meta, params)); +} + +TEST(PqInt8Quantizer, TrainAndEncode) { + const size_t DIM = 16; + const size_t NSQ = 4; + const size_t COUNT = 1000; + + auto quantizer = make_pq_quantizer(DIM, NSQ); + ASSERT_TRUE(quantizer); + EXPECT_TRUE(quantizer->require_train()); + + auto holder = make_random_holder(COUNT, DIM); + ASSERT_EQ(0, quantizer->train(holder)); + + // Quantize a few vectors and check code length. + auto iter = holder->create_iterator(); + size_t checked = 0; + std::vector code(quantizer->quantized_datapoint_vector_length()); + for (; iter->is_valid() && checked < 10; iter->next(), ++checked) { + quantizer->quantize_data(iter->data(), code.data()); + // Each code byte should be in [0, 255]. + for (size_t m = 0; m < NSQ; ++m) { + EXPECT_LE(code[m], 255u); + } + } + EXPECT_EQ(10u, checked); +} + +TEST(PqInt8Quantizer, AdcDistance) { + const size_t DIM = 32; + const size_t NSQ = 8; + const size_t COUNT = 2000; + + auto quantizer = make_pq_quantizer(DIM, NSQ); + ASSERT_TRUE(quantizer); + + auto holder = make_random_holder(COUNT, DIM); + ASSERT_EQ(0, quantizer->train(holder)); + + // Collect raw vectors and PQ codes. + std::vector> raw_vecs(COUNT); + std::vector> pq_codes(COUNT); + size_t code_len = quantizer->quantized_datapoint_vector_length(); + size_t lut_len = quantizer->quantized_query_vector_length(); + + auto iter = holder->create_iterator(); + for (size_t i = 0; iter->is_valid(); iter->next(), ++i) { + const float *v = reinterpret_cast(iter->data()); + raw_vecs[i].assign(v, v + DIM); + pq_codes[i].resize(code_len); + quantizer->quantize_data(iter->data(), pq_codes[i].data()); + } + + // Build LUT for query = raw_vecs[0] + std::vector lut(lut_len / sizeof(float)); + quantizer->quantize_query(raw_vecs[0].data(), lut.data()); + + // ADC distances should be a reasonable approximation of true distance. + // With 8 sub-quantizers and 32 dims (sub_dim=4), PQ error is non-trivial + // but should be bounded. + float max_rel_error = 0.0f; + for (size_t i = 1; i < COUNT; ++i) { + float adc_dist = + quantizer->calc_distance_dp_query(pq_codes[i].data(), lut.data()); + float true_dist = + reference_sq_euclidean(raw_vecs[i].data(), raw_vecs[0].data(), DIM); + if (true_dist > 1e-6f) { + float rel = std::fabs(adc_dist - true_dist) / true_dist; + max_rel_error = std::max(max_rel_error, rel); + } + // ADC distance must be non-negative. + EXPECT_GE(adc_dist, 0.0f) << "i=" << i; + } + // With 8 subs and 2000 training points, max relative error should be + // well below 100% (generous bound; actual error is typically <30%). + EXPECT_LT(max_rel_error, 1.0f) << "max_rel_error=" << max_rel_error; +} + +TEST(PqInt8Quantizer, SdcDistance) { + const size_t DIM = 16; + const size_t NSQ = 4; + const size_t COUNT = 2000; + + auto quantizer = make_pq_quantizer(DIM, NSQ); + ASSERT_TRUE(quantizer); + + auto holder = make_random_holder(COUNT, DIM); + ASSERT_EQ(0, quantizer->train(holder)); + + // Encode two vectors and compute SDC distance. + auto iter = holder->create_iterator(); + std::vector code1(quantizer->quantized_datapoint_vector_length()); + std::vector code2(quantizer->quantized_datapoint_vector_length()); + + iter->is_valid(); + quantizer->quantize_data(iter->data(), code1.data()); + iter->next(); + iter->is_valid(); + quantizer->quantize_data(iter->data(), code2.data()); + + float sdc_dist = quantizer->calc_distance_dp_dp(code1.data(), code2.data()); + EXPECT_GE(sdc_dist, 0.0f); +} + +TEST(PqInt8Quantizer, DistanceImplAdcAndSdc) { + const size_t DIM = 16; + const size_t NSQ = 4; + const size_t COUNT = 1000; + + auto quantizer = make_pq_quantizer(DIM, NSQ); + ASSERT_TRUE(quantizer); + + auto holder = make_random_holder(COUNT, DIM); + ASSERT_EQ(0, quantizer->train(holder)); + + // Quantize query[0] as LUT. + auto iter = holder->create_iterator(); + iter->is_valid(); + const float *query_raw = reinterpret_cast(iter->data()); + + size_t lut_bytes = quantizer->quantized_query_vector_length(); + std::string lut_storage(lut_bytes, '\0'); + quantizer->quantize_query(query_raw, &lut_storage[0]); + + IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, DIM); + auto dist_impl = quantizer->distance(lut_storage.data(), qmeta); + ASSERT_TRUE(dist_impl.valid()); + + // func() and sym_func() should both be set. + EXPECT_TRUE(static_cast(dist_impl.func())); + EXPECT_TRUE(static_cast(dist_impl.sym_func())); + + // Encode a candidate and compute distance via DistanceImpl (ADC path). + iter->next(); + iter->is_valid(); + std::vector code(quantizer->quantized_datapoint_vector_length()); + quantizer->quantize_data(iter->data(), code.data()); + + float d = dist_impl(code.data()); + EXPECT_GE(d, 0.0f); + + // SDC (symmetric) via sym_func() directly. + std::vector code2(quantizer->quantized_datapoint_vector_length()); + iter->next(); + iter->is_valid(); + quantizer->quantize_data(iter->data(), code2.data()); + + float sdc_d = 0.0f; + dist_impl.sym_func()(code.data(), code2.data(), NSQ, &sdc_d); + EXPECT_GE(sdc_d, 0.0f); +} + +TEST(PqInt8Quantizer, SerializeDeserialize) { + const size_t DIM = 16; + const size_t NSQ = 4; + const size_t COUNT = 500; + + auto quantizer = make_pq_quantizer(DIM, NSQ); + ASSERT_TRUE(quantizer); + + auto holder = make_random_holder(COUNT, DIM); + ASSERT_EQ(0, quantizer->train(holder)); + + // Serialize. + std::string blob; + ASSERT_EQ(0, quantizer->serialize(&blob)); + EXPECT_GT(blob.size(), sizeof(zvec::turbo::QuantizerSerHeader)); + + // Deserialize into a fresh quantizer. + auto q2 = IndexFactory::CreateQuantizer("PqInt8Quantizer"); + ASSERT_TRUE(q2); + ASSERT_EQ(0, q2->deserialize(blob)); + + // Encode the same vector with both and compare codes. + auto iter = holder->create_iterator(); + iter->is_valid(); + std::vector code1(quantizer->quantized_datapoint_vector_length()); + std::vector code2(q2->quantized_datapoint_vector_length()); + quantizer->quantize_data(iter->data(), code1.data()); + q2->quantize_data(iter->data(), code2.data()); + + for (size_t m = 0; m < NSQ; ++m) { + EXPECT_EQ(code1[m], code2[m]) << "m=" << m; + } + + // ADC distances should also match (same codebook → same LUT → same ADC). + size_t lut_len = quantizer->quantized_query_vector_length(); + std::vector lut1(lut_len / sizeof(float)); + std::vector lut2(lut_len / sizeof(float)); + quantizer->quantize_query(iter->data(), lut1.data()); + q2->quantize_query(iter->data(), lut2.data()); + + float adc1 = quantizer->calc_distance_dp_query(code1.data(), lut1.data()); + float adc2 = q2->calc_distance_dp_query(code2.data(), lut2.data()); + EXPECT_NEAR(adc1, adc2, 1e-6f); + + // Note: SDC (calc_distance_dp_dp) is intentionally NOT tested after + // deserialization because dist_table_ is a build-phase-only structure + // and is not persisted. +} + +// --------------------------------------------------------------------------- +// SIMD Consistency Tests +// --------------------------------------------------------------------------- + +namespace { + +// Helper to generate random uint8 codes +void fill_random_codes(uint8_t *codes, size_t len, std::mt19937 &gen) { + std::uniform_int_distribution dist(0, 255); + for (size_t i = 0; i < len; ++i) { + codes[i] = static_cast(dist(gen)); + } +} + +// Helper to generate random LUT (ADC) +void fill_random_lut(float *lut, size_t num_subquantizers, std::mt19937 &gen) { + constexpr size_t kNumCentroids = 256; + std::uniform_real_distribution dist(0.0f, 1.0f); + for (size_t m = 0; m < num_subquantizers; ++m) { + for (size_t c = 0; c < kNumCentroids; ++c) { + lut[m * kNumCentroids + c] = dist(gen); + } + } +} + +// Helper to generate random dist_table (SDC) +void fill_random_sdc_table(float *table, size_t num_subquantizers, + std::mt19937 &gen) { + constexpr size_t kTablePerSub = 256 * 256; + std::uniform_real_distribution dist(0.0f, 1.0f); + for (size_t m = 0; m < num_subquantizers; ++m) { + for (size_t i = 0; i < kTablePerSub; ++i) { + table[m * kTablePerSub + i] = dist(gen); + } + } +} + +} // anonymous namespace + +// Test ADC SIMD consistency across multiple M values +TEST(PqInt8SimdConsistency, AdcDistance) { + std::mt19937 gen(2024); + + // Test various M values including boundary cases + // M=4,8: exact multiples of AVX2 chunk (8) + // M=12: not multiple of 8, has leftover + // M=16: exact multiple of AVX512 chunk (16) + for (size_t num_sq : {4, 8, 12, 16}) { + constexpr size_t kNumCentroids = 256; + std::vector codes(num_sq); + std::vector lut(num_sq * kNumCentroids); + + fill_random_codes(codes.data(), num_sq, gen); + fill_random_lut(lut.data(), num_sq, gen); + + // Compute reference (scalar) + float scalar_result = 0.0f; + zvec::turbo::scalar::pq_adc_int8_distance(codes.data(), lut.data(), num_sq, + &scalar_result); + +#if defined(__AVX2__) + { + float avx2_result = 0.0f; + zvec::turbo::avx2::pq_adc_int8_distance_avx2(codes.data(), lut.data(), + num_sq, &avx2_result); + EXPECT_NEAR(scalar_result, avx2_result, 1e-5f) + << "AVX2 ADC mismatch for M=" << num_sq; + } +#endif + +#if defined(__AVX512F__) + { + float avx512_result = 0.0f; + zvec::turbo::avx512::pq_adc_int8_distance_avx512(codes.data(), lut.data(), + num_sq, &avx512_result); + EXPECT_NEAR(scalar_result, avx512_result, 1e-5f) + << "AVX512 ADC mismatch for M=" << num_sq; + } +#endif + } +} + +// Test SDC SIMD consistency across multiple M values +TEST(PqInt8SimdConsistency, SdcDistance) { + std::mt19937 gen(2025); + + for (size_t num_sq : {4, 8, 12, 16}) { + constexpr size_t kTablePerSub = 256 * 256; + std::vector codes_a(num_sq); + std::vector codes_b(num_sq); + std::vector dist_table(num_sq * kTablePerSub); + + fill_random_codes(codes_a.data(), num_sq, gen); + fill_random_codes(codes_b.data(), num_sq, gen); + fill_random_sdc_table(dist_table.data(), num_sq, gen); + + // Compute reference (scalar) + float scalar_result = 0.0f; + zvec::turbo::scalar::pq_sdc_int8_distance(codes_a.data(), codes_b.data(), + dist_table.data(), num_sq, + &scalar_result); + +#if defined(__AVX2__) + { + float avx2_result = 0.0f; + zvec::turbo::avx2::pq_sdc_int8_distance_avx2( + codes_a.data(), codes_b.data(), dist_table.data(), num_sq, + &avx2_result); + EXPECT_NEAR(scalar_result, avx2_result, 1e-5f) + << "AVX2 SDC mismatch for M=" << num_sq; + } +#endif + +#if defined(__AVX512F__) + { + float avx512_result = 0.0f; + zvec::turbo::avx512::pq_sdc_int8_distance_avx512( + codes_a.data(), codes_b.data(), dist_table.data(), num_sq, + &avx512_result); + EXPECT_NEAR(scalar_result, avx512_result, 1e-5f) + << "AVX512 SDC mismatch for M=" << num_sq; + } +#endif + } +} + +// Test edge case: M=1 (minimum valid value) +TEST(PqInt8SimdConsistency, AdcDistanceM1) { + std::mt19937 gen(123); + constexpr size_t kNumCentroids = 256; + constexpr size_t num_sq = 1; + + std::vector codes(num_sq); + std::vector lut(num_sq * kNumCentroids); + + fill_random_codes(codes.data(), num_sq, gen); + fill_random_lut(lut.data(), num_sq, gen); + + float scalar_result = 0.0f; + zvec::turbo::scalar::pq_adc_int8_distance(codes.data(), lut.data(), num_sq, + &scalar_result); + +#if defined(__AVX2__) + { + float avx2_result = 0.0f; + zvec::turbo::avx2::pq_adc_int8_distance_avx2(codes.data(), lut.data(), + num_sq, &avx2_result); + EXPECT_NEAR(scalar_result, avx2_result, 1e-5f); + } +#endif + +#if defined(__AVX512F__) + { + float avx512_result = 0.0f; + zvec::turbo::avx512::pq_adc_int8_distance_avx512(codes.data(), lut.data(), + num_sq, &avx512_result); + EXPECT_NEAR(scalar_result, avx512_result, 1e-5f); + } +#endif +} + +// --------------------------------------------------------------------------- +// Cosine Metric Tests +// --------------------------------------------------------------------------- + +// Helper to create a PqInt8Quantizer with Cosine metric. +static std::shared_ptr make_pq_cosine_quantizer( + size_t dim, size_t num_subquantizers) { + auto q = IndexFactory::CreateQuantizer("PqInt8Quantizer"); + if (!q) return nullptr; + + IndexMeta meta; + meta.set_meta(IndexMeta::DataType::DT_FP32, dim); + meta.set_metric("Cosine", 0, Params()); + + Params params; + params.set("num_subquantizers", static_cast(num_subquantizers)); + if (q->init(meta, params) != 0) return nullptr; + return q; +} + +// Reference cosine distance: 1 - (a·b) / (||a|| * ||b||). +static float reference_cosine_distance(const float *a, const float *b, + size_t dim) { + float dot = 0.0f, norm_a = 0.0f, norm_b = 0.0f; + for (size_t i = 0; i < dim; ++i) { + dot += a[i] * b[i]; + norm_a += a[i] * a[i]; + norm_b += b[i] * b[i]; + } + float denom = std::sqrt(norm_a) * std::sqrt(norm_b); + if (denom < 1e-12f) return 1.0f; + return 1.0f - dot / denom; +} + +// Helper: generate random vectors with varying norms (not unit length). +static std::shared_ptr> +make_cosine_holder(size_t count, size_t dim, uint32_t seed = 42) { + auto holder = + std::make_shared>(dim); + std::mt19937 gen(seed); + std::uniform_real_distribution dist(-1.0f, 1.0f); + std::uniform_real_distribution scale(0.5f, 5.0f); + for (size_t i = 0; i < count; ++i) { + NumericalVector vec(dim); + float s = scale(gen); + for (size_t j = 0; j < dim; ++j) vec[j] = dist(gen) * s; + holder->emplace(i + 1, vec); + } + return holder; +} + +// Verify that quantize_data stores the correct L2 norm after the PQ code. +TEST(PqInt8Quantizer, CosineNormStorage) { + const size_t DIM = 32; + const size_t NSQ = 8; + const size_t COUNT = 500; + + auto quantizer = make_pq_cosine_quantizer(DIM, NSQ); + ASSERT_TRUE(quantizer); + + // extra_meta_size should be sizeof(float) for Cosine. + EXPECT_EQ(quantizer->quantized_datapoint_vector_length(), + NSQ + sizeof(float)); + + auto holder = make_cosine_holder(COUNT, DIM); + ASSERT_EQ(0, quantizer->train(holder)); + + auto iter = holder->create_iterator(); + std::vector code(quantizer->quantized_datapoint_vector_length()); + + for (size_t checked = 0; iter->is_valid() && checked < 10; + iter->next(), ++checked) { + const float *v = reinterpret_cast(iter->data()); + + // Compute expected norm. + float expected_norm_sq = 0.0f; + for (size_t j = 0; j < DIM; ++j) expected_norm_sq += v[j] * v[j]; + float expected_norm = std::sqrt(expected_norm_sq); + + quantizer->quantize_data(iter->data(), code.data()); + + // Read stored norm from after the PQ code. + float stored_norm = 0.0f; + std::memcpy(&stored_norm, code.data() + NSQ, sizeof(float)); + + EXPECT_NEAR(stored_norm, expected_norm, expected_norm * 1e-5f) + << "Norm mismatch at vector " << checked; + } +} + +// Verify that dequantize reconstructs a vector with approximately correct +// direction (cosine similarity close to 1) and magnitude. +TEST(PqInt8Quantizer, CosineDequantize) { + const size_t DIM = 32; + const size_t NSQ = 8; + const size_t COUNT = 2000; + + auto quantizer = make_pq_cosine_quantizer(DIM, NSQ); + ASSERT_TRUE(quantizer); + + auto holder = make_cosine_holder(COUNT, DIM); + ASSERT_EQ(0, quantizer->train(holder)); + + auto iter = holder->create_iterator(); + size_t code_len = quantizer->quantized_datapoint_vector_length(); + + float max_cos_dist = 0.0f; + float max_norm_rel_error = 0.0f; + + for (size_t i = 0; iter->is_valid() && i < 50; iter->next(), ++i) { + const float *v = reinterpret_cast(iter->data()); + + // Compute original norm. + float orig_norm_sq = 0.0f; + for (size_t j = 0; j < DIM; ++j) orig_norm_sq += v[j] * v[j]; + float orig_norm = std::sqrt(orig_norm_sq); + + // Encode. + std::vector code(code_len); + quantizer->quantize_data(iter->data(), code.data()); + + // Decode. + IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, DIM); + std::string decoded; + ASSERT_EQ(0, quantizer->dequantize(code.data(), qmeta, &decoded)); + ASSERT_EQ(decoded.size(), DIM * sizeof(float)); + + const float *recon = reinterpret_cast(decoded.data()); + + // Check cosine similarity between original and reconstructed. + float cos_dist = reference_cosine_distance(v, recon, DIM); + max_cos_dist = std::max(max_cos_dist, cos_dist); + + // Check norm of reconstructed vector ≈ original norm. + float recon_norm_sq = 0.0f; + for (size_t j = 0; j < DIM; ++j) recon_norm_sq += recon[j] * recon[j]; + float recon_norm = std::sqrt(recon_norm_sq); + + if (orig_norm > 1e-6f) { + float rel_err = std::fabs(recon_norm - orig_norm) / orig_norm; + max_norm_rel_error = std::max(max_norm_rel_error, rel_err); + } + } + + // PQ with 8 subs and 2000 training vectors: cosine distance should be + // small (< 0.3 is generous; typically < 0.1). + EXPECT_LT(max_cos_dist, 0.3f) << "max_cos_dist=" << max_cos_dist; + + // Reconstructed norm should closely match original. PQ centroid + // concatenation in normalized space is not exactly unit-length, so + // the norm carries the centroid approximation error (~10-15% for + // 8 subs with sub_dim=4 and 2000 training vectors). + EXPECT_LT(max_norm_rel_error, 0.2f) + << "max_norm_rel_error=" << max_norm_rel_error; +} + +// Verify Cosine search distances via ADC are consistent with true cosine +// distance and fall in the expected range [0, 2]. +TEST(PqInt8Quantizer, CosineAdcDistance) { + const size_t DIM = 32; + const size_t NSQ = 8; + const size_t COUNT = 2000; + + auto quantizer = make_pq_cosine_quantizer(DIM, NSQ); + ASSERT_TRUE(quantizer); + + auto holder = make_cosine_holder(COUNT, DIM); + ASSERT_EQ(0, quantizer->train(holder)); + + // Collect raw vectors and PQ codes. + std::vector> raw_vecs(COUNT); + std::vector> pq_codes(COUNT); + size_t code_len = quantizer->quantized_datapoint_vector_length(); + size_t lut_len = quantizer->quantized_query_vector_length(); + + auto iter = holder->create_iterator(); + for (size_t i = 0; iter->is_valid(); iter->next(), ++i) { + const float *v = reinterpret_cast(iter->data()); + raw_vecs[i].assign(v, v + DIM); + pq_codes[i].resize(code_len); + quantizer->quantize_data(iter->data(), pq_codes[i].data()); + } + + // Build LUT for query = raw_vecs[0]. + std::vector lut(lut_len / sizeof(float)); + quantizer->quantize_query(raw_vecs[0].data(), lut.data()); + + for (size_t i = 1; i < COUNT; ++i) { + float adc_dist = + quantizer->calc_distance_dp_query(pq_codes[i].data(), lut.data()); + float true_dist = + reference_cosine_distance(raw_vecs[i].data(), raw_vecs[0].data(), DIM); + + // Cosine distance should be in [0, 2]. + EXPECT_GE(adc_dist, -0.01f) << "i=" << i; + EXPECT_LE(adc_dist, 2.01f) << "i=" << i; + + // PQ approximation: should be roughly correlated (within 0.5). + EXPECT_LT(std::fabs(adc_dist - true_dist), 0.5f) + << "i=" << i << " adc=" << adc_dist << " true=" << true_dist; + } +} + +// Verify that dequantize for L2 metric (no norm storage) still works. +TEST(PqInt8Quantizer, L2Dequantize) { + const size_t DIM = 16; + const size_t NSQ = 4; + const size_t COUNT = 1000; + + auto quantizer = make_pq_quantizer(DIM, NSQ); + ASSERT_TRUE(quantizer); + + // L2 metric: no extra meta. + EXPECT_EQ(quantizer->quantized_datapoint_vector_length(), NSQ); + + auto holder = make_random_holder(COUNT, DIM); + ASSERT_EQ(0, quantizer->train(holder)); + + auto iter = holder->create_iterator(); + iter->is_valid(); + + std::vector code(quantizer->quantized_datapoint_vector_length()); + quantizer->quantize_data(iter->data(), code.data()); + + IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, DIM); + std::string decoded; + ASSERT_EQ(0, quantizer->dequantize(code.data(), qmeta, &decoded)); + ASSERT_EQ(decoded.size(), DIM * sizeof(float)); + + const float *recon = reinterpret_cast(decoded.data()); + const float *orig = reinterpret_cast(iter->data()); + + // L2 PQ reconstruction should be a reasonable approximation. + float recon_err = reference_sq_euclidean(orig, recon, DIM); + float orig_norm = reference_sq_euclidean(orig, orig, DIM); + // Relative reconstruction error should be bounded. + if (orig_norm > 1e-6f) { + EXPECT_LT(recon_err / orig_norm, 1.0f) + << "recon_err=" << recon_err << " orig_norm=" << orig_norm; + } +} + +// --------------------------------------------------------------------------- +// InnerProduct Metric Tests +// --------------------------------------------------------------------------- + +// Helper to create a PqInt8Quantizer with InnerProduct metric. +static std::shared_ptr make_pq_ip_quantizer( + size_t dim, size_t num_subquantizers) { + auto q = IndexFactory::CreateQuantizer("PqInt8Quantizer"); + if (!q) return nullptr; + + IndexMeta meta; + meta.set_meta(IndexMeta::DataType::DT_FP32, dim); + meta.set_metric("InnerProduct", 0, Params()); + + Params params; + params.set("num_subquantizers", static_cast(num_subquantizers)); + if (q->init(meta, params) != 0) return nullptr; + return q; +} + +// Reference inner-product distance: -dot(a, b). +static float reference_ip_distance(const float *a, const float *b, size_t dim) { + float dot = 0.0f; + for (size_t i = 0; i < dim; ++i) { + dot += a[i] * b[i]; + } + return -dot; +} + +// Verify IP metric: no extra meta, ADC distances approximate true IP. +TEST(PqInt8Quantizer, InnerProductAdcDistance) { + const size_t DIM = 32; + const size_t NSQ = 8; + const size_t COUNT = 2000; + + auto quantizer = make_pq_ip_quantizer(DIM, NSQ); + ASSERT_TRUE(quantizer); + + // IP metric should NOT add extra meta (unlike Cosine). + EXPECT_EQ(quantizer->quantized_datapoint_vector_length(), NSQ); + + auto holder = make_random_holder(COUNT, DIM); + ASSERT_EQ(0, quantizer->train(holder)); + + // Collect raw vectors and PQ codes. + std::vector> raw_vecs(COUNT); + std::vector> pq_codes(COUNT); + size_t code_len = quantizer->quantized_datapoint_vector_length(); + size_t lut_len = quantizer->quantized_query_vector_length(); + + auto iter = holder->create_iterator(); + for (size_t i = 0; iter->is_valid(); iter->next(), ++i) { + const float *v = reinterpret_cast(iter->data()); + raw_vecs[i].assign(v, v + DIM); + pq_codes[i].resize(code_len); + quantizer->quantize_data(iter->data(), pq_codes[i].data()); + } + + // Build LUT for query = raw_vecs[0]. + std::vector lut(lut_len / sizeof(float)); + quantizer->quantize_query(raw_vecs[0].data(), lut.data()); + + float max_abs_error = 0.0f; + for (size_t i = 1; i < COUNT; ++i) { + float adc_dist = + quantizer->calc_distance_dp_query(pq_codes[i].data(), lut.data()); + float true_dist = + reference_ip_distance(raw_vecs[i].data(), raw_vecs[0].data(), DIM); + + // IP distance can be positive or negative; relative error is + // meaningless near zero crossings. Use absolute error instead. + float abs_err = std::fabs(adc_dist - true_dist); + max_abs_error = std::max(max_abs_error, abs_err); + } + // PQ IP distance: absolute error should be bounded. + // With 8 subs and dim=32, typical max abs error is a few units. + EXPECT_LT(max_abs_error, static_cast(DIM) * 0.5f) + << "max_abs_error=" << max_abs_error; +} + +// Verify IP dequantize works (same as L2: centroid concat, no norm rescale). +TEST(PqInt8Quantizer, InnerProductDequantize) { + const size_t DIM = 16; + const size_t NSQ = 4; + const size_t COUNT = 1000; + + auto quantizer = make_pq_ip_quantizer(DIM, NSQ); + ASSERT_TRUE(quantizer); + + EXPECT_EQ(quantizer->quantized_datapoint_vector_length(), NSQ); + + auto holder = make_random_holder(COUNT, DIM); + ASSERT_EQ(0, quantizer->train(holder)); + + auto iter = holder->create_iterator(); + iter->is_valid(); + + std::vector code(quantizer->quantized_datapoint_vector_length()); + quantizer->quantize_data(iter->data(), code.data()); + + IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, DIM); + std::string decoded; + ASSERT_EQ(0, quantizer->dequantize(code.data(), qmeta, &decoded)); + ASSERT_EQ(decoded.size(), DIM * sizeof(float)); + + const float *recon = reinterpret_cast(decoded.data()); + const float *orig = reinterpret_cast(iter->data()); + + // IP PQ reconstruction: centroid concat, same as L2. + float recon_err = reference_sq_euclidean(orig, recon, DIM); + float orig_norm = reference_sq_euclidean(orig, orig, DIM); + if (orig_norm > 1e-6f) { + EXPECT_LT(recon_err / orig_norm, 1.0f) + << "recon_err=" << recon_err << " orig_norm=" << orig_norm; + } +}