From b1313a466a883f440976ac60c13805078f7d0ac6 Mon Sep 17 00:00:00 2001 From: ray Date: Mon, 29 Jun 2026 19:47:27 +0800 Subject: [PATCH 01/42] refactor: add quantizer and distance --- src/turbo/quantizer/distance.h | 107 +++++++++++++++++++ src/turbo/quantizer/quantizer.h | 180 ++++++++++++++++++++++++++++++++ 2 files changed, 287 insertions(+) create mode 100644 src/turbo/quantizer/distance.h create mode 100644 src/turbo/quantizer/quantizer.h diff --git a/src/turbo/quantizer/distance.h b/src/turbo/quantizer/distance.h new file mode 100644 index 000000000..bc8af6c1a --- /dev/null +++ b/src/turbo/quantizer/distance.h @@ -0,0 +1,107 @@ +// 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)), + 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)), + 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). Useful for pairwise node-vs-node distance where no + //! stored query is involved. + const DistanceFunc &func() const { + return func_; + } + + //! Raw batch distance function. + const BatchDistanceFunc &batch_func() const { + return batch_func_; + } + + private: + DistanceFunc func_{}; + BatchDistanceFunc batch_func_{}; + std::string query_storage_{}; + size_t dim_{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..c6d1db2bc --- /dev/null +++ b/src/turbo/quantizer/quantizer.h @@ -0,0 +1,180 @@ +// 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 +#include "distance.h" + +using namespace zvec::core; + +namespace zvec { +namespace turbo { + +//! 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 IndexError_NotImplemented; + } + + //! 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 IndexError_NotImplemented; + } + + //! 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 IndexError_NotImplemented; + } + + //! Dequantize a result vector back to original format + virtual int dequantize(const void * /*in*/, const IndexQueryMeta & /*qmeta*/, + std::string * /*out*/) const { + return IndexError_NotImplemented; + } + + virtual DistanceImpl distance(const void * /*query*/, + const IndexQueryMeta & /*qmeta*/) const { + return DistanceImpl{}; + } + + //! Serialize quantizer parameters + virtual int serialize(std::string * /*out*/) const { + return IndexError_NotImplemented; + } + + //! Deserialize quantizer parameters + virtual int deserialize(std::string & /*in*/) { + return IndexError_NotImplemented; + } + + //! 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 IndexError_NotImplemented; + } + + 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 From 45d9b54a10259cfe988dfbadc39bc516dbef5d19 Mon Sep 17 00:00:00 2001 From: ray Date: Tue, 30 Jun 2026 14:16:10 +0800 Subject: [PATCH 02/42] fix: add turo quantizer --- src/turbo/CMakeLists.txt | 4 +- .../record_quantized_int8/common.h | 0 .../record_quantized_int8/cosine.cc | 0 .../record_quantized_int8/cosine.h | 0 .../squared_euclidean.cc | 0 .../record_quantized_int8/squared_euclidean.h | 0 .../avx512_vnni/uniform_int8/quantize.cc | 0 .../avx512_vnni/uniform_int8/quantize.h | 0 .../uniform_int8/squared_euclidean.cc | 0 .../uniform_int8/squared_euclidean.h | 0 tests/turbo/turbo_fp32_quantizer_test.cc | 83 +++++++++++++++++++ 11 files changed, 85 insertions(+), 2 deletions(-) rename src/turbo/{ => distance}/avx512_vnni/record_quantized_int8/common.h (100%) rename src/turbo/{ => distance}/avx512_vnni/record_quantized_int8/cosine.cc (100%) rename src/turbo/{ => distance}/avx512_vnni/record_quantized_int8/cosine.h (100%) rename src/turbo/{ => distance}/avx512_vnni/record_quantized_int8/squared_euclidean.cc (100%) rename src/turbo/{ => distance}/avx512_vnni/record_quantized_int8/squared_euclidean.h (100%) rename src/turbo/{ => distance}/avx512_vnni/uniform_int8/quantize.cc (100%) rename src/turbo/{ => distance}/avx512_vnni/uniform_int8/quantize.h (100%) rename src/turbo/{ => distance}/avx512_vnni/uniform_int8/squared_euclidean.cc (100%) rename src/turbo/{ => distance}/avx512_vnni/uniform_int8/squared_euclidean.h (100%) create mode 100644 tests/turbo/turbo_fp32_quantizer_test.cc diff --git a/src/turbo/CMakeLists.txt b/src/turbo/CMakeLists.txt index 9cbb2fac7..fcf9dc431 100644 --- a/src/turbo/CMakeLists.txt +++ b/src/turbo/CMakeLists.txt @@ -19,7 +19,7 @@ file(GLOB_RECURSE ALL_SRCS *.cc *.c *.h) # subdirectory). if(NOT ANDROID AND AUTO_DETECT_ARCH) if (HOST_ARCH MATCHES "^(x86|x64)$") - file(GLOB_RECURSE AVX512_VNNI_SRCS ${CMAKE_CURRENT_SOURCE_DIR}/avx512_vnni/*.cc) + file(GLOB_RECURSE AVX512_VNNI_SRCS ${CMAKE_CURRENT_SOURCE_DIR}/distance/avx512_vnni/*.cc) set_source_files_properties( ${AVX512_VNNI_SRCS} PROPERTIES @@ -32,5 +32,5 @@ cc_library( NAME zvec_turbo STATIC STRICT 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/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/tests/turbo/turbo_fp32_quantizer_test.cc b/tests/turbo/turbo_fp32_quantizer_test.cc new file mode 100644 index 000000000..40165a5d3 --- /dev/null +++ b/tests/turbo/turbo_fp32_quantizer_test.cc @@ -0,0 +1,83 @@ +// 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 "zvec/core/framework/index_factory.h" + +using namespace zvec; +using namespace zvec::core; +using namespace zvec::ailego; + +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); + } + } +} \ No newline at end of file From e82d73f343e75680ba7c66710d91f022c462dd3d Mon Sep 17 00:00:00 2001 From: ray Date: Tue, 30 Jun 2026 15:26:47 +0800 Subject: [PATCH 03/42] refactor: add quantizer --- src/core/framework/index_factory.cc | 13 ++++++++++++ .../zvec/core/framework/index_factory.h | 19 ++++++++++++++++++ src/include/zvec/turbo/turbo.h | 20 +++++++++++++++++++ tests/CMakeLists.txt | 1 + tests/turbo/CMakeLists.txt | 14 +++++++++++++ 5 files changed, 67 insertions(+) create mode 100644 tests/turbo/CMakeLists.txt 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/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/turbo/turbo.h b/src/include/zvec/turbo/turbo.h index 2fbf6d680..df353bded 100644 --- a/src/include/zvec/turbo/turbo.h +++ b/src/include/zvec/turbo/turbo.h @@ -36,18 +36,38 @@ using UniformQuantizeFunc = void (*)(const float *in, size_t dim, float scale, enum class MetricType { kSquaredEuclidean, kCosine, + kInnerProduct, kMipsSquaredEuclidean, kUnknown, }; enum class DataType { + kInt4, kInt8, + kFp16, + kFp32, kUnknown, }; enum class QuantizeType { kDefault, kUniform, + kRecord, + kFp16, + kFp32, + kPQ, + kRabit +}; + +enum class CpuArchType { + kAuto, + kScalar, + kSSE, + kAVX, + kAVX2, + kAVX512, + kAVX512VNNI, + kAVX512FP16 }; DistanceFunc get_distance_func(MetricType metric_type, DataType data_type, 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/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 From 829f53e198aa092f0a62ee6f6da2e56b11e4b894 Mon Sep 17 00:00:00 2001 From: ray Date: Tue, 30 Jun 2026 15:29:46 +0800 Subject: [PATCH 04/42] refactor: update meta --- src/core/framework/index_meta.cc | 4 +- src/include/zvec/core/framework/index_meta.h | 50 +++++++++++++++++++- 2 files changed, 51 insertions(+), 3 deletions(-) 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/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}; }; From f08463011b0af38ecde1d987d9279a78f247aefe Mon Sep 17 00:00:00 2001 From: ray Date: Tue, 30 Jun 2026 15:46:13 +0800 Subject: [PATCH 05/42] refactor: update meta --- src/turbo/quantizer/quantizer.h | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/turbo/quantizer/quantizer.h b/src/turbo/quantizer/quantizer.h index c6d1db2bc..572d372b5 100644 --- a/src/turbo/quantizer/quantizer.h +++ b/src/turbo/quantizer/quantizer.h @@ -24,11 +24,11 @@ #include #include "distance.h" -using namespace zvec::core; - namespace zvec { namespace turbo { +using namespace zvec::core; + //! Magic number ('QTZR') stamped at the start of a serialized quantizer blob. constexpr uint32_t kQuantizerMagic = 0x52545A51u; //! Current quantizer serialization format version. From c3c4ca35c0d1c7389c9c96206a7527be97bf3596 Mon Sep 17 00:00:00 2001 From: ray Date: Tue, 30 Jun 2026 16:22:51 +0800 Subject: [PATCH 06/42] refactor: add fp32 quantizer --- src/include/zvec/turbo/turbo.h | 15 +- .../fp32_quantizer/fp32_quantizer.cc | 161 ++++++++++++++++++ .../quantizer/fp32_quantizer/fp32_quantizer.h | 127 ++++++++++++++ src/turbo/turbo.cc | 22 ++- 4 files changed, 312 insertions(+), 13 deletions(-) create mode 100644 src/turbo/quantizer/fp32_quantizer/fp32_quantizer.cc create mode 100644 src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h diff --git a/src/include/zvec/turbo/turbo.h b/src/include/zvec/turbo/turbo.h index df353bded..2eff572d5 100644 --- a/src/include/zvec/turbo/turbo.h +++ b/src/include/zvec/turbo/turbo.h @@ -71,15 +71,16 @@ enum class CpuArchType { }; 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 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..efd4e951e --- /dev/null +++ b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.cc @@ -0,0 +1,161 @@ +// 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()); + + 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_); + } + + // `meta.dimension()` may be either the raw dim (when the caller passes a + // bare meta, e.g. unit tests) or the inflated storage dim (data + extras) + // when an upstream converter such as CosineConverter has already appended + // a norm float and set meta.extra_meta_size(). Distinguish via the input + // meta's extra_meta_size: if it is already set, the dim is inflated and + // we strip the extras to recover the raw dim; otherwise the dim is raw. + if (meta.extra_meta_size() > 0) { + original_dim_ = meta.dimension() - meta.extra_meta_size() / sizeof(float); + } else { + original_dim_ = meta.dimension(); + } + + // 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 IndexError_Unsupported; + } + + // 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); + out->resize(byte_size); + 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 byte_size = qmeta.dimension() * sizeof(float); + 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 { + std::memcpy(output, input, + static_cast(original_dim_) * sizeof(float)); +} + +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..998a19d68 --- /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). + size_t quantized_length() const { + return static_cast(original_dim_) * sizeof(float); + } + + //! 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}; + IndexMeta::DataType data_type_{}; + + //! Cached distance dispatch (bound in init()). + DistanceFunc dp_query_func_{}; + BatchDistanceFunc dp_query_batch_func_{}; +}; + + +} // namespace turbo +} // namespace zvec diff --git a/src/turbo/turbo.cc b/src/turbo/turbo.cc index adc9b785e..11a7a43cc 100644 --- a/src/turbo/turbo.cc +++ b/src/turbo/turbo.cc @@ -22,10 +22,13 @@ 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::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 +50,13 @@ 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::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 +73,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; } From ec4ea1fdfbfd6b5240cef161c1820b15f070c1f3 Mon Sep 17 00:00:00 2001 From: ray Date: Tue, 30 Jun 2026 16:35:47 +0800 Subject: [PATCH 07/42] refactor: add quantizer --- src/turbo/CMakeLists.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/turbo/CMakeLists.txt b/src/turbo/CMakeLists.txt index fcf9dc431..7ec9c0e5d 100644 --- a/src/turbo/CMakeLists.txt +++ b/src/turbo/CMakeLists.txt @@ -29,7 +29,7 @@ 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} ${CMAKE_CURRENT_SOURCE_DIR}/distance ${PROJECT_ROOT_DIR}/src/include From e36f0c5495735fba1ba04e82149e9e39fcbac122 Mon Sep 17 00:00:00 2001 From: ray Date: Tue, 30 Jun 2026 17:34:23 +0800 Subject: [PATCH 08/42] refactor: quantizer --- src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h index 998a19d68..79a7caa09 100644 --- a/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h +++ b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h @@ -57,7 +57,7 @@ class Fp32Quantizer : public Quantizer { } int train(core::IndexHolder::Pointer /*holder*/) override { - return 0; + return IndexError_NotImplemented; } size_t quantized_datapoint_vector_length() const override { From c1cdd0b8c0fbbf337a85cc17bc4391cf234540e4 Mon Sep 17 00:00:00 2001 From: ray Date: Tue, 30 Jun 2026 19:08:56 +0800 Subject: [PATCH 09/42] refactor: add scalar --- src/turbo/distance/scalar/fp32/cosine.cc | 47 +++++++ src/turbo/distance/scalar/fp32/cosine.h | 30 +++++ .../distance/scalar/fp32/inner_product.cc | 43 +++++++ .../distance/scalar/fp32/inner_product.h | 31 +++++ .../distance/scalar/fp32/squared_euclidean.cc | 41 ++++++ .../distance/scalar/fp32/squared_euclidean.h | 31 +++++ src/turbo/turbo.cc | 33 +++++ tests/turbo/turbo_fp32_quantizer_test.cc | 119 ++++++++++++++++++ 8 files changed, 375 insertions(+) create mode 100644 src/turbo/distance/scalar/fp32/cosine.cc create mode 100644 src/turbo/distance/scalar/fp32/cosine.h create mode 100644 src/turbo/distance/scalar/fp32/inner_product.cc create mode 100644 src/turbo/distance/scalar/fp32/inner_product.h create mode 100644 src/turbo/distance/scalar/fp32/squared_euclidean.cc create mode 100644 src/turbo/distance/scalar/fp32/squared_euclidean.h diff --git a/src/turbo/distance/scalar/fp32/cosine.cc b/src/turbo/distance/scalar/fp32/cosine.cc new file mode 100644 index 000000000..112a6fc80 --- /dev/null +++ b/src/turbo/distance/scalar/fp32/cosine.cc @@ -0,0 +1,47 @@ +// 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 + +namespace zvec::turbo::scalar { + +void cosine_fp32_distance(const void *a, const void *b, size_t dim, + float *distance) { + const float *fa = static_cast(a); + const float *fb = static_cast(b); + float dot = 0.0f; + float na = 0.0f; + float nb = 0.0f; + for (size_t i = 0; i < dim; ++i) { + dot += fa[i] * fb[i]; + na += fa[i] * fa[i]; + nb += fb[i] * fb[i]; + } + float denom = std::sqrt(na) * std::sqrt(nb); + if (denom < 1e-12f) { + *distance = 1.0f; + } else { + *distance = 1.0f - dot / denom; + } +} + +void cosine_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) { + cosine_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/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..63a0575a3 --- /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 squared Euclidean distance between a single quantized FP32 +// vector pair. +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/turbo.cc b/src/turbo/turbo.cc index 11a7a43cc..80f2ee77d 100644 --- a/src/turbo/turbo.cc +++ b/src/turbo/turbo.cc @@ -18,12 +18,30 @@ #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/fp32/cosine.h" +#include "scalar/fp32/inner_product.h" +#include "scalar/fp32/squared_euclidean.h" namespace zvec::turbo { DistanceFunc get_distance_func(MetricType metric_type, DataType data_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 && @@ -52,6 +70,21 @@ BatchDistanceFunc get_batch_distance_func(MetricType metric_type, DataType data_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 && diff --git a/tests/turbo/turbo_fp32_quantizer_test.cc b/tests/turbo/turbo_fp32_quantizer_test.cc index 40165a5d3..c180de4a4 100644 --- a/tests/turbo/turbo_fp32_quantizer_test.cc +++ b/tests/turbo/turbo_fp32_quantizer_test.cc @@ -12,7 +12,9 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include #include +#include #include #include #include @@ -22,6 +24,18 @@ 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); @@ -80,4 +94,109 @@ TEST(Fp32Quantizer, General) { 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 From 628982db5738d6e431e37d362e9f85816d93e09c Mon Sep 17 00:00:00 2001 From: ray Date: Wed, 1 Jul 2026 11:53:21 +0800 Subject: [PATCH 10/42] fix: fix cosine --- src/turbo/distance/scalar/fp32/cosine.cc | 31 +++++++------------ .../quantizer/fp32_quantizer/fp32_quantizer.h | 2 +- 2 files changed, 13 insertions(+), 20 deletions(-) diff --git a/src/turbo/distance/scalar/fp32/cosine.cc b/src/turbo/distance/scalar/fp32/cosine.cc index 112a6fc80..6dd9fc532 100644 --- a/src/turbo/distance/scalar/fp32/cosine.cc +++ b/src/turbo/distance/scalar/fp32/cosine.cc @@ -13,34 +13,27 @@ // limitations under the License. #include "scalar/fp32/cosine.h" -#include +#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) { - const float *fa = static_cast(a); - const float *fb = static_cast(b); - float dot = 0.0f; - float na = 0.0f; - float nb = 0.0f; - for (size_t i = 0; i < dim; ++i) { - dot += fa[i] * fb[i]; - na += fa[i] * fa[i]; - nb += fb[i] * fb[i]; - } - float denom = std::sqrt(na) * std::sqrt(nb); - if (denom < 1e-12f) { - *distance = 1.0f; - } else { - *distance = 1.0f - dot / denom; - } + // 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) { - for (size_t i = 0; i < n; ++i) { - cosine_fp32_distance(vectors[i], query, dim, &distances[i]); + 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]; } } diff --git a/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h index 79a7caa09..998a19d68 100644 --- a/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h +++ b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h @@ -57,7 +57,7 @@ class Fp32Quantizer : public Quantizer { } int train(core::IndexHolder::Pointer /*holder*/) override { - return IndexError_NotImplemented; + return 0; } size_t quantized_datapoint_vector_length() const override { From de3322042a2b7f300d99b6403c47666135676f91 Mon Sep 17 00:00:00 2001 From: ray Date: Wed, 1 Jul 2026 14:17:05 +0800 Subject: [PATCH 11/42] fix: add quantizer --- src/turbo/distance/scalar/fp32/inner_product.cc | 4 ++-- src/turbo/quantizer/fp32_quantizer/fp32_quantizer.cc | 6 ------ 2 files changed, 2 insertions(+), 8 deletions(-) diff --git a/src/turbo/distance/scalar/fp32/inner_product.cc b/src/turbo/distance/scalar/fp32/inner_product.cc index 63a0575a3..1f966d788 100644 --- a/src/turbo/distance/scalar/fp32/inner_product.cc +++ b/src/turbo/distance/scalar/fp32/inner_product.cc @@ -16,8 +16,8 @@ namespace zvec::turbo::scalar { -// Compute squared Euclidean distance between a single quantized FP32 -// vector pair. +// 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); diff --git a/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.cc b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.cc index efd4e951e..6f8853c5c 100644 --- a/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.cc +++ b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.cc @@ -36,12 +36,6 @@ int Fp32Quantizer::init(const IndexMeta &meta, meta_.set_extra_meta_size(extra_meta_size_); } - // `meta.dimension()` may be either the raw dim (when the caller passes a - // bare meta, e.g. unit tests) or the inflated storage dim (data + extras) - // when an upstream converter such as CosineConverter has already appended - // a norm float and set meta.extra_meta_size(). Distinguish via the input - // meta's extra_meta_size: if it is already set, the dim is inflated and - // we strip the extras to recover the raw dim; otherwise the dim is raw. if (meta.extra_meta_size() > 0) { original_dim_ = meta.dimension() - meta.extra_meta_size() / sizeof(float); } else { From 49780f897918d2006d0e9b1b1fd95eddf0801786 Mon Sep 17 00:00:00 2001 From: ray Date: Wed, 1 Jul 2026 14:59:26 +0800 Subject: [PATCH 12/42] fix: fix quantizer --- .../fp32_quantizer/fp32_quantizer.cc | 64 +++++++++++++++---- .../quantizer/fp32_quantizer/fp32_quantizer.h | 5 +- 2 files changed, 54 insertions(+), 15 deletions(-) diff --git a/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.cc b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.cc index 6f8853c5c..f5c29b76b 100644 --- a/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.cc +++ b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.cc @@ -16,6 +16,7 @@ #include #include #include +#include #include #include #include @@ -30,18 +31,14 @@ int Fp32Quantizer::init(const IndexMeta &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_); } - if (meta.extra_meta_size() > 0) { - original_dim_ = meta.dimension() - meta.extra_meta_size() / sizeof(float); - } else { - original_dim_ = meta.dimension(); - } - // Cache the distance dispatch for the new Quantizer interface. dp_query_func_ = get_distance_func(metric_from_name(metric_name), DataType::kFp32, @@ -65,9 +62,22 @@ int Fp32Quantizer::quantize(const void *query, const IndexQueryMeta &qmeta, size_t raw_dim = (original_dim_ != 0 && qmeta.dimension() >= original_dim_) ? original_dim_ : qmeta.dimension(); - size_t byte_size = raw_dim * sizeof(float); + size_t byte_size = raw_dim * sizeof(float) + extra_meta_size_; out->resize(byte_size); - std::memcpy(&(*out)[0], query, 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, @@ -78,9 +88,26 @@ int Fp32Quantizer::quantize(const void *query, const IndexQueryMeta &qmeta, int Fp32Quantizer::dequantize(const void *in, const IndexQueryMeta &qmeta, std::string *out) const { - size_t byte_size = qmeta.dimension() * sizeof(float); - out->resize(byte_size); - std::memcpy(out->data(), in, byte_size); + 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; } @@ -103,8 +130,19 @@ DistanceImpl Fp32Quantizer::distance(const void *query, } void Fp32Quantizer::quantize_one(const void *input, void *output) const { - std::memcpy(output, input, - static_cast(original_dim_) * sizeof(float)); + 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, diff --git a/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h index 998a19d68..328a71e29 100644 --- a/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h +++ b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h @@ -102,9 +102,10 @@ class Fp32Quantizer : public Quantizer { const core::IndexQueryMeta &qmeta) const override; private: - //! Byte length of a quantized vector (raw fp32 data). + //! Byte length of a quantized vector (raw fp32 data + extra meta). size_t quantized_length() const { - return static_cast(original_dim_) * sizeof(float); + return static_cast(original_dim_) * sizeof(float) + + extra_meta_size_; } //! Quantize a single fp32 vector into a caller-provided buffer of From b1d484956fa425aaba5c3c9726dd7046e4eeb32e Mon Sep 17 00:00:00 2001 From: ray Date: Thu, 2 Jul 2026 10:58:18 +0800 Subject: [PATCH 13/42] fix: remove unused var --- src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h | 1 - 1 file changed, 1 deletion(-) diff --git a/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h index 328a71e29..885390639 100644 --- a/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h +++ b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h @@ -116,7 +116,6 @@ class Fp32Quantizer : public Quantizer { IndexMeta meta_{}; uint32_t original_dim_{0}; - IndexMeta::DataType data_type_{}; //! Cached distance dispatch (bound in init()). DistanceFunc dp_query_func_{}; From f0fdf5d65fe8663ed89b91b43d8b25635b346e25 Mon Sep 17 00:00:00 2001 From: ray Date: Thu, 2 Jul 2026 14:08:42 +0800 Subject: [PATCH 14/42] fix: header --- examples/c++/CMakeLists.txt | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) 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 From 0c4cec659b4e285b668411ab829f1bc00e07311e Mon Sep 17 00:00:00 2001 From: ray Date: Thu, 2 Jul 2026 15:05:02 +0800 Subject: [PATCH 15/42] fix: header --- src/turbo/quantizer/quantizer.h | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/src/turbo/quantizer/quantizer.h b/src/turbo/quantizer/quantizer.h index 572d372b5..14e889787 100644 --- a/src/turbo/quantizer/quantizer.h +++ b/src/turbo/quantizer/quantizer.h @@ -18,7 +18,6 @@ #include #include #include -#include #include #include #include @@ -31,6 +30,7 @@ using namespace zvec::core; //! 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; @@ -75,7 +75,7 @@ class Quantizer { //! Train the quantizer with a contiguous batch of data virtual int train(const void * /*data*/, size_t /*num*/, size_t /*stride*/) { - return IndexError_NotImplemented; + return 0; } //! Whether the quantizer requires training before use @@ -83,7 +83,7 @@ class Quantizer { //! Train the quantizer with data from an IndexHolder virtual int train(IndexHolder::Pointer /*holder*/) { - return IndexError_NotImplemented; + return 0; } //! Byte length of a quantized datapoint vector @@ -123,13 +123,13 @@ class Quantizer { virtual int quantize(const void * /*query*/, const IndexQueryMeta & /*qmeta*/, std::string * /*out*/, IndexQueryMeta * /*ometa*/) const { - return IndexError_NotImplemented; + return 0; } //! Dequantize a result vector back to original format virtual int dequantize(const void * /*in*/, const IndexQueryMeta & /*qmeta*/, std::string * /*out*/) const { - return IndexError_NotImplemented; + return 0; } virtual DistanceImpl distance(const void * /*query*/, @@ -139,18 +139,18 @@ class Quantizer { //! Serialize quantizer parameters virtual int serialize(std::string * /*out*/) const { - return IndexError_NotImplemented; + return 0 } //! Deserialize quantizer parameters virtual int deserialize(std::string & /*in*/) { - return IndexError_NotImplemented; + 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 IndexError_NotImplemented; + return 0; } protected: From 73f5bc3bab498315bac70d8c3142a7ecb530020c Mon Sep 17 00:00:00 2001 From: ray Date: Thu, 2 Jul 2026 15:36:15 +0800 Subject: [PATCH 16/42] fix: header --- src/turbo/quantizer/quantizer.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/turbo/quantizer/quantizer.h b/src/turbo/quantizer/quantizer.h index 14e889787..3e5e65ec6 100644 --- a/src/turbo/quantizer/quantizer.h +++ b/src/turbo/quantizer/quantizer.h @@ -139,7 +139,7 @@ class Quantizer { //! Serialize quantizer parameters virtual int serialize(std::string * /*out*/) const { - return 0 + return 0; } //! Deserialize quantizer parameters From 89843dd92a2131921ad2bcfdefaf93d73e827010 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Thu, 2 Jul 2026 15:42:47 +0800 Subject: [PATCH 17/42] add preprocessor --- ...26\345\231\250\350\256\276\350\256\241.md" | 731 ++++++++++++++++++ doc/pq.md | 94 +++ src/include/zvec/turbo/turbo.h | 20 + src/turbo/CMakeLists.txt | 23 +- src/turbo/distance/avx2/fht/fht.h | 37 + src/turbo/distance/avx2/fht/fht_avx2.cc | 138 ++++ src/turbo/distance/avx512/fht/fht.h | 37 + src/turbo/distance/avx512/fht/fht_avx512.cc | 143 ++++ src/turbo/distance/scalar/fht/fht.h | 38 + src/turbo/distance/scalar/fht/fht_scalar.cc | 80 ++ src/turbo/distance/sse/fht/fht.h | 34 + src/turbo/distance/sse/fht/fht_sse.cc | 119 +++ .../quantizer/preprocessor/fht_rotator.cc | 266 +++++++ .../quantizer/preprocessor/fht_rotator.h | 97 +++ .../quantizer/preprocessor/preprocessor.h | 93 +++ src/turbo/turbo.cc | 54 ++ tests/turbo/turbo_fht_rotator_test.cc | 204 +++++ 17 files changed, 2205 insertions(+), 3 deletions(-) create mode 100644 "doc/Zvec Turbo \351\207\217\345\214\226\345\231\250\350\256\276\350\256\241.md" create mode 100644 doc/pq.md create mode 100644 src/turbo/distance/avx2/fht/fht.h create mode 100644 src/turbo/distance/avx2/fht/fht_avx2.cc create mode 100644 src/turbo/distance/avx512/fht/fht.h create mode 100644 src/turbo/distance/avx512/fht/fht_avx512.cc create mode 100644 src/turbo/distance/scalar/fht/fht.h create mode 100644 src/turbo/distance/scalar/fht/fht_scalar.cc create mode 100644 src/turbo/distance/sse/fht/fht.h create mode 100644 src/turbo/distance/sse/fht/fht_sse.cc create mode 100644 src/turbo/quantizer/preprocessor/fht_rotator.cc create mode 100644 src/turbo/quantizer/preprocessor/fht_rotator.h create mode 100644 src/turbo/quantizer/preprocessor/preprocessor.h create mode 100644 tests/turbo/turbo_fht_rotator_test.cc diff --git "a/doc/Zvec Turbo \351\207\217\345\214\226\345\231\250\350\256\276\350\256\241.md" "b/doc/Zvec Turbo \351\207\217\345\214\226\345\231\250\350\256\276\350\256\241.md" new file mode 100644 index 000000000..66f02415e --- /dev/null +++ "b/doc/Zvec Turbo \351\207\217\345\214\226\345\231\250\350\256\276\350\256\241.md" @@ -0,0 +1,731 @@ +# 背景 +当前，距离度量沿用了 Proxima 的风格，采用模板方式进行定义。系统通过模块化设计实现了度量、转换器及指令集等组件的解耦，检索过程由Metric（距离度量）、Quantizer（量化器）、Algorithm（索引算法）等多个模块协作完成。然而，一些现有的组件存在职责边界模糊、实现方式耦合等问题，计划通过改为turbo的形式，统一切面通过输入多个维度信息组合以获得对应的函数集合，从而使得不同逻辑判断统一。 + +# 设计现状 +当前的系统架构设计存在着一些问题，下面将进行简要说明。 + +## Converter和Reformer功能耦合 +当前的量化器由Converter和Reformer两个组件构成，其中，Converter主要负责离线构建，而Reformer主要负责在线查询，不过由于当前设计未区分流式插入与批量合并，Converter和Reformer在离线阶段存在一定的功能耦合： + +![画板](https://intranetproxy.alipay.com/skylark/lark/0/2026/jpeg/226557004/1782716498906-f0957db9-24f8-466b-bfb2-351e65d1b5c5.jpeg) + +这个实现最初的想法是将离线的一次性操作`build`、`train`等放到Converter中，而可以反复进行的操作`add`和`search`等放到Reformer，但是，这在本身也引入了一些额外的限制，例如，如果量化器需要训练，但索引是流式插入的，例如hnsw+kUniformInt8，就会存在问题。以下对当前的使用接口进行说明： + +| **Converter接口** | **功能** | **Reformer接口** | **功能** | +| --- | --- | --- | --- | +| init() | 声明量化后的 meta 信息 | init() | 从 Converter 声明的 reformer_params 中加载运行时参数 | +| train() | 在全量数据上执行训练，收集统计信息 | load() | 从持久化存储中加载/卸载 Reformer 运行时状态 | +| transform() | 将 全量向量从原始格式转换为量化格式 | transform() | 将查询向量转换为与索引一致的量化格式 | +| dump() | 将 Converter 内部状态（如量化参数）持久化到 dumper | normalize() | 将量化空间计算的距离/分数映射回原始距离尺度 | +| meta() | 返回 Converter 修改后的 IndexMeta | revert() | 将量化向量反量化回 FP32 | + + +## Quantizer与Metric存在耦合 +当前设计存在不同功能分布在同一文件/同一功能分布在不同文件的问题，例如： + +| 注册文件 | 工厂注册名 | Metric | 备注 | +| --- | --- | --- | --- | +| **integer_quantizer_converter.cc** | Int8StreamingConverter | L2 / IP / MIPS | 基于单条数据的INT8 | +| | Int4StreamingConverter | L2 / IP / MIPS | 基于单条数据的INT4 | +| | Int8QuantizerConverter | - | 全局数据的INT8 | +| | Int4QuantizerConverter | - | 全局数据的INT4 | +| **cosine_converter.cc** | CosineInt8Converter | Cosine | 基于单条数据的INT8 | +| | CosineInt4Converter | Cosine | 基于单条数据的INT4 | +| | CosineFp16Converter | Cosine | FP16 | +| | ... | ... | ... | + + +由于这样的设计，一些量化器内部需要判断度量类型、数据类型、架构类型，以合理分配量化类型。 + +## 其他问题 +1. Builder、Searcher、Streamer职责边界模糊： + 1. Builder：全量数据的构建，负责操作 `build` + 2. Searcher：只被 IndexFlow（CLI 工具）和单元测试使用 + 3. Streamer：主路径组件负责读写，负责操作 `add` 和`search` +2. 部分索引实现与量化器耦合： + 1. 与Quantizer完全解耦：HNSW / Flat / Vamana + 2. 硬编码与Quantizer耦合：IVF + 3. 当前自成一体：diskann/hnsw rabitq + +# Turbo设计 +Turbo 旨在通过统一接口收敛度量（Metric）与量化（Quantization）逻辑，提高代码内聚性，从而降低维护成本，提升代码的可读性与简洁性。 + +## 使用场景 +![画板](https://intranetproxy.alipay.com/skylark/lark/0/2026/jpeg/95883/1782788752754-8c017988-c7ed-469d-8025-b1ab5ce94281.jpeg) + +## 模块 +Turbo系统包含以下几个模块： + ++ 量化器 ++ 预处理器 ++ 距离计算 + +都以量化器为核心进行调度： + +![画板](https://intranetproxy.alipay.com/skylark/lark/0/2026/jpeg/95883/1782802278889-9e04f816-ec15-42f6-9ff6-3e53ceb7c83e.jpeg) + +## 设计目标 +设计的目标主要包含以下六个维度的： + +| **维度** | **说明** | +| :---: | :---: | +| 源数据类型 | 量化前的数据类型 | +| 目标数据类型 | 量化后的数据类型 | +| 距离类型 | 支持的距离度量方式 | +| 量化器类型 | 具体量化算法及参数管理方式 | +| 距离计算模式 | 单查询单点、单查询多点连续、单查询多点离散、点对点 | +| 指令架构 | 该量化器支持的指令架构 | +| 量化器使用场景 | 训练、流式插入、批量插入/合并、序列化、反序列化 | + + +## 调用场景 +以下为构建、查询的调用场景： + +### 构建 +```plain + 源向量 + │ + ▼ +┌───────────────────────┐ +│ Quantizer │ +│ .quantize_datapoint │ +└───────────────────────┘ + │ + ▼ +[量化向量 + Extra Meta(如有)] → 写入 Streamer +``` + +要求： + ++ Uniform：无需训练即可编码（全局参数已确定）。 ++ Record：逐向量独立计算 scale/bias，无需全局训练。 ++ PQ/RabitQ：不支持纯流式插入（需先训练）。 + +### 查询 +```plain + 源向量 + │ + ▼ +┌─────────────────────┐ +│ Quantizer │ +│ .quantize_query() │ +└─────────────────────┘ + │ + ▼ +[量化向量 + Extra Meta(如有)] → 检索 Streamer +``` + +### 合并 +```plain +多个源 Segment (FP32) + │ + ▼ +┌─────────────┐ +│ write_holder│ ← 聚合数据到Holder +└─────────────┘ + │ + ▼ +┌─────────────┐ +│ Quantizer │ ← 训练全局参数 +│ .train │ +└─────────────┘ + │ + ▼ +┌───────────────────────┐ +│ Quantizer │ ← 用训练后的Quantizer编码所有向量 +│ .quantize_datapoint | +└───────────────────────┘ + │ + ▼ +[量化向量 + Extra Meta(如有)] → 写入 Streamer +``` + +# 支持范围 +## 源数据类型 → 目标数据类型 +考虑到 db 层有不保留 Flat FP32 的需求 + ++ 源数据类型需要支持 4 种通用类型： + - FP32，P0 + - FP16，P0 + - INT8，P1 + - INT4，P2 ++ 目标数据类型支持范围： + +| **目标类型** | **说明** | **优先级** | +| :---: | :---: | :---: | +| FP32 | 原始数据不做量化，仅做度量转换（如 Cosine） | P0 | +| FP16 | 半精度浮点 | P0 | +| INT8 | 8-bit 整数量化 | P0 | +| INT4 | 4-bit 整数量化 | P0 | +| X-Bits | 自定义 bit 宽度的量化（如 5-bit、6-bit） | P1 | + + +## 量化器类型 +量化器类型是核心分类维度，决定了端到端的数据布局、训练方式和距离计算方式。 + +| **量化器类型** | **特征** | **是否带Extra Meta** | **距离计算方式** | **优先级** | +| :---: | :---: | :---: | :---: | :---: | +| **Record** | 逐向量保存 scale/bias，按维度量化 | 是
（Extra Meta 存 scale/bias） | 非对称距离，需反量化或查表 | P0 | +| **Uniform** | 全局统一scale & bias，按维度量化 | 否 | 对称距离，可 SIMD 优化 | P0 | +| **PQ**
** (Product Quantization)** | 分段聚类码本量化 | 是
（码本参数） | 非对称距离，ADC/SDC | P1 | +| **OPQ** | 优化的 PQ，含正交变换 | 是
（变换矩阵 + 码本） | 非对称距离，ADC | P2 | +| **Raw** | 不做数值量化，仅改变类型（FP32→FP16）或做度量预处理 | 否 | 原始距离计算 | P0 | +| ~~**RabitQ**~~ | ~~旋转 + 二进制量化~~ | ~~是（旋转矩阵等）~~ | ~~汉明距离近似~~ | ~~P1~~ | + + +_**注意**__：量化器类型与目标数据类型是正交但有关联的概念。例如 Record 和 Uniform 都可以输出到 INT8 或 INT4；Raw 可以输出到 FP32、FP16、BF16。_ + +## 距离类型支持 +| 距离类型 | 支持说明 | 优先级 | +| --- | --- | --- | +| **L2 (Squared Euclidean)** | 所有量化方案的基础支持 | P0 | +| **Cosine** | 所有量化方案支持，可通过预处理归一化转化为 IP 或 L2 | P0 | +| **Inner Product (IP)** | 所有量化方案支持 | P0 | +| **MIPS** | 基于 IP 的变形，需做转换 | P1 | + + +# 量化器 +## 接口语义 +每个具体的量化器实现以下统一接口，但内部行为因方案而异： + +```cpp +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 IndexError_NotImplemented; + } + + //! 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 IndexError_NotImplemented; + } + + //! 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 IndexError_NotImplemented; + } + + //! Dequantize a result vector back to original format + virtual int dequantize(const void * /*in*/, const IndexQueryMeta & /*qmeta*/, + std::string * /*out*/) const { + return IndexError_NotImplemented; + } + + virtual DistanceImpl distance(const void * /*query*/, + const IndexQueryMeta & /*qmeta*/) const { + return DistanceImpl{}; + } + + //! Serialize quantizer parameters + virtual int serialize(std::string * /*out*/) const { + return IndexError_NotImplemented; + } + + //! Deserialize quantizer parameters + virtual int deserialize(std::string & /*in*/) { + return IndexError_NotImplemented; + } + + //! 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 IndexError_NotImplemented; + } + + 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}; +}; +``` + +dp * dp batch for IVF -> P1 + +## 量化方案 +#### Raw 方案（类型转换/度量预处理） ++ **Extra Meta**：对于Cosine向量增加归一化值。 ++ **编码流程**： + - FP32 → FP16：数值截断/舍入。 + - FP32 → FP32 + Cosine 预处理：向量归一化。 ++ **距离计算**：使用对应精度的原始距离计算（如 FP16 L2、Cosine 转 IP）。 ++ **使用场景**：无需训练，实时转换。不需要训练，支持流式 + +#### Record 方案（逐向量 scale/bias） ++ **Extra Meta**：每个向量后附加 scale、bias相关的数组（或其他压缩形式）。 ++ **编码流程**：输入 FP32 → 计算向量的scale/bias参数 → 量化到 INT8/INT4 → 输出 [codes + scale + bias]。 ++ **距离计算**：非对称，查询向量需经过相同 scale/bias 反量化或预处理到同一域。 ++ **使用场景**：流式插入友好（每条向量独立计算 meta），但 Extra Meta 增加存储。 + +#### Uniform 方案（全局 scale/bias） ++ **Extra Meta**：无 per-vector meta，全局参数存于 Quantizer 内部（通过 serialize/deserialize 持久化）。 ++ **编码流程**：输入 FP32 → 使用全局 scale/bias → 量化到 INT8/INT4 → 纯码本输出。 ++ **距离计算**：对称距离，查询向量和库向量使用同一组全局参数，可高度 SIMD 优化。 ++ **使用场景**：批量训练确定全局参数后，流式插入和检索效率最高。 + +#### PQ / OPQ 方案 ++ **Extra Meta**：无 per-vector meta，全局码本参数（centroids）存于 Quantizer 内部（通过 serialize/deserialize 持久化）。 ++ **预处理：**OPQ有随机旋转预处理 ++ **编码流程**：输入 FP32 → （OPQ 先做正交变换）→ 分段聚类 → 输出码本索引。 ++ **距离计算**：ADC（Asymmetric Distance Computation）~~或 SDC~~。 ++ **使用场景**：高维向量压缩，适合批量构建，流式插入需支持码本增量更新（P2）。 ++ **实现计划**：[PQ in Zvec Turbo](https://yuque.alibaba-inc.com/proxima/aby36h/dt2x3o15c1mbqgxz) + +#### TurboQuant方案 ++ **Extra Meta：**无per-vector meta， ++ **预处理：**需要加随机旋转 ++ **距离计算：**SIMD， INT计算 ++ **使用场景：**极致压缩，保持一定召回。不需要训练，支持流式 + +#### ~~RabitQ 方案~~ ++ **Extra Meta**：旋转矩阵等参数。 ++ **编码流程**：输入 FP32 → 旋转 → 二进制量化。 ++ **距离计算**：基于汉明距离的近似计算。 ++ **使用场景**：极高压缩比，精度要求可容忍的场景。 + + + +## 训练 +量化器训练过程会在合并阶段被调用。量化器根据各自实现，来确定是否需要训练。量化器的训练过程，通过传入一组连续内存的数据，并用步长来指定采样方式，进行训练。 + +# 距离计算器 +为了对距离方式进行封装，设计了距离计算器DistanceImpl类： + +![画板](https://intranetproxy.alipay.com/skylark/lark/0/2026/jpeg/95883/1782875740300-e562d52a-5bfc-485d-b6cf-61e523e3edaf.jpeg) + +# 距离表 +基于上述分类，Turbo 的函数分发不再通过分散的模板特化，而是通过**量化方案 + 目标数据类型 + 距离类型 + SIMD 架构**四个维度组合查询函数表： + +类型定义： + ++ 度量类型MetricType定义： + - kSquaredEuclidean + - kCosine + - kInnerProduct + - kMipsSquaredEuclidean + - kUnknown ++ 数据类型DataType定义 + - kInt4 + - kInt8 + - kFp16 + - kFp32 + - kUnknown ++ 量化类型QuantizeType定义 + - kDefault + - kUniform + - kRecord + - kFp16 + - kFp32 + - kPQ + - kRabit ++ 架构类型CpuArchType定义： + - kAuto + - kScalar + - kSSE + - kAVX + - kAVX2 + - kAVX512 + - kAVX512VNNI + - kAVX512FP16 + +相应的，获取函数的接口也需要枚举对应的参数类型： + ++ DistanceFunc **get_distance_func**(MetricType metric_type, DataType data_type, QuantizeType quantize_type, CpuArchType cpu_arch_type = CpuArchType::kAuto); ++ 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, CpuArchType cpu_arch_type = CpuArchType::kAuto); + +**说明**： + ++ 不存在的组合在编译期报错（static_assert），避免运行时意外。 ++ 运行时通过 Quantizer 返回的 DistanceCalculator 绑定已确定的 scheme/dst_type/metric，仅需March进行动态分发。 + +# 原值构建支持 +由于使用量化值进行图索引构建，会出现不够精确的情况，进一步支持使用原值而非量化值来进行。streamer 和 builder 在实现时严格区分 rdp * rdp （fp32 raw quantizer）; query*dp ，并提供设置原值索引 provider 的接口（之后可考虑和 external_storage 需求中设计的VectorSource 合并） + +术语：query / dp / r(eference)dp + +# 预处理器 +区分概念：预处理操作和预处理器操作，预处理操作包括：旋转、normalize，数据类型偏移等，但并不是所有的预处理操作都需要预处理器来完成，只有那些有状态依赖的操作（例如：旋转）需要使用旋转完成，而normalize，数据类型偏移等操作使用函数指针来完成更加便捷。 + +预处理器是一种量化器的内部可选组件，主要负责在进行数据检索前对数据进行预处理，主要包括旋转功能，可以对量化前的数据进行预处理，以帮助量化器发挥更好的效果。其中，在需要进行预处理时，量化器**主动**调用预处理器，在进行量化前对向量数据进行处理；需要保存时，预处理器需要的数据同量化器中的参数一起进行保存。该组件需要实现的功能包括： + ++ 随机旋转（optional）+ Record INT8/INT4 ++ OPQ旋转（optional）+ PQ ++ 通过某种方式，对数据进行降维（optional） + +职责边界：哪些功能属于预处理范畴，与索引/量化的边界在何处？预处理器是一种**可选**的量化方法的优化组件，缺少也不影响量化方法的正常运行，由量化方案创建的组件属于量化器。 + +预处理器中，处理数据同样是一个费时的操作，尤其是在在线处理的情况下，需要使用**距离表**中的方案进行加速，在获取相应的预处理函数的接口时，也需要枚举对应的参数类型，比如： + ++ FwhtFunc get_fwht_func(CpuArchType cpu_arch_type = CpuArchType::kAuto); ++ MatVecFunc get_matvec_func(CpuArchType cpu_arch_type = CpuArchType::kAuto); + +通过分发的方式可以快速获得对应的预处理函数。具体的实现如下： + +![画板](https://intranetproxy.alipay.com/skylark/lark/0/2026/jpeg/226557004/1782812444145-e4b10b66-8bd1-4f6d-aa94-c205ca57235d.jpeg) + +由于当前FHT可以适用于任意维度的向量，不必再将FHT和Matrix作为rotator下的子类，去除rotator，将FHT和Matrix作为两个独立的类。对于INT8/INT4使用FHT，而对于OPQ使用Matrix，未来还可能有降维等功能。 + +# RabitQ量化接入挑战 +针对当前HnswRabitq索引接入RabitQ量化器，主要有两个方面的挑战： + +1. 距离计算是渐进进行，先计算粗略距离，再进行精确距离计算，这个涉及到在索引遍历时的判断； +2. 返回结果时距离组合，对于流程中的距离返回需要进行改造，涉及到TopK队列的结构体定义； + +结论：暂时在第一期不做接入。 + +# IndexMeta调整 +为支持上述设计，IndexMeta 需增加以下字段： + +| 字段 | 类型 | 说明 | +| --- | --- | --- | +| quantize_type | QuantizeType | 量化方案（Record/Uniform/PQ...） | +| extra_meta_size | uint32_t | 每条向量的 Extra Meta 大小（0 表示无） | + + +兼容性处理： + ++ 存量数据无 quantize_type 字段时，按旧逻辑推断（如根据数据类型和是否有 Extra Meta 推断为 Record 或 Raw）。 ++ 新增字段以扩展字段方式存入 IndexMeta，保证向后兼容。 + +# 代码说明 +## 距离计算器定义 +```cpp +class DistanceImpl { + public: + DistanceImpl() = default; + + DistanceImpl(DistanceFunc func, std::string quantized_query, size_t dim) + : func_(std::move(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)), + 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). Useful for pairwise node-vs-node distance where no + //! stored query is involved. + const DistanceFunc &func() const { + return func_; + } + + //! Raw batch distance function. + const BatchDistanceFunc &batch_func() const { + return batch_func_; + } + + private: + DistanceFunc func_{}; + BatchDistanceFunc batch_func_{}; + std::string query_storage_{}; + size_t dim_{0}; +}; +``` + +## 索引对接示例 +在索引中，通过接入量化器来实现距离对象的接入： + +```cpp +class HnswDistCalculator { + ... .. + private: + zvec::turbo::DistanceImpl dist_impl_{}; +}; +``` + +初始化： + +```cpp +inline void reset_query(const void *query) { + if (quantizer_) { + dist_impl_ = quantizer_->distance(query, qmeta_); + } + ... ... +} +``` + + 距离计算调用： + +```cpp +const auto &func = dist_impl_.func(); +if (func) { + func(vec_lhs, vec_rhs, dist_impl_.dim(), &score); + return score; +} +``` + +## 量化器实现示例 +对于有状态的距离计算，可以在量化器中，通过lambda来捕获对应的状态数据。 + +```cpp + +DistanceImpl MyQuantizer::distance(const void *query, + const IndexQueryMeta &qmeta) const { + // Create a lambda-based DistanceFunc that captures this quantizer + auto dim = padded_dim_; + auto metric = metric_type_; + auto cb_val = cb(); + + DistanceFunc func = [dim, metric, cb_val](const void *m, const void *q, + size_t /*dim_arg*/, float *out) { + const auto *dp_bytes = reinterpret_cast(m); + const auto *q_floats = reinterpret_cast(q); + + const auto *code = reinterpret_cast(dp_bytes); + const float *factors = reinterpret_cast(dp_bytes + dim); + float f_add = factors[0]; + float f_rescale = factors[1]; + + const float *rotated_q = q_floats; + float sum_q = q_floats[dim]; + float norm_sq_q = q_floats[dim + 1]; + + float ip = ip_float_code(rotated_q, code, dim); + float est_dist = f_add + f_rescale * (ip + cb_val * sum_q); + if (metric == MetricType::kSquaredEuclidean) { + est_dist += norm_sq_q; + } + *out = est_dist; + }; + + return DistanceImpl(std::move(func), std::move(qbuf), padded_dim_); +} +``` + +## 预处理示例 +增加预处理来支持向量的前置处理过程，其中旋转也是预处理的一部分。定义如下： + +```cpp +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). + 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 the preprocessor from a contiguous batch of training data. + //! \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 (header + 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; +}; +``` + +对于旋转定义： + +```cpp +class Rotator : public Preprocessor { + public: + using Pointer = std::shared_ptr; + + //! Kind of rotator. + virtual RotatorType type() const = 0; + + // in_dim(), out_dim(), apply(), apply_inverse(), train(), serialize(), + // and deserialize() are inherited from Preprocessor. +}; + +//! Create an untrained rotator of the given type for in_dim dimensions. +Rotator::Pointer CreateRotator(RotatorType type, int in_dim); + +//! Create and restore a rotator from a serialized blob (reads the type from +//! the embedded RotatorSerHeader). Returns nullptr on malformed input. +Rotator::Pointer CreateRotatorFromBlob(const void *data, size_t len); + +} // namespace turbo +``` + +预处理器作为量化器的成员，进行接入。 + +# 实现优先级 +## 一期（P0） ++ **量化方案**： + - Raw（FP32→FP16, FP32→FP32 Cosine） + - Uniform（FP32→INT8, FP32→INT4） + - Record（FP32→INT8, FP32→INT4） ++ **距离类型**：L2, Cosine, IP ++ **微架构类型：**SSE，AVX，AVX512，AVX512FP16 ++ **距离计算模式**：所有单点查询、部分批量、 ++ **索引支持：**HNSW，IVF，Vamana，DiskAnn ++ **使用场景**：流式插入、Merge 训练、流式检索、序列化/反序列化 + +## 二期（P1） ++ **量化方案**： + - PQ（FP32→PQ） ++ **距离类型**：L2, Cosine, IP ++ **微架构类型：**SSE，AVX，AVX512，AVX512FP16 ++ **距离计算模式**：所有批量查询 + +# 模块分工 +在基础PR的基础上，其它模块可以分PR进行提交、验证和合并。 + +| **模块** | **内容** | **人员** | **开始日期** | **结束日期** | **人日** | **优先级** | +| :---: | :---: | :---: | :---: | :---: | :---: | :---: | +| **量化器** | 基础实现 | 弃疾 | 6月22日 | 6月25日 | 3 | P0 | +| | 预处理 | 莊霖 | | | 5 | P0 | +| | turbo分发 | 弃疾 | 6月26日 | 6月29日 | 2 | P0 | +| | FP32量化器 | 荐宁 | | | 3 | P0 | +| | FP16量化器 | 荐宁 | | | 3 | P0 | +| | Int8量化器 | 荐宁 | | | 3 | P0 | +| | Int4量化器 | 荐宁 | | | 3 | P0 | +| | PQ量化器 | 莊霖 | | | 10 | P1 | +| **索引对接** | HNSW | 荐宁 | | | 2 | P0 | +| | IVF | 荐宁 | | | 2 | P0 | +| | Vamana | 荐宁 | | | 2 | P0 | +| | DiskAnn | 弃疾 | 7月6日 | 7月8日 | 2 | P1 | +| **单元测试** | 单元测试 | ALL | | | 5 | P0 | +| **集成测试** | 集成测试 | ALL | | | 5 | P0 | + + +# 结语 ++ 本设计以**量化方案**为核心分类维度，将原本分散在量化器、度量、数据类型中的判断逻辑统一到 Turbo 量化器中。 ++ 确保一期聚焦高频使用场景的 Uniform/Record + INT8/INT4/FP16 + L2/Cosine/IP 组合，再逐步扩展 PQ量化器等复杂方案。 + + + +# 参考 ++ [距离计算能力](https://yuque.alibaba-inc.com/proxima/aby36h/lr69q2egk5ob18cw) ++ [CI&CD 支持集合](https://yuque.alibaba-inc.com/proxima/aby36h/pyvb301f3fs8s3ua) ++ [TurboQuant-近似最优失真率的在线量化方案](https://yuque.alibaba-inc.com/proxima/aby36h/lgllif6v3hht61be) diff --git a/doc/pq.md b/doc/pq.md new file mode 100644 index 000000000..58612e75d --- /dev/null +++ b/doc/pq.md @@ -0,0 +1,94 @@ +## 在turbo中实现preprocessor + +### 1. Preprocessor + +```cpp +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). + 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 the preprocessor from a contiguous batch of training data. + //! \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 (header + 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; +}; +``` +定位：属于/src/turbo/quantizer的一部分，定义在/src/turbo/quantizer/preprocessor.h + +### 2. FhtRotator + +```cpp +class FhtRotator : public Preprocessor { + public: + using Pointer = std::shared_ptr; + + //! Kind of rotator. + virtual RotatorType type() const = 0; + + // in_dim(), out_dim(), apply(), apply_inverse(), train(), serialize(), + // and deserialize() are inherited from Preprocessor. +} + +//! Create an untrained rotator of the given type for in_dim dimensions. +FhtRotator::Pointer CreateFhtRotator(int in_dim); + +//! Create and restore a rotator from a serialized blob (reads the type from +//! the embedded RotatorSerHeader). Returns nullptr on malformed input. +FhtRotator::Pointer CreateFhtRotatorFromBlob(const void *data, size_t len); +``` + +### 3. 数据保存 + +#### 保存格式 +```cpp +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, +``` +定义写在preprocessor.h中 + +在FHT中， +1. rotator_type设置为1，表示FHT旋转 +2. in_dim和out_dim设置为相同， +4. 具体实现参考/root/code/zvec/中fht的实现 + +#### 保存位置 +RotatorSerHeader放到当前已经存在的QuantizerSerHeader中，和他一起保存。具体来说： +1. 量化器调用serialize -> 其中调用preprocessor的serialize +2. serialize后的数据保存到Indexmeta中 +3. 当前方案已经有了保存Indexmeta，这样就可以自然保存 + + +### 4. 分发函数 +在 turbo.h/.cc 中新增函数，将fht的分发功能，也放在其中 diff --git a/src/include/zvec/turbo/turbo.h b/src/include/zvec/turbo/turbo.h index 2eff572d5..c3506f68f 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,6 +35,21 @@ 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); + +// 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; +}; + enum class MetricType { kSquaredEuclidean, kCosine, @@ -90,4 +107,7 @@ QueryPreprocessFunc get_query_preprocess_func( // 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(); + } // namespace zvec::turbo diff --git a/src/turbo/CMakeLists.txt b/src/turbo/CMakeLists.txt index 7ec9c0e5d..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}/distance/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}" ) 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/fht/fht_avx2.cc b/src/turbo/distance/avx2/fht/fht_avx2.cc new file mode 100644 index 000000000..903305931 --- /dev/null +++ b/src/turbo/distance/avx2/fht/fht_avx2.cc @@ -0,0 +1,138 @@ +// 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/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/fht/fht_avx512.cc b/src/turbo/distance/avx512/fht/fht_avx512.cc new file mode 100644 index 000000000..b7ea37d63 --- /dev/null +++ b/src/turbo/distance/avx512/fht/fht_avx512.cc @@ -0,0 +1,143 @@ +// 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/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/fht/fht_scalar.cc b/src/turbo/distance/scalar/fht/fht_scalar.cc new file mode 100644 index 000000000..008c6a2a6 --- /dev/null +++ b/src/turbo/distance/scalar/fht/fht_scalar.cc @@ -0,0 +1,80 @@ +// 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/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/distance/sse/fht/fht_sse.cc b/src/turbo/distance/sse/fht/fht_sse.cc new file mode 100644 index 000000000..7dbbf1951 --- /dev/null +++ b/src/turbo/distance/sse/fht/fht_sse.cc @@ -0,0 +1,119 @@ +// 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/quantizer/preprocessor/fht_rotator.cc b/src/turbo/quantizer/preprocessor/fht_rotator.cc new file mode 100644 index 000000000..94a19a916 --- /dev/null +++ b/src/turbo/quantizer/preprocessor/fht_rotator.cc @@ -0,0 +1,266 @@ +// 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 + +namespace zvec { +namespace turbo { + +using core::IndexError_InvalidArgument; +using core::IndexError_Runtime; +using core::IndexError_Unsupported; + +// ============================================================================ +// 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; + // flip_ stays empty until train() is called. + 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*/) { + if (in_dim_ <= 0) return; + + flip_offset_ = (static_cast(in_dim_) + kByteLen - 1) / kByteLen; + flip_.resize(4 * flip_offset_); + + std::random_device rd; + std::mt19937 gen(rd()); + std::uniform_int_distribution dist(0, 255); + for (auto &b : flip_) b = static_cast(dist(gen)); +} + +// --------------------------------------------------------------------------- +// 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 IndexError_InvalidArgument; + if (flip_.empty()) return IndexError_Runtime; + + 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 IndexError_InvalidArgument; + + const auto *hdr = reinterpret_cast(data); + if (hdr->magic != kRotatorMagic) return IndexError_Unsupported; + if (hdr->version != kRotatorSerVersion) return IndexError_Unsupported; + if (static_cast(hdr->rotator_type) != RotatorType::kFht) { + return IndexError_Unsupported; + } + + const size_t total = sizeof(RotatorSerHeader) + hdr->payload_size; + if (len < total) return IndexError_InvalidArgument; + + 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..15dfbb0c4 --- /dev/null +++ b/src/turbo/quantizer/preprocessor/fht_rotator.h @@ -0,0 +1,97 @@ +// 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 "preprocessor.h" +#include + +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 an untrained rotator for \p in_dim dimensions. + //! Call train() afterwards to generate the random flip-sign arrays. + 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; + + //! Generate 4 rounds of random flip-sign arrays. + //! For FhtRotator the training data is ignored; only \p in_dim() matters. + 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/turbo.cc b/src/turbo/turbo.cc index 80f2ee77d..8e2e1b617 100644 --- a/src/turbo/turbo.cc +++ b/src/turbo/turbo.cc @@ -18,10 +18,21 @@ #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" +#if defined(__SSE2__) +#include "sse/fht/fht.h" +#endif +#if defined(__AVX2__) +#include "avx2/fht/fht.h" +#endif +#if defined(__AVX512F__) +#include "avx512/fht/fht.h" +#endif + namespace zvec::turbo { DistanceFunc get_distance_func(MetricType metric_type, DataType data_type, @@ -143,4 +154,47 @@ UniformQuantizeFunc get_uniform_quantize_func(DataType data_type) { return nullptr; } +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 +} + } // namespace zvec::turbo diff --git a/tests/turbo/turbo_fht_rotator_test.cc b/tests/turbo/turbo_fht_rotator_test.cc new file mode 100644 index 000000000..178dad3fd --- /dev/null +++ b/tests/turbo/turbo_fht_rotator_test.cc @@ -0,0 +1,204 @@ +// 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; + rot->train(nullptr, 0, 0); + + 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; + rot->train(nullptr, 0, 0); + + 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 and train original rotator. + auto rot = FhtRotator::create(dim); + ASSERT_TRUE(rot); + rot->train(nullptr, 0, 0); + + // 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, TrainGeneratesFlip) { + for (int dim : {8, 64, 97}) { + auto rot = FhtRotator::create(dim); + ASSERT_TRUE(rot); + + // Before train, apply should not crash but flip is empty — we skip + // calling apply before train. After train, serialize must succeed + // (which requires flip to be populated). + rot->train(nullptr, 0, 0); + + std::string blob; + EXPECT_EQ(0, rot->serialize(&blob)) + << "serialize failed after train 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); +} From 437a7546c7df926ab748ce129f2cfa24083b0edc Mon Sep 17 00:00:00 2001 From: ray Date: Thu, 2 Jul 2026 16:39:13 +0800 Subject: [PATCH 18/42] fix: fix symbol --- .../quantizer/fp32_quantizer/fp32_quantizer.cc | 3 +-- src/turbo/quantizer/quantizer.h | 14 ++++++++++++++ 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.cc b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.cc index f5c29b76b..839137e6d 100644 --- a/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.cc +++ b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.cc @@ -17,7 +17,6 @@ #include #include #include -#include #include #include #include "core/quantizer/record_quantizer.h" @@ -53,7 +52,7 @@ int Fp32Quantizer::init(const IndexMeta &meta, int Fp32Quantizer::quantize(const void *query, const IndexQueryMeta &qmeta, std::string *out, IndexQueryMeta *ometa) const { if (qmeta.unit_size() != sizeof(float)) { - return IndexError_Unsupported; + return kErrUnsupported; } // qmeta.dimension() may be the inflated (data + extras) dimension when the diff --git a/src/turbo/quantizer/quantizer.h b/src/turbo/quantizer/quantizer.h index 3e5e65ec6..3b7a7d944 100644 --- a/src/turbo/quantizer/quantizer.h +++ b/src/turbo/quantizer/quantizer.h @@ -28,6 +28,20 @@ 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 kErrNotImplemented = -11; +constexpr int kErrUnsupported = -12; + //! Magic number ('QTZR') stamped at the start of a serialized quantizer blob. constexpr uint32_t kQuantizerMagic = 0x52545A51u; From e16cda561e2630b10bd805ccca3aa3c9e54c8089 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Thu, 2 Jul 2026 16:47:42 +0800 Subject: [PATCH 19/42] add untest --- tests/turbo/turbo_fht_rotator_test.cc | 255 ++++++++++++++++++++++++++ 1 file changed, 255 insertions(+) diff --git a/tests/turbo/turbo_fht_rotator_test.cc b/tests/turbo/turbo_fht_rotator_test.cc index 178dad3fd..6acdaba80 100644 --- a/tests/turbo/turbo_fht_rotator_test.cc +++ b/tests/turbo/turbo_fht_rotator_test.cc @@ -202,3 +202,258 @@ TEST(FhtRotator, FromBlobMalformed) { 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); + rot->train(nullptr, 0, 0); + + 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); + rot->train(nullptr, 0, 0); + + 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); + rot->train(nullptr, 0, 0); + + 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); + rot->train(nullptr, 0, 0); + + 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 and train original. + auto rot1 = FhtRotator::create(dim); + ASSERT_TRUE(rot1); + rot1->train(nullptr, 0, 0); + + 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); + rot->train(nullptr, 0, 0); + + 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; + rot->train(nullptr, 0, 0); + + 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); + rot->train(nullptr, 0, 0); + + 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; + } +} From f3a67820ee31c4288614ff7d0ed3beaa9d9b7b60 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Thu, 2 Jul 2026 17:34:32 +0800 Subject: [PATCH 20/42] remove extra doc --- .gitignore | 1 + ...26\345\231\250\350\256\276\350\256\241.md" | 731 ------------------ doc/pq.md | 94 --- 3 files changed, 1 insertion(+), 825 deletions(-) delete mode 100644 "doc/Zvec Turbo \351\207\217\345\214\226\345\231\250\350\256\276\350\256\241.md" delete mode 100644 doc/pq.md diff --git a/.gitignore b/.gitignore index b2fd51657..7df6fa939 100644 --- a/.gitignore +++ b/.gitignore @@ -60,3 +60,4 @@ allure-* !build_android.sh !build_ios.sh +doc/ \ No newline at end of file diff --git "a/doc/Zvec Turbo \351\207\217\345\214\226\345\231\250\350\256\276\350\256\241.md" "b/doc/Zvec Turbo \351\207\217\345\214\226\345\231\250\350\256\276\350\256\241.md" deleted file mode 100644 index 66f02415e..000000000 --- "a/doc/Zvec Turbo \351\207\217\345\214\226\345\231\250\350\256\276\350\256\241.md" +++ /dev/null @@ -1,731 +0,0 @@ -# 背景 -当前，距离度量沿用了 Proxima 的风格，采用模板方式进行定义。系统通过模块化设计实现了度量、转换器及指令集等组件的解耦，检索过程由Metric（距离度量）、Quantizer（量化器）、Algorithm（索引算法）等多个模块协作完成。然而，一些现有的组件存在职责边界模糊、实现方式耦合等问题，计划通过改为turbo的形式，统一切面通过输入多个维度信息组合以获得对应的函数集合，从而使得不同逻辑判断统一。 - -# 设计现状 -当前的系统架构设计存在着一些问题，下面将进行简要说明。 - -## Converter和Reformer功能耦合 -当前的量化器由Converter和Reformer两个组件构成，其中，Converter主要负责离线构建，而Reformer主要负责在线查询，不过由于当前设计未区分流式插入与批量合并，Converter和Reformer在离线阶段存在一定的功能耦合： - -![画板](https://intranetproxy.alipay.com/skylark/lark/0/2026/jpeg/226557004/1782716498906-f0957db9-24f8-466b-bfb2-351e65d1b5c5.jpeg) - -这个实现最初的想法是将离线的一次性操作`build`、`train`等放到Converter中，而可以反复进行的操作`add`和`search`等放到Reformer，但是，这在本身也引入了一些额外的限制，例如，如果量化器需要训练，但索引是流式插入的，例如hnsw+kUniformInt8，就会存在问题。以下对当前的使用接口进行说明： - -| **Converter接口** | **功能** | **Reformer接口** | **功能** | -| --- | --- | --- | --- | -| init() | 声明量化后的 meta 信息 | init() | 从 Converter 声明的 reformer_params 中加载运行时参数 | -| train() | 在全量数据上执行训练，收集统计信息 | load() | 从持久化存储中加载/卸载 Reformer 运行时状态 | -| transform() | 将 全量向量从原始格式转换为量化格式 | transform() | 将查询向量转换为与索引一致的量化格式 | -| dump() | 将 Converter 内部状态（如量化参数）持久化到 dumper | normalize() | 将量化空间计算的距离/分数映射回原始距离尺度 | -| meta() | 返回 Converter 修改后的 IndexMeta | revert() | 将量化向量反量化回 FP32 | - - -## Quantizer与Metric存在耦合 -当前设计存在不同功能分布在同一文件/同一功能分布在不同文件的问题，例如： - -| 注册文件 | 工厂注册名 | Metric | 备注 | -| --- | --- | --- | --- | -| **integer_quantizer_converter.cc** | Int8StreamingConverter | L2 / IP / MIPS | 基于单条数据的INT8 | -| | Int4StreamingConverter | L2 / IP / MIPS | 基于单条数据的INT4 | -| | Int8QuantizerConverter | - | 全局数据的INT8 | -| | Int4QuantizerConverter | - | 全局数据的INT4 | -| **cosine_converter.cc** | CosineInt8Converter | Cosine | 基于单条数据的INT8 | -| | CosineInt4Converter | Cosine | 基于单条数据的INT4 | -| | CosineFp16Converter | Cosine | FP16 | -| | ... | ... | ... | - - -由于这样的设计，一些量化器内部需要判断度量类型、数据类型、架构类型，以合理分配量化类型。 - -## 其他问题 -1. Builder、Searcher、Streamer职责边界模糊： - 1. Builder：全量数据的构建，负责操作 `build` - 2. Searcher：只被 IndexFlow（CLI 工具）和单元测试使用 - 3. Streamer：主路径组件负责读写，负责操作 `add` 和`search` -2. 部分索引实现与量化器耦合： - 1. 与Quantizer完全解耦：HNSW / Flat / Vamana - 2. 硬编码与Quantizer耦合：IVF - 3. 当前自成一体：diskann/hnsw rabitq - -# Turbo设计 -Turbo 旨在通过统一接口收敛度量（Metric）与量化（Quantization）逻辑，提高代码内聚性，从而降低维护成本，提升代码的可读性与简洁性。 - -## 使用场景 -![画板](https://intranetproxy.alipay.com/skylark/lark/0/2026/jpeg/95883/1782788752754-8c017988-c7ed-469d-8025-b1ab5ce94281.jpeg) - -## 模块 -Turbo系统包含以下几个模块： - -+ 量化器 -+ 预处理器 -+ 距离计算 - -都以量化器为核心进行调度： - -![画板](https://intranetproxy.alipay.com/skylark/lark/0/2026/jpeg/95883/1782802278889-9e04f816-ec15-42f6-9ff6-3e53ceb7c83e.jpeg) - -## 设计目标 -设计的目标主要包含以下六个维度的： - -| **维度** | **说明** | -| :---: | :---: | -| 源数据类型 | 量化前的数据类型 | -| 目标数据类型 | 量化后的数据类型 | -| 距离类型 | 支持的距离度量方式 | -| 量化器类型 | 具体量化算法及参数管理方式 | -| 距离计算模式 | 单查询单点、单查询多点连续、单查询多点离散、点对点 | -| 指令架构 | 该量化器支持的指令架构 | -| 量化器使用场景 | 训练、流式插入、批量插入/合并、序列化、反序列化 | - - -## 调用场景 -以下为构建、查询的调用场景： - -### 构建 -```plain - 源向量 - │ - ▼ -┌───────────────────────┐ -│ Quantizer │ -│ .quantize_datapoint │ -└───────────────────────┘ - │ - ▼ -[量化向量 + Extra Meta(如有)] → 写入 Streamer -``` - -要求： - -+ Uniform：无需训练即可编码（全局参数已确定）。 -+ Record：逐向量独立计算 scale/bias，无需全局训练。 -+ PQ/RabitQ：不支持纯流式插入（需先训练）。 - -### 查询 -```plain - 源向量 - │ - ▼ -┌─────────────────────┐ -│ Quantizer │ -│ .quantize_query() │ -└─────────────────────┘ - │ - ▼ -[量化向量 + Extra Meta(如有)] → 检索 Streamer -``` - -### 合并 -```plain -多个源 Segment (FP32) - │ - ▼ -┌─────────────┐ -│ write_holder│ ← 聚合数据到Holder -└─────────────┘ - │ - ▼ -┌─────────────┐ -│ Quantizer │ ← 训练全局参数 -│ .train │ -└─────────────┘ - │ - ▼ -┌───────────────────────┐ -│ Quantizer │ ← 用训练后的Quantizer编码所有向量 -│ .quantize_datapoint | -└───────────────────────┘ - │ - ▼ -[量化向量 + Extra Meta(如有)] → 写入 Streamer -``` - -# 支持范围 -## 源数据类型 → 目标数据类型 -考虑到 db 层有不保留 Flat FP32 的需求 - -+ 源数据类型需要支持 4 种通用类型： - - FP32，P0 - - FP16，P0 - - INT8，P1 - - INT4，P2 -+ 目标数据类型支持范围： - -| **目标类型** | **说明** | **优先级** | -| :---: | :---: | :---: | -| FP32 | 原始数据不做量化，仅做度量转换（如 Cosine） | P0 | -| FP16 | 半精度浮点 | P0 | -| INT8 | 8-bit 整数量化 | P0 | -| INT4 | 4-bit 整数量化 | P0 | -| X-Bits | 自定义 bit 宽度的量化（如 5-bit、6-bit） | P1 | - - -## 量化器类型 -量化器类型是核心分类维度，决定了端到端的数据布局、训练方式和距离计算方式。 - -| **量化器类型** | **特征** | **是否带Extra Meta** | **距离计算方式** | **优先级** | -| :---: | :---: | :---: | :---: | :---: | -| **Record** | 逐向量保存 scale/bias，按维度量化 | 是
（Extra Meta 存 scale/bias） | 非对称距离，需反量化或查表 | P0 | -| **Uniform** | 全局统一scale & bias，按维度量化 | 否 | 对称距离，可 SIMD 优化 | P0 | -| **PQ**
** (Product Quantization)** | 分段聚类码本量化 | 是
（码本参数） | 非对称距离，ADC/SDC | P1 | -| **OPQ** | 优化的 PQ，含正交变换 | 是
（变换矩阵 + 码本） | 非对称距离，ADC | P2 | -| **Raw** | 不做数值量化，仅改变类型（FP32→FP16）或做度量预处理 | 否 | 原始距离计算 | P0 | -| ~~**RabitQ**~~ | ~~旋转 + 二进制量化~~ | ~~是（旋转矩阵等）~~ | ~~汉明距离近似~~ | ~~P1~~ | - - -_**注意**__：量化器类型与目标数据类型是正交但有关联的概念。例如 Record 和 Uniform 都可以输出到 INT8 或 INT4；Raw 可以输出到 FP32、FP16、BF16。_ - -## 距离类型支持 -| 距离类型 | 支持说明 | 优先级 | -| --- | --- | --- | -| **L2 (Squared Euclidean)** | 所有量化方案的基础支持 | P0 | -| **Cosine** | 所有量化方案支持，可通过预处理归一化转化为 IP 或 L2 | P0 | -| **Inner Product (IP)** | 所有量化方案支持 | P0 | -| **MIPS** | 基于 IP 的变形，需做转换 | P1 | - - -# 量化器 -## 接口语义 -每个具体的量化器实现以下统一接口，但内部行为因方案而异： - -```cpp -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 IndexError_NotImplemented; - } - - //! 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 IndexError_NotImplemented; - } - - //! 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 IndexError_NotImplemented; - } - - //! Dequantize a result vector back to original format - virtual int dequantize(const void * /*in*/, const IndexQueryMeta & /*qmeta*/, - std::string * /*out*/) const { - return IndexError_NotImplemented; - } - - virtual DistanceImpl distance(const void * /*query*/, - const IndexQueryMeta & /*qmeta*/) const { - return DistanceImpl{}; - } - - //! Serialize quantizer parameters - virtual int serialize(std::string * /*out*/) const { - return IndexError_NotImplemented; - } - - //! Deserialize quantizer parameters - virtual int deserialize(std::string & /*in*/) { - return IndexError_NotImplemented; - } - - //! 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 IndexError_NotImplemented; - } - - 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}; -}; -``` - -dp * dp batch for IVF -> P1 - -## 量化方案 -#### Raw 方案（类型转换/度量预处理） -+ **Extra Meta**：对于Cosine向量增加归一化值。 -+ **编码流程**： - - FP32 → FP16：数值截断/舍入。 - - FP32 → FP32 + Cosine 预处理：向量归一化。 -+ **距离计算**：使用对应精度的原始距离计算（如 FP16 L2、Cosine 转 IP）。 -+ **使用场景**：无需训练，实时转换。不需要训练，支持流式 - -#### Record 方案（逐向量 scale/bias） -+ **Extra Meta**：每个向量后附加 scale、bias相关的数组（或其他压缩形式）。 -+ **编码流程**：输入 FP32 → 计算向量的scale/bias参数 → 量化到 INT8/INT4 → 输出 [codes + scale + bias]。 -+ **距离计算**：非对称，查询向量需经过相同 scale/bias 反量化或预处理到同一域。 -+ **使用场景**：流式插入友好（每条向量独立计算 meta），但 Extra Meta 增加存储。 - -#### Uniform 方案（全局 scale/bias） -+ **Extra Meta**：无 per-vector meta，全局参数存于 Quantizer 内部（通过 serialize/deserialize 持久化）。 -+ **编码流程**：输入 FP32 → 使用全局 scale/bias → 量化到 INT8/INT4 → 纯码本输出。 -+ **距离计算**：对称距离，查询向量和库向量使用同一组全局参数，可高度 SIMD 优化。 -+ **使用场景**：批量训练确定全局参数后，流式插入和检索效率最高。 - -#### PQ / OPQ 方案 -+ **Extra Meta**：无 per-vector meta，全局码本参数（centroids）存于 Quantizer 内部（通过 serialize/deserialize 持久化）。 -+ **预处理：**OPQ有随机旋转预处理 -+ **编码流程**：输入 FP32 → （OPQ 先做正交变换）→ 分段聚类 → 输出码本索引。 -+ **距离计算**：ADC（Asymmetric Distance Computation）~~或 SDC~~。 -+ **使用场景**：高维向量压缩，适合批量构建，流式插入需支持码本增量更新（P2）。 -+ **实现计划**：[PQ in Zvec Turbo](https://yuque.alibaba-inc.com/proxima/aby36h/dt2x3o15c1mbqgxz) - -#### TurboQuant方案 -+ **Extra Meta：**无per-vector meta， -+ **预处理：**需要加随机旋转 -+ **距离计算：**SIMD， INT计算 -+ **使用场景：**极致压缩，保持一定召回。不需要训练，支持流式 - -#### ~~RabitQ 方案~~ -+ **Extra Meta**：旋转矩阵等参数。 -+ **编码流程**：输入 FP32 → 旋转 → 二进制量化。 -+ **距离计算**：基于汉明距离的近似计算。 -+ **使用场景**：极高压缩比，精度要求可容忍的场景。 - - - -## 训练 -量化器训练过程会在合并阶段被调用。量化器根据各自实现，来确定是否需要训练。量化器的训练过程，通过传入一组连续内存的数据，并用步长来指定采样方式，进行训练。 - -# 距离计算器 -为了对距离方式进行封装，设计了距离计算器DistanceImpl类： - -![画板](https://intranetproxy.alipay.com/skylark/lark/0/2026/jpeg/95883/1782875740300-e562d52a-5bfc-485d-b6cf-61e523e3edaf.jpeg) - -# 距离表 -基于上述分类，Turbo 的函数分发不再通过分散的模板特化，而是通过**量化方案 + 目标数据类型 + 距离类型 + SIMD 架构**四个维度组合查询函数表： - -类型定义： - -+ 度量类型MetricType定义： - - kSquaredEuclidean - - kCosine - - kInnerProduct - - kMipsSquaredEuclidean - - kUnknown -+ 数据类型DataType定义 - - kInt4 - - kInt8 - - kFp16 - - kFp32 - - kUnknown -+ 量化类型QuantizeType定义 - - kDefault - - kUniform - - kRecord - - kFp16 - - kFp32 - - kPQ - - kRabit -+ 架构类型CpuArchType定义： - - kAuto - - kScalar - - kSSE - - kAVX - - kAVX2 - - kAVX512 - - kAVX512VNNI - - kAVX512FP16 - -相应的，获取函数的接口也需要枚举对应的参数类型： - -+ DistanceFunc **get_distance_func**(MetricType metric_type, DataType data_type, QuantizeType quantize_type, CpuArchType cpu_arch_type = CpuArchType::kAuto); -+ 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, CpuArchType cpu_arch_type = CpuArchType::kAuto); - -**说明**： - -+ 不存在的组合在编译期报错（static_assert），避免运行时意外。 -+ 运行时通过 Quantizer 返回的 DistanceCalculator 绑定已确定的 scheme/dst_type/metric，仅需March进行动态分发。 - -# 原值构建支持 -由于使用量化值进行图索引构建，会出现不够精确的情况，进一步支持使用原值而非量化值来进行。streamer 和 builder 在实现时严格区分 rdp * rdp （fp32 raw quantizer）; query*dp ，并提供设置原值索引 provider 的接口（之后可考虑和 external_storage 需求中设计的VectorSource 合并） - -术语：query / dp / r(eference)dp - -# 预处理器 -区分概念：预处理操作和预处理器操作，预处理操作包括：旋转、normalize，数据类型偏移等，但并不是所有的预处理操作都需要预处理器来完成，只有那些有状态依赖的操作（例如：旋转）需要使用旋转完成，而normalize，数据类型偏移等操作使用函数指针来完成更加便捷。 - -预处理器是一种量化器的内部可选组件，主要负责在进行数据检索前对数据进行预处理，主要包括旋转功能，可以对量化前的数据进行预处理，以帮助量化器发挥更好的效果。其中，在需要进行预处理时，量化器**主动**调用预处理器，在进行量化前对向量数据进行处理；需要保存时，预处理器需要的数据同量化器中的参数一起进行保存。该组件需要实现的功能包括： - -+ 随机旋转（optional）+ Record INT8/INT4 -+ OPQ旋转（optional）+ PQ -+ 通过某种方式，对数据进行降维（optional） - -职责边界：哪些功能属于预处理范畴，与索引/量化的边界在何处？预处理器是一种**可选**的量化方法的优化组件，缺少也不影响量化方法的正常运行，由量化方案创建的组件属于量化器。 - -预处理器中，处理数据同样是一个费时的操作，尤其是在在线处理的情况下，需要使用**距离表**中的方案进行加速，在获取相应的预处理函数的接口时，也需要枚举对应的参数类型，比如： - -+ FwhtFunc get_fwht_func(CpuArchType cpu_arch_type = CpuArchType::kAuto); -+ MatVecFunc get_matvec_func(CpuArchType cpu_arch_type = CpuArchType::kAuto); - -通过分发的方式可以快速获得对应的预处理函数。具体的实现如下： - -![画板](https://intranetproxy.alipay.com/skylark/lark/0/2026/jpeg/226557004/1782812444145-e4b10b66-8bd1-4f6d-aa94-c205ca57235d.jpeg) - -由于当前FHT可以适用于任意维度的向量，不必再将FHT和Matrix作为rotator下的子类，去除rotator，将FHT和Matrix作为两个独立的类。对于INT8/INT4使用FHT，而对于OPQ使用Matrix，未来还可能有降维等功能。 - -# RabitQ量化接入挑战 -针对当前HnswRabitq索引接入RabitQ量化器，主要有两个方面的挑战： - -1. 距离计算是渐进进行，先计算粗略距离，再进行精确距离计算，这个涉及到在索引遍历时的判断； -2. 返回结果时距离组合，对于流程中的距离返回需要进行改造，涉及到TopK队列的结构体定义； - -结论：暂时在第一期不做接入。 - -# IndexMeta调整 -为支持上述设计，IndexMeta 需增加以下字段： - -| 字段 | 类型 | 说明 | -| --- | --- | --- | -| quantize_type | QuantizeType | 量化方案（Record/Uniform/PQ...） | -| extra_meta_size | uint32_t | 每条向量的 Extra Meta 大小（0 表示无） | - - -兼容性处理： - -+ 存量数据无 quantize_type 字段时，按旧逻辑推断（如根据数据类型和是否有 Extra Meta 推断为 Record 或 Raw）。 -+ 新增字段以扩展字段方式存入 IndexMeta，保证向后兼容。 - -# 代码说明 -## 距离计算器定义 -```cpp -class DistanceImpl { - public: - DistanceImpl() = default; - - DistanceImpl(DistanceFunc func, std::string quantized_query, size_t dim) - : func_(std::move(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)), - 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). Useful for pairwise node-vs-node distance where no - //! stored query is involved. - const DistanceFunc &func() const { - return func_; - } - - //! Raw batch distance function. - const BatchDistanceFunc &batch_func() const { - return batch_func_; - } - - private: - DistanceFunc func_{}; - BatchDistanceFunc batch_func_{}; - std::string query_storage_{}; - size_t dim_{0}; -}; -``` - -## 索引对接示例 -在索引中，通过接入量化器来实现距离对象的接入： - -```cpp -class HnswDistCalculator { - ... .. - private: - zvec::turbo::DistanceImpl dist_impl_{}; -}; -``` - -初始化： - -```cpp -inline void reset_query(const void *query) { - if (quantizer_) { - dist_impl_ = quantizer_->distance(query, qmeta_); - } - ... ... -} -``` - - 距离计算调用： - -```cpp -const auto &func = dist_impl_.func(); -if (func) { - func(vec_lhs, vec_rhs, dist_impl_.dim(), &score); - return score; -} -``` - -## 量化器实现示例 -对于有状态的距离计算，可以在量化器中，通过lambda来捕获对应的状态数据。 - -```cpp - -DistanceImpl MyQuantizer::distance(const void *query, - const IndexQueryMeta &qmeta) const { - // Create a lambda-based DistanceFunc that captures this quantizer - auto dim = padded_dim_; - auto metric = metric_type_; - auto cb_val = cb(); - - DistanceFunc func = [dim, metric, cb_val](const void *m, const void *q, - size_t /*dim_arg*/, float *out) { - const auto *dp_bytes = reinterpret_cast(m); - const auto *q_floats = reinterpret_cast(q); - - const auto *code = reinterpret_cast(dp_bytes); - const float *factors = reinterpret_cast(dp_bytes + dim); - float f_add = factors[0]; - float f_rescale = factors[1]; - - const float *rotated_q = q_floats; - float sum_q = q_floats[dim]; - float norm_sq_q = q_floats[dim + 1]; - - float ip = ip_float_code(rotated_q, code, dim); - float est_dist = f_add + f_rescale * (ip + cb_val * sum_q); - if (metric == MetricType::kSquaredEuclidean) { - est_dist += norm_sq_q; - } - *out = est_dist; - }; - - return DistanceImpl(std::move(func), std::move(qbuf), padded_dim_); -} -``` - -## 预处理示例 -增加预处理来支持向量的前置处理过程，其中旋转也是预处理的一部分。定义如下： - -```cpp -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). - 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 the preprocessor from a contiguous batch of training data. - //! \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 (header + 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; -}; -``` - -对于旋转定义： - -```cpp -class Rotator : public Preprocessor { - public: - using Pointer = std::shared_ptr; - - //! Kind of rotator. - virtual RotatorType type() const = 0; - - // in_dim(), out_dim(), apply(), apply_inverse(), train(), serialize(), - // and deserialize() are inherited from Preprocessor. -}; - -//! Create an untrained rotator of the given type for in_dim dimensions. -Rotator::Pointer CreateRotator(RotatorType type, int in_dim); - -//! Create and restore a rotator from a serialized blob (reads the type from -//! the embedded RotatorSerHeader). Returns nullptr on malformed input. -Rotator::Pointer CreateRotatorFromBlob(const void *data, size_t len); - -} // namespace turbo -``` - -预处理器作为量化器的成员，进行接入。 - -# 实现优先级 -## 一期（P0） -+ **量化方案**： - - Raw（FP32→FP16, FP32→FP32 Cosine） - - Uniform（FP32→INT8, FP32→INT4） - - Record（FP32→INT8, FP32→INT4） -+ **距离类型**：L2, Cosine, IP -+ **微架构类型：**SSE，AVX，AVX512，AVX512FP16 -+ **距离计算模式**：所有单点查询、部分批量、 -+ **索引支持：**HNSW，IVF，Vamana，DiskAnn -+ **使用场景**：流式插入、Merge 训练、流式检索、序列化/反序列化 - -## 二期（P1） -+ **量化方案**： - - PQ（FP32→PQ） -+ **距离类型**：L2, Cosine, IP -+ **微架构类型：**SSE，AVX，AVX512，AVX512FP16 -+ **距离计算模式**：所有批量查询 - -# 模块分工 -在基础PR的基础上，其它模块可以分PR进行提交、验证和合并。 - -| **模块** | **内容** | **人员** | **开始日期** | **结束日期** | **人日** | **优先级** | -| :---: | :---: | :---: | :---: | :---: | :---: | :---: | -| **量化器** | 基础实现 | 弃疾 | 6月22日 | 6月25日 | 3 | P0 | -| | 预处理 | 莊霖 | | | 5 | P0 | -| | turbo分发 | 弃疾 | 6月26日 | 6月29日 | 2 | P0 | -| | FP32量化器 | 荐宁 | | | 3 | P0 | -| | FP16量化器 | 荐宁 | | | 3 | P0 | -| | Int8量化器 | 荐宁 | | | 3 | P0 | -| | Int4量化器 | 荐宁 | | | 3 | P0 | -| | PQ量化器 | 莊霖 | | | 10 | P1 | -| **索引对接** | HNSW | 荐宁 | | | 2 | P0 | -| | IVF | 荐宁 | | | 2 | P0 | -| | Vamana | 荐宁 | | | 2 | P0 | -| | DiskAnn | 弃疾 | 7月6日 | 7月8日 | 2 | P1 | -| **单元测试** | 单元测试 | ALL | | | 5 | P0 | -| **集成测试** | 集成测试 | ALL | | | 5 | P0 | - - -# 结语 -+ 本设计以**量化方案**为核心分类维度，将原本分散在量化器、度量、数据类型中的判断逻辑统一到 Turbo 量化器中。 -+ 确保一期聚焦高频使用场景的 Uniform/Record + INT8/INT4/FP16 + L2/Cosine/IP 组合，再逐步扩展 PQ量化器等复杂方案。 - - - -# 参考 -+ [距离计算能力](https://yuque.alibaba-inc.com/proxima/aby36h/lr69q2egk5ob18cw) -+ [CI&CD 支持集合](https://yuque.alibaba-inc.com/proxima/aby36h/pyvb301f3fs8s3ua) -+ [TurboQuant-近似最优失真率的在线量化方案](https://yuque.alibaba-inc.com/proxima/aby36h/lgllif6v3hht61be) diff --git a/doc/pq.md b/doc/pq.md deleted file mode 100644 index 58612e75d..000000000 --- a/doc/pq.md +++ /dev/null @@ -1,94 +0,0 @@ -## 在turbo中实现preprocessor - -### 1. Preprocessor - -```cpp -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). - 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 the preprocessor from a contiguous batch of training data. - //! \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 (header + 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; -}; -``` -定位：属于/src/turbo/quantizer的一部分，定义在/src/turbo/quantizer/preprocessor.h - -### 2. FhtRotator - -```cpp -class FhtRotator : public Preprocessor { - public: - using Pointer = std::shared_ptr; - - //! Kind of rotator. - virtual RotatorType type() const = 0; - - // in_dim(), out_dim(), apply(), apply_inverse(), train(), serialize(), - // and deserialize() are inherited from Preprocessor. -} - -//! Create an untrained rotator of the given type for in_dim dimensions. -FhtRotator::Pointer CreateFhtRotator(int in_dim); - -//! Create and restore a rotator from a serialized blob (reads the type from -//! the embedded RotatorSerHeader). Returns nullptr on malformed input. -FhtRotator::Pointer CreateFhtRotatorFromBlob(const void *data, size_t len); -``` - -### 3. 数据保存 - -#### 保存格式 -```cpp -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, -``` -定义写在preprocessor.h中 - -在FHT中， -1. rotator_type设置为1，表示FHT旋转 -2. in_dim和out_dim设置为相同， -4. 具体实现参考/root/code/zvec/中fht的实现 - -#### 保存位置 -RotatorSerHeader放到当前已经存在的QuantizerSerHeader中，和他一起保存。具体来说： -1. 量化器调用serialize -> 其中调用preprocessor的serialize -2. serialize后的数据保存到Indexmeta中 -3. 当前方案已经有了保存Indexmeta，这样就可以自然保存 - - -### 4. 分发函数 -在 turbo.h/.cc 中新增函数，将fht的分发功能，也放在其中 From 81d8eb5bb3fd4d55308e54f5ea536665e94bc151 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Thu, 2 Jul 2026 19:49:11 +0800 Subject: [PATCH 21/42] tmp --- .../algorithm/hnsw/hnsw_dist_calculator.h | 19 +- src/include/zvec/turbo/turbo.h | 28 + src/turbo/distance/scalar/pq/pq_scalar.cc | 44 ++ src/turbo/distance/scalar/pq/pq_scalar.h | 43 ++ src/turbo/quantizer/distance.h | 20 + .../pq_int8_quantizer/pq_int8_quantizer.cc | 491 ++++++++++++++++++ .../pq_int8_quantizer/pq_int8_quantizer.h | 146 ++++++ src/turbo/turbo.cc | 15 +- tests/turbo/turbo_pq_int8_quantizer_test.cc | 288 ++++++++++ 9 files changed, 1092 insertions(+), 2 deletions(-) create mode 100644 src/turbo/distance/scalar/pq/pq_scalar.cc create mode 100644 src/turbo/distance/scalar/pq/pq_scalar.h create mode 100644 src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc create mode 100644 src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h create mode 100644 tests/turbo/turbo_pq_int8_quantizer_test.cc diff --git a/src/core/algorithm/hnsw/hnsw_dist_calculator.h b/src/core/algorithm/hnsw/hnsw_dist_calculator.h index 2e4b22d1f..f11b064c4 100644 --- a/src/core/algorithm/hnsw/hnsw_dist_calculator.h +++ b/src/core/algorithm/hnsw/hnsw_dist_calculator.h @@ -84,6 +84,16 @@ class HnswDistCalculator { batch_distance_ = batch_distance; } + //! Set the pairwise (SDC) distance function used for node-vs-node + //! comparisons (e.g. HNSW neighbor pruning). When set, dist(lhs, rhs) + //! will use this function instead of the ADC distance_. For non-PQ + //! quantizers this is left unset (nullptr) and dist(lhs, rhs) falls + //! back to distance_ as before. + inline void set_pairwise_distance( + const IndexMetric::MatrixDistance &pairwise_distance) { + pairwise_distance_ = pairwise_distance; + } + //! Reset query vector data inline void reset_query(const void *query) { error_ = false; @@ -100,7 +110,13 @@ class HnswDistCalculator { float score{0.0f}; - distance_(vec_lhs, vec_rhs, dim_, &score); + // Use pairwise (SDC) distance when available (PQ node-vs-node path); + // otherwise fall back to the standard ADC distance_. + if (pairwise_distance_) { + pairwise_distance_(vec_lhs, vec_rhs, dim_, &score); + } else { + distance_(vec_lhs, vec_rhs, dim_, &score); + } return score; } @@ -233,6 +249,7 @@ class HnswDistCalculator { const HnswEntity *entity_; IndexMetric::MatrixDistance distance_; + IndexMetric::MatrixDistance pairwise_distance_; // SDC for PQ (optional) IndexMetric::MatrixBatchDistance batch_distance_; const void *query_; diff --git a/src/include/zvec/turbo/turbo.h b/src/include/zvec/turbo/turbo.h index c3506f68f..463273c7a 100644 --- a/src/include/zvec/turbo/turbo.h +++ b/src/include/zvec/turbo/turbo.h @@ -41,6 +41,21 @@ 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 +using PqAdcDistanceFunc = void (*)(const uint8_t *pq_code, const float *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 +using PqSdcKernelFunc = void (*)(const uint8_t *a, const uint8_t *b, + const float *dist_table, + size_t num_subquantizers, float *out); + // Aggregate of all FHT kernels needed by FhtRotator, dispatched by ISA. struct FhtKernels { FhtFlipSignFunc flip_sign; @@ -50,6 +65,15 @@ struct FhtKernels { FhtVecRescaleFunc rescale; }; +// Aggregate of all PQ-specific kernels needed by PqInt8Quantizer, dispatched +// by ISA. 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. +struct PqKernels { + PqAdcDistanceFunc adc_distance; + PqSdcKernelFunc sdc_distance; +}; + enum class MetricType { kSquaredEuclidean, kCosine, @@ -110,4 +134,8 @@ 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 quantize_type and CPU arch. +PqKernels get_pq_kernels(QuantizeType quantize_type, + CpuArchType cpu_arch_type = CpuArchType::kAuto); + } // namespace zvec::turbo diff --git a/src/turbo/distance/scalar/pq/pq_scalar.cc b/src/turbo/distance/scalar/pq/pq_scalar.cc new file mode 100644 index 000000000..28e8d9756 --- /dev/null +++ b/src/turbo/distance/scalar/pq/pq_scalar.cc @@ -0,0 +1,44 @@ +// 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/pq_scalar.h" + +namespace zvec::turbo::scalar { + +void pq_adc_distance(const uint8_t *pq_code, const float *lut, + size_t num_subquantizers, float *out) { + constexpr size_t kNumCentroids = 256; + 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_distance(const uint8_t *a, const uint8_t *b, + const float *dist_table, size_t num_subquantizers, + float *out) { + constexpr size_t kNumCentroids = 256; + constexpr size_t kTablePerSub = kNumCentroids * kNumCentroids; // 65536 + 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; +} + +} // namespace zvec::turbo::scalar diff --git a/src/turbo/distance/scalar/pq/pq_scalar.h b/src/turbo/distance/scalar/pq/pq_scalar.h new file mode 100644 index 000000000..a0a576f4b --- /dev/null +++ b/src/turbo/distance/scalar/pq/pq_scalar.h @@ -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. + +#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_distance(const uint8_t *pq_code, const float *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_distance(const uint8_t *a, const uint8_t *b, + const float *dist_table, size_t num_subquantizers, + float *out); + +} // namespace zvec::turbo::scalar diff --git a/src/turbo/quantizer/distance.h b/src/turbo/quantizer/distance.h index bc8af6c1a..03e119dbe 100644 --- a/src/turbo/quantizer/distance.h +++ b/src/turbo/quantizer/distance.h @@ -44,6 +44,18 @@ class DistanceImpl { query_storage_(std::move(quantized_query)), dim_(dim) {} + //! PQ-specific constructor: pairwise_func_ is the SDC kernel (distinct + //! from func_ which is the ADC kernel). Non-PQ quantizers should use + //! the constructors above where pairwise_func_ is left as nullptr. + DistanceImpl(DistanceFunc func, DistanceFunc pairwise_func, + BatchDistanceFunc batch_func, + std::string quantized_query, size_t dim) + : func_(std::move(func)), + pairwise_func_(std::move(pairwise_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_); @@ -91,6 +103,13 @@ class DistanceImpl { return func_; } + //! Pairwise (SDC) distance function. Returns nullptr for non-PQ + //! quantizers; for PQ quantizers it points to the SDC kernel that + //! operates on two PQ codes via a centroid-to-centroid table. + const DistanceFunc &pairwise_func() const { + return pairwise_func_; + } + //! Raw batch distance function. const BatchDistanceFunc &batch_func() const { return batch_func_; @@ -98,6 +117,7 @@ class DistanceImpl { private: DistanceFunc func_{}; + DistanceFunc pairwise_func_{}; // nullptr for non-PQ; set explicitly for PQ (SDC) BatchDistanceFunc batch_func_{}; std::string query_storage_{}; size_t dim_{0}; 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..f67b6c329 --- /dev/null +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc @@ -0,0 +1,491 @@ +// 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 + +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(QuantizeType::kPQ); + adc_fn_ = pq_k.adc_distance; + sdc_fn_ = pq_k.sdc_distance; + + // LUT / dist_table computation reuses the metric-aware fp32 batch distance + // function (already SIMD-optimized), no need for a hand-written kernel. + fp32_batch_fn_ = get_batch_distance_func( + metric_from_name(meta_.metric_name()), DataType::kFp32, + QuantizeType::kDefault, CpuArchType::kAuto); + + meta_.set_meta(IndexMeta::DataType::DT_FP32, d); + return 0; +} + +// --------------------------------------------------------------------------- +// Simple Lloyd's KMeans for one sub-quantizer. +// --------------------------------------------------------------------------- +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); + + for (uint32_t iter = 0; iter < kMaxKmeansIters; ++iter) { + bool changed = false; + + // Assignment step. + for (size_t i = 0; i < num; ++i) { + const float *sub_vec = + reinterpret_cast( + reinterpret_cast(data) + i * stride) + + sub_idx * d; + + float best_dist = std::numeric_limits::max(); + uint32_t best_idx = 0; + for (size_t c = 0; c < k; ++c) { + float dist = 0.0f; + const float *cent = centroids_m + c * d; + for (size_t j = 0; j < d; ++j) { + float diff = sub_vec[j] - cent[j]; + dist += diff * diff; + } + if (dist < best_dist) { + best_dist = dist; + best_idx = static_cast(c); + } + } + if (assignments[i] != best_idx) { + changed = true; + assignments[i] = best_idx; + } + } + + if (!changed) 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) { + if (!holder) { + return kErrUnsupported; + } + + size_t num = holder->count(); + size_t stride = holder->element_size(); + + // Collect all data into a contiguous buffer (stride may include extras). + 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)); + } + + size_t data_stride = original_dim_ * sizeof(float); + + // Train each sub-quantizer independently. + for (uint32_t m = 0; m < num_subquantizers_; ++m) { + train_subquantizer(all_data.data(), num, data_stride, m); + } + + // Pre-compute SDC dist_table. + compute_dist_table(); + + 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. + std::vector centroid_ptrs(k); + 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; + + for (uint32_t c = 0; c < k; ++c) { + centroid_ptrs[c] = centroids_m + c * d; + } + for (uint32_t i = 0; i < k; ++i) { + fp32_batch_fn_(centroid_ptrs.data(), + reinterpret_cast(centroids_m + i * d), + k, d, table_m + i * k); + } + } +} + +size_t PqInt8Quantizer::find_nearest_centroid(const float *sub_vec, + size_t sub_idx) const { + const size_t k = kNumCentroids; + const size_t d = sub_dim_; + const float *centroids_m = + centroids_.data() + sub_idx * k * d; + + float best_dist = std::numeric_limits::max(); + size_t best_idx = 0; + for (size_t c = 0; c < k; ++c) { + float dist = 0.0f; + const float *cent = centroids_m + c * d; + for (size_t j = 0; j < d; ++j) { + float diff = sub_vec[j] - cent[j]; + dist += diff * diff; + } + if (dist < best_dist) { + best_dist = dist; + best_idx = c; + } + } + return best_idx; +} + +void PqInt8Quantizer::quantize_data(const void *input, void *output) const { + const float *vec = reinterpret_cast(input); + uint8_t *code = reinterpret_cast(output); + + for (uint32_t m = 0; m < num_subquantizers_; ++m) { + code[m] = static_cast( + find_nearest_centroid(vec + m * sub_dim_, m)); + } +} + +void PqInt8Quantizer::quantize_query(const void *input, void *output) const { + const float *query = reinterpret_cast(input); + float *lut = reinterpret_cast(output); + + // For each sub-quantizer, compute distance from query sub-vector to all + // 256 centroids via the metric-aware fp32 batch distance function. + std::vector centroid_ptrs(kNumCentroids); + for (uint32_t m = 0; m < num_subquantizers_; ++m) { + const float *centroids_m = + centroids_.data() + static_cast(m) * kNumCentroids * sub_dim_; + for (uint32_t k = 0; k < kNumCentroids; ++k) { + centroid_ptrs[k] = centroids_m + k * sub_dim_; + } + fp32_batch_fn_(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); + return d; +} + +void PqInt8Quantizer::calc_distance_dp_query_batch( + const void *const *dp_list, int dp_num, const void *query, + float *dist_list) const { + const uint8_t *lut = reinterpret_cast(query); + for (int i = 0; i < dp_num; ++i) { + float d = 0.0f; + adc_fn_(reinterpret_cast(dp_list[i]), + reinterpret_cast(query), num_subquantizers_, &d); + dist_list[i] = d; + } +} + +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); + 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()); + for (int i = 0; i < dp_num; ++i) { + float d = 0.0f; + adc_fn_(reinterpret_cast(dp_list[i]), lut.data(), + num_subquantizers_, &d); + dist_list[i] = d; + } +} + +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 { + 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; +} + +DistanceImpl PqInt8Quantizer::distance(const void *query, + const IndexQueryMeta &qmeta) const { + // Build the LUT from the (already quantized) query. + size_t lut_bytes = quantized_query_vector_length(); + std::string lut_storage(static_cast(query), lut_bytes); + + // ADC function: DistanceFunc signature + // (const void *candidate_pq_code, const void *lut, size_t dim, float *out) + // dim here is num_subquantizers (passed through DistanceImpl::dim()). + auto nsq = num_subquantizers_; + auto adc = adc_fn_; + DistanceFunc adc_func = [nsq, adc](const void *cand, const void *lut, + size_t /*dim*/, float *out) { + adc(reinterpret_cast(cand), + reinterpret_cast(lut), nsq, out); + }; + + // SDC pairwise function: captures dist_table + sdc kernel. + auto sdc = sdc_fn_; + const float *dt = dist_table_.data(); + DistanceFunc sdc_func = [nsq, sdc, dt](const void *a, const void *b, + size_t /*dim*/, float *out) { + sdc(reinterpret_cast(a), + reinterpret_cast(b), dt, nsq, out); + }; + + // Batch ADC: iterate over candidates. + BatchDistanceFunc batch_func = + [nsq, adc](const void **candidates, const void *lut, size_t num, + size_t /*dim*/, float *out) { + for (size_t i = 0; i < num; ++i) { + adc(reinterpret_cast(candidates[i]), + reinterpret_cast(lut), nsq, out + i); + } + }; + + 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); + size_t dist_table_bytes = dist_table_.size() * sizeof(float); + hdr.payload_size = static_cast( + sizeof(payload) + centroids_bytes + dist_table_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); + out->append(reinterpret_cast(dist_table_.data()), + dist_table_bytes); + 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); + size_t dist_table_bytes = + static_cast(num_subquantizers_) * kNumCentroids * kNumCentroids * + sizeof(float); + + centroids_.resize(centroids_bytes / sizeof(float)); + std::memcpy(centroids_.data(), ptr, centroids_bytes); + ptr += centroids_bytes; + + dist_table_.resize(dist_table_bytes / sizeof(float)); + std::memcpy(dist_table_.data(), ptr, dist_table_bytes); + + // Re-dispatch kernels. + auto pq_k = get_pq_kernels(QuantizeType::kPQ); + adc_fn_ = pq_k.adc_distance; + sdc_fn_ = pq_k.sdc_distance; + + fp32_batch_fn_ = get_batch_distance_func( + metric_from_name(meta_.metric_name()), DataType::kFp32, + QuantizeType::kDefault, CpuArchType::kAuto); + + 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..af5d111f2 --- /dev/null +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h @@ -0,0 +1,146 @@ +// 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; + + size_t quantized_datapoint_vector_length() const override { + return num_subquantizers_; + } + + 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; + + 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(); + + //! Find the nearest centroid for a single sub-vector. + size_t find_nearest_centroid(const float *sub_vec, size_t sub_idx) const; + + static constexpr uint32_t kNumCentroids = 256; + static constexpr uint32_t kMaxKmeansIters = 25; + + 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_; + + //! ISA-dispatched kernel function pointers (ADC / SDC). + PqAdcDistanceFunc adc_fn_{nullptr}; + PqSdcKernelFunc sdc_fn_{nullptr}; + + //! Reused fp32 batch distance function for LUT computation and SDC + //! dist_table precomputation. Obtained from get_batch_distance_func() + //! which is metric-aware and already SIMD-optimized. + BatchDistanceFunc fp32_batch_fn_{}; +}; + +} // namespace turbo +} // namespace zvec diff --git a/src/turbo/turbo.cc b/src/turbo/turbo.cc index 8e2e1b617..9f14f6017 100644 --- a/src/turbo/turbo.cc +++ b/src/turbo/turbo.cc @@ -22,6 +22,7 @@ #include "scalar/fp32/cosine.h" #include "scalar/fp32/inner_product.h" #include "scalar/fp32/squared_euclidean.h" +#include "scalar/pq/pq_scalar.h" #if defined(__SSE2__) #include "sse/fht/fht.h" @@ -197,4 +198,16 @@ FhtKernels get_fht_kernels() { return k; // scalar } -} // namespace zvec::turbo +PqKernels get_pq_kernels(QuantizeType quantize_type, + CpuArchType cpu_arch_type) { + PqKernels k{}; + if (quantize_type == QuantizeType::kPQ) { + // scalar is the only implementation for now; future SIMD paths gate on + // cpu_arch_type here (same pattern as get_distance_func). + k.adc_distance = scalar::pq_adc_distance; + k.sdc_distance = scalar::pq_sdc_distance; + } + return k; +} + +} // namespace zvec::turbo \ No newline at end of file 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..74cb94b48 --- /dev/null +++ b/tests/turbo/turbo_pq_int8_quantizer_test.cc @@ -0,0 +1,288 @@ +// 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 "quantizer/pq_int8_quantizer/pq_int8_quantizer.h" +#include "zvec/core/framework/index_factory.h" + +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< + MultiPassIndexHolder> +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 pairwise_func() should both be set. + EXPECT_TRUE(static_cast(dist_impl.func())); + EXPECT_TRUE(static_cast(dist_impl.pairwise_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); + + // Pairwise (SDC) via pairwise_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.pairwise_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; + } + + // SDC distances should match. + iter->next(); + iter->is_valid(); + std::vector code3(quantizer->quantized_datapoint_vector_length()); + std::vector code4(q2->quantized_datapoint_vector_length()); + quantizer->quantize_data(iter->data(), code3.data()); + q2->quantize_data(iter->data(), code4.data()); + + float d1 = quantizer->calc_distance_dp_dp(code1.data(), code3.data()); + float d2 = q2->calc_distance_dp_dp(code2.data(), code4.data()); + EXPECT_NEAR(d1, d2, 1e-6f); +} From f9b43f1c643eb8d6da1a7df8f86fc493cd6c8eb5 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Fri, 3 Jul 2026 10:08:55 +0800 Subject: [PATCH 22/42] tmp --- .../algorithm/hnsw/hnsw_dist_calculator.h | 7 ++++++ .../{pq => pq_quantizer_int8}/pq_scalar.cc | 12 ++++----- .../{pq => pq_quantizer_int8}/pq_scalar.h | 10 ++++---- src/turbo/quantizer/distance.h | 25 ++++++++++--------- .../pq_int8_quantizer/pq_int8_quantizer.cc | 2 +- src/turbo/turbo.cc | 6 ++--- tests/turbo/turbo_pq_int8_quantizer_test.cc | 8 +++--- 7 files changed, 39 insertions(+), 31 deletions(-) rename src/turbo/distance/scalar/{pq => pq_quantizer_int8}/pq_scalar.cc (77%) rename src/turbo/distance/scalar/{pq => pq_quantizer_int8}/pq_scalar.h (81%) diff --git a/src/core/algorithm/hnsw/hnsw_dist_calculator.h b/src/core/algorithm/hnsw/hnsw_dist_calculator.h index f11b064c4..eb5871f74 100644 --- a/src/core/algorithm/hnsw/hnsw_dist_calculator.h +++ b/src/core/algorithm/hnsw/hnsw_dist_calculator.h @@ -16,6 +16,13 @@ #include #include "hnsw_entity.h" +// TODO(HNSW+PQ): Adapt this file to the origin architecture: +// 1. Integrate turbo::Quantizer + turbo::DistanceImpl (see distance.h) +// 2. Use dist_impl_.sym_func() for zero-branch dist(void*,void*) dispatch +// 3. HnswContext needs dual calculators: add_dc_ + search_dc_ with mode_ +// 4. reset_query() should call quantizer_->distance() to build DistanceImpl +// Reference: origin/turbo/src/core/algorithm/hnsw/hnsw_dist_calculator.h + namespace zvec { namespace core { diff --git a/src/turbo/distance/scalar/pq/pq_scalar.cc b/src/turbo/distance/scalar/pq_quantizer_int8/pq_scalar.cc similarity index 77% rename from src/turbo/distance/scalar/pq/pq_scalar.cc rename to src/turbo/distance/scalar/pq_quantizer_int8/pq_scalar.cc index 28e8d9756..12c88afcb 100644 --- a/src/turbo/distance/scalar/pq/pq_scalar.cc +++ b/src/turbo/distance/scalar/pq_quantizer_int8/pq_scalar.cc @@ -12,12 +12,12 @@ // See the License for the specific language governing permissions and // limitations under the License. -#include "scalar/pq/pq_scalar.h" +#include "scalar/pq_quantizer_int8/pq_scalar.h" namespace zvec::turbo::scalar { -void pq_adc_distance(const uint8_t *pq_code, const float *lut, - size_t num_subquantizers, float *out) { +void pq_adc_int8_distance(const uint8_t *pq_code, const float *lut, + size_t num_subquantizers, float *out) { constexpr size_t kNumCentroids = 256; float sum = 0.0f; for (size_t m = 0; m < num_subquantizers; ++m) { @@ -26,9 +26,9 @@ void pq_adc_distance(const uint8_t *pq_code, const float *lut, *out = sum; } -void pq_sdc_distance(const uint8_t *a, const uint8_t *b, - const float *dist_table, size_t num_subquantizers, - float *out) { +void pq_sdc_int8_distance(const uint8_t *a, const uint8_t *b, + const float *dist_table, size_t num_subquantizers, + float *out) { constexpr size_t kNumCentroids = 256; constexpr size_t kTablePerSub = kNumCentroids * kNumCentroids; // 65536 float sum = 0.0f; diff --git a/src/turbo/distance/scalar/pq/pq_scalar.h b/src/turbo/distance/scalar/pq_quantizer_int8/pq_scalar.h similarity index 81% rename from src/turbo/distance/scalar/pq/pq_scalar.h rename to src/turbo/distance/scalar/pq_quantizer_int8/pq_scalar.h index a0a576f4b..80657a4e3 100644 --- a/src/turbo/distance/scalar/pq/pq_scalar.h +++ b/src/turbo/distance/scalar/pq_quantizer_int8/pq_scalar.h @@ -23,8 +23,8 @@ namespace zvec::turbo::scalar { // 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_distance(const uint8_t *pq_code, const float *lut, - size_t num_subquantizers, float *out); +void pq_adc_int8_distance(const uint8_t *pq_code, const float *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 @@ -36,8 +36,8 @@ void pq_adc_distance(const uint8_t *pq_code, const float *lut, // // distance = sum_{m=0}^{num_subquantizers-1} // dist_table[m * 65536 + a[m] * 256 + b[m]] -void pq_sdc_distance(const uint8_t *a, const uint8_t *b, - const float *dist_table, size_t num_subquantizers, - float *out); +void pq_sdc_int8_distance(const uint8_t *a, const uint8_t *b, + const float *dist_table, size_t num_subquantizers, + float *out); } // namespace zvec::turbo::scalar diff --git a/src/turbo/quantizer/distance.h b/src/turbo/quantizer/distance.h index 03e119dbe..47f191890 100644 --- a/src/turbo/quantizer/distance.h +++ b/src/turbo/quantizer/distance.h @@ -34,24 +34,26 @@ class DistanceImpl { 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: pairwise_func_ is the SDC kernel (distinct + //! 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 pairwise_func_ is left as nullptr. - DistanceImpl(DistanceFunc func, DistanceFunc pairwise_func, + //! 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)), - pairwise_func_(std::move(pairwise_func)), + sym_func_(std::move(sym_func)), batch_func_(std::move(batch_func)), query_storage_(std::move(quantized_query)), dim_(dim) {} @@ -97,17 +99,16 @@ class DistanceImpl { } //! Raw scalar distance function (operates on already-quantized - //! candidates). Useful for pairwise node-vs-node distance where no - //! stored query is involved. + //! candidates). This is the ADC path for PQ. const DistanceFunc &func() const { return func_; } - //! Pairwise (SDC) distance function. Returns nullptr for non-PQ - //! quantizers; for PQ quantizers it points to the SDC kernel that - //! operates on two PQ codes via a centroid-to-centroid table. - const DistanceFunc &pairwise_func() const { - return pairwise_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. @@ -117,7 +118,7 @@ class DistanceImpl { private: DistanceFunc func_{}; - DistanceFunc pairwise_func_{}; // nullptr for non-PQ; set explicitly for PQ (SDC) + DistanceFunc sym_func_{}; // defaults to func_; PQ sets SDC kernel BatchDistanceFunc batch_func_{}; std::string query_storage_{}; size_t dim_{0}; diff --git a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc index f67b6c329..3d8312a42 100644 --- a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc @@ -375,7 +375,7 @@ DistanceImpl PqInt8Quantizer::distance(const void *query, reinterpret_cast(lut), nsq, out); }; - // SDC pairwise function: captures dist_table + sdc kernel. + // SDC (symmetric) function: captures dist_table + sdc kernel. auto sdc = sdc_fn_; const float *dt = dist_table_.data(); DistanceFunc sdc_func = [nsq, sdc, dt](const void *a, const void *b, diff --git a/src/turbo/turbo.cc b/src/turbo/turbo.cc index 9f14f6017..c90c1e67f 100644 --- a/src/turbo/turbo.cc +++ b/src/turbo/turbo.cc @@ -22,7 +22,7 @@ #include "scalar/fp32/cosine.h" #include "scalar/fp32/inner_product.h" #include "scalar/fp32/squared_euclidean.h" -#include "scalar/pq/pq_scalar.h" +#include "scalar/pq_quantizer_int8/pq_scalar.h" #if defined(__SSE2__) #include "sse/fht/fht.h" @@ -204,8 +204,8 @@ PqKernels get_pq_kernels(QuantizeType quantize_type, if (quantize_type == QuantizeType::kPQ) { // scalar is the only implementation for now; future SIMD paths gate on // cpu_arch_type here (same pattern as get_distance_func). - k.adc_distance = scalar::pq_adc_distance; - k.sdc_distance = scalar::pq_sdc_distance; + k.adc_distance = scalar::pq_adc_int8_distance; + k.sdc_distance = scalar::pq_sdc_int8_distance; } return k; } diff --git a/tests/turbo/turbo_pq_int8_quantizer_test.cc b/tests/turbo/turbo_pq_int8_quantizer_test.cc index 74cb94b48..fdc0a2d0c 100644 --- a/tests/turbo/turbo_pq_int8_quantizer_test.cc +++ b/tests/turbo/turbo_pq_int8_quantizer_test.cc @@ -217,9 +217,9 @@ TEST(PqInt8Quantizer, DistanceImplAdcAndSdc) { auto dist_impl = quantizer->distance(lut_storage.data(), qmeta); ASSERT_TRUE(dist_impl.valid()); - // func() and pairwise_func() should both be set. + // func() and sym_func() should both be set. EXPECT_TRUE(static_cast(dist_impl.func())); - EXPECT_TRUE(static_cast(dist_impl.pairwise_func())); + EXPECT_TRUE(static_cast(dist_impl.sym_func())); // Encode a candidate and compute distance via DistanceImpl (ADC path). iter->next(); @@ -230,14 +230,14 @@ TEST(PqInt8Quantizer, DistanceImplAdcAndSdc) { float d = dist_impl(code.data()); EXPECT_GE(d, 0.0f); - // Pairwise (SDC) via pairwise_func() directly. + // 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.pairwise_func()(code.data(), code2.data(), NSQ, &sdc_d); + dist_impl.sym_func()(code.data(), code2.data(), NSQ, &sdc_d); EXPECT_GE(sdc_d, 0.0f); } From fc7efd86c93ed219b48336cd2f25e838f0cdfec3 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Fri, 3 Jul 2026 11:06:09 +0800 Subject: [PATCH 23/42] tmp --- src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc | 5 ++--- src/turbo/turbo.cc | 1 + 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc index 3d8312a42..199727e9a 100644 --- a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc @@ -186,9 +186,8 @@ int PqInt8Quantizer::train(IndexHolder::Pointer holder) { } size_t num = holder->count(); - size_t stride = holder->element_size(); - // Collect all data into a contiguous buffer (stride may include extras). + // Collect all data into a contiguous buffer. auto iter = holder->create_iterator(); std::vector all_data(num * original_dim_); size_t row = 0; @@ -301,7 +300,6 @@ float PqInt8Quantizer::calc_distance_dp_query(const void *dp, void PqInt8Quantizer::calc_distance_dp_query_batch( const void *const *dp_list, int dp_num, const void *query, float *dist_list) const { - const uint8_t *lut = reinterpret_cast(query); for (int i = 0; i < dp_num; ++i) { float d = 0.0f; adc_fn_(reinterpret_cast(dp_list[i]), @@ -360,6 +358,7 @@ int PqInt8Quantizer::quantize(const void *query, const IndexQueryMeta &qmeta, DistanceImpl PqInt8Quantizer::distance(const void *query, const IndexQueryMeta &qmeta) const { + (void)qmeta; // reserved for future use (e.g. data-type dispatch) // Build the LUT from the (already quantized) query. size_t lut_bytes = quantized_query_vector_length(); std::string lut_storage(static_cast(query), lut_bytes); diff --git a/src/turbo/turbo.cc b/src/turbo/turbo.cc index c90c1e67f..b96d10f05 100644 --- a/src/turbo/turbo.cc +++ b/src/turbo/turbo.cc @@ -200,6 +200,7 @@ FhtKernels get_fht_kernels() { PqKernels get_pq_kernels(QuantizeType quantize_type, CpuArchType cpu_arch_type) { + (void)cpu_arch_type; // reserved for future SIMD dispatch PqKernels k{}; if (quantize_type == QuantizeType::kPQ) { // scalar is the only implementation for now; future SIMD paths gate on From 79a08eb16ccd15e28f2b6404c32bcdf99684e945 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Fri, 3 Jul 2026 11:15:01 +0800 Subject: [PATCH 24/42] style: fix clang-format issues from upstream merge --- tests/db/collection_test.cc | 8 ++++---- tests/db/index/common/query_params_test.cc | 1 - 2 files changed, 4 insertions(+), 5 deletions(-) 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/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) { From 78ace2bf41e6708ab31f329641561164a5e182f5 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Fri, 3 Jul 2026 11:18:16 +0800 Subject: [PATCH 25/42] style: fix clang-format issues in turbo module --- src/turbo/distance/avx2/fht/fht_avx2.cc | 3 +-- src/turbo/distance/avx512/fht/fht_avx512.cc | 3 +-- src/turbo/distance/scalar/fht/fht_scalar.cc | 3 +-- src/turbo/distance/sse/fht/fht_sse.cc | 3 +-- src/turbo/quantizer/preprocessor/fht_rotator.cc | 4 ++-- src/turbo/quantizer/preprocessor/fht_rotator.h | 14 ++++++++++---- tests/turbo/turbo_fht_rotator_test.cc | 12 +++++++----- 7 files changed, 23 insertions(+), 19 deletions(-) diff --git a/src/turbo/distance/avx2/fht/fht_avx2.cc b/src/turbo/distance/avx2/fht/fht_avx2.cc index 903305931..1ed941e54 100644 --- a/src/turbo/distance/avx2/fht/fht_avx2.cc +++ b/src/turbo/distance/avx2/fht/fht_avx2.cc @@ -14,13 +14,12 @@ #if defined(__AVX2__) -#include "fht.h" - #include #include #include #include #include +#include "fht.h" namespace zvec::turbo::avx2 { diff --git a/src/turbo/distance/avx512/fht/fht_avx512.cc b/src/turbo/distance/avx512/fht/fht_avx512.cc index b7ea37d63..12aec2410 100644 --- a/src/turbo/distance/avx512/fht/fht_avx512.cc +++ b/src/turbo/distance/avx512/fht/fht_avx512.cc @@ -14,13 +14,12 @@ #if defined(__AVX512F__) -#include "fht.h" - #include #include #include #include #include +#include "fht.h" namespace zvec::turbo::avx512 { diff --git a/src/turbo/distance/scalar/fht/fht_scalar.cc b/src/turbo/distance/scalar/fht/fht_scalar.cc index 008c6a2a6..9e01a536b 100644 --- a/src/turbo/distance/scalar/fht/fht_scalar.cc +++ b/src/turbo/distance/scalar/fht/fht_scalar.cc @@ -12,11 +12,10 @@ // See the License for the specific language governing permissions and // limitations under the License. -#include "fht.h" - #include #include #include +#include "fht.h" namespace zvec::turbo::scalar { diff --git a/src/turbo/distance/sse/fht/fht_sse.cc b/src/turbo/distance/sse/fht/fht_sse.cc index 7dbbf1951..fc6b45f3c 100644 --- a/src/turbo/distance/sse/fht/fht_sse.cc +++ b/src/turbo/distance/sse/fht/fht_sse.cc @@ -14,13 +14,12 @@ #if defined(__SSE2__) -#include "fht.h" - #include #include #include #include #include +#include "fht.h" namespace zvec::turbo::sse { diff --git a/src/turbo/quantizer/preprocessor/fht_rotator.cc b/src/turbo/quantizer/preprocessor/fht_rotator.cc index 94a19a916..41d318d77 100644 --- a/src/turbo/quantizer/preprocessor/fht_rotator.cc +++ b/src/turbo/quantizer/preprocessor/fht_rotator.cc @@ -13,7 +13,6 @@ // limitations under the License. #include "fht_rotator.h" - #include #include #include @@ -230,7 +229,8 @@ int FhtRotator::serialize(std::string *out) const { } int FhtRotator::deserialize(const void *data, size_t len) { - if (!data || len < sizeof(RotatorSerHeader)) return IndexError_InvalidArgument; + if (!data || len < sizeof(RotatorSerHeader)) + return IndexError_InvalidArgument; const auto *hdr = reinterpret_cast(data); if (hdr->magic != kRotatorMagic) return IndexError_Unsupported; diff --git a/src/turbo/quantizer/preprocessor/fht_rotator.h b/src/turbo/quantizer/preprocessor/fht_rotator.h index 15dfbb0c4..91c670f69 100644 --- a/src/turbo/quantizer/preprocessor/fht_rotator.h +++ b/src/turbo/quantizer/preprocessor/fht_rotator.h @@ -18,8 +18,8 @@ #include #include #include -#include "preprocessor.h" #include +#include "preprocessor.h" namespace zvec { namespace turbo { @@ -46,8 +46,12 @@ class FhtRotator : public Preprocessor { // -- Preprocessor interface ------------------------------------------------ - int in_dim() const override { return in_dim_; } - int out_dim() const override { return out_dim_; } + 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; @@ -60,7 +64,9 @@ class FhtRotator : public Preprocessor { int deserialize(const void *data, size_t len) override; //! Rotator type tag (kFht = 1). - RotatorType rotator_type() const { return RotatorType::kFht; } + RotatorType rotator_type() const { + return RotatorType::kFht; + } private: FhtRotator() = default; diff --git a/tests/turbo/turbo_fht_rotator_test.cc b/tests/turbo/turbo_fht_rotator_test.cc index 6acdaba80..8d7f06edc 100644 --- a/tests/turbo/turbo_fht_rotator_test.cc +++ b/tests/turbo/turbo_fht_rotator_test.cc @@ -164,11 +164,14 @@ TEST(FhtRotator, TrainGeneratesFlip) { 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); + 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; } + if (payload[i] != 0) { + any_nonzero = true; + break; + } } EXPECT_TRUE(any_nonzero) << "flip payload is all-zero for dim=" << dim; } @@ -424,8 +427,7 @@ TEST(FhtRotator, LargeDimension) { 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; + EXPECT_FLOAT_EQ(r1[i], r2[i]) << "mismatch at i=" << i << " dim=" << dim; } } } From b80bc4d34723b48fd57ef90824a9af6a25b5d8a6 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Fri, 3 Jul 2026 13:56:44 +0800 Subject: [PATCH 26/42] fix: LNK1120 --- .../quantizer/preprocessor/fht_rotator.cc | 22 +++++++------------ src/turbo/quantizer/quantizer.h | 2 ++ 2 files changed, 10 insertions(+), 14 deletions(-) diff --git a/src/turbo/quantizer/preprocessor/fht_rotator.cc b/src/turbo/quantizer/preprocessor/fht_rotator.cc index 41d318d77..c600ee163 100644 --- a/src/turbo/quantizer/preprocessor/fht_rotator.cc +++ b/src/turbo/quantizer/preprocessor/fht_rotator.cc @@ -16,15 +16,10 @@ #include #include #include -#include - +#include "../quantizer.h" namespace zvec { namespace turbo { -using core::IndexError_InvalidArgument; -using core::IndexError_Runtime; -using core::IndexError_Unsupported; - // ============================================================================ // FhtRotator method implementations // ============================================================================ @@ -210,8 +205,8 @@ void FhtRotator::apply_inverse(const float *in, float *out) const { // --------------------------------------------------------------------------- int FhtRotator::serialize(std::string *out) const { - if (!out) return IndexError_InvalidArgument; - if (flip_.empty()) return IndexError_Runtime; + if (!out) return kErrInvalidArgument; + if (flip_.empty()) return kErrRuntime; RotatorSerHeader hdr{}; hdr.magic = kRotatorMagic; @@ -229,18 +224,17 @@ int FhtRotator::serialize(std::string *out) const { } int FhtRotator::deserialize(const void *data, size_t len) { - if (!data || len < sizeof(RotatorSerHeader)) - return IndexError_InvalidArgument; + if (!data || len < sizeof(RotatorSerHeader)) return kErrInvalidArgument; const auto *hdr = reinterpret_cast(data); - if (hdr->magic != kRotatorMagic) return IndexError_Unsupported; - if (hdr->version != kRotatorSerVersion) return IndexError_Unsupported; + if (hdr->magic != kRotatorMagic) return kErrUnsupported; + if (hdr->version != kRotatorSerVersion) return kErrUnsupported; if (static_cast(hdr->rotator_type) != RotatorType::kFht) { - return IndexError_Unsupported; + return kErrUnsupported; } const size_t total = sizeof(RotatorSerHeader) + hdr->payload_size; - if (len < total) return IndexError_InvalidArgument; + if (len < total) return kErrInvalidArgument; in_dim_ = static_cast(hdr->in_dim); out_dim_ = static_cast(hdr->out_dim); diff --git a/src/turbo/quantizer/quantizer.h b/src/turbo/quantizer/quantizer.h index 3b7a7d944..2a9d54cb5 100644 --- a/src/turbo/quantizer/quantizer.h +++ b/src/turbo/quantizer/quantizer.h @@ -39,8 +39,10 @@ using namespace zvec::core; //! //! 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; From 6ae749c8a17fe93a37b19e1a4970b3da317a2c49 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Fri, 3 Jul 2026 14:07:11 +0800 Subject: [PATCH 27/42] tmp --- .../scalar/pq_quantizer_int8/{pq_scalar.cc => pq_distance.cc} | 2 +- .../scalar/pq_quantizer_int8/{pq_scalar.h => pq_distance.h} | 0 src/turbo/turbo.cc | 2 +- 3 files changed, 2 insertions(+), 2 deletions(-) rename src/turbo/distance/scalar/pq_quantizer_int8/{pq_scalar.cc => pq_distance.cc} (96%) rename src/turbo/distance/scalar/pq_quantizer_int8/{pq_scalar.h => pq_distance.h} (100%) diff --git a/src/turbo/distance/scalar/pq_quantizer_int8/pq_scalar.cc b/src/turbo/distance/scalar/pq_quantizer_int8/pq_distance.cc similarity index 96% rename from src/turbo/distance/scalar/pq_quantizer_int8/pq_scalar.cc rename to src/turbo/distance/scalar/pq_quantizer_int8/pq_distance.cc index 12c88afcb..2d56c3c7a 100644 --- a/src/turbo/distance/scalar/pq_quantizer_int8/pq_scalar.cc +++ b/src/turbo/distance/scalar/pq_quantizer_int8/pq_distance.cc @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -#include "scalar/pq_quantizer_int8/pq_scalar.h" +#include "scalar/pq_quantizer_int8/pq_distance.h" namespace zvec::turbo::scalar { diff --git a/src/turbo/distance/scalar/pq_quantizer_int8/pq_scalar.h b/src/turbo/distance/scalar/pq_quantizer_int8/pq_distance.h similarity index 100% rename from src/turbo/distance/scalar/pq_quantizer_int8/pq_scalar.h rename to src/turbo/distance/scalar/pq_quantizer_int8/pq_distance.h diff --git a/src/turbo/turbo.cc b/src/turbo/turbo.cc index b96d10f05..205903498 100644 --- a/src/turbo/turbo.cc +++ b/src/turbo/turbo.cc @@ -22,7 +22,7 @@ #include "scalar/fp32/cosine.h" #include "scalar/fp32/inner_product.h" #include "scalar/fp32/squared_euclidean.h" -#include "scalar/pq_quantizer_int8/pq_scalar.h" +#include "scalar/pq_quantizer_int8/pq_distance.h" #if defined(__SSE2__) #include "sse/fht/fht.h" From 281fdb40b9fc72eba9aabdae939b3f6cdc15f20d Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Fri, 3 Jul 2026 14:16:23 +0800 Subject: [PATCH 28/42] Rename to conform to the style guide --- src/turbo/distance/avx2/fht/{fht_avx2.cc => fht.cc} | 0 src/turbo/distance/avx512/fht/{fht_avx512.cc => fht.cc} | 0 src/turbo/distance/scalar/fht/{fht_scalar.cc => fht.cc} | 0 src/turbo/distance/sse/fht/{fht_sse.cc => fht.cc} | 0 4 files changed, 0 insertions(+), 0 deletions(-) rename src/turbo/distance/avx2/fht/{fht_avx2.cc => fht.cc} (100%) rename src/turbo/distance/avx512/fht/{fht_avx512.cc => fht.cc} (100%) rename src/turbo/distance/scalar/fht/{fht_scalar.cc => fht.cc} (100%) rename src/turbo/distance/sse/fht/{fht_sse.cc => fht.cc} (100%) diff --git a/src/turbo/distance/avx2/fht/fht_avx2.cc b/src/turbo/distance/avx2/fht/fht.cc similarity index 100% rename from src/turbo/distance/avx2/fht/fht_avx2.cc rename to src/turbo/distance/avx2/fht/fht.cc diff --git a/src/turbo/distance/avx512/fht/fht_avx512.cc b/src/turbo/distance/avx512/fht/fht.cc similarity index 100% rename from src/turbo/distance/avx512/fht/fht_avx512.cc rename to src/turbo/distance/avx512/fht/fht.cc diff --git a/src/turbo/distance/scalar/fht/fht_scalar.cc b/src/turbo/distance/scalar/fht/fht.cc similarity index 100% rename from src/turbo/distance/scalar/fht/fht_scalar.cc rename to src/turbo/distance/scalar/fht/fht.cc diff --git a/src/turbo/distance/sse/fht/fht_sse.cc b/src/turbo/distance/sse/fht/fht.cc similarity index 100% rename from src/turbo/distance/sse/fht/fht_sse.cc rename to src/turbo/distance/sse/fht/fht.cc From 10cbe03e2b2773494d7c01908c1cd66ff34d30a3 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Fri, 3 Jul 2026 14:20:13 +0800 Subject: [PATCH 29/42] Rename to conform to the style guide --- src/turbo/distance/avx2/fht/fht.cc | 2 +- src/turbo/distance/avx512/fht/fht.cc | 2 +- src/turbo/distance/scalar/fht/fht.cc | 2 +- src/turbo/distance/sse/fht/fht.cc | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/turbo/distance/avx2/fht/fht.cc b/src/turbo/distance/avx2/fht/fht.cc index 1ed941e54..a6d3edc36 100644 --- a/src/turbo/distance/avx2/fht/fht.cc +++ b/src/turbo/distance/avx2/fht/fht.cc @@ -14,12 +14,12 @@ #if defined(__AVX2__) +#include "fht.h" #include #include #include #include #include -#include "fht.h" namespace zvec::turbo::avx2 { diff --git a/src/turbo/distance/avx512/fht/fht.cc b/src/turbo/distance/avx512/fht/fht.cc index 12aec2410..2ba9e4b3b 100644 --- a/src/turbo/distance/avx512/fht/fht.cc +++ b/src/turbo/distance/avx512/fht/fht.cc @@ -14,12 +14,12 @@ #if defined(__AVX512F__) +#include "fht.h" #include #include #include #include #include -#include "fht.h" namespace zvec::turbo::avx512 { diff --git a/src/turbo/distance/scalar/fht/fht.cc b/src/turbo/distance/scalar/fht/fht.cc index 9e01a536b..f1cdffd78 100644 --- a/src/turbo/distance/scalar/fht/fht.cc +++ b/src/turbo/distance/scalar/fht/fht.cc @@ -12,10 +12,10 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include "fht.h" #include #include #include -#include "fht.h" namespace zvec::turbo::scalar { diff --git a/src/turbo/distance/sse/fht/fht.cc b/src/turbo/distance/sse/fht/fht.cc index fc6b45f3c..18b12c545 100644 --- a/src/turbo/distance/sse/fht/fht.cc +++ b/src/turbo/distance/sse/fht/fht.cc @@ -14,12 +14,12 @@ #if defined(__SSE2__) +#include "fht.h" #include #include #include #include #include -#include "fht.h" namespace zvec::turbo::sse { From 8734866e71455db747c7a96d17b868ba7b4ca66e Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Fri, 3 Jul 2026 15:03:07 +0800 Subject: [PATCH 30/42] pq_int8 --- src/turbo/distance/avx2/pq/pq_distance.cc | 139 ++++++++++++++++ src/turbo/distance/avx2/pq/pq_distance.h | 35 ++++ src/turbo/distance/avx512/pq/pq_distance.cc | 133 +++++++++++++++ src/turbo/distance/avx512/pq/pq_distance.h | 33 ++++ src/turbo/turbo.cc | 21 ++- tests/turbo/turbo_pq_int8_quantizer_test.cc | 174 ++++++++++++++++++++ 6 files changed, 532 insertions(+), 3 deletions(-) create mode 100644 src/turbo/distance/avx2/pq/pq_distance.cc create mode 100644 src/turbo/distance/avx2/pq/pq_distance.h create mode 100644 src/turbo/distance/avx512/pq/pq_distance.cc create mode 100644 src/turbo/distance/avx512/pq/pq_distance.h diff --git a/src/turbo/distance/avx2/pq/pq_distance.cc b/src/turbo/distance/avx2/pq/pq_distance.cc new file mode 100644 index 000000000..73f92be9d --- /dev/null +++ b/src/turbo/distance/avx2/pq/pq_distance.cc @@ -0,0 +1,139 @@ +// 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/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 uint8_t *pq_code, const float *lut, + size_t num_subquantizers, float *out) { + constexpr int kNumCentroids = 256; + constexpr int kChunkSize = 8; // AVX2 processes 8 floats at once + + __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 uint8_t *a, const uint8_t *b, + const float *dist_table, + size_t num_subquantizers, float *out) { + constexpr int kNumCentroids = 256; + constexpr int kTablePerSub = kNumCentroids * kNumCentroids; // 65536 + constexpr int kChunkSize = 8; + + __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; +} + +} // namespace zvec::turbo::avx2 diff --git a/src/turbo/distance/avx2/pq/pq_distance.h b/src/turbo/distance/avx2/pq/pq_distance.h new file mode 100644 index 000000000..bde9aec75 --- /dev/null +++ b/src/turbo/distance/avx2/pq/pq_distance.h @@ -0,0 +1,35 @@ +// 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 uint8_t *pq_code, const float *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 uint8_t *a, const uint8_t *b, + const float *dist_table, + size_t num_subquantizers, float *out); + +} // namespace zvec::turbo::avx2 diff --git a/src/turbo/distance/avx512/pq/pq_distance.cc b/src/turbo/distance/avx512/pq/pq_distance.cc new file mode 100644 index 000000000..cf136f459 --- /dev/null +++ b/src/turbo/distance/avx512/pq/pq_distance.cc @@ -0,0 +1,133 @@ +// 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/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 uint8_t *pq_code, const float *lut, + size_t num_subquantizers, float *out) { + constexpr int kNumCentroids = 256; + constexpr int kChunkSize = 16; // AVX512 processes 16 floats at once + + __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 uint8_t *a, const uint8_t *b, + const float *dist_table, + size_t num_subquantizers, float *out) { + constexpr int kNumCentroids = 256; + constexpr int kTablePerSub = kNumCentroids * kNumCentroids; // 65536 + constexpr int kChunkSize = 16; + + __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; +} + +} // namespace zvec::turbo::avx512 diff --git a/src/turbo/distance/avx512/pq/pq_distance.h b/src/turbo/distance/avx512/pq/pq_distance.h new file mode 100644 index 000000000..e67e8210f --- /dev/null +++ b/src/turbo/distance/avx512/pq/pq_distance.h @@ -0,0 +1,33 @@ +// 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 uint8_t *pq_code, const float *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 uint8_t *a, const uint8_t *b, + const float *dist_table, + size_t num_subquantizers, float *out); + +} // namespace zvec::turbo::avx512 diff --git a/src/turbo/turbo.cc b/src/turbo/turbo.cc index 205903498..55187257c 100644 --- a/src/turbo/turbo.cc +++ b/src/turbo/turbo.cc @@ -29,9 +29,11 @@ #endif #if defined(__AVX2__) #include "avx2/fht/fht.h" +#include "avx2/pq/pq_distance.h" #endif #if defined(__AVX512F__) #include "avx512/fht/fht.h" +#include "avx512/pq/pq_distance.h" #endif namespace zvec::turbo { @@ -200,13 +202,26 @@ FhtKernels get_fht_kernels() { PqKernels get_pq_kernels(QuantizeType quantize_type, CpuArchType cpu_arch_type) { - (void)cpu_arch_type; // reserved for future SIMD dispatch + (void)cpu_arch_type; // currently unused, reserved for future use PqKernels k{}; if (quantize_type == QuantizeType::kPQ) { - // scalar is the only implementation for now; future SIMD paths gate on - // cpu_arch_type here (same pattern as get_distance_func). + // Default: scalar fallback k.adc_distance = scalar::pq_adc_int8_distance; k.sdc_distance = scalar::pq_sdc_int8_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; + 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; + } +#endif } return k; } diff --git a/tests/turbo/turbo_pq_int8_quantizer_test.cc b/tests/turbo/turbo_pq_int8_quantizer_test.cc index fdc0a2d0c..2a28f2efd 100644 --- a/tests/turbo/turbo_pq_int8_quantizer_test.cc +++ b/tests/turbo/turbo_pq_int8_quantizer_test.cc @@ -15,13 +15,22 @@ #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/pq_distance.h" +#endif +#if defined(__AVX512F__) +#include "distance/avx512/pq/pq_distance.h" +#endif + using namespace zvec; using namespace zvec::core; using namespace zvec::ailego; @@ -286,3 +295,168 @@ TEST(PqInt8Quantizer, SerializeDeserialize) { float d2 = q2->calc_distance_dp_dp(code2.data(), code4.data()); EXPECT_NEAR(d1, d2, 1e-6f); } + +// --------------------------------------------------------------------------- +// 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 +} From 74e9f930cb1ca3729cf50f1a3d56b5660ffe969b Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Fri, 3 Jul 2026 17:20:42 +0800 Subject: [PATCH 31/42] tmp --- .../pq_int8_quantizer/pq_int8_quantizer.cc | 19 ++++++-------- tests/turbo/turbo_pq_int8_quantizer_test.cc | 25 +++++++++++-------- 2 files changed, 21 insertions(+), 23 deletions(-) diff --git a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc index 199727e9a..acc7ceaf9 100644 --- a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc @@ -418,17 +418,17 @@ int PqInt8Quantizer::serialize(std::string *out) const { payload.num_centroids = kNumCentroids; size_t centroids_bytes = centroids_.size() * sizeof(float); - size_t dist_table_bytes = dist_table_.size() * sizeof(float); - hdr.payload_size = static_cast( - sizeof(payload) + centroids_bytes + dist_table_bytes); + 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); - out->append(reinterpret_cast(dist_table_.data()), - dist_table_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; } @@ -461,16 +461,11 @@ int PqInt8Quantizer::deserialize(const void *data, size_t len) { size_t centroids_bytes = static_cast(num_subquantizers_) * kNumCentroids * sub_dim_ * sizeof(float); - size_t dist_table_bytes = - static_cast(num_subquantizers_) * kNumCentroids * kNumCentroids * - sizeof(float); centroids_.resize(centroids_bytes / sizeof(float)); std::memcpy(centroids_.data(), ptr, centroids_bytes); - ptr += centroids_bytes; - - dist_table_.resize(dist_table_bytes / sizeof(float)); - std::memcpy(dist_table_.data(), ptr, dist_table_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(QuantizeType::kPQ); diff --git a/tests/turbo/turbo_pq_int8_quantizer_test.cc b/tests/turbo/turbo_pq_int8_quantizer_test.cc index 2a28f2efd..5650911b8 100644 --- a/tests/turbo/turbo_pq_int8_quantizer_test.cc +++ b/tests/turbo/turbo_pq_int8_quantizer_test.cc @@ -283,17 +283,20 @@ TEST(PqInt8Quantizer, SerializeDeserialize) { EXPECT_EQ(code1[m], code2[m]) << "m=" << m; } - // SDC distances should match. - iter->next(); - iter->is_valid(); - std::vector code3(quantizer->quantized_datapoint_vector_length()); - std::vector code4(q2->quantized_datapoint_vector_length()); - quantizer->quantize_data(iter->data(), code3.data()); - q2->quantize_data(iter->data(), code4.data()); - - float d1 = quantizer->calc_distance_dp_dp(code1.data(), code3.data()); - float d2 = q2->calc_distance_dp_dp(code2.data(), code4.data()); - EXPECT_NEAR(d1, d2, 1e-6f); + // 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. } // --------------------------------------------------------------------------- From 7c5d299c1fa10093b879637df4619aea39b9805b Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Mon, 6 Jul 2026 14:04:20 +0800 Subject: [PATCH 32/42] tmp --- src/include/zvec/turbo/turbo.h | 15 +- src/turbo/distance/avx2/pq/pq_distance.cc | 86 ++++++++++- src/turbo/distance/avx2/pq/pq_distance.h | 13 +- src/turbo/distance/avx512/pq/pq_distance.cc | 94 +++++++++++- src/turbo/distance/avx512/pq/pq_distance.h | 13 +- .../scalar/pq_quantizer_int8/pq_distance.cc | 54 ++++++- .../scalar/pq_quantizer_int8/pq_distance.h | 14 +- .../pq_int8_quantizer/pq_int8_quantizer.cc | 145 ++++++++++++------ .../pq_int8_quantizer/pq_int8_quantizer.h | 5 +- src/turbo/quantizer/quantizer.h | 6 + src/turbo/turbo.cc | 3 + 11 files changed, 378 insertions(+), 70 deletions(-) diff --git a/src/include/zvec/turbo/turbo.h b/src/include/zvec/turbo/turbo.h index 463273c7a..4eea20eec 100644 --- a/src/include/zvec/turbo/turbo.h +++ b/src/include/zvec/turbo/turbo.h @@ -46,16 +46,24 @@ using FhtVecRescaleFunc = void (*)(float *data, size_t n, float factor); // 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 -using PqAdcDistanceFunc = void (*)(const uint8_t *pq_code, const float *lut, +// 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 -using PqSdcKernelFunc = void (*)(const uint8_t *a, const uint8_t *b, - const float *dist_table, +// 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; @@ -72,6 +80,7 @@ struct FhtKernels { struct PqKernels { PqAdcDistanceFunc adc_distance; PqSdcKernelFunc sdc_distance; + PqBatchAdcFunc batch_adc_distance; }; enum class MetricType { diff --git a/src/turbo/distance/avx2/pq/pq_distance.cc b/src/turbo/distance/avx2/pq/pq_distance.cc index 73f92be9d..072ec65d1 100644 --- a/src/turbo/distance/avx2/pq/pq_distance.cc +++ b/src/turbo/distance/avx2/pq/pq_distance.cc @@ -37,10 +37,12 @@ inline float horizontal_sum_avx2(__m256 v) { } // namespace -void pq_adc_int8_distance_avx2(const uint8_t *pq_code, const float *lut, +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(); @@ -82,12 +84,15 @@ void pq_adc_int8_distance_avx2(const uint8_t *pq_code, const float *lut, *out = sum; } -void pq_sdc_int8_distance_avx2(const uint8_t *a, const uint8_t *b, - const float *dist_table, +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(); @@ -136,4 +141,79 @@ void pq_sdc_int8_distance_avx2(const uint8_t *a, const uint8_t *b, *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/pq_distance.h b/src/turbo/distance/avx2/pq/pq_distance.h index bde9aec75..f1da57f64 100644 --- a/src/turbo/distance/avx2/pq/pq_distance.h +++ b/src/turbo/distance/avx2/pq/pq_distance.h @@ -22,14 +22,21 @@ 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 uint8_t *pq_code, const float *lut, +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 uint8_t *a, const uint8_t *b, - const float *dist_table, +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/pq/pq_distance.cc b/src/turbo/distance/avx512/pq/pq_distance.cc index cf136f459..91ba21cdd 100644 --- a/src/turbo/distance/avx512/pq/pq_distance.cc +++ b/src/turbo/distance/avx512/pq/pq_distance.cc @@ -28,10 +28,12 @@ inline float horizontal_sum_avx512(__m512 v) { } // namespace -void pq_adc_int8_distance_avx512(const uint8_t *pq_code, const float *lut, +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(); @@ -74,12 +76,15 @@ void pq_adc_int8_distance_avx512(const uint8_t *pq_code, const float *lut, *out = sum; } -void pq_sdc_int8_distance_avx512(const uint8_t *a, const uint8_t *b, - const float *dist_table, +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(); @@ -130,4 +135,87 @@ void pq_sdc_int8_distance_avx512(const uint8_t *a, const uint8_t *b, *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/pq_distance.h b/src/turbo/distance/avx512/pq/pq_distance.h index e67e8210f..5026b4779 100644 --- a/src/turbo/distance/avx512/pq/pq_distance.h +++ b/src/turbo/distance/avx512/pq/pq_distance.h @@ -21,13 +21,20 @@ 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 uint8_t *pq_code, const float *lut, +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 uint8_t *a, const uint8_t *b, - const float *dist_table, +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/distance/scalar/pq_quantizer_int8/pq_distance.cc b/src/turbo/distance/scalar/pq_quantizer_int8/pq_distance.cc index 2d56c3c7a..16362f54b 100644 --- a/src/turbo/distance/scalar/pq_quantizer_int8/pq_distance.cc +++ b/src/turbo/distance/scalar/pq_quantizer_int8/pq_distance.cc @@ -16,9 +16,11 @@ namespace zvec::turbo::scalar { -void pq_adc_int8_distance(const uint8_t *pq_code, const float *lut, +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]]; @@ -26,11 +28,14 @@ void pq_adc_int8_distance(const uint8_t *pq_code, const float *lut, *out = sum; } -void pq_sdc_int8_distance(const uint8_t *a, const uint8_t *b, - const float *dist_table, size_t num_subquantizers, +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 + @@ -41,4 +46,47 @@ void pq_sdc_int8_distance(const uint8_t *a, const uint8_t *b, *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 index 80657a4e3..1d0c8d810 100644 --- a/src/turbo/distance/scalar/pq_quantizer_int8/pq_distance.h +++ b/src/turbo/distance/scalar/pq_quantizer_int8/pq_distance.h @@ -23,7 +23,7 @@ namespace zvec::turbo::scalar { // 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 uint8_t *pq_code, const float *lut, +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 @@ -36,8 +36,16 @@ void pq_adc_int8_distance(const uint8_t *pq_code, const float *lut, // // distance = sum_{m=0}^{num_subquantizers-1} // dist_table[m * 65536 + a[m] * 256 + b[m]] -void pq_sdc_int8_distance(const uint8_t *a, const uint8_t *b, - const float *dist_table, size_t num_subquantizers, +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/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc index acc7ceaf9..48d693722 100644 --- a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc @@ -17,6 +17,7 @@ #include #include #include +#include #include #include #include @@ -66,6 +67,7 @@ int PqInt8Quantizer::init(const IndexMeta &meta, auto pq_k = get_pq_kernels(QuantizeType::kPQ); adc_fn_ = pq_k.adc_distance; sdc_fn_ = pq_k.sdc_distance; + batch_adc_fn_ = pq_k.batch_adc_distance; // LUT / dist_table computation reuses the metric-aware fp32 batch distance // function (already SIMD-optimized), no need for a hand-written kernel. @@ -112,37 +114,41 @@ void PqInt8Quantizer::train_subquantizer(const float *data, size_t 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. + // Assignment step — use SIMD-accelerated fp32_batch_fn_ instead of + // scalar L2 loop. Each call computes distances from one sub-vector + // to all k centroids in one batch. for (size_t i = 0; i < num; ++i) { const float *sub_vec = reinterpret_cast( reinterpret_cast(data) + i * stride) + sub_idx * d; - float best_dist = std::numeric_limits::max(); - uint32_t best_idx = 0; - for (size_t c = 0; c < k; ++c) { - float dist = 0.0f; - const float *cent = centroids_m + c * d; - for (size_t j = 0; j < d; ++j) { - float diff = sub_vec[j] - cent[j]; - dist += diff * diff; - } - if (dist < best_dist) { - best_dist = dist; - best_idx = static_cast(c); - } - } + fp32_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) break; + 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); @@ -181,6 +187,10 @@ void PqInt8Quantizer::train_subquantizer(const float *data, size_t num, } int PqInt8Quantizer::train(IndexHolder::Pointer holder) { + return train(holder, 1); +} + +int PqInt8Quantizer::train(IndexHolder::Pointer holder, int thread_count) { if (!holder) { return kErrUnsupported; } @@ -199,14 +209,56 @@ int PqInt8Quantizer::train(IndexHolder::Pointer holder) { size_t data_stride = original_dim_ * sizeof(float); - // Train each sub-quantizer independently. - for (uint32_t m = 0; m < num_subquantizers_; ++m) { - train_subquantizer(all_data.data(), num, data_stride, m); + // 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-compute SDC dist_table. + LOG_INFO("Computing SDC dist_table ..."); compute_dist_table(); + LOG_INFO("PQ training complete."); return 0; } @@ -262,6 +314,10 @@ void PqInt8Quantizer::quantize_data(const void *input, void *output) const { const float *vec = reinterpret_cast(input); uint8_t *code = reinterpret_cast(output); + // using _mm256_cmp_ps + _mm256_blendv_ps for branchless min+argmin + // without an intermediate distance buffer. Consider adopting the same + // fused approach for small sub_dim to eliminate the scalar argmin loop + // in find_nearest_centroid. for (uint32_t m = 0; m < num_subquantizers_; ++m) { code[m] = static_cast( find_nearest_centroid(vec + m * sub_dim_, m)); @@ -346,6 +402,10 @@ float PqInt8Quantizer::calc_distance_dp_dp(const void *dp1, 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]); @@ -358,40 +418,28 @@ int PqInt8Quantizer::quantize(const void *query, const IndexQueryMeta &qmeta, DistanceImpl PqInt8Quantizer::distance(const void *query, const IndexQueryMeta &qmeta) const { - (void)qmeta; // reserved for future use (e.g. data-type dispatch) - // Build the LUT from the (already quantized) query. - size_t lut_bytes = quantized_query_vector_length(); - std::string lut_storage(static_cast(query), lut_bytes); + (void)qmeta; - // ADC function: DistanceFunc signature - // (const void *candidate_pq_code, const void *lut, size_t dim, float *out) - // dim here is num_subquantizers (passed through DistanceImpl::dim()). - auto nsq = num_subquantizers_; - auto adc = adc_fn_; - DistanceFunc adc_func = [nsq, adc](const void *cand, const void *lut, - size_t /*dim*/, float *out) { - adc(reinterpret_cast(cand), - reinterpret_cast(lut), nsq, out); - }; + // ADC: PqAdcDistanceFunc signature now matches DistanceFunc directly + // (both use void*), so we can assign without a lambda wrapper. + DistanceFunc adc_func = adc_fn_; - // SDC (symmetric) function: captures dist_table + sdc kernel. + // SDC: still needs a lambda because the kernel has 5 parameters + // (extra dist_table pointer), vs DistanceFunc's 4. auto sdc = sdc_fn_; - const float *dt = dist_table_.data(); - DistanceFunc sdc_func = [nsq, sdc, dt](const void *a, const void *b, - size_t /*dim*/, float *out) { - sdc(reinterpret_cast(a), - reinterpret_cast(b), dt, nsq, out); + 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: iterate over candidates. - BatchDistanceFunc batch_func = - [nsq, adc](const void **candidates, const void *lut, size_t num, - size_t /*dim*/, float *out) { - for (size_t i = 0; i < num; ++i) { - adc(reinterpret_cast(candidates[i]), - reinterpret_cast(lut), nsq, out + i); - } - }; + // 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), @@ -471,6 +519,7 @@ int PqInt8Quantizer::deserialize(const void *data, size_t len) { auto pq_k = get_pq_kernels(QuantizeType::kPQ); adc_fn_ = pq_k.adc_distance; sdc_fn_ = pq_k.sdc_distance; + batch_adc_fn_ = pq_k.batch_adc_distance; fp32_batch_fn_ = get_batch_distance_func( metric_from_name(meta_.metric_name()), DataType::kFp32, diff --git a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h index af5d111f2..933297a19 100644 --- a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h @@ -64,6 +64,8 @@ class PqInt8Quantizer : public Quantizer { 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_; } @@ -132,9 +134,10 @@ class PqInt8Quantizer : public Quantizer { //! [num_subquantizers * kNumCentroids * kNumCentroids] std::vector dist_table_; - //! ISA-dispatched kernel function pointers (ADC / SDC). + //! ISA-dispatched kernel function pointers (ADC / SDC / Batch ADC). PqAdcDistanceFunc adc_fn_{nullptr}; PqSdcKernelFunc sdc_fn_{nullptr}; + PqBatchAdcFunc batch_adc_fn_{nullptr}; //! Reused fp32 batch distance function for LUT computation and SDC //! dist_table precomputation. Obtained from get_batch_distance_func() diff --git a/src/turbo/quantizer/quantizer.h b/src/turbo/quantizer/quantizer.h index 2a9d54cb5..4ef239bd5 100644 --- a/src/turbo/quantizer/quantizer.h +++ b/src/turbo/quantizer/quantizer.h @@ -102,6 +102,12 @@ class Quantizer { 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; diff --git a/src/turbo/turbo.cc b/src/turbo/turbo.cc index 55187257c..db04b8e12 100644 --- a/src/turbo/turbo.cc +++ b/src/turbo/turbo.cc @@ -208,11 +208,13 @@ PqKernels get_pq_kernels(QuantizeType quantize_type, // Default: scalar fallback 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 @@ -220,6 +222,7 @@ PqKernels get_pq_kernels(QuantizeType quantize_type, 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 } From 8ecdd17b318628651a05056be390737be72c13c1 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Mon, 6 Jul 2026 14:30:25 +0800 Subject: [PATCH 33/42] uniform func sign --- .../pq_int8_quantizer/pq_int8_quantizer.cc | 105 +++++++++--------- .../pq_int8_quantizer/pq_int8_quantizer.h | 12 +- 2 files changed, 62 insertions(+), 55 deletions(-) diff --git a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc index 48d693722..5db04d994 100644 --- a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc @@ -82,6 +82,21 @@ int PqInt8Quantizer::init(const IndexMeta &meta, // --------------------------------------------------------------------------- // 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; @@ -254,6 +269,9 @@ int PqInt8Quantizer::train(IndexHolder::Pointer holder, int thread_count) { 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(); @@ -270,57 +288,38 @@ void PqInt8Quantizer::compute_dist_table() { // For each sub-quantizer, compute centroid-to-centroid distances via the // metric-aware fp32 batch distance function. - std::vector centroid_ptrs(k); 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; - for (uint32_t c = 0; c < k; ++c) { - centroid_ptrs[c] = centroids_m + c * d; - } + // 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_(centroid_ptrs.data(), + fp32_batch_fn_(const_cast(centroid_ptrs.data()), reinterpret_cast(centroids_m + i * d), k, d, table_m + i * k); } } } -size_t PqInt8Quantizer::find_nearest_centroid(const float *sub_vec, - size_t sub_idx) const { - const size_t k = kNumCentroids; - const size_t d = sub_dim_; - const float *centroids_m = - centroids_.data() + sub_idx * k * d; - - float best_dist = std::numeric_limits::max(); - size_t best_idx = 0; - for (size_t c = 0; c < k; ++c) { - float dist = 0.0f; - const float *cent = centroids_m + c * d; - for (size_t j = 0; j < d; ++j) { - float diff = sub_vec[j] - cent[j]; - dist += diff * diff; - } - if (dist < best_dist) { - best_dist = dist; - best_idx = c; - } - } - return best_idx; -} - void PqInt8Quantizer::quantize_data(const void *input, void *output) const { const float *vec = reinterpret_cast(input); uint8_t *code = reinterpret_cast(output); - // using _mm256_cmp_ps + _mm256_blendv_ps for branchless min+argmin - // without an intermediate distance buffer. Consider adopting the same - // fused approach for small sub_dim to eliminate the scalar argmin loop - // in find_nearest_centroid. + // Reuse quantize_query to build LUT (distances from each sub-vector to + // all 256 centroids), then argmin each sub to get the PQ code. + // This shares the same fp32_batch_fn_ + centroid_ptrs_cache_ path as + // LUT computation, eliminating the separate find_nearest_centroid. + const size_t lut_size = static_cast(num_subquantizers_) * kNumCentroids; + std::vector lut(lut_size); + quantize_query(vec, lut.data()); + for (uint32_t m = 0; m < num_subquantizers_; ++m) { + const float *row = lut.data() + m * kNumCentroids; code[m] = static_cast( - find_nearest_centroid(vec + m * sub_dim_, m)); + std::min_element(row, row + kNumCentroids) - row); } } @@ -330,14 +329,11 @@ void PqInt8Quantizer::quantize_query(const void *input, void *output) const { // For each sub-quantizer, compute distance from query sub-vector to all // 256 centroids via the metric-aware fp32 batch distance function. - std::vector centroid_ptrs(kNumCentroids); + // Uses pre-built centroid pointer cache to avoid per-sub allocations. + // const_cast: see compute_dist_table for rationale. for (uint32_t m = 0; m < num_subquantizers_; ++m) { - const float *centroids_m = - centroids_.data() + static_cast(m) * kNumCentroids * sub_dim_; - for (uint32_t k = 0; k < kNumCentroids; ++k) { - centroid_ptrs[k] = centroids_m + k * sub_dim_; - } - fp32_batch_fn_(centroid_ptrs.data(), + 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); } @@ -356,12 +352,13 @@ float PqInt8Quantizer::calc_distance_dp_query(const void *dp, void PqInt8Quantizer::calc_distance_dp_query_batch( const void *const *dp_list, int dp_num, const void *query, float *dist_list) const { - for (int i = 0; i < dp_num; ++i) { - float d = 0.0f; - adc_fn_(reinterpret_cast(dp_list[i]), - reinterpret_cast(query), num_subquantizers_, &d); - dist_list[i] = d; - } + // 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); } float PqInt8Quantizer::calc_distance_dp_query_unquantized( @@ -382,12 +379,11 @@ void PqInt8Quantizer::calc_distance_dp_query_batch_unquantized( std::vector lut(static_cast(num_subquantizers_) * kNumCentroids); quantize_query(query, lut.data()); - for (int i = 0; i < dp_num; ++i) { - float d = 0.0f; - adc_fn_(reinterpret_cast(dp_list[i]), lut.data(), - num_subquantizers_, &d); - dist_list[i] = d; - } + // 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); } float PqInt8Quantizer::calc_distance_dp_dp(const void *dp1, @@ -525,6 +521,9 @@ int PqInt8Quantizer::deserialize(const void *data, size_t len) { metric_from_name(meta_.metric_name()), DataType::kFp32, QuantizeType::kDefault, CpuArchType::kAuto); + // Pre-build centroid pointer cache for fast encode/search. + build_centroid_ptrs_cache(); + return 0; } diff --git a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h index 933297a19..304b6b9bb 100644 --- a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h @@ -116,8 +116,9 @@ class PqInt8Quantizer : public Quantizer { //! Compute the centroid-to-centroid distance table for SDC. void compute_dist_table(); - //! Find the nearest centroid for a single sub-vector. - size_t find_nearest_centroid(const float *sub_vec, size_t sub_idx) const; + //! 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; @@ -134,11 +135,18 @@ class PqInt8Quantizer : public Quantizer { //! [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}; + //! Reused fp32 batch distance function for LUT computation and SDC //! dist_table precomputation. Obtained from get_batch_distance_func() //! which is metric-aware and already SIMD-optimized. From f1a838e0099930b051c6ec59f1438741dd035781 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Mon, 6 Jul 2026 16:08:41 +0800 Subject: [PATCH 34/42] support ip/cosine --- .../{pq => pq_quantizer_int8}/pq_distance.cc | 2 +- .../{pq => pq_quantizer_int8}/pq_distance.h | 0 .../{pq => pq_quantizer_int8}/pq_distance.cc | 2 +- .../{pq => pq_quantizer_int8}/pq_distance.h | 0 .../pq_int8_quantizer/pq_int8_quantizer.cc | 220 ++++++++++++++-- .../pq_int8_quantizer/pq_int8_quantizer.h | 21 +- src/turbo/turbo.cc | 4 +- tests/turbo/turbo_pq_int8_quantizer_test.cc | 249 +++++++++++++++++- 8 files changed, 465 insertions(+), 33 deletions(-) rename src/turbo/distance/avx2/{pq => pq_quantizer_int8}/pq_distance.cc (99%) rename src/turbo/distance/avx2/{pq => pq_quantizer_int8}/pq_distance.h (100%) rename src/turbo/distance/avx512/{pq => pq_quantizer_int8}/pq_distance.cc (99%) rename src/turbo/distance/avx512/{pq => pq_quantizer_int8}/pq_distance.h (100%) diff --git a/src/turbo/distance/avx2/pq/pq_distance.cc b/src/turbo/distance/avx2/pq_quantizer_int8/pq_distance.cc similarity index 99% rename from src/turbo/distance/avx2/pq/pq_distance.cc rename to src/turbo/distance/avx2/pq_quantizer_int8/pq_distance.cc index 072ec65d1..3a74fb84a 100644 --- a/src/turbo/distance/avx2/pq/pq_distance.cc +++ b/src/turbo/distance/avx2/pq_quantizer_int8/pq_distance.cc @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -#include "avx2/pq/pq_distance.h" +#include "avx2/pq_quantizer_int8/pq_distance.h" #include diff --git a/src/turbo/distance/avx2/pq/pq_distance.h b/src/turbo/distance/avx2/pq_quantizer_int8/pq_distance.h similarity index 100% rename from src/turbo/distance/avx2/pq/pq_distance.h rename to src/turbo/distance/avx2/pq_quantizer_int8/pq_distance.h diff --git a/src/turbo/distance/avx512/pq/pq_distance.cc b/src/turbo/distance/avx512/pq_quantizer_int8/pq_distance.cc similarity index 99% rename from src/turbo/distance/avx512/pq/pq_distance.cc rename to src/turbo/distance/avx512/pq_quantizer_int8/pq_distance.cc index 91ba21cdd..2fe2e4204 100644 --- a/src/turbo/distance/avx512/pq/pq_distance.cc +++ b/src/turbo/distance/avx512/pq_quantizer_int8/pq_distance.cc @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -#include "avx512/pq/pq_distance.h" +#include "avx512/pq_quantizer_int8/pq_distance.h" #include diff --git a/src/turbo/distance/avx512/pq/pq_distance.h b/src/turbo/distance/avx512/pq_quantizer_int8/pq_distance.h similarity index 100% rename from src/turbo/distance/avx512/pq/pq_distance.h rename to src/turbo/distance/avx512/pq_quantizer_int8/pq_distance.h diff --git a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc index 5db04d994..87288b508 100644 --- a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc @@ -14,6 +14,7 @@ #include "quantizer/pq_int8_quantizer/pq_int8_quantizer.h" #include +#include #include #include #include @@ -69,12 +70,40 @@ int PqInt8Quantizer::init(const IndexMeta &meta, sdc_fn_ = pq_k.sdc_distance; batch_adc_fn_ = pq_k.batch_adc_distance; - // LUT / dist_table computation reuses the metric-aware fp32 batch distance - // function (already SIMD-optimized), no need for a hand-written kernel. + // 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( - metric_from_name(meta_.metric_name()), DataType::kFp32, + 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; } @@ -139,18 +168,19 @@ void PqInt8Quantizer::train_subquantizer(const float *data, size_t num, for (uint32_t iter = 0; iter < kMaxKmeansIters; ++iter) { bool changed = false; - // Assignment step — use SIMD-accelerated fp32_batch_fn_ instead of - // scalar L2 loop. Each call computes distances from one sub-vector - // to all k centroids in one batch. + // 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_batch_fn_(centroid_ptrs.data(), - reinterpret_cast(sub_vec), - k, d, dists.data()); + 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()); @@ -224,6 +254,25 @@ int PqInt8Quantizer::train(IndexHolder::Pointer holder, int thread_count) { 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_sq = 0.0f; + for (uint32_t j = 0; j < original_dim_; ++j) { + norm_sq += v[j] * v[j]; + } + if (norm_sq > 0.0f) { + const float inv_norm = 1.0f / std::sqrt(norm_sq); + for (uint32_t j = 0; j < original_dim_; ++j) { + v[j] *= inv_norm; + } + } + } + } + // Clamp thread count to [1, num_subquantizers_]. thread_count = std::max(1, std::min(thread_count, static_cast(num_subquantizers_))); @@ -286,8 +335,11 @@ void PqInt8Quantizer::compute_dist_table() { 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. + // 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; @@ -308,18 +360,58 @@ void PqInt8Quantizer::quantize_data(const void *input, void *output) const { const float *vec = reinterpret_cast(input); uint8_t *code = reinterpret_cast(output); - // Reuse quantize_query to build LUT (distances from each sub-vector to - // all 256 centroids), then argmin each sub to get the PQ code. - // This shares the same fp32_batch_fn_ + centroid_ptrs_cache_ path as - // LUT computation, eliminating the separate find_nearest_centroid. - const size_t lut_size = static_cast(num_subquantizers_) * kNumCentroids; - std::vector lut(lut_size); - quantize_query(vec, lut.data()); + // 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_); + float norm_sq = 0.0f; + for (uint32_t j = 0; j < original_dim_; ++j) { + norm_sq += norm_vec_storage[j] * norm_vec_storage[j]; + } + vec_norm = (norm_sq > 0.0f) ? std::sqrt(norm_sq) : 0.0f; + if (vec_norm > 0.0f) { + const float inv_norm = 1.0f / vec_norm; + for (uint32_t j = 0; j < original_dim_; ++j) { + norm_vec_storage[j] *= inv_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 *row = lut.data() + m * kNumCentroids; - code[m] = static_cast( - std::min_element(row, row + kNumCentroids) - row); + 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; } } @@ -327,9 +419,31 @@ 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_sq = 0.0f; + for (uint32_t i = 0; i < original_dim_; ++i) { + norm_sq += norm_query_storage[i] * norm_query_storage[i]; + } + const float inv_norm = (norm_sq > 0.0f) + ? 1.0f / std::sqrt(norm_sq) + : 0.0f; + for (uint32_t i = 0; i < original_dim_; ++i) { + norm_query_storage[i] *= inv_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. - // Uses pre-built centroid pointer cache to avoid per-sub allocations. + // 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]; @@ -346,6 +460,10 @@ float PqInt8Quantizer::calc_distance_dp_query(const void *dp, 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; } @@ -359,6 +477,12 @@ void PqInt8Quantizer::calc_distance_dp_query_batch( 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( @@ -370,6 +494,9 @@ float PqInt8Quantizer::calc_distance_dp_query_unquantized( 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; } @@ -384,6 +511,11 @@ void PqInt8Quantizer::calc_distance_dp_query_batch_unquantized( 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, @@ -412,6 +544,38 @@ int PqInt8Quantizer::quantize(const void *query, const IndexQueryMeta &qmeta, 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; @@ -517,9 +681,21 @@ int PqInt8Quantizer::deserialize(const void *data, size_t len) { 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(); diff --git a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h index 304b6b9bb..d117bec10 100644 --- a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h @@ -67,7 +67,7 @@ class PqInt8Quantizer : public Quantizer { int train(IndexHolder::Pointer holder, int thread_count) override; size_t quantized_datapoint_vector_length() const override { - return num_subquantizers_; + return num_subquantizers_ + extra_meta_size_; } size_t quantized_query_vector_length() const override { @@ -98,6 +98,9 @@ class PqInt8Quantizer : public Quantizer { 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; @@ -122,6 +125,7 @@ class PqInt8Quantizer : public Quantizer { static constexpr uint32_t kNumCentroids = 256; static constexpr uint32_t kMaxKmeansIters = 25; + static constexpr uint32_t kExtraMetaSizeCosine = sizeof(float); IndexMeta meta_{}; uint32_t original_dim_{0}; @@ -146,11 +150,18 @@ class PqInt8Quantizer : public Quantizer { PqSdcKernelFunc sdc_fn_{nullptr}; PqBatchAdcFunc batch_adc_fn_{nullptr}; - - //! Reused fp32 batch distance function for LUT computation and SDC - //! dist_table precomputation. Obtained from get_batch_distance_func() - //! which is metric-aware and already SIMD-optimized. + //! 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 diff --git a/src/turbo/turbo.cc b/src/turbo/turbo.cc index db04b8e12..d3d0ab7ad 100644 --- a/src/turbo/turbo.cc +++ b/src/turbo/turbo.cc @@ -29,11 +29,11 @@ #endif #if defined(__AVX2__) #include "avx2/fht/fht.h" -#include "avx2/pq/pq_distance.h" +#include "avx2/pq_quantizer_int8/pq_distance.h" #endif #if defined(__AVX512F__) #include "avx512/fht/fht.h" -#include "avx512/pq/pq_distance.h" +#include "avx512/pq_quantizer_int8/pq_distance.h" #endif namespace zvec::turbo { diff --git a/tests/turbo/turbo_pq_int8_quantizer_test.cc b/tests/turbo/turbo_pq_int8_quantizer_test.cc index 5650911b8..55737445e 100644 --- a/tests/turbo/turbo_pq_int8_quantizer_test.cc +++ b/tests/turbo/turbo_pq_int8_quantizer_test.cc @@ -25,10 +25,10 @@ #include "zvec/core/framework/index_factory.h" #if defined(__AVX2__) -#include "distance/avx2/pq/pq_distance.h" +#include "distance/avx2/pq_quantizer_int8/pq_distance.h" #endif #if defined(__AVX512F__) -#include "distance/avx512/pq/pq_distance.h" +#include "distance/avx512/pq_quantizer_int8/pq_distance.h" #endif using namespace zvec; @@ -463,3 +463,248 @@ TEST(PqInt8SimdConsistency, AdcDistanceM1) { } #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< + MultiPassIndexHolder> +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; + } +} From 58ea18caa5e4b198a814b4ca864c8887da7888f3 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Mon, 6 Jul 2026 16:16:09 +0800 Subject: [PATCH 35/42] support ip/cosine --- tests/turbo/turbo_pq_int8_quantizer_test.cc | 119 ++++++++++++++++++++ 1 file changed, 119 insertions(+) diff --git a/tests/turbo/turbo_pq_int8_quantizer_test.cc b/tests/turbo/turbo_pq_int8_quantizer_test.cc index 55737445e..06c6fd7bb 100644 --- a/tests/turbo/turbo_pq_int8_quantizer_test.cc +++ b/tests/turbo/turbo_pq_int8_quantizer_test.cc @@ -708,3 +708,122 @@ TEST(PqInt8Quantizer, L2Dequantize) { << "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; + } +} From 3122b1745b6edacbdb6fe23ab68af1b63319b426 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Mon, 6 Jul 2026 17:11:08 +0800 Subject: [PATCH 36/42] set upper --- .../pq_int8_quantizer/pq_int8_quantizer.cc | 24 +++++++++++++++++++ .../pq_int8_quantizer/pq_int8_quantizer.h | 1 + 2 files changed, 25 insertions(+) diff --git a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc index 87288b508..a6d47f4ce 100644 --- a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc @@ -252,6 +252,30 @@ int PqInt8Quantizer::train(IndexHolder::Pointer holder, int thread_count) { 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 diff --git a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h index d117bec10..6417170c2 100644 --- a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h @@ -125,6 +125,7 @@ class PqInt8Quantizer : public Quantizer { static constexpr uint32_t kNumCentroids = 256; static constexpr uint32_t kMaxKmeansIters = 25; + static constexpr size_t kMaxTrainVectors = 65535; static constexpr uint32_t kExtraMetaSizeCosine = sizeof(float); IndexMeta meta_{}; From 64595d1de829d54511c56b11b2e2dab79805451f Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Mon, 6 Jul 2026 17:31:48 +0800 Subject: [PATCH 37/42] norm match --- .../pq_int8_quantizer/pq_int8_quantizer.cc | 39 ++++--------------- 1 file changed, 8 insertions(+), 31 deletions(-) diff --git a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc index a6d47f4ce..d7be9055e 100644 --- a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc @@ -22,6 +22,7 @@ #include #include #include +#include namespace zvec { namespace turbo { @@ -284,16 +285,8 @@ int PqInt8Quantizer::train(IndexHolder::Pointer holder, int thread_count) { if (meta_.metric_name() == "Cosine") { for (size_t i = 0; i < num; ++i) { float *v = all_data.data() + i * original_dim_; - float norm_sq = 0.0f; - for (uint32_t j = 0; j < original_dim_; ++j) { - norm_sq += v[j] * v[j]; - } - if (norm_sq > 0.0f) { - const float inv_norm = 1.0f / std::sqrt(norm_sq); - for (uint32_t j = 0; j < original_dim_; ++j) { - v[j] *= inv_norm; - } - } + float norm = 0.0f; + ailego::Normalizer::L2(v, original_dim_, &norm); } } @@ -391,17 +384,8 @@ void PqInt8Quantizer::quantize_data(const void *input, void *output) const { float vec_norm = 0.0f; if (meta_.metric_name() == "Cosine") { norm_vec_storage.assign(vec, vec + original_dim_); - float norm_sq = 0.0f; - for (uint32_t j = 0; j < original_dim_; ++j) { - norm_sq += norm_vec_storage[j] * norm_vec_storage[j]; - } - vec_norm = (norm_sq > 0.0f) ? std::sqrt(norm_sq) : 0.0f; - if (vec_norm > 0.0f) { - const float inv_norm = 1.0f / vec_norm; - for (uint32_t j = 0; j < original_dim_; ++j) { - norm_vec_storage[j] *= inv_norm; - } - } + ailego::Normalizer::L2(norm_vec_storage.data(), original_dim_, + &vec_norm); vec = norm_vec_storage.data(); } @@ -450,16 +434,9 @@ void PqInt8Quantizer::quantize_query(const void *input, void *output) const { std::vector norm_query_storage; if (meta_.metric_name() == "Cosine") { norm_query_storage.assign(query, query + original_dim_); - float norm_sq = 0.0f; - for (uint32_t i = 0; i < original_dim_; ++i) { - norm_sq += norm_query_storage[i] * norm_query_storage[i]; - } - const float inv_norm = (norm_sq > 0.0f) - ? 1.0f / std::sqrt(norm_sq) - : 0.0f; - for (uint32_t i = 0; i < original_dim_; ++i) { - norm_query_storage[i] *= inv_norm; - } + float norm = 0.0f; + ailego::Normalizer::L2(norm_query_storage.data(), original_dim_, + &norm); query = norm_query_storage.data(); } From 17e5a0dfd1191243618b973709ee32a251fd30d3 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Mon, 6 Jul 2026 17:45:39 +0800 Subject: [PATCH 38/42] rollback --- .../algorithm/hnsw/hnsw_dist_calculator.h | 26 +------------------ 1 file changed, 1 insertion(+), 25 deletions(-) diff --git a/src/core/algorithm/hnsw/hnsw_dist_calculator.h b/src/core/algorithm/hnsw/hnsw_dist_calculator.h index eb5871f74..2e4b22d1f 100644 --- a/src/core/algorithm/hnsw/hnsw_dist_calculator.h +++ b/src/core/algorithm/hnsw/hnsw_dist_calculator.h @@ -16,13 +16,6 @@ #include #include "hnsw_entity.h" -// TODO(HNSW+PQ): Adapt this file to the origin architecture: -// 1. Integrate turbo::Quantizer + turbo::DistanceImpl (see distance.h) -// 2. Use dist_impl_.sym_func() for zero-branch dist(void*,void*) dispatch -// 3. HnswContext needs dual calculators: add_dc_ + search_dc_ with mode_ -// 4. reset_query() should call quantizer_->distance() to build DistanceImpl -// Reference: origin/turbo/src/core/algorithm/hnsw/hnsw_dist_calculator.h - namespace zvec { namespace core { @@ -91,16 +84,6 @@ class HnswDistCalculator { batch_distance_ = batch_distance; } - //! Set the pairwise (SDC) distance function used for node-vs-node - //! comparisons (e.g. HNSW neighbor pruning). When set, dist(lhs, rhs) - //! will use this function instead of the ADC distance_. For non-PQ - //! quantizers this is left unset (nullptr) and dist(lhs, rhs) falls - //! back to distance_ as before. - inline void set_pairwise_distance( - const IndexMetric::MatrixDistance &pairwise_distance) { - pairwise_distance_ = pairwise_distance; - } - //! Reset query vector data inline void reset_query(const void *query) { error_ = false; @@ -117,13 +100,7 @@ class HnswDistCalculator { float score{0.0f}; - // Use pairwise (SDC) distance when available (PQ node-vs-node path); - // otherwise fall back to the standard ADC distance_. - if (pairwise_distance_) { - pairwise_distance_(vec_lhs, vec_rhs, dim_, &score); - } else { - distance_(vec_lhs, vec_rhs, dim_, &score); - } + distance_(vec_lhs, vec_rhs, dim_, &score); return score; } @@ -256,7 +233,6 @@ class HnswDistCalculator { const HnswEntity *entity_; IndexMetric::MatrixDistance distance_; - IndexMetric::MatrixDistance pairwise_distance_; // SDC for PQ (optional) IndexMetric::MatrixBatchDistance batch_distance_; const void *query_; From 3e8abad3adcc6a99d42873913bab197810d51007 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Tue, 7 Jul 2026 13:32:33 +0800 Subject: [PATCH 39/42] add int4 --- src/include/zvec/turbo/turbo.h | 19 +- .../avx2/pq_quantizer_int4/pq_distance.cc | 226 ++++++ .../avx2/pq_quantizer_int4/pq_distance.h | 41 ++ .../avx512/pq_quantizer_int4/pq_distance.cc | 212 ++++++ .../avx512/pq_quantizer_int4/pq_distance.h | 40 ++ .../scalar/pq_quantizer_int4/pq_distance.cc | 109 +++ .../scalar/pq_quantizer_int4/pq_distance.h | 52 ++ .../pq_int4_quantizer/pq_int4_quantizer.cc | 651 ++++++++++++++++++ .../pq_int4_quantizer/pq_int4_quantizer.h | 163 +++++ .../pq_int8_quantizer/pq_int8_quantizer.cc | 4 +- src/turbo/turbo.cc | 30 +- tests/turbo/turbo_pq_int4_quantizer_test.cc | 592 ++++++++++++++++ 12 files changed, 2129 insertions(+), 10 deletions(-) create mode 100644 src/turbo/distance/avx2/pq_quantizer_int4/pq_distance.cc create mode 100644 src/turbo/distance/avx2/pq_quantizer_int4/pq_distance.h create mode 100644 src/turbo/distance/avx512/pq_quantizer_int4/pq_distance.cc create mode 100644 src/turbo/distance/avx512/pq_quantizer_int4/pq_distance.h create mode 100644 src/turbo/distance/scalar/pq_quantizer_int4/pq_distance.cc create mode 100644 src/turbo/distance/scalar/pq_quantizer_int4/pq_distance.h create mode 100644 src/turbo/quantizer/pq_int4_quantizer/pq_int4_quantizer.cc create mode 100644 src/turbo/quantizer/pq_int4_quantizer/pq_int4_quantizer.h create mode 100644 tests/turbo/turbo_pq_int4_quantizer_test.cc diff --git a/src/include/zvec/turbo/turbo.h b/src/include/zvec/turbo/turbo.h index 4eea20eec..5feaf06fd 100644 --- a/src/include/zvec/turbo/turbo.h +++ b/src/include/zvec/turbo/turbo.h @@ -73,10 +73,15 @@ struct FhtKernels { FhtVecRescaleFunc rescale; }; -// Aggregate of all PQ-specific kernels needed by PqInt8Quantizer, dispatched -// by ISA. 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. +// 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; @@ -143,8 +148,10 @@ 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 quantize_type and CPU arch. -PqKernels get_pq_kernels(QuantizeType quantize_type, +// 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/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..c68c245e3 --- /dev/null +++ b/src/turbo/distance/avx2/pq_quantizer_int4/pq_distance.cc @@ -0,0 +1,226 @@ +// 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..603dddadd --- /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/avx512/pq_quantizer_int4/pq_distance.cc b/src/turbo/distance/avx512/pq_quantizer_int4/pq_distance.cc new file mode 100644 index 000000000..8ce143f1c --- /dev/null +++ b/src/turbo/distance/avx512/pq_quantizer_int4/pq_distance.cc @@ -0,0 +1,212 @@ +// 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..e3fce63d2 --- /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/scalar/pq_quantizer_int4/pq_distance.cc b/src/turbo/distance/scalar/pq_quantizer_int4/pq_distance.cc new file mode 100644 index 000000000..a0a25b13d --- /dev/null +++ b/src/turbo/distance/scalar/pq_quantizer_int4/pq_distance.cc @@ -0,0 +1,109 @@ +// 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..6644ba97e --- /dev/null +++ b/src/turbo/distance/scalar/pq_quantizer_int4/pq_distance.h @@ -0,0 +1,52 @@ +// 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/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..54314e7cb --- /dev/null +++ b/src/turbo/quantizer/pq_int4_quantizer/pq_int4_quantizer.cc @@ -0,0 +1,651 @@ +// 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..709b3a2fc --- /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 = 65535; + 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 index d7be9055e..dfe157d31 100644 --- a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc @@ -66,7 +66,7 @@ int PqInt8Quantizer::init(const IndexMeta &meta, 0.0f); // Dispatch ISA kernels (scalar only for now). - auto pq_k = get_pq_kernels(QuantizeType::kPQ); + 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; @@ -677,7 +677,7 @@ int PqInt8Quantizer::deserialize(const void *data, size_t len) { // offline build, not after deserialization (search uses ADC). // Re-dispatch kernels. - auto pq_k = get_pq_kernels(QuantizeType::kPQ); + 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; diff --git a/src/turbo/turbo.cc b/src/turbo/turbo.cc index d3d0ab7ad..fdb00e30f 100644 --- a/src/turbo/turbo.cc +++ b/src/turbo/turbo.cc @@ -22,6 +22,7 @@ #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__) @@ -29,10 +30,12 @@ #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 @@ -200,12 +203,35 @@ FhtKernels get_fht_kernels() { return k; // scalar } -PqKernels get_pq_kernels(QuantizeType quantize_type, +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) { - // Default: scalar fallback + 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; 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..fa444de9c --- /dev/null +++ b/tests/turbo/turbo_pq_int4_quantizer_test.cc @@ -0,0 +1,592 @@ +// 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< + MultiPassIndexHolder> +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 +} From 5b23f09c200a801970afb48a64d1a3d3de93bbe0 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Tue, 7 Jul 2026 14:25:01 +0800 Subject: [PATCH 40/42] clang-format --- .../algorithm/diskann/diskann_visit_filter.h | 4 +- .../hnsw_sparse/hnsw_sparse_entity.h | 2 +- src/core/utility/visit_filter.h | 4 +- .../inverted_column/inverted_column_indexer.h | 2 +- .../column/inverted_column/inverted_indexer.h | 2 +- .../inverted_column/inverted_search_result.h | 2 +- src/db/index/common/delete_store.h | 2 +- src/db/index/common/id_map.h | 2 +- src/db/index/common/index_params.cc | 3 +- src/db/index/common/meta.h | 9 +- src/db/index/common/schema.cc | 6 +- src/db/index/common/stats.cc | 4 +- src/db/sqlengine/parser/zvec_parser.h | 2 +- src/include/zvec/ailego/pattern/closure.h | 2 +- src/include/zvec/turbo/turbo.h | 14 +- .../avx2/pq_quantizer_int4/pq_distance.cc | 46 ++-- .../avx2/pq_quantizer_int4/pq_distance.h | 10 +- .../avx2/pq_quantizer_int8/pq_distance.cc | 56 ++--- .../avx2/pq_quantizer_int8/pq_distance.h | 10 +- .../avx512/pq_quantizer_int4/pq_distance.cc | 72 +++--- .../avx512/pq_quantizer_int4/pq_distance.h | 10 +- .../avx512/pq_quantizer_int8/pq_distance.cc | 32 +-- .../avx512/pq_quantizer_int8/pq_distance.h | 10 +- .../scalar/pq_quantizer_int4/pq_distance.cc | 13 +- .../scalar/pq_quantizer_int4/pq_distance.h | 11 +- .../scalar/pq_quantizer_int8/pq_distance.cc | 6 +- .../scalar/pq_quantizer_int8/pq_distance.h | 5 +- src/turbo/quantizer/distance.h | 4 +- .../pq_int4_quantizer/pq_int4_quantizer.cc | 109 ++++----- .../pq_int4_quantizer/pq_int4_quantizer.h | 4 +- .../pq_int8_quantizer/pq_int8_quantizer.cc | 119 +++++---- .../pq_int8_quantizer/pq_int8_quantizer.h | 5 +- src/turbo/turbo.cc | 2 +- .../ailego/parallel/multi_thread_list_test.cc | 4 +- tests/ailego/parallel/thread_queue_test.cc | 2 +- tests/android_gmock_main.cc | 2 +- tests/android_gtest_main.cc | 2 +- .../hnsw/hnsw_streamer_buffer_test.cc | 3 +- .../core/utility/buffer_storage_write_test.cc | 30 ++- tests/db/crash_recovery/utility.h | 2 - ...erted_column_indexer_array_numbers_test.cc | 2 +- .../inverted_column_indexer_bool_test.cc | 5 +- ...rted_column_indexer_cyclic_numbers_test.cc | 2 +- ..._column_indexer_sequential_numbers_test.cc | 5 +- .../inverted_column_indexer_string_test.cc | 2 +- tests/db/index/segment/segment_test.cc | 229 +++++++++--------- tests/ios_test_sandbox.cc | 55 +++-- tests/turbo/turbo_pq_int4_quantizer_test.cc | 49 ++-- tests/turbo/turbo_pq_int8_quantizer_test.cc | 74 +++--- 49 files changed, 498 insertions(+), 554 deletions(-) diff --git a/src/core/algorithm/diskann/diskann_visit_filter.h b/src/core/algorithm/diskann/diskann_visit_filter.h index b0b879955..f0ba76cad 100644 --- a/src/core/algorithm/diskann/diskann_visit_filter.h +++ b/src/core/algorithm/diskann/diskann_visit_filter.h @@ -42,7 +42,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}; @@ -358,7 +358,7 @@ class VisitFilter { ByteMap = VisitByteMap::mode }; - VisitFilter() : mode_(0), ctx_(nullptr) {}; + VisitFilter() : mode_(0), ctx_(nullptr){}; inline bool visited(id_t idx) { PROXIMA_DISKANN_VISITFILTER_CALL_IMPL(visited, idx); 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/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/turbo/turbo.h b/src/include/zvec/turbo/turbo.h index 5feaf06fd..36dd8aae1 100644 --- a/src/include/zvec/turbo/turbo.h +++ b/src/include/zvec/turbo/turbo.h @@ -48,21 +48,21 @@ using FhtVecRescaleFunc = void (*)(float *data, size_t n, float factor); // 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); + 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); + 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); + size_t num, size_t num_subquantizers, + float *out); // Aggregate of all FHT kernels needed by FhtRotator, dispatched by ISA. struct FhtKernels { @@ -151,7 +151,7 @@ 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); + QuantizeType quantize_type = QuantizeType::kPQ, + CpuArchType cpu_arch_type = CpuArchType::kAuto); } // namespace zvec::turbo diff --git a/src/turbo/distance/avx2/pq_quantizer_int4/pq_distance.cc b/src/turbo/distance/avx2/pq_quantizer_int4/pq_distance.cc index c68c245e3..657c335d3 100644 --- a/src/turbo/distance/avx2/pq_quantizer_int4/pq_distance.cc +++ b/src/turbo/distance/avx2/pq_quantizer_int4/pq_distance.cc @@ -13,7 +13,6 @@ // limitations under the License. #include "avx2/pq_quantizer_int4/pq_distance.h" - #include namespace zvec::turbo::avx2 { @@ -45,7 +44,7 @@ inline float horizontal_sum_avx2(__m256 v) { // 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 &hi_nib) { __m128i packed128 = _mm_loadl_epi64(reinterpret_cast(packed)); __m256i packed256 = _mm256_castsi128_si256(packed128); @@ -56,11 +55,11 @@ inline void unpack_nibbles(const uint8_t *packed, __m128i &lo_nib, // 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); + __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. @@ -72,13 +71,11 @@ inline void unpack_nibbles(const uint8_t *packed, __m128i &lo_nib, // 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) { + __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); + 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); @@ -88,13 +85,14 @@ inline void accumulate_adc(__m256 &acc, const float *lut, __m128i lo_nib, // 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)); + 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) { + 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(); @@ -117,16 +115,15 @@ void pq_adc_int4_distance_avx2(const void *pq_code_v, const void *lut_v, } 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 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 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(); @@ -167,8 +164,8 @@ void pq_sdc_int4_distance_avx2(const void *a_v, const void *b_v, } 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 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); @@ -206,7 +203,8 @@ void pq_adc_int4_batch_distance_avx2(const void **candidates_v, // 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]; + 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)]; diff --git a/src/turbo/distance/avx2/pq_quantizer_int4/pq_distance.h b/src/turbo/distance/avx2/pq_quantizer_int4/pq_distance.h index 603dddadd..1d4164685 100644 --- a/src/turbo/distance/avx2/pq_quantizer_int4/pq_distance.h +++ b/src/turbo/distance/avx2/pq_quantizer_int4/pq_distance.h @@ -23,19 +23,19 @@ namespace zvec::turbo::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); + 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); + 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); + 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 index 3a74fb84a..18c177069 100644 --- a/src/turbo/distance/avx2/pq_quantizer_int8/pq_distance.cc +++ b/src/turbo/distance/avx2/pq_quantizer_int8/pq_distance.cc @@ -13,7 +13,6 @@ // limitations under the License. #include "avx2/pq_quantizer_int8/pq_distance.h" - #include namespace zvec::turbo::avx2 { @@ -38,7 +37,7 @@ inline float horizontal_sum_avx2(__m256 v) { } // namespace void pq_adc_int8_distance_avx2(const void *pq_code_v, const void *lut_v, - size_t num_subquantizers, float *out) { + 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); @@ -48,10 +47,9 @@ void pq_adc_int8_distance_avx2(const void *pq_code_v, const void *lut_v, // 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); + 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; @@ -59,8 +57,8 @@ void pq_adc_int8_distance_avx2(const void *pq_code_v, const void *lut_v, 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)); + __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] @@ -68,8 +66,7 @@ void pq_adc_int8_distance_avx2(const void *pq_code_v, const void *lut_v, // 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); + __m256 gathered = _mm256_i32gather_ps(lut + m * kNumCentroids, indices, 4); acc = _mm256_add_ps(acc, gathered); } @@ -85,8 +82,8 @@ void pq_adc_int8_distance_avx2(const void *pq_code_v, const void *lut_v, } 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) { + 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; @@ -102,18 +99,15 @@ void pq_sdc_int8_distance_avx2(const void *a_v, const void *b_v, 5 * kTablePerSub, 6 * kTablePerSub, 7 * kTablePerSub); // Multiplier for a[m] * 256 - const __m256i a_multiplier = - _mm256_set1_epi32(kNumCentroids); + 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)); + __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); @@ -132,8 +126,7 @@ void pq_sdc_int8_distance_avx2(const void *a_v, const void *b_v, // Scalar leftover for (; m < num_subquantizers; ++m) { - size_t idx = m * kTablePerSub + - static_cast(a[m]) * kNumCentroids + + size_t idx = m * kTablePerSub + static_cast(a[m]) * kNumCentroids + static_cast(b[m]); sum += dist_table[idx]; } @@ -142,8 +135,8 @@ void pq_sdc_int8_distance_avx2(const void *a_v, const void *b_v, } void pq_adc_int8_batch_distance_avx2(const void **candidates_v, - const void *lut_v, size_t num, - size_t num_subquantizers, float *out) { + 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; @@ -152,10 +145,9 @@ void pq_adc_int8_batch_distance_avx2(const void **candidates_v, 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); + 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) { @@ -181,10 +173,14 @@ void pq_adc_int8_batch_distance_avx2(const void **candidates_v, __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); + __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)); diff --git a/src/turbo/distance/avx2/pq_quantizer_int8/pq_distance.h b/src/turbo/distance/avx2/pq_quantizer_int8/pq_distance.h index f1da57f64..1d1b412cd 100644 --- a/src/turbo/distance/avx2/pq_quantizer_int8/pq_distance.h +++ b/src/turbo/distance/avx2/pq_quantizer_int8/pq_distance.h @@ -23,20 +23,20 @@ namespace zvec::turbo::avx2 { // 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); + 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); + 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); + size_t num, size_t num_subquantizers, + float *out); } // namespace zvec::turbo::avx2 diff --git a/src/turbo/distance/avx512/pq_quantizer_int4/pq_distance.cc b/src/turbo/distance/avx512/pq_quantizer_int4/pq_distance.cc index 8ce143f1c..e93c4f92c 100644 --- a/src/turbo/distance/avx512/pq_quantizer_int4/pq_distance.cc +++ b/src/turbo/distance/avx512/pq_quantizer_int4/pq_distance.cc @@ -13,7 +13,6 @@ // limitations under the License. #include "avx512/pq_quantizer_int4/pq_distance.h" - #include namespace zvec::turbo::avx512 { @@ -43,10 +42,9 @@ inline __m512i unpack_nibbles_512(const uint8_t *packed) { // 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); + 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); @@ -54,7 +52,7 @@ inline __m512i unpack_nibbles_512(const uint8_t *packed) { // 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)); + _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). @@ -64,17 +62,17 @@ inline __m512i unpack_nibbles_512(const uint8_t *packed) { } // namespace void pq_adc_int4_distance_avx512(const void *pq_code_v, const void *lut_v, - size_t num_subquantizers, float *out) { + 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); + 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; @@ -82,8 +80,7 @@ void pq_adc_int4_distance_avx512(const void *pq_code_v, const void *lut_v, 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); + __m512 gathered = _mm512_i32gather_ps(indices, lut + m * kNumCentroids, 4); acc = _mm512_add_ps(acc, gathered); } @@ -98,17 +95,16 @@ void pq_adc_int4_distance_avx512(const void *pq_code_v, const void *lut_v, } 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 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, + 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); @@ -138,18 +134,18 @@ void pq_sdc_int4_distance_avx512(const void *a_v, const void *b_v, } 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 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); + 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) { @@ -170,17 +166,17 @@ void pq_adc_int4_batch_distance_avx512(const void **candidates_v, __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)); + 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)); + 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)); + 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)); + acc3, _mm512_i32gather_ps(_mm512_add_epi32(nib3, base_offsets), + lut_base, 4)); } float s0 = _mm512_reduce_add_ps(acc0); @@ -191,7 +187,8 @@ void pq_adc_int4_batch_distance_avx512(const void **candidates_v, // 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]; + 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)]; @@ -204,8 +201,7 @@ void pq_adc_int4_batch_distance_avx512(const void **candidates_v, } // Remaining candidates: use single ADC kernel. for (; i < num; ++i) { - pq_adc_int4_distance_avx512(candidates[i], lut, num_subquantizers, - out + i); + pq_adc_int4_distance_avx512(candidates[i], lut, num_subquantizers, out + i); } } diff --git a/src/turbo/distance/avx512/pq_quantizer_int4/pq_distance.h b/src/turbo/distance/avx512/pq_quantizer_int4/pq_distance.h index e3fce63d2..6a84bfe5b 100644 --- a/src/turbo/distance/avx512/pq_quantizer_int4/pq_distance.h +++ b/src/turbo/distance/avx512/pq_quantizer_int4/pq_distance.h @@ -23,18 +23,18 @@ namespace zvec::turbo::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); + 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); + 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); + 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 index 2fe2e4204..8ad40e1e4 100644 --- a/src/turbo/distance/avx512/pq_quantizer_int8/pq_distance.cc +++ b/src/turbo/distance/avx512/pq_quantizer_int8/pq_distance.cc @@ -13,7 +13,6 @@ // limitations under the License. #include "avx512/pq_quantizer_int8/pq_distance.h" - #include namespace zvec::turbo::avx512 { @@ -29,7 +28,7 @@ inline float horizontal_sum_avx512(__m512 v) { } // namespace void pq_adc_int8_distance_avx512(const void *pq_code_v, const void *lut_v, - size_t num_subquantizers, float *out) { + 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); @@ -60,8 +59,7 @@ void pq_adc_int8_distance_avx512(const void *pq_code_v, const void *lut_v, __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); + __m512 gathered = _mm512_i32gather_ps(indices, lut + m * kNumCentroids, 4); acc = _mm512_add_ps(acc, gathered); } @@ -77,8 +75,8 @@ void pq_adc_int8_distance_avx512(const void *pq_code_v, const void *lut_v, } 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) { + 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; @@ -92,9 +90,9 @@ void pq_sdc_int8_distance_avx512(const void *a_v, const void *b_v, 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); + 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); @@ -104,10 +102,8 @@ void pq_sdc_int8_distance_avx512(const void *a_v, const void *b_v, // 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)); + __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); @@ -126,8 +122,7 @@ void pq_sdc_int8_distance_avx512(const void *a_v, const void *b_v, // Scalar leftover for (; m < num_subquantizers; ++m) { - size_t idx = m * kTablePerSub + - static_cast(a[m]) * kNumCentroids + + size_t idx = m * kTablePerSub + static_cast(a[m]) * kNumCentroids + static_cast(b[m]); sum += dist_table[idx]; } @@ -136,8 +131,8 @@ void pq_sdc_int8_distance_avx512(const void *a_v, const void *b_v, } void pq_adc_int8_batch_distance_avx512(const void **candidates_v, - const void *lut_v, size_t num, - size_t num_subquantizers, float *out) { + 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; @@ -213,8 +208,7 @@ void pq_adc_int8_batch_distance_avx512(const void **candidates_v, } // Remaining candidates: use single ADC kernel. for (; i < num; ++i) { - pq_adc_int8_distance_avx512(candidates[i], lut, num_subquantizers, - out + i); + pq_adc_int8_distance_avx512(candidates[i], lut, num_subquantizers, out + i); } } diff --git a/src/turbo/distance/avx512/pq_quantizer_int8/pq_distance.h b/src/turbo/distance/avx512/pq_quantizer_int8/pq_distance.h index 5026b4779..2c2ea92d5 100644 --- a/src/turbo/distance/avx512/pq_quantizer_int8/pq_distance.h +++ b/src/turbo/distance/avx512/pq_quantizer_int8/pq_distance.h @@ -22,19 +22,19 @@ 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); + 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); + 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); + size_t num, size_t num_subquantizers, + float *out); } // namespace zvec::turbo::avx512 diff --git a/src/turbo/distance/scalar/pq_quantizer_int4/pq_distance.cc b/src/turbo/distance/scalar/pq_quantizer_int4/pq_distance.cc index a0a25b13d..4684fe67c 100644 --- a/src/turbo/distance/scalar/pq_quantizer_int4/pq_distance.cc +++ b/src/turbo/distance/scalar/pq_quantizer_int4/pq_distance.cc @@ -30,7 +30,7 @@ inline uint8_t decode_nibble(const uint8_t *code, size_t m) { } // namespace void pq_adc_int4_distance(const void *pq_code_v, const void *lut_v, - size_t num_subquantizers, float *out) { + 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); @@ -43,8 +43,8 @@ void pq_adc_int4_distance(const void *pq_code_v, const void *lut_v, } void pq_sdc_int4_distance(const void *a_v, const void *b_v, - const void *dist_table_v, - size_t num_subquantizers, float *out) { + 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); @@ -54,8 +54,7 @@ void pq_sdc_int4_distance(const void *a_v, const void *b_v, 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 + + size_t idx = m * kTablePerSub + static_cast(ai) * kNumCentroids + static_cast(bi); sum += dist_table[idx]; } @@ -63,8 +62,8 @@ void pq_sdc_int4_distance(const void *a_v, const void *b_v, } void pq_adc_int4_batch_distance(const void **candidates_v, const void *lut_v, - size_t num, size_t num_subquantizers, - float *out) { + size_t num, size_t num_subquantizers, + float *out) { constexpr size_t kNumCentroids = 16; const auto *lut = reinterpret_cast(lut_v); const auto *candidates = diff --git a/src/turbo/distance/scalar/pq_quantizer_int4/pq_distance.h b/src/turbo/distance/scalar/pq_quantizer_int4/pq_distance.h index 6644ba97e..561df9b6b 100644 --- a/src/turbo/distance/scalar/pq_quantizer_int4/pq_distance.h +++ b/src/turbo/distance/scalar/pq_quantizer_int4/pq_distance.h @@ -30,7 +30,7 @@ namespace zvec::turbo::scalar { // // 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); + size_t num_subquantizers, float *out); // SDC (Symmetric Distance Computation) for int4 PQ codes. // @@ -38,15 +38,14 @@ void pq_adc_int4_distance(const void *pq_code, const void *lut, // 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); +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); + 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 index 16362f54b..3fb46f1a8 100644 --- a/src/turbo/distance/scalar/pq_quantizer_int8/pq_distance.cc +++ b/src/turbo/distance/scalar/pq_quantizer_int8/pq_distance.cc @@ -38,8 +38,7 @@ void pq_sdc_int8_distance(const void *a_v, const void *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 + + size_t idx = m * kTablePerSub + static_cast(a[m]) * kNumCentroids + static_cast(b[m]); sum += dist_table[idx]; } @@ -53,8 +52,7 @@ void pq_adc_int8_batch_distance(const void **candidates_v, const void *lut_v, 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); + auto candidates = reinterpret_cast(candidates_v); size_t i = 0; // Main loop: process 4 candidates per iteration. diff --git a/src/turbo/distance/scalar/pq_quantizer_int8/pq_distance.h b/src/turbo/distance/scalar/pq_quantizer_int8/pq_distance.h index 1d0c8d810..d968aad7d 100644 --- a/src/turbo/distance/scalar/pq_quantizer_int8/pq_distance.h +++ b/src/turbo/distance/scalar/pq_quantizer_int8/pq_distance.h @@ -36,9 +36,8 @@ void pq_adc_int8_distance(const void *pq_code, const void *lut, // // 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); +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 diff --git a/src/turbo/quantizer/distance.h b/src/turbo/quantizer/distance.h index 47f191890..ecc43f21a 100644 --- a/src/turbo/quantizer/distance.h +++ b/src/turbo/quantizer/distance.h @@ -50,8 +50,8 @@ class DistanceImpl { //! 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) + 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)), diff --git a/src/turbo/quantizer/pq_int4_quantizer/pq_int4_quantizer.cc b/src/turbo/quantizer/pq_int4_quantizer/pq_int4_quantizer.cc index 54314e7cb..c86492a79 100644 --- a/src/turbo/quantizer/pq_int4_quantizer/pq_int4_quantizer.cc +++ b/src/turbo/quantizer/pq_int4_quantizer/pq_int4_quantizer.cc @@ -20,9 +20,9 @@ #include #include #include +#include #include #include -#include namespace zvec { namespace turbo { @@ -37,8 +37,7 @@ struct PqInt4SerPayload { uint32_t num_centroids; // always 16 for int4 }; -int PqInt4Quantizer::init(const IndexMeta &meta, - const ailego::Params ¶ms) { +int PqInt4Quantizer::init(const IndexMeta &meta, const ailego::Params ¶ms) { meta_ = meta; uint32_t d = meta.dimension(); @@ -71,8 +70,7 @@ int PqInt4Quantizer::init(const IndexMeta &meta, // Pre-allocate centroids (filled by train()). centroids_.resize( - static_cast(num_subquantizers_) * kNumCentroids * sub_dim_, - 0.0f); + 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); @@ -85,20 +83,19 @@ int PqInt4Quantizer::init(const IndexMeta &meta, // 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); + 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); + 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); + 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). @@ -134,8 +131,7 @@ 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; + 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)); @@ -177,8 +173,8 @@ void PqInt4Quantizer::train_subquantizer(const float *data, size_t num, sub_idx * d; fp32_l2_batch_fn_(centroid_ptrs.data(), - reinterpret_cast(sub_vec), - k, d, dists.data()); + reinterpret_cast(sub_vec), k, d, + dists.data()); uint32_t best_idx = static_cast( std::min_element(dists.begin(), dists.end()) - dists.begin()); @@ -279,13 +275,14 @@ int PqInt4Quantizer::train(IndexHolder::Pointer holder, int thread_count) { } } - thread_count = std::max(1, std::min(thread_count, - static_cast(num_subquantizers_))); + 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); + 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) { @@ -302,20 +299,20 @@ int PqInt4Quantizer::train(IndexHolder::Pointer holder, int 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); - } - }); + 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); + LOG_INFO(" all %u sub-quantizers trained (%d threads)", num_subquantizers_, + thread_count); } build_centroid_ptrs_cache(); @@ -330,8 +327,7 @@ int PqInt4Quantizer::train(IndexHolder::Pointer holder, int thread_count) { 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); + 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; @@ -340,8 +336,8 @@ void PqInt4Quantizer::compute_dist_table() { 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); + reinterpret_cast(centroids_m + i * d), k, d, + table_m + i * k); } } } @@ -372,8 +368,8 @@ void PqInt4Quantizer::quantize_data(const void *input, void *output) const { 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); + reinterpret_cast(sub_vec), kNumCentroids, + sub_dim_, dists); // Argmin: find nearest centroid (0..15). float best_dist = dists[0]; @@ -434,12 +430,12 @@ float PqInt4Quantizer::calc_distance_dp_query(const void *dp, return d; } -void PqInt4Quantizer::calc_distance_dp_query_batch( - const void *const *dp_list, int dp_num, const void *query, - float *dist_list) const { +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); + 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]; @@ -467,8 +463,7 @@ void PqInt4Quantizer::calc_distance_dp_query_batch_unquantized( kNumCentroids); quantize_query(query, lut.data()); batch_adc_fn_(const_cast(dp_list), lut.data(), - static_cast(dp_num), - num_subquantizers_, dist_list); + 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]; @@ -484,8 +479,7 @@ float PqInt4Quantizer::calc_distance_dp_dp(const void *dp1, } int PqInt4Quantizer::quantize(const void *query, const IndexQueryMeta &qmeta, - std::string *out, - IndexQueryMeta *ometa) const { + std::string *out, IndexQueryMeta *ometa) const { if (qmeta.unit_size() != sizeof(float)) { return kErrUnsupported; } @@ -543,10 +537,8 @@ DistanceImpl PqInt4Quantizer::distance(const void *query, 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); - }; + 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_; @@ -614,9 +606,8 @@ int PqInt4Quantizer::deserialize(const void *data, size_t len) { meta_.set_meta(IndexMeta::DataType::DT_FP32, original_dim_); - size_t centroids_bytes = - static_cast(num_subquantizers_) * kNumCentroids * sub_dim_ * - sizeof(float); + 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); @@ -627,17 +618,17 @@ int PqInt4Quantizer::deserialize(const void *data, size_t len) { 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_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); + fp32_batch_fn_ = + get_batch_distance_func(MetricType::kInnerProduct, DataType::kFp32, + QuantizeType::kDefault, CpuArchType::kAuto); } build_centroid_ptrs_cache(); diff --git a/src/turbo/quantizer/pq_int4_quantizer/pq_int4_quantizer.h b/src/turbo/quantizer/pq_int4_quantizer/pq_int4_quantizer.h index 709b3a2fc..1aa4b8f78 100644 --- a/src/turbo/quantizer/pq_int4_quantizer/pq_int4_quantizer.h +++ b/src/turbo/quantizer/pq_int4_quantizer/pq_int4_quantizer.h @@ -100,8 +100,8 @@ class PqInt4Quantizer : public Quantizer { 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 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; diff --git a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc index dfe157d31..2dd724ccf 100644 --- a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc @@ -20,9 +20,9 @@ #include #include #include +#include #include #include -#include namespace zvec { namespace turbo { @@ -37,8 +37,7 @@ struct PqInt8SerPayload { uint32_t num_centroids; // always 256 for int8 }; -int PqInt8Quantizer::init(const IndexMeta &meta, - const ailego::Params ¶ms) { +int PqInt8Quantizer::init(const IndexMeta &meta, const ailego::Params ¶ms) { meta_ = meta; uint32_t d = meta.dimension(); @@ -62,8 +61,7 @@ int PqInt8Quantizer::init(const IndexMeta &meta, // Pre-allocate centroids (filled by train()). centroids_.resize( - static_cast(num_subquantizers_) * kNumCentroids * sub_dim_, - 0.0f); + static_cast(num_subquantizers_) * kNumCentroids * sub_dim_, 0.0f); // Dispatch ISA kernels (scalar only for now). auto pq_k = get_pq_kernels(DataType::kInt8); @@ -78,24 +76,23 @@ int PqInt8Quantizer::init(const IndexMeta &meta, // 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); + 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); + 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); + 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) @@ -118,7 +115,8 @@ void PqInt8Quantizer::build_centroid_ptrs_cache() { 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; + 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) { @@ -131,8 +129,7 @@ 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; + 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)); @@ -180,8 +177,8 @@ void PqInt8Quantizer::train_subquantizer(const float *data, size_t num, sub_idx * d; fp32_l2_batch_fn_(centroid_ptrs.data(), - reinterpret_cast(sub_vec), - k, d, dists.data()); + reinterpret_cast(sub_vec), k, d, + dists.data()); uint32_t best_idx = static_cast( std::min_element(dists.begin(), dists.end()) - dists.begin()); @@ -291,20 +288,20 @@ int PqInt8Quantizer::train(IndexHolder::Pointer holder, int thread_count) { } // Clamp thread count to [1, num_subquantizers_]. - thread_count = std::max(1, std::min(thread_count, - static_cast(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); + 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_); + LOG_INFO(" sub-quantizer [%u/%u] done", m + 1, num_subquantizers_); } } else { // Multi-threaded path: each thread handles a contiguous range of @@ -319,20 +316,20 @@ int PqInt8Quantizer::train(IndexHolder::Pointer holder, int 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); - } - }); + 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); + LOG_INFO(" all %u sub-quantizers trained (%d threads)", num_subquantizers_, + thread_count); } // Pre-build centroid pointer cache (needed by compute_dist_table). @@ -349,8 +346,7 @@ int PqInt8Quantizer::train(IndexHolder::Pointer holder, int thread_count) { 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); + 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_). @@ -367,8 +363,8 @@ void PqInt8Quantizer::compute_dist_table() { 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); + reinterpret_cast(centroids_m + i * d), k, d, + table_m + i * k); } } } @@ -401,8 +397,8 @@ void PqInt8Quantizer::quantize_data(const void *input, void *output) const { // 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); + reinterpret_cast(sub_vec), kNumCentroids, + sub_dim_, dists); // Argmin: find nearest centroid. float best_dist = dists[0]; @@ -468,16 +464,16 @@ float PqInt8Quantizer::calc_distance_dp_query(const void *dp, return d; } -void PqInt8Quantizer::calc_distance_dp_query_batch( - const void *const *dp_list, int dp_num, const void *query, - float *dist_list) const { +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); + 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) { @@ -493,8 +489,8 @@ float PqInt8Quantizer::calc_distance_dp_query_unquantized( kNumCentroids); quantize_query(query, lut.data()); float d = 0.0f; - adc_fn_(reinterpret_cast(dp), lut.data(), - num_subquantizers_, &d); + adc_fn_(reinterpret_cast(dp), lut.data(), num_subquantizers_, + &d); if (meta_.metric_name() == "Cosine") { d = 1.0f + d; } @@ -510,8 +506,7 @@ void PqInt8Quantizer::calc_distance_dp_query_batch_unquantized( // 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); + 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]; @@ -529,8 +524,7 @@ float PqInt8Quantizer::calc_distance_dp_dp(const void *dp1, } int PqInt8Quantizer::quantize(const void *query, const IndexQueryMeta &qmeta, - std::string *out, - IndexQueryMeta *ometa) const { + std::string *out, IndexQueryMeta *ometa) const { if (qmeta.unit_size() != sizeof(float)) { return kErrUnsupported; } @@ -589,10 +583,8 @@ DistanceImpl PqInt8Quantizer::distance(const void *query, // (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); - }; + 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_; @@ -667,9 +659,8 @@ int PqInt8Quantizer::deserialize(const void *data, size_t len) { meta_.set_meta(IndexMeta::DataType::DT_FP32, original_dim_); - size_t centroids_bytes = - static_cast(num_subquantizers_) * kNumCentroids * sub_dim_ * - sizeof(float); + 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); @@ -683,9 +674,9 @@ int PqInt8Quantizer::deserialize(const void *data, size_t len) { 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); + 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). @@ -693,9 +684,9 @@ int PqInt8Quantizer::deserialize(const void *data, size_t len) { 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); + fp32_batch_fn_ = + get_batch_distance_func(MetricType::kInnerProduct, DataType::kFp32, + QuantizeType::kDefault, CpuArchType::kAuto); } // Pre-build centroid pointer cache for fast encode/search. diff --git a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h index 6417170c2..c1e33e28b 100644 --- a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h @@ -95,8 +95,8 @@ class PqInt8Quantizer : public Quantizer { 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 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; @@ -162,7 +162,6 @@ class PqInt8Quantizer : public Quantizer { //! This separation ensures encoding and search LUT use the correct //! distance semantics independently. BatchDistanceFunc fp32_l2_batch_fn_{}; - }; } // namespace turbo diff --git a/src/turbo/turbo.cc b/src/turbo/turbo.cc index fdb00e30f..c7655a4dc 100644 --- a/src/turbo/turbo.cc +++ b/src/turbo/turbo.cc @@ -204,7 +204,7 @@ FhtKernels get_fht_kernels() { } PqKernels get_pq_kernels(DataType data_type, QuantizeType quantize_type, - CpuArchType cpu_arch_type) { + CpuArchType cpu_arch_type) { (void)cpu_arch_type; // currently unused, reserved for future use PqKernels k{}; if (quantize_type == QuantizeType::kPQ) { 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/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/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/turbo_pq_int4_quantizer_test.cc b/tests/turbo/turbo_pq_int4_quantizer_test.cc index fa444de9c..86249f156 100644 --- a/tests/turbo/turbo_pq_int4_quantizer_test.cc +++ b/tests/turbo/turbo_pq_int4_quantizer_test.cc @@ -57,19 +57,16 @@ static std::shared_ptr make_pq4_quantizer( meta.set_metric("SquaredEuclidean", 0, Params()); Params params; - params.set("num_subquantizers", - static_cast(num_subquantizers)); + 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< - MultiPassIndexHolder> +static std::shared_ptr> make_random_holder(size_t count, size_t dim, uint32_t seed = 42) { auto holder = - std::make_shared>( - dim); + 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) { @@ -172,8 +169,8 @@ TEST(PqInt4Quantizer, AdcDistance) { 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 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) { @@ -314,7 +311,7 @@ inline uint8_t test_decode_nibble(const uint8_t *code, size_t m) { } void fill_random_int4_codes(uint8_t *codes, size_t num_subquantizers, - std::mt19937 &gen) { + std::mt19937 &gen) { std::uniform_int_distribution dist(0, 15); size_t code_bytes = num_subquantizers / 2; std::memset(codes, 0, code_bytes); @@ -329,7 +326,7 @@ void fill_random_int4_codes(uint8_t *codes, size_t num_subquantizers, } void fill_random_int4_lut(float *lut, size_t num_subquantizers, - std::mt19937 &gen) { + 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) { @@ -361,8 +358,8 @@ TEST(PqInt4Kernel, AdcDistance) { } float result = 0.0f; - zvec::turbo::scalar::pq_adc_int4_distance(codes.data(), lut.data(), - nsq, &result); + zvec::turbo::scalar::pq_adc_int4_distance(codes.data(), lut.data(), nsq, + &result); EXPECT_NEAR(ref, result, 1e-5f) << "ADC mismatch for nsq=" << nsq; } } @@ -389,16 +386,14 @@ TEST(PqInt4Kernel, SdcDistance) { 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 + + 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); + dist_table.data(), nsq, &result); EXPECT_NEAR(ref, result, 1e-5f) << "SDC mismatch for nsq=" << nsq; } } @@ -411,7 +406,7 @@ TEST(PqInt4Kernel, BatchAdcDistance) { const size_t num_candidates = 7; // not multiple of 4 to test leftover std::vector> codes(num_candidates, - std::vector(nsq / 2)); + std::vector(nsq / 2)); std::vector lut(nsq * kNumCentroids); for (auto &c : codes) fill_random_int4_codes(c.data(), nsq, gen); @@ -448,7 +443,7 @@ TEST(PqInt4Kernel, BatchAdcDistance) { // 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) { + 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) { @@ -475,14 +470,14 @@ TEST(PqInt4SimdConsistency, AdcDistance) { 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); + 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); + num_sq, &avx2_result); EXPECT_NEAR(scalar_result, avx2_result, 1e-5f) << "AVX2 ADC mismatch for M=" << num_sq; } @@ -491,8 +486,8 @@ TEST(PqInt4SimdConsistency, AdcDistance) { #if defined(__AVX512F__) { float avx512_result = 0.0f; - zvec::turbo::avx512::pq_adc_int4_distance_avx512( - codes.data(), lut.data(), num_sq, &avx512_result); + 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; } @@ -516,8 +511,8 @@ TEST(PqInt4SimdConsistency, SdcDistance) { 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); + dist_table.data(), num_sq, + &scalar_result); #if defined(__AVX2__) { @@ -547,11 +542,11 @@ TEST(PqInt4SimdConsistency, SdcDistance) { TEST(PqInt4SimdConsistency, BatchAdcDistance) { std::mt19937 gen(2026); constexpr size_t kNumCentroids = 16; - const size_t nsq = 16; // chunk-aligned + 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(nsq / 2)); std::vector lut(nsq * kNumCentroids); for (auto &c : codes) fill_random_int4_codes(c.data(), nsq, gen); diff --git a/tests/turbo/turbo_pq_int8_quantizer_test.cc b/tests/turbo/turbo_pq_int8_quantizer_test.cc index 06c6fd7bb..438b19c96 100644 --- a/tests/turbo/turbo_pq_int8_quantizer_test.cc +++ b/tests/turbo/turbo_pq_int8_quantizer_test.cc @@ -57,19 +57,16 @@ static std::shared_ptr make_pq_quantizer( meta.set_metric("SquaredEuclidean", 0, Params()); Params params; - params.set("num_subquantizers", - static_cast(num_subquantizers)); + 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< - MultiPassIndexHolder> +static std::shared_ptr> make_random_holder(size_t count, size_t dim, uint32_t seed = 42) { auto holder = - std::make_shared>( - dim); + 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) { @@ -160,8 +157,8 @@ TEST(PqInt8Quantizer, AdcDistance) { // 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 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) { @@ -314,8 +311,7 @@ void fill_random_codes(uint8_t *codes, size_t len, std::mt19937 &gen) { } // Helper to generate random LUT (ADC) -void fill_random_lut(float *lut, size_t num_subquantizers, - std::mt19937 &gen) { +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) { @@ -357,14 +353,14 @@ TEST(PqInt8SimdConsistency, AdcDistance) { // Compute reference (scalar) float scalar_result = 0.0f; - zvec::turbo::scalar::pq_adc_int8_distance(codes.data(), lut.data(), - num_sq, &scalar_result); + 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); + num_sq, &avx2_result); EXPECT_NEAR(scalar_result, avx2_result, 1e-5f) << "AVX2 ADC mismatch for M=" << num_sq; } @@ -373,9 +369,8 @@ TEST(PqInt8SimdConsistency, AdcDistance) { #if defined(__AVX512F__) { float avx512_result = 0.0f; - zvec::turbo::avx512::pq_adc_int8_distance_avx512(codes.data(), - lut.data(), num_sq, - &avx512_result); + 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; } @@ -400,16 +395,15 @@ TEST(PqInt8SimdConsistency, SdcDistance) { // 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); + 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); + 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; } @@ -418,10 +412,9 @@ TEST(PqInt8SimdConsistency, SdcDistance) { #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); + 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; } @@ -443,13 +436,13 @@ TEST(PqInt8SimdConsistency, AdcDistanceM1) { float scalar_result = 0.0f; zvec::turbo::scalar::pq_adc_int8_distance(codes.data(), lut.data(), num_sq, - &scalar_result); + &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); + num_sq, &avx2_result); EXPECT_NEAR(scalar_result, avx2_result, 1e-5f); } #endif @@ -458,7 +451,7 @@ TEST(PqInt8SimdConsistency, AdcDistanceM1) { { float avx512_result = 0.0f; zvec::turbo::avx512::pq_adc_int8_distance_avx512(codes.data(), lut.data(), - num_sq, &avx512_result); + num_sq, &avx512_result); EXPECT_NEAR(scalar_result, avx512_result, 1e-5f); } #endif @@ -479,8 +472,7 @@ static std::shared_ptr make_pq_cosine_quantizer( meta.set_metric("Cosine", 0, Params()); Params params; - params.set("num_subquantizers", - static_cast(num_subquantizers)); + params.set("num_subquantizers", static_cast(num_subquantizers)); if (q->init(meta, params) != 0) return nullptr; return q; } @@ -500,12 +492,10 @@ static float reference_cosine_distance(const float *a, const float *b, } // Helper: generate random vectors with varying norms (not unit length). -static std::shared_ptr< - MultiPassIndexHolder> +static std::shared_ptr> make_cosine_holder(size_t count, size_t dim, uint32_t seed = 42) { auto holder = - std::make_shared>( - dim); + 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); @@ -655,8 +645,8 @@ TEST(PqInt8Quantizer, CosineAdcDistance) { 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 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); @@ -724,15 +714,13 @@ static std::shared_ptr make_pq_ip_quantizer( meta.set_metric("InnerProduct", 0, Params()); Params params; - params.set("num_subquantizers", - static_cast(num_subquantizers)); + 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) { +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]; @@ -775,8 +763,8 @@ TEST(PqInt8Quantizer, InnerProductAdcDistance) { 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 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); From 06da5d7505bb2b639afe976295f9d1a7acb4cd4f Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Tue, 7 Jul 2026 16:16:39 +0800 Subject: [PATCH 41/42] debug --- src/turbo/quantizer/pq_int4_quantizer/pq_int4_quantizer.h | 2 +- src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/turbo/quantizer/pq_int4_quantizer/pq_int4_quantizer.h b/src/turbo/quantizer/pq_int4_quantizer/pq_int4_quantizer.h index 1aa4b8f78..2a1649cfc 100644 --- a/src/turbo/quantizer/pq_int4_quantizer/pq_int4_quantizer.h +++ b/src/turbo/quantizer/pq_int4_quantizer/pq_int4_quantizer.h @@ -128,7 +128,7 @@ class PqInt4Quantizer : public Quantizer { static constexpr uint32_t kNumCentroids = 16; static constexpr uint32_t kMaxKmeansIters = 25; - static constexpr size_t kMaxTrainVectors = 65535; + static constexpr size_t kMaxTrainVectors = 65536; static constexpr uint32_t kExtraMetaSizeCosine = sizeof(float); IndexMeta meta_{}; diff --git a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h index c1e33e28b..deb661941 100644 --- a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h @@ -125,7 +125,7 @@ class PqInt8Quantizer : public Quantizer { static constexpr uint32_t kNumCentroids = 256; static constexpr uint32_t kMaxKmeansIters = 25; - static constexpr size_t kMaxTrainVectors = 65535; + static constexpr size_t kMaxTrainVectors = 65536; static constexpr uint32_t kExtraMetaSizeCosine = sizeof(float); IndexMeta meta_{}; From abc7a4b90087c6d5d2a8db92451bb91fc805d7d8 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Tue, 7 Jul 2026 20:30:43 +0800 Subject: [PATCH 42/42] remove train to create int fht --- .../quantizer/preprocessor/fht_rotator.cc | 20 ++++++++-------- .../quantizer/preprocessor/fht_rotator.h | 9 +++---- tests/turbo/turbo_fht_rotator_test.cc | 24 ++++--------------- 3 files changed, 20 insertions(+), 33 deletions(-) diff --git a/src/turbo/quantizer/preprocessor/fht_rotator.cc b/src/turbo/quantizer/preprocessor/fht_rotator.cc index c600ee163..4d31f6949 100644 --- a/src/turbo/quantizer/preprocessor/fht_rotator.cc +++ b/src/turbo/quantizer/preprocessor/fht_rotator.cc @@ -46,7 +46,15 @@ FhtRotator::Pointer FhtRotator::create(int dim) { r->inv_kacs_walk_fn_ = k.inv_kacs_walk; r->inplace_fn_ = k.inplace; r->rescale_fn_ = k.rescale; - // flip_ stays empty until train() is called. + + // 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; } @@ -72,15 +80,7 @@ FhtRotator::Pointer FhtRotator::from_blob(const void *data, size_t len) { void FhtRotator::train(const void * /*data*/, size_t /*num*/, size_t /*stride*/) { - if (in_dim_ <= 0) return; - - flip_offset_ = (static_cast(in_dim_) + kByteLen - 1) / kByteLen; - flip_.resize(4 * flip_offset_); - - std::random_device rd; - std::mt19937 gen(rd()); - std::uniform_int_distribution dist(0, 255); - for (auto &b : flip_) b = static_cast(dist(gen)); + // No-op: flip-sign arrays are generated in create(). } // --------------------------------------------------------------------------- diff --git a/src/turbo/quantizer/preprocessor/fht_rotator.h b/src/turbo/quantizer/preprocessor/fht_rotator.h index 91c670f69..1bad6e924 100644 --- a/src/turbo/quantizer/preprocessor/fht_rotator.h +++ b/src/turbo/quantizer/preprocessor/fht_rotator.h @@ -36,8 +36,9 @@ class FhtRotator : public Preprocessor { public: using Pointer = std::shared_ptr; - //! Create an untrained rotator for \p in_dim dimensions. - //! Call train() afterwards to generate the random flip-sign arrays. + //! 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 @@ -56,8 +57,8 @@ class FhtRotator : public Preprocessor { void apply(const float *in, float *out) const override; void apply_inverse(const float *in, float *out) const override; - //! Generate 4 rounds of random flip-sign arrays. - //! For FhtRotator the training data is ignored; only \p in_dim() matters. + //! 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; diff --git a/tests/turbo/turbo_fht_rotator_test.cc b/tests/turbo/turbo_fht_rotator_test.cc index 8d7f06edc..4d5ca82a6 100644 --- a/tests/turbo/turbo_fht_rotator_test.cc +++ b/tests/turbo/turbo_fht_rotator_test.cc @@ -56,7 +56,6 @@ TEST(FhtRotator, PowerOf2RoundTrip) { 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; - rot->train(nullptr, 0, 0); std::vector input(dim); fill_random(input.data(), dim, gen); @@ -73,7 +72,6 @@ TEST(FhtRotator, NonPowerOf2RoundTrip) { 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; - rot->train(nullptr, 0, 0); std::vector input(dim); fill_random(input.data(), dim, gen); @@ -88,10 +86,9 @@ TEST(FhtRotator, NonPowerOf2RoundTrip) { TEST(FhtRotator, SerializeDeserialize) { std::mt19937 gen(999); for (int dim : {32, 97, 128}) { - // Build and train original rotator. + // Build original rotator. auto rot = FhtRotator::create(dim); ASSERT_TRUE(rot); - rot->train(nullptr, 0, 0); // Serialize. std::string blob; @@ -140,19 +137,16 @@ TEST(FhtRotator, DimensionPreserved) { // Train generates non-zero flip signs // --------------------------------------------------------------------------- -TEST(FhtRotator, TrainGeneratesFlip) { +TEST(FhtRotator, CreateGeneratesFlip) { for (int dim : {8, 64, 97}) { auto rot = FhtRotator::create(dim); ASSERT_TRUE(rot); - // Before train, apply should not crash but flip is empty — we skip - // calling apply before train. After train, serialize must succeed - // (which requires flip to be populated). - rot->train(nullptr, 0, 0); + // After create, flip is already populated, so serialize must succeed. std::string blob; EXPECT_EQ(0, rot->serialize(&blob)) - << "serialize failed after train for dim=" << dim; + << "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()); @@ -217,7 +211,6 @@ TEST(FhtRotator, L2DistancePreserved) { for (int dim : {32, 64, 97, 128}) { auto rot = FhtRotator::create(dim); ASSERT_TRUE(rot); - rot->train(nullptr, 0, 0); const int N = 50; std::vector> raw(N, std::vector(dim)); @@ -264,7 +257,6 @@ TEST(FhtRotator, CosineDistancePreserved) { for (int dim : {32, 97, 128}) { auto rot = FhtRotator::create(dim); ASSERT_TRUE(rot); - rot->train(nullptr, 0, 0); const int N = 50; std::vector> raw(N, std::vector(dim)); @@ -292,7 +284,6 @@ TEST(FhtRotator, ApplyIsNonTrivial) { for (int dim : {32, 97, 128}) { auto rot = FhtRotator::create(dim); ASSERT_TRUE(rot); - rot->train(nullptr, 0, 0); std::vector input(dim); fill_random(input.data(), dim, gen); @@ -321,7 +312,6 @@ TEST(FhtRotator, ApplyDeterministic) { for (int dim : {32, 97, 128}) { auto rot = FhtRotator::create(dim); ASSERT_TRUE(rot); - rot->train(nullptr, 0, 0); std::vector input(dim); fill_random(input.data(), dim, gen); @@ -344,10 +334,9 @@ TEST(FhtRotator, ApplyDeterministic) { TEST(FhtRotator, DeserializeOnExistingObject) { std::mt19937 gen(314); for (int dim : {32, 97, 128}) { - // Build and train original. + // Build original. auto rot1 = FhtRotator::create(dim); ASSERT_TRUE(rot1); - rot1->train(nullptr, 0, 0); std::string blob; ASSERT_EQ(0, rot1->serialize(&blob)); @@ -385,7 +374,6 @@ TEST(FhtRotator, DeserializeTruncatedPayload) { std::mt19937 gen(42); auto rot = FhtRotator::create(64); ASSERT_TRUE(rot); - rot->train(nullptr, 0, 0); std::string blob; ASSERT_EQ(0, rot->serialize(&blob)); @@ -411,7 +399,6 @@ TEST(FhtRotator, LargeDimension) { for (int dim : {1024, 2048, 4096}) { auto rot = FhtRotator::create(dim); ASSERT_TRUE(rot) << "create failed for dim=" << dim; - rot->train(nullptr, 0, 0); std::vector input(dim); fill_random(input.data(), dim, gen); @@ -441,7 +428,6 @@ TEST(FhtRotator, NormPreserved) { for (int dim : {32, 97, 128}) { auto rot = FhtRotator::create(dim); ASSERT_TRUE(rot); - rot->train(nullptr, 0, 0); std::vector input(dim); fill_random(input.data(), dim, gen);