From b1313a466a883f440976ac60c13805078f7d0ac6 Mon Sep 17 00:00:00 2001 From: ray Date: Mon, 29 Jun 2026 19:47:27 +0800 Subject: [PATCH 001/100] 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 002/100] 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 003/100] 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 004/100] 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 005/100] 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 006/100] 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 007/100] 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 008/100] 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 009/100] 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 010/100] 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 011/100] 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 012/100] 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 013/100] 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 014/100] 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 015/100] 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 016/100] 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 017/100] 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 018/100] 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 019/100] 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 020/100] 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 021/100] 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 022/100] 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 023/100] 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 024/100] 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 025/100] 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 026/100] 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 027/100] 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 028/100] 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 029/100] 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 030/100] 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 031/100] 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 032/100] 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 033/100] 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 034/100] 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 035/100] 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 036/100] 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 037/100] 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 038/100] 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 039/100] 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 040/100] 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 041/100] 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 8406e9fdc34e45e0990daba6f527faff0f7022d4 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Tue, 7 Jul 2026 20:30:43 +0800 Subject: [PATCH 042/100] 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); 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 043/100] 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); From 8edecca420b934f258ba9405444403d93a6068fc Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Thu, 9 Jul 2026 10:24:44 +0800 Subject: [PATCH 044/100] clang-format --- src/core/algorithm/hnsw_sparse/hnsw_sparse_entity.h | 2 +- src/core/utility/visit_filter.h | 4 ++-- .../column/inverted_column/inverted_column_indexer.h | 2 +- src/db/index/column/inverted_column/inverted_indexer.h | 2 +- .../column/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/meta.h | 9 +++++---- src/db/sqlengine/parser/zvec_parser.h | 2 +- src/include/zvec/ailego/pattern/closure.h | 2 +- 10 files changed, 15 insertions(+), 14 deletions(-) diff --git a/src/core/algorithm/hnsw_sparse/hnsw_sparse_entity.h b/src/core/algorithm/hnsw_sparse/hnsw_sparse_entity.h index ee4730a7e..7e6f814b1 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 853110ac1..959a0b450 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 56b8b6f57..94e3c05b1 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 13884ab43..19c7b392d 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 3bf391f9f..3c9720a6e 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 56ac8733c..e6b742758 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 67b0e8ce0..59e6b3ecc 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/meta.h b/src/db/index/common/meta.h index 81a5f0fc9..272e831ef 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,7 +368,8 @@ 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/sqlengine/parser/zvec_parser.h b/src/db/sqlengine/parser/zvec_parser.h index 2e1152b76..44102d0aa 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 554e8a77e..df624ac92 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 From d600917de422d38658c2e31a3ffc36b7a8d61616 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Thu, 9 Jul 2026 12:00:41 +0800 Subject: [PATCH 045/100] clang-format --- src/include/zvec/turbo/turbo.h | 6 ------ 1 file changed, 6 deletions(-) diff --git a/src/include/zvec/turbo/turbo.h b/src/include/zvec/turbo/turbo.h index 36dd8aae1..9a78a8349 100644 --- a/src/include/zvec/turbo/turbo.h +++ b/src/include/zvec/turbo/turbo.h @@ -73,12 +73,6 @@ struct FhtKernels { FhtVecRescaleFunc rescale; }; -// Aggregate of all PQ-specific kernels needed by PqInt8Quantizer / -// PqInt4Quantizer, dispatched by ISA and data_type. LUT computation -// (compute_distance_table) is NOT included here — it reuses the existing -// fp32 BatchDistanceFunc from get_batch_distance_func() which is metric- -// aware and already SIMD-optimized. -// // data_type selects the code packing layout: // kInt8: one uint8 per sub-quantizer (256 centroids, stride=256) // kInt4: two sub-quantizers packed into one uint8 (16 centroids, stride=16) From ac17ee916c1d0851e0b18a9d541b26fbf6888f87 Mon Sep 17 00:00:00 2001 From: ray Date: Fri, 10 Jul 2026 12:45:00 +0800 Subject: [PATCH 046/100] fix: add neon --- src/include/zvec/turbo/turbo.h | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/include/zvec/turbo/turbo.h b/src/include/zvec/turbo/turbo.h index 2eff572d5..0b0e03c2a 100644 --- a/src/include/zvec/turbo/turbo.h +++ b/src/include/zvec/turbo/turbo.h @@ -62,12 +62,14 @@ enum class QuantizeType { enum class CpuArchType { kAuto, kScalar, + // x86 SIMD kSSE, kAVX, kAVX2, kAVX512, kAVX512VNNI, - kAVX512FP16 + kAVX512FP16, + kNEON }; DistanceFunc get_distance_func(MetricType metric_type, DataType data_type, From 7bb01c35f7c809a1ab873071beb4c151ab3c9066 Mon Sep 17 00:00:00 2001 From: ray Date: Fri, 10 Jul 2026 12:46:01 +0800 Subject: [PATCH 047/100] fix: add arm arch --- src/include/zvec/turbo/turbo.h | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/include/zvec/turbo/turbo.h b/src/include/zvec/turbo/turbo.h index 0b0e03c2a..e6f1587ac 100644 --- a/src/include/zvec/turbo/turbo.h +++ b/src/include/zvec/turbo/turbo.h @@ -69,7 +69,10 @@ enum class CpuArchType { kAVX512, kAVX512VNNI, kAVX512FP16, - kNEON + // ARM SIMD + kNEON, + kSVE, + kSVE2 }; DistanceFunc get_distance_func(MetricType metric_type, DataType data_type, From ea71d23b8412d2b53e626600f9d3c5d68803b0a4 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Fri, 10 Jul 2026 13:50:49 +0800 Subject: [PATCH 048/100] CpuArchType --- src/include/zvec/turbo/turbo.h | 2 +- src/turbo/turbo.cc | 35 +++++++++++++++++++++++++--------- 2 files changed, 27 insertions(+), 10 deletions(-) diff --git a/src/include/zvec/turbo/turbo.h b/src/include/zvec/turbo/turbo.h index 9a78a8349..9d5ba8a5b 100644 --- a/src/include/zvec/turbo/turbo.h +++ b/src/include/zvec/turbo/turbo.h @@ -140,7 +140,7 @@ QueryPreprocessFunc get_query_preprocess_func( UniformQuantizeFunc get_uniform_quantize_func(DataType data_type); // Returns all FHT kernels dispatched for the current CPU. -FhtKernels get_fht_kernels(); +FhtKernels get_fht_kernels(CpuArchType cpu_arch_type = CpuArchType::kAuto); // 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). diff --git a/src/turbo/turbo.cc b/src/turbo/turbo.cc index c7655a4dc..9c9ffd7de 100644 --- a/src/turbo/turbo.cc +++ b/src/turbo/turbo.cc @@ -160,7 +160,9 @@ UniformQuantizeFunc get_uniform_quantize_func(DataType data_type) { return nullptr; } -FhtKernels get_fht_kernels() { +FhtKernels get_fht_kernels(CpuArchType cpu_arch_type) { + // Suppress unused-parameter warning when no SIMD #if blocks are compiled in. + (void)cpu_arch_type; FhtKernels k; // Default: scalar fallback for all k.flip_sign = scalar::fht_flip_sign; @@ -171,7 +173,9 @@ FhtKernels get_fht_kernels() { #if defined(__AVX512F__) if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F && - zvec::ailego::internal::CpuFeatures::static_flags_.AVX512DQ) { + zvec::ailego::internal::CpuFeatures::static_flags_.AVX512DQ && + (cpu_arch_type == CpuArchType::kAuto || + cpu_arch_type == CpuArchType::kAVX512)) { 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; @@ -181,7 +185,9 @@ FhtKernels get_fht_kernels() { } #endif #if defined(__AVX2__) - if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) { + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2 && + (cpu_arch_type == CpuArchType::kAuto || + cpu_arch_type == CpuArchType::kAVX2)) { 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; @@ -191,7 +197,9 @@ FhtKernels get_fht_kernels() { } #endif #if defined(__SSE2__) - if (zvec::ailego::internal::CpuFeatures::static_flags_.SSE2) { + if (zvec::ailego::internal::CpuFeatures::static_flags_.SSE2 && + (cpu_arch_type == CpuArchType::kAuto || + cpu_arch_type == CpuArchType::kSSE)) { 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; @@ -205,7 +213,8 @@ FhtKernels get_fht_kernels() { PqKernels get_pq_kernels(DataType data_type, QuantizeType quantize_type, CpuArchType cpu_arch_type) { - (void)cpu_arch_type; // currently unused, reserved for future use + // Suppress unused-parameter warning when no SIMD #if blocks are compiled in. + (void)cpu_arch_type; PqKernels k{}; if (quantize_type == QuantizeType::kPQ) { if (data_type == DataType::kInt4) { @@ -215,7 +224,9 @@ PqKernels get_pq_kernels(DataType data_type, QuantizeType quantize_type, k.batch_adc_distance = scalar::pq_adc_int4_batch_distance; #if defined(__AVX512F__) - if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) { + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F && + (cpu_arch_type == CpuArchType::kAuto || + cpu_arch_type == CpuArchType::kAVX512)) { 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; @@ -223,7 +234,9 @@ PqKernels get_pq_kernels(DataType data_type, QuantizeType quantize_type, } #endif #if defined(__AVX2__) - if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) { + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2 && + (cpu_arch_type == CpuArchType::kAuto || + cpu_arch_type == CpuArchType::kAVX2)) { 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; @@ -237,7 +250,9 @@ PqKernels get_pq_kernels(DataType data_type, QuantizeType quantize_type, k.batch_adc_distance = scalar::pq_adc_int8_batch_distance; #if defined(__AVX512F__) - if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) { + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F && + (cpu_arch_type == CpuArchType::kAuto || + cpu_arch_type == CpuArchType::kAVX512)) { 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; @@ -245,7 +260,9 @@ PqKernels get_pq_kernels(DataType data_type, QuantizeType quantize_type, } #endif #if defined(__AVX2__) - if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) { + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2 && + (cpu_arch_type == CpuArchType::kAuto || + cpu_arch_type == CpuArchType::kAVX2)) { 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; From 6abf1b8b25e6ccb9ff87343b20511eb9785859e7 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Fri, 10 Jul 2026 15:30:34 +0800 Subject: [PATCH 049/100] CpuArchType FHT --- src/include/zvec/turbo/turbo.h | 2 +- src/turbo/turbo.cc | 16 ++++++++++++---- 2 files changed, 13 insertions(+), 5 deletions(-) diff --git a/src/include/zvec/turbo/turbo.h b/src/include/zvec/turbo/turbo.h index c3506f68f..8e8d687dd 100644 --- a/src/include/zvec/turbo/turbo.h +++ b/src/include/zvec/turbo/turbo.h @@ -108,6 +108,6 @@ QueryPreprocessFunc get_query_preprocess_func( UniformQuantizeFunc get_uniform_quantize_func(DataType data_type); // Returns all FHT kernels dispatched for the current CPU. -FhtKernels get_fht_kernels(); +FhtKernels get_fht_kernels(CpuArchType cpu_arch_type = CpuArchType::kAuto); } // namespace zvec::turbo diff --git a/src/turbo/turbo.cc b/src/turbo/turbo.cc index 8e2e1b617..b5b3e9e51 100644 --- a/src/turbo/turbo.cc +++ b/src/turbo/turbo.cc @@ -154,7 +154,9 @@ UniformQuantizeFunc get_uniform_quantize_func(DataType data_type) { return nullptr; } -FhtKernels get_fht_kernels() { +FhtKernels get_fht_kernels(CpuArchType cpu_arch_type) { + // Suppress unused-parameter warning when no SIMD #if blocks are compiled in. + (void)cpu_arch_type; FhtKernels k; // Default: scalar fallback for all k.flip_sign = scalar::fht_flip_sign; @@ -165,7 +167,9 @@ FhtKernels get_fht_kernels() { #if defined(__AVX512F__) if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F && - zvec::ailego::internal::CpuFeatures::static_flags_.AVX512DQ) { + zvec::ailego::internal::CpuFeatures::static_flags_.AVX512DQ && + (cpu_arch_type == CpuArchType::kAuto || + cpu_arch_type == CpuArchType::kAVX512)) { 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; @@ -175,7 +179,9 @@ FhtKernels get_fht_kernels() { } #endif #if defined(__AVX2__) - if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) { + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2 && + (cpu_arch_type == CpuArchType::kAuto || + cpu_arch_type == CpuArchType::kAVX2)) { 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; @@ -185,7 +191,9 @@ FhtKernels get_fht_kernels() { } #endif #if defined(__SSE2__) - if (zvec::ailego::internal::CpuFeatures::static_flags_.SSE2) { + if (zvec::ailego::internal::CpuFeatures::static_flags_.SSE2 && + (cpu_arch_type == CpuArchType::kAuto || + cpu_arch_type == CpuArchType::kSSE)) { 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; From ea816ccb6200acbf00825a67a7ba840b65f79172 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Fri, 10 Jul 2026 15:46:59 +0800 Subject: [PATCH 050/100] reset preprocessor --- .../preprocessor => preprocessor/fht_rotator}/fht_rotator.cc | 2 +- .../preprocessor => preprocessor/fht_rotator}/fht_rotator.h | 2 +- src/turbo/{quantizer => }/preprocessor/preprocessor.h | 0 tests/turbo/turbo_fht_rotator_test.cc | 2 +- 4 files changed, 3 insertions(+), 3 deletions(-) rename src/turbo/{quantizer/preprocessor => preprocessor/fht_rotator}/fht_rotator.cc (99%) rename src/turbo/{quantizer/preprocessor => preprocessor/fht_rotator}/fht_rotator.h (98%) rename src/turbo/{quantizer => }/preprocessor/preprocessor.h (100%) diff --git a/src/turbo/quantizer/preprocessor/fht_rotator.cc b/src/turbo/preprocessor/fht_rotator/fht_rotator.cc similarity index 99% rename from src/turbo/quantizer/preprocessor/fht_rotator.cc rename to src/turbo/preprocessor/fht_rotator/fht_rotator.cc index 4d31f6949..979112305 100644 --- a/src/turbo/quantizer/preprocessor/fht_rotator.cc +++ b/src/turbo/preprocessor/fht_rotator/fht_rotator.cc @@ -16,7 +16,7 @@ #include #include #include -#include "../quantizer.h" +#include "quantizer/quantizer.h" namespace zvec { namespace turbo { diff --git a/src/turbo/quantizer/preprocessor/fht_rotator.h b/src/turbo/preprocessor/fht_rotator/fht_rotator.h similarity index 98% rename from src/turbo/quantizer/preprocessor/fht_rotator.h rename to src/turbo/preprocessor/fht_rotator/fht_rotator.h index 1bad6e924..450420276 100644 --- a/src/turbo/quantizer/preprocessor/fht_rotator.h +++ b/src/turbo/preprocessor/fht_rotator/fht_rotator.h @@ -19,7 +19,7 @@ #include #include #include -#include "preprocessor.h" +#include "preprocessor/preprocessor.h" namespace zvec { namespace turbo { diff --git a/src/turbo/quantizer/preprocessor/preprocessor.h b/src/turbo/preprocessor/preprocessor.h similarity index 100% rename from src/turbo/quantizer/preprocessor/preprocessor.h rename to src/turbo/preprocessor/preprocessor.h diff --git a/tests/turbo/turbo_fht_rotator_test.cc b/tests/turbo/turbo_fht_rotator_test.cc index 4d5ca82a6..f1ad4877b 100644 --- a/tests/turbo/turbo_fht_rotator_test.cc +++ b/tests/turbo/turbo_fht_rotator_test.cc @@ -17,7 +17,7 @@ #include #include #include -#include "quantizer/preprocessor/fht_rotator.h" +#include "preprocessor/fht_rotator/fht_rotator.h" using namespace zvec::turbo; From b909b4db74a1547e7f8ff836995b1ea32f93b803 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Fri, 10 Jul 2026 15:56:01 +0800 Subject: [PATCH 051/100] add neon --- src/turbo/CMakeLists.txt | 6 +- src/turbo/distance/neon/fht/fht.cc | 119 +++++++++++++++++++++++++++++ src/turbo/distance/neon/fht/fht.h | 34 +++++++++ src/turbo/turbo.cc | 14 ++++ 4 files changed, 170 insertions(+), 3 deletions(-) create mode 100644 src/turbo/distance/neon/fht/fht.cc create mode 100644 src/turbo/distance/neon/fht/fht.h diff --git a/src/turbo/CMakeLists.txt b/src/turbo/CMakeLists.txt index f34f24e66..2ce56d40e 100644 --- a/src/turbo/CMakeLists.txt +++ b/src/turbo/CMakeLists.txt @@ -5,9 +5,9 @@ if(NOT ANDROID AND AUTO_DETECT_ARCH) if (HOST_ARCH MATCHES "^(x86|x64)$") setup_compiler_march_for_x86(TURBO_MARCH_FLAG_SSE TURBO_MARCH_FLAG_AVX2 TURBO_MARCH_FLAG_AVX512 TURBO_MARCH_FLAG_AVX512FP16) elseif (HOST_ARCH MATCHES "^(arm|arm64)$") - # ARM64 architecture - no special march flags needed for now - # NEON implementations can be added here if needed - message(STATUS "turbo: ARM64 detected, skipping x86-specific optimizations") + # ARM64 architecture - NEON is enabled by default on aarch64, + # no special march flags needed. + message(STATUS "turbo: ARM64 detected, NEON enabled by default") endif() endif() diff --git a/src/turbo/distance/neon/fht/fht.cc b/src/turbo/distance/neon/fht/fht.cc new file mode 100644 index 000000000..8ba9c9f1e --- /dev/null +++ b/src/turbo/distance/neon/fht/fht.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(__ARM_NEON) && defined(__aarch64__) + +#include "fht.h" +#include +#include +#include +#include +#include + +namespace zvec::turbo::neon { + +void fht_flip_sign_neon(const uint8_t *flip, float *data, size_t dim) { + const uint32x4_t sign_bit = vdupq_n_u32(0x80000000u); + 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; + uint32x4_t bit_mask = {b0, b1, b2, b3}; + uint32x4_t sign_mask = vmulq_u32(bit_mask, sign_bit); + float32x4_t v = vld1q_f32(&data[i]); + v = vreinterpretq_f32_u32(veorq_u32(vreinterpretq_u32_f32(v), sign_mask)); + vst1q_f32(&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_neon(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) { + float32x4_t x = vld1q_f32(&data[i]); + float32x4_t y = vld1q_f32(&data[i + offset]); + vst1q_f32(&data[i], vaddq_f32(x, y)); + vst1q_f32(&data[i + offset], vsubq_f32(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_neon(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 float32x4_t half_fac = vdupq_n_f32(0.5f); + for (size_t i = 0; i < half_end; i += 4) { + float32x4_t a = vld1q_f32(&data[i]); + float32x4_t b = vld1q_f32(&data[i + offset]); + vst1q_f32(&data[i], vmulq_f32(vaddq_f32(a, b), half_fac)); + vst1q_f32(&data[i + offset], vmulq_f32(vsubq_f32(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_neon(float *data, size_t n, float factor) { + const float32x4_t fac = vdupq_n_f32(factor); + size_t simd_end = n & ~3u; + for (size_t i = 0; i < simd_end; i += 4) { + float32x4_t v = vld1q_f32(&data[i]); + vst1q_f32(&data[i], vmulq_f32(v, fac)); + } + // Scalar tail + for (size_t i = simd_end; i < n; ++i) { + data[i] *= factor; + } +} + +} // namespace zvec::turbo::neon + +#endif // __ARM_NEON && __aarch64__ diff --git a/src/turbo/distance/neon/fht/fht.h b/src/turbo/distance/neon/fht/fht.h new file mode 100644 index 000000000..409e8b875 --- /dev/null +++ b/src/turbo/distance/neon/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::neon { + +//! Apply bitwise sign-flip mask to a float vector (NEON). +void fht_flip_sign_neon(const uint8_t *flip, float *data, size_t dim); + +//! Apply KacsWalk butterfly operation (NEON). +void fht_kacs_walk_neon(float *data, size_t len); + +//! Inverse KacsWalk butterfly operation (NEON). +void fht_inv_kacs_walk_neon(float *data, size_t len); + +//! Element-wise rescale: data[i] *= factor (NEON). +void fht_vec_rescale_neon(float *data, size_t n, float factor); + +} // namespace zvec::turbo::neon diff --git a/src/turbo/turbo.cc b/src/turbo/turbo.cc index b5b3e9e51..87826f12f 100644 --- a/src/turbo/turbo.cc +++ b/src/turbo/turbo.cc @@ -32,6 +32,9 @@ #if defined(__AVX512F__) #include "avx512/fht/fht.h" #endif +#if defined(__ARM_NEON) && defined(__aarch64__) +#include "neon/fht/fht.h" +#endif namespace zvec::turbo { @@ -201,6 +204,17 @@ FhtKernels get_fht_kernels(CpuArchType cpu_arch_type) { // inplace fallback to scalar (SSE has no fht_inplace) return k; } +#endif +#if defined(__ARM_NEON) && defined(__aarch64__) + if (cpu_arch_type == CpuArchType::kAuto || + cpu_arch_type == CpuArchType::kNEON) { + k.flip_sign = neon::fht_flip_sign_neon; + k.kacs_walk = neon::fht_kacs_walk_neon; + k.inv_kacs_walk = neon::fht_inv_kacs_walk_neon; + k.rescale = neon::fht_vec_rescale_neon; + // inplace fallback to scalar (NEON has no fht_inplace) + return k; + } #endif return k; // scalar } From 403c1234b9c876d29ee5de1d142c7a49090e975c Mon Sep 17 00:00:00 2001 From: ray Date: Fri, 10 Jul 2026 18:38:23 +0800 Subject: [PATCH 052/100] fix: add meta --- src/core/framework/index_meta.cc | 1 + 1 file changed, 1 insertion(+) diff --git a/src/core/framework/index_meta.cc b/src/core/framework/index_meta.cc index d0eadb02d..0cdf4815a 100644 --- a/src/core/framework/index_meta.cc +++ b/src/core/framework/index_meta.cc @@ -145,6 +145,7 @@ bool IndexMeta::deserialize(const void *data, size_t len) { unit_size_ = format->unit_size; element_size_ = IndexMeta::ElementSizeof(data_type_, unit_size_, dimension_); space_id_ = format->space_id; + extra_meta_size_ = format->extra_meta_size; // Read attachment ailego::Params attachment; From b3e6a56a149f3487230485180c01877b572233a2 Mon Sep 17 00:00:00 2001 From: ray Date: Fri, 10 Jul 2026 18:42:00 +0800 Subject: [PATCH 053/100] fix: query meta --- src/core/framework/index_meta.cc | 5 ++-- src/include/zvec/core/framework/index_meta.h | 25 ++++++++++++++++++-- 2 files changed, 26 insertions(+), 4 deletions(-) diff --git a/src/core/framework/index_meta.cc b/src/core/framework/index_meta.cc index 0cdf4815a..4a0f13e90 100644 --- a/src/core/framework/index_meta.cc +++ b/src/core/framework/index_meta.cc @@ -143,9 +143,10 @@ bool IndexMeta::deserialize(const void *data, size_t len) { data_type_ = static_cast(format->data_type); dimension_ = format->dimension; unit_size_ = format->unit_size; - element_size_ = IndexMeta::ElementSizeof(data_type_, unit_size_, dimension_); - space_id_ = format->space_id; extra_meta_size_ = format->extra_meta_size; + element_size_ = IndexMeta::ElementSizeof(data_type_, unit_size_, dimension_) + + extra_meta_size_; + space_id_ = format->space_id; // Read attachment ailego::Params attachment; diff --git a/src/include/zvec/core/framework/index_meta.h b/src/include/zvec/core/framework/index_meta.h index c6c5137db..010fca857 100644 --- a/src/include/zvec/core/framework/index_meta.h +++ b/src/include/zvec/core/framework/index_meta.h @@ -440,7 +440,8 @@ class IndexMeta { //! Set dimension of feature void set_dimension(uint32_t dim) { dimension_ = dim; - element_size_ = IndexMeta::ElementSizeof(data_type_, unit_size_, dim); + element_size_ = IndexMeta::ElementSizeof(data_type_, unit_size_, dim) + + extra_meta_size_; } //! Set meta information of feature @@ -454,7 +455,7 @@ class IndexMeta { data_type_ = data_type; dimension_ = dim; unit_size_ = unit; - element_size_ = ElementSizeof(data_type, unit, dim); + element_size_ = ElementSizeof(data_type, unit, dim) + extra_meta_size_; } //! Set meta information of feature @@ -465,6 +466,8 @@ class IndexMeta { //! Set extra meta size void set_extra_meta_size(uint32_t size) { extra_meta_size_ = size; + element_size_ = + ElementSizeof(data_type_, unit_size_, dimension_) + extra_meta_size_; } //! Set information of metric @@ -703,6 +706,16 @@ class IndexQueryMeta { return element_size_; } + //! Retrieve quantize type + uint32_t quantize_type(void) const { + return quantize_type_; + } + + //! Retrieve extra meta size in bytes + uint32_t extra_meta_size(void) const { + return extra_meta_size_; + } + //! Set dimension of feature void set_dimension(uint32_t dim) { dimension_ = dim; @@ -720,6 +733,14 @@ class IndexQueryMeta { data_type_ = data_type; } + //! Set extra meta size + void set_extra_meta_size(uint32_t size) { + extra_meta_size_ = size; + element_size_ = + IndexMeta::ElementSizeof(data_type_, unit_size_, dimension_) + + extra_meta_size_; + } + //! Set meta information of feature void set_meta(IndexMeta::DataType data_type, uint32_t unit, uint32_t dim) { data_type_ = data_type; From 3cd090c03b5b3ff1bfce6bfe871bef699de30ee3 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Mon, 13 Jul 2026 10:38:03 +0800 Subject: [PATCH 054/100] clang-format --- src/turbo/distance/neon/fht/fht.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/turbo/distance/neon/fht/fht.cc b/src/turbo/distance/neon/fht/fht.cc index 8ba9c9f1e..8a94e7e2f 100644 --- a/src/turbo/distance/neon/fht/fht.cc +++ b/src/turbo/distance/neon/fht/fht.cc @@ -15,11 +15,11 @@ #if defined(__ARM_NEON) && defined(__aarch64__) #include "fht.h" -#include #include #include #include #include +#include namespace zvec::turbo::neon { From aa4a4f9b5112dd2ea0d46ad8434796a0d2567d90 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Tue, 14 Jul 2026 14:15:04 +0800 Subject: [PATCH 055/100] reuse k-mean --- src/include/zvec/turbo/turbo.h | 3 +- .../avx2/pq_quantizer_int4/pq_distance.cc | 224 ------ .../avx2/pq_quantizer_int4/pq_distance.h | 41 -- .../avx512/pq_quantizer_int4/pq_distance.cc | 208 ------ .../avx512/pq_quantizer_int4/pq_distance.h | 40 -- .../scalar/pq_quantizer_int4/pq_distance.cc | 108 --- .../scalar/pq_quantizer_int4/pq_distance.h | 51 -- .../pq_int4_quantizer/pq_int4_quantizer.cc | 642 ------------------ .../pq_int4_quantizer/pq_int4_quantizer.h | 163 ----- .../pq_int8_quantizer/pq_int8_quantizer.cc | 189 ++---- .../pq_int8_quantizer/pq_int8_quantizer.h | 11 + src/turbo/turbo.cc | 32 +- tests/turbo/turbo_pq_int4_quantizer_test.cc | 587 ---------------- 13 files changed, 78 insertions(+), 2221 deletions(-) delete mode 100644 src/turbo/distance/avx2/pq_quantizer_int4/pq_distance.cc delete mode 100644 src/turbo/distance/avx2/pq_quantizer_int4/pq_distance.h delete mode 100644 src/turbo/distance/avx512/pq_quantizer_int4/pq_distance.cc delete mode 100644 src/turbo/distance/avx512/pq_quantizer_int4/pq_distance.h delete mode 100644 src/turbo/distance/scalar/pq_quantizer_int4/pq_distance.cc delete mode 100644 src/turbo/distance/scalar/pq_quantizer_int4/pq_distance.h delete mode 100644 src/turbo/quantizer/pq_int4_quantizer/pq_int4_quantizer.cc delete mode 100644 src/turbo/quantizer/pq_int4_quantizer/pq_int4_quantizer.h delete 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 2a81cafba..50732ed31 100644 --- a/src/include/zvec/turbo/turbo.h +++ b/src/include/zvec/turbo/turbo.h @@ -75,7 +75,6 @@ struct FhtKernels { // 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; @@ -148,7 +147,7 @@ UniformQuantizeFunc get_uniform_quantize_func(DataType data_type); FhtKernels get_fht_kernels(CpuArchType cpu_arch_type = CpuArchType::kAuto); // 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). +// and CPU arch. PqKernels get_pq_kernels(DataType data_type, QuantizeType quantize_type = QuantizeType::kPQ, CpuArchType cpu_arch_type = CpuArchType::kAuto); diff --git a/src/turbo/distance/avx2/pq_quantizer_int4/pq_distance.cc b/src/turbo/distance/avx2/pq_quantizer_int4/pq_distance.cc deleted file mode 100644 index 657c335d3..000000000 --- a/src/turbo/distance/avx2/pq_quantizer_int4/pq_distance.cc +++ /dev/null @@ -1,224 +0,0 @@ -// 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 deleted file mode 100644 index 1d4164685..000000000 --- a/src/turbo/distance/avx2/pq_quantizer_int4/pq_distance.h +++ /dev/null @@ -1,41 +0,0 @@ -// 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 deleted file mode 100644 index e93c4f92c..000000000 --- a/src/turbo/distance/avx512/pq_quantizer_int4/pq_distance.cc +++ /dev/null @@ -1,208 +0,0 @@ -// 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 deleted file mode 100644 index 6a84bfe5b..000000000 --- a/src/turbo/distance/avx512/pq_quantizer_int4/pq_distance.h +++ /dev/null @@ -1,40 +0,0 @@ -// 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 deleted file mode 100644 index 4684fe67c..000000000 --- a/src/turbo/distance/scalar/pq_quantizer_int4/pq_distance.cc +++ /dev/null @@ -1,108 +0,0 @@ -// 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 deleted file mode 100644 index 561df9b6b..000000000 --- a/src/turbo/distance/scalar/pq_quantizer_int4/pq_distance.h +++ /dev/null @@ -1,51 +0,0 @@ -// 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 deleted file mode 100644 index c86492a79..000000000 --- a/src/turbo/quantizer/pq_int4_quantizer/pq_int4_quantizer.cc +++ /dev/null @@ -1,642 +0,0 @@ -// 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 deleted file mode 100644 index 2a1649cfc..000000000 --- a/src/turbo/quantizer/pq_int4_quantizer/pq_int4_quantizer.h +++ /dev/null @@ -1,163 +0,0 @@ -// Copyright 2025-present the zvec project -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -#pragma once - -#include -#include -#include -#include -#include "quantizer/quantizer.h" - -namespace zvec { -namespace turbo { - -using namespace zvec::core; - -//! Product Quantizer with 4-bit sub-codes (num_bits=4, 16 centroids). -//! -//! Datapoints are encoded as packed nibbles: two 4-bit indices per byte. -//! byte[i] = (code[2*i+1] << 4) | code[2*i] -//! num_subquantizers MUST be even for clean byte packing. -//! -//! Queries are encoded as a float LUT of size [num_subquantizers * 16] -//! via quantize_query(). Distance between a PQ code and a query uses -//! ADC (LUT look-up); distance between two PQ codes uses SDC -//! (centroid-to-centroid distance table of size [nsq * 16 * 16]). -class PqInt4Quantizer : public Quantizer { - public: - PqInt4Quantizer() { - type_ = QuantizeType::kPQ; // shares kPQ with int8; DataType distinguishes - } - - ~PqInt4Quantizer() override = default; - - int init(const IndexMeta &meta, const ailego::Params ¶ms) override; - - const IndexMeta &meta() const override { - return meta_; - } - - DataType input_data_type() const override { - return DataType::kFp32; - } - - QuantizeType type() const override { - return type_; - } - - int dim() const override { - return static_cast(original_dim_); - } - - bool require_train() const override { - return true; - } - - int train(IndexHolder::Pointer holder) override; - - int train(IndexHolder::Pointer holder, int thread_count) override; - - // Packed nibbles: num_subquantizers / 2 bytes + optional Cosine norm. - size_t quantized_datapoint_vector_length() const override { - return static_cast(num_subquantizers_) / 2 + extra_meta_size_; - } - - // LUT: [num_subquantizers * 16] floats. - size_t quantized_query_vector_length() const override { - return static_cast(num_subquantizers_) * kNumCentroids * - sizeof(float); - } - - void quantize_data(const void *input, void *output) const override; - - void quantize_query(const void *input, void *output) const override; - - float calc_distance_dp_query(const void *dp, - const void *query) const override; - - void calc_distance_dp_query_batch(const void *const *dp_list, int dp_num, - const void *query, - float *dist_list) const override; - - float calc_distance_dp_query_unquantized(const void *dp, - const void *query) const override; - - void calc_distance_dp_query_batch_unquantized( - const void *const *dp_list, int dp_num, const void *query, - float *dist_list) const override; - - float calc_distance_dp_dp(const void *dp1, const void *dp2) const override; - - int quantize(const void *query, const IndexQueryMeta &qmeta, std::string *out, - IndexQueryMeta *ometa) const override; - - int dequantize(const void *in, const IndexQueryMeta &qmeta, - std::string *out) const override; - - DistanceImpl distance(const void *query, - const IndexQueryMeta &qmeta) const override; - - int serialize(std::string *out) const override; - - int deserialize(std::string &in) override; - - int deserialize(const void *data, size_t len) override; - - private: - //! Train a single sub-quantizer (KMeans, k=16) on the sub-vectors. - void train_subquantizer(const float *data, size_t num, size_t stride, - size_t sub_idx); - - //! Compute the centroid-to-centroid distance table for SDC. - void compute_dist_table(); - - //! Build centroid_ptrs_cache_ from current centroids_. - void build_centroid_ptrs_cache(); - - static constexpr uint32_t kNumCentroids = 16; - static constexpr uint32_t kMaxKmeansIters = 25; - static constexpr size_t kMaxTrainVectors = 65536; - static constexpr uint32_t kExtraMetaSizeCosine = sizeof(float); - - IndexMeta meta_{}; - uint32_t original_dim_{0}; - uint32_t num_subquantizers_{0}; - uint32_t sub_dim_{0}; - - //! Centroids: [num_subquantizers * kNumCentroids * sub_dim] - std::vector centroids_; - - //! Centroid-to-centroid distance table for SDC: - //! [num_subquantizers * kNumCentroids * kNumCentroids] - std::vector dist_table_; - - //! Pre-built centroid pointer arrays for each sub-quantizer. - std::vector> centroid_ptrs_cache_; - - //! ISA-dispatched kernel function pointers (ADC / SDC / Batch ADC). - //! These point to int4-specific kernels (packed nibble codes). - PqAdcDistanceFunc adc_fn_{nullptr}; - PqSdcKernelFunc sdc_fn_{nullptr}; - PqBatchAdcFunc batch_adc_fn_{nullptr}; - - //! Metric-aware fp32 batch distance for search-side LUT and SDC table. - BatchDistanceFunc fp32_batch_fn_{}; - - //! L2-only fp32 batch distance for encoding and KMeans training. - BatchDistanceFunc fp32_l2_batch_fn_{}; -}; - -} // namespace turbo -} // namespace zvec diff --git a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc index 2dd724ccf..f2ea03348 100644 --- a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc @@ -14,15 +14,17 @@ #include "quantizer/pq_int8_quantizer/pq_int8_quantizer.h" #include +#include #include #include #include #include -#include #include +#include #include #include #include +#include namespace zvec { namespace turbo { @@ -102,6 +104,11 @@ int PqInt8Quantizer::init(const IndexMeta &meta, const ailego::Params ¶ms) { meta_.set_extra_meta_size(extra_meta_size_); } + // Read optional training params (aligned with multi_chunk_cluster) + params.get("thread_count", &thread_count_); + params.get("markov_chain_length", &markov_chain_length_); + params.get("epsilon", &epsilon_); + meta_.set_meta(IndexMeta::DataType::DT_FP32, d); return 0; } @@ -131,106 +138,55 @@ void PqInt8Quantizer::train_subquantizer(const float *data, size_t num, 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); + // Create KMeans algorithm (L2 distance via NumericalKmeansContext) + ailego::NumericalKmeans algorithm(k, d); - size_t n_init = std::min(static_cast(k), num); - for (size_t i = 0; i < n_init; ++i) { - const float *src = + // Append sub-vectors (NumericalKmeans handles transpose internally) + for (size_t i = 0; i < num; ++i) { + const float *sub_vec = reinterpret_cast( - reinterpret_cast(data) + perm[i] * stride) + + reinterpret_cast(data) + i * stride) + sub_idx * d; - std::memcpy(centroids_m + i * d, src, d * sizeof(float)); - } - // Duplicate if num < k (rare edge case). - for (size_t i = n_init; i < k; ++i) { - std::memcpy(centroids_m + i * d, centroids_m + (i % n_init) * d, - d * sizeof(float)); - } - - // -- Lloyd iterations --------------------------------------------------- - std::vector assignments(num); - std::vector new_centroids(k * d); - std::vector counts(k); - - // Pre-build centroid pointer array for SIMD batch distance. - std::vector centroid_ptrs(k); - for (size_t c = 0; c < k; ++c) { - centroid_ptrs[c] = centroids_m + c * d; + algorithm.append(sub_vec, d); } - std::vector dists(k); + // Single-threaded pool — parallelism is at sub-quantizer level + // (aligned with multi_chunk_cluster.h L184-185) + auto local_threads = std::make_shared(1, false); + + // KMC2 centroid initialization (aligned with multi_chunk_cluster.h L191-196) + ailego::Kmc2CentroidsGenerator< + ailego::NumericalKmeans, + SingleQueueIndexThreads> + gen; + gen.set_chain_length(markov_chain_length_); + gen.set_assumption_free(false); + algorithm.init_centroids(*local_threads, gen); + + // Lloyd iterations (aligned with multi_chunk_cluster.h L200-216) + double cost = 0.0; for (uint32_t iter = 0; iter < kMaxKmeansIters; ++iter) { - bool changed = false; - - // Assignment step — always use L2 distance (fp32_l2_batch_fn_) to - // minimize L2 quantization error, regardless of the search metric. - // KMeans operates in L2 space; the search metric only affects LUT - // construction in quantize_query. - for (size_t i = 0; i < num; ++i) { - const float *sub_vec = - reinterpret_cast( - reinterpret_cast(data) + i * stride) + - sub_idx * d; - - fp32_l2_batch_fn_(centroid_ptrs.data(), - reinterpret_cast(sub_vec), k, d, - dists.data()); - uint32_t best_idx = static_cast( - std::min_element(dists.begin(), dists.end()) - dists.begin()); - - if (assignments[i] != best_idx) { - changed = true; - assignments[i] = best_idx; - } - } - - if (!changed) { - LOG_INFO(" sub[%zu] converged at iter %u", sub_idx, iter + 1); + double old_cost = cost; + bool result = algorithm.cluster_once(*local_threads, &cost); + if (!result) { + LOG_ERROR("sub[%zu] cluster_once failed at iter %u", sub_idx, iter); 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; - } + double new_epsilon = std::abs(cost - old_cost); + if (new_epsilon < epsilon_) { + break; } + } - // 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)); - } - } + // Extract centroids into the flat centroids_ buffer + const auto ¢s = algorithm.centroids(); + for (size_t c = 0; c < cents.count(); ++c) { + std::memcpy(centroids_m + c * d, cents[c], d * sizeof(float)); } } int PqInt8Quantizer::train(IndexHolder::Pointer holder) { - return train(holder, 1); + return train(holder, static_cast(thread_count_)); } int PqInt8Quantizer::train(IndexHolder::Pointer holder, int thread_count) { @@ -287,50 +243,35 @@ 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_))); - 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); + // Create thread pool (aligned with multi_chunk_cluster L183-185) + auto threads = std::make_shared( + static_cast(thread_count), false); + auto task_group = threads->make_group(); + + // Distribute sub-quantizers across threads (aligned with L198-201) + std::atomic finished{0}; + size_t pool_count = threads->count(); + for (size_t i = 0; i < pool_count; ++i) { + task_group->submit(ailego::Closure::New( + [this, &all_data, num, data_stride, i, pool_count, &finished]() { + for (uint32_t m = static_cast(i); + m < num_subquantizers_; + m += static_cast(pool_count)) { + train_subquantizer(all_data.data(), num, data_stride, m); + finished++; + } + })); } + task_group->wait_finish(); + + LOG_INFO(" all %u sub-quantizers trained (%zu threads)", + num_subquantizers_, pool_count); // Pre-build centroid pointer cache (needed by compute_dist_table). 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 deb661941..83a25c4e6 100644 --- a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h @@ -15,6 +15,7 @@ #pragma once #include +#include #include #include #include @@ -128,6 +129,16 @@ class PqInt8Quantizer : public Quantizer { static constexpr size_t kMaxTrainVectors = 65536; static constexpr uint32_t kExtraMetaSizeCosine = sizeof(float); + //! Thread count for KMeans training (0 = hardware_concurrency). + //! Read from params in init(), aligned with multi_chunk_cluster. + uint32_t thread_count_{0}; + + //! KMC2 Markov chain length (aligned with multi_chunk_cluster default 32). + uint32_t markov_chain_length_{32}; + + //! Cost-based convergence threshold (aligned with multi_chunk_cluster). + double epsilon_{std::numeric_limits::epsilon()}; + IndexMeta meta_{}; uint32_t original_dim_{0}; uint32_t num_subquantizers_{0}; diff --git a/src/turbo/turbo.cc b/src/turbo/turbo.cc index 51f1bdeed..3673ec696 100644 --- a/src/turbo/turbo.cc +++ b/src/turbo/turbo.cc @@ -22,7 +22,6 @@ #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__) @@ -30,12 +29,10 @@ #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 #if defined(__ARM_NEON) && defined(__aarch64__) @@ -231,34 +228,7 @@ PqKernels get_pq_kernels(DataType data_type, QuantizeType quantize_type, (void)cpu_arch_type; PqKernels k{}; if (quantize_type == QuantizeType::kPQ) { - if (data_type == DataType::kInt4) { - // int4 packed nibble path — scalar by default. - k.adc_distance = scalar::pq_adc_int4_distance; - k.sdc_distance = scalar::pq_sdc_int4_distance; - k.batch_adc_distance = scalar::pq_adc_int4_batch_distance; - -#if defined(__AVX512F__) - if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F && - (cpu_arch_type == CpuArchType::kAuto || - cpu_arch_type == CpuArchType::kAVX512)) { - 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 && - (cpu_arch_type == CpuArchType::kAuto || - cpu_arch_type == CpuArchType::kAVX2)) { - 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. + // Default (kInt8): 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 deleted file mode 100644 index 86249f156..000000000 --- a/tests/turbo/turbo_pq_int4_quantizer_test.cc +++ /dev/null @@ -1,587 +0,0 @@ -// Copyright 2025-present the zvec project -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -#include -#include -#include -#include -#include -#include -#include -#include -#include "distance/scalar/pq_quantizer_int4/pq_distance.h" -#include "quantizer/pq_int4_quantizer/pq_int4_quantizer.h" -#include "zvec/core/framework/index_factory.h" - -#if defined(__AVX2__) -#include "distance/avx2/pq_quantizer_int4/pq_distance.h" -#endif -#if defined(__AVX512F__) -#include "distance/avx512/pq_quantizer_int4/pq_distance.h" -#endif - -using namespace zvec; -using namespace zvec::core; -using namespace zvec::ailego; - -// Reference squared Euclidean distance between two raw fp32 vectors. -static float reference_sq_euclidean(const float *a, const float *b, - size_t dim) { - float sum = 0.0f; - for (size_t i = 0; i < dim; ++i) { - float diff = a[i] - b[i]; - sum += diff * diff; - } - return sum; -} - -// Helper to create a PqInt4Quantizer via the factory. -static std::shared_ptr make_pq4_quantizer( - size_t dim, size_t num_subquantizers) { - auto q = IndexFactory::CreateQuantizer("PqInt4Quantizer"); - if (!q) return nullptr; - - IndexMeta meta; - meta.set_meta(IndexMeta::DataType::DT_FP32, dim); - meta.set_metric("SquaredEuclidean", 0, Params()); - - Params params; - params.set("num_subquantizers", static_cast(num_subquantizers)); - if (q->init(meta, params) != 0) return nullptr; - return q; -} - -// Helper: build a holder with random fp32 vectors. -static std::shared_ptr> -make_random_holder(size_t count, size_t dim, uint32_t seed = 42) { - auto holder = - std::make_shared>(dim); - std::mt19937 gen(seed); - std::uniform_real_distribution dist(-1.0f, 1.0f); - for (size_t i = 0; i < count; ++i) { - NumericalVector vec(dim); - for (size_t j = 0; j < dim; ++j) vec[j] = dist(gen); - holder->emplace(i + 1, vec); - } - return holder; -} - -// --------------------------------------------------------------------------- -// Tests -// --------------------------------------------------------------------------- - -TEST(PqInt4Quantizer, InitInvalidParams) { - // dim not divisible by num_subquantizers - auto q = make_pq4_quantizer(10, 4); - EXPECT_EQ(q, nullptr); - - // num_subquantizers = 0 - auto q2 = IndexFactory::CreateQuantizer("PqInt4Quantizer"); - ASSERT_TRUE(q2); - IndexMeta meta; - meta.set_meta(IndexMeta::DataType::DT_FP32, 16); - meta.set_metric("SquaredEuclidean", 0, Params()); - Params params; - params.set("num_subquantizers", static_cast(0)); - EXPECT_NE(0, q2->init(meta, params)); - - // Odd num_subquantizers: must be rejected for int4 packed nibble. - auto q3 = IndexFactory::CreateQuantizer("PqInt4Quantizer"); - ASSERT_TRUE(q3); - IndexMeta meta3; - meta3.set_meta(IndexMeta::DataType::DT_FP32, 12); - meta3.set_metric("SquaredEuclidean", 0, Params()); - Params params3; - params3.set("num_subquantizers", static_cast(3)); - EXPECT_NE(0, q3->init(meta3, params3)); -} - -TEST(PqInt4Quantizer, TrainAndEncode) { - const size_t DIM = 16; - const size_t NSQ = 4; // must be even - const size_t COUNT = 1000; - - auto quantizer = make_pq4_quantizer(DIM, NSQ); - ASSERT_TRUE(quantizer); - EXPECT_TRUE(quantizer->require_train()); - - // Packed nibble: code length = NSQ/2 bytes (no extra meta for L2). - EXPECT_EQ(quantizer->quantized_datapoint_vector_length(), NSQ / 2); - - auto holder = make_random_holder(COUNT, DIM); - ASSERT_EQ(0, quantizer->train(holder)); - - // Quantize a few vectors and check code length + nibble range. - auto iter = holder->create_iterator(); - size_t checked = 0; - std::vector code(quantizer->quantized_datapoint_vector_length()); - for (; iter->is_valid() && checked < 10; iter->next(), ++checked) { - quantizer->quantize_data(iter->data(), code.data()); - // Each nibble (low and high) should be in [0, 15]. - for (size_t b = 0; b < NSQ / 2; ++b) { - uint8_t lo = code[b] & 0x0Fu; - uint8_t hi = (code[b] >> 4) & 0x0Fu; - EXPECT_LE(lo, 15u); - EXPECT_LE(hi, 15u); - } - } - EXPECT_EQ(10u, checked); -} - -TEST(PqInt4Quantizer, AdcDistance) { - const size_t DIM = 32; - const size_t NSQ = 8; // even - const size_t COUNT = 2000; - - auto quantizer = make_pq4_quantizer(DIM, NSQ); - ASSERT_TRUE(quantizer); - - auto holder = make_random_holder(COUNT, DIM); - ASSERT_EQ(0, quantizer->train(holder)); - - std::vector> raw_vecs(COUNT); - std::vector> pq_codes(COUNT); - size_t code_len = quantizer->quantized_datapoint_vector_length(); - size_t lut_len = quantizer->quantized_query_vector_length(); - - auto iter = holder->create_iterator(); - for (size_t i = 0; iter->is_valid(); iter->next(), ++i) { - const float *v = reinterpret_cast(iter->data()); - raw_vecs[i].assign(v, v + DIM); - pq_codes[i].resize(code_len); - quantizer->quantize_data(iter->data(), pq_codes[i].data()); - } - - // Build LUT for query = raw_vecs[0]. - std::vector lut(lut_len / sizeof(float)); - quantizer->quantize_query(raw_vecs[0].data(), lut.data()); - - float max_rel_error = 0.0f; - for (size_t i = 1; i < COUNT; ++i) { - float adc_dist = - quantizer->calc_distance_dp_query(pq_codes[i].data(), lut.data()); - float true_dist = - reference_sq_euclidean(raw_vecs[i].data(), raw_vecs[0].data(), DIM); - if (true_dist > 1e-6f) { - float rel = std::fabs(adc_dist - true_dist) / true_dist; - max_rel_error = std::max(max_rel_error, rel); - } - EXPECT_GE(adc_dist, 0.0f) << "i=" << i; - } - // int4 PQ has higher quantization error than int8 (only 16 centroids). - // Generous bound: relative error < 2.0 (200%). - EXPECT_LT(max_rel_error, 2.0f) << "max_rel_error=" << max_rel_error; -} - -TEST(PqInt4Quantizer, SdcDistance) { - const size_t DIM = 16; - const size_t NSQ = 4; - const size_t COUNT = 2000; - - auto quantizer = make_pq4_quantizer(DIM, NSQ); - ASSERT_TRUE(quantizer); - - auto holder = make_random_holder(COUNT, DIM); - ASSERT_EQ(0, quantizer->train(holder)); - - auto iter = holder->create_iterator(); - std::vector code1(quantizer->quantized_datapoint_vector_length()); - std::vector code2(quantizer->quantized_datapoint_vector_length()); - - iter->is_valid(); - quantizer->quantize_data(iter->data(), code1.data()); - iter->next(); - iter->is_valid(); - quantizer->quantize_data(iter->data(), code2.data()); - - float sdc_dist = quantizer->calc_distance_dp_dp(code1.data(), code2.data()); - EXPECT_GE(sdc_dist, 0.0f); -} - -TEST(PqInt4Quantizer, DistanceImplAdcAndSdc) { - const size_t DIM = 16; - const size_t NSQ = 4; - const size_t COUNT = 1000; - - auto quantizer = make_pq4_quantizer(DIM, NSQ); - ASSERT_TRUE(quantizer); - - auto holder = make_random_holder(COUNT, DIM); - ASSERT_EQ(0, quantizer->train(holder)); - - auto iter = holder->create_iterator(); - iter->is_valid(); - const float *query_raw = reinterpret_cast(iter->data()); - - size_t lut_bytes = quantizer->quantized_query_vector_length(); - std::string lut_storage(lut_bytes, '\0'); - quantizer->quantize_query(query_raw, &lut_storage[0]); - - IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, DIM); - auto dist_impl = quantizer->distance(lut_storage.data(), qmeta); - ASSERT_TRUE(dist_impl.valid()); - EXPECT_TRUE(static_cast(dist_impl.func())); - EXPECT_TRUE(static_cast(dist_impl.sym_func())); - - // Encode a candidate and compute distance via ADC. - iter->next(); - iter->is_valid(); - std::vector code(quantizer->quantized_datapoint_vector_length()); - quantizer->quantize_data(iter->data(), code.data()); - - float d = dist_impl(code.data()); - EXPECT_GE(d, 0.0f); - - // SDC (symmetric) via sym_func(). - std::vector code2(quantizer->quantized_datapoint_vector_length()); - iter->next(); - iter->is_valid(); - quantizer->quantize_data(iter->data(), code2.data()); - - float sdc_d = 0.0f; - dist_impl.sym_func()(code.data(), code2.data(), NSQ, &sdc_d); - EXPECT_GE(sdc_d, 0.0f); -} - -TEST(PqInt4Quantizer, SerializeDeserialize) { - const size_t DIM = 16; - const size_t NSQ = 4; - const size_t COUNT = 500; - - auto quantizer = make_pq4_quantizer(DIM, NSQ); - ASSERT_TRUE(quantizer); - - auto holder = make_random_holder(COUNT, DIM); - ASSERT_EQ(0, quantizer->train(holder)); - - std::string blob; - ASSERT_EQ(0, quantizer->serialize(&blob)); - EXPECT_GT(blob.size(), sizeof(zvec::turbo::QuantizerSerHeader)); - - auto q2 = IndexFactory::CreateQuantizer("PqInt4Quantizer"); - ASSERT_TRUE(q2); - ASSERT_EQ(0, q2->deserialize(blob)); - - auto iter = holder->create_iterator(); - iter->is_valid(); - std::vector code1(quantizer->quantized_datapoint_vector_length()); - std::vector code2(q2->quantized_datapoint_vector_length()); - quantizer->quantize_data(iter->data(), code1.data()); - q2->quantize_data(iter->data(), code2.data()); - - size_t code_bytes = NSQ / 2; - for (size_t b = 0; b < code_bytes; ++b) { - EXPECT_EQ(code1[b], code2[b]) << "b=" << b; - } - - // ADC distances should also match. - size_t lut_len = quantizer->quantized_query_vector_length(); - std::vector lut1(lut_len / sizeof(float)); - std::vector lut2(lut_len / sizeof(float)); - quantizer->quantize_query(iter->data(), lut1.data()); - q2->quantize_query(iter->data(), lut2.data()); - - float adc1 = quantizer->calc_distance_dp_query(code1.data(), lut1.data()); - float adc2 = q2->calc_distance_dp_query(code2.data(), lut2.data()); - EXPECT_NEAR(adc1, adc2, 1e-6f); -} - -// --------------------------------------------------------------------------- -// Scalar int4 kernel direct tests -// --------------------------------------------------------------------------- - -namespace { - -// Decode nibble helper for test verification. -inline uint8_t test_decode_nibble(const uint8_t *code, size_t m) { - uint8_t byte = code[m >> 1]; - return (m & 1u) ? static_cast(byte >> 4) - : static_cast(byte & 0x0Fu); -} - -void fill_random_int4_codes(uint8_t *codes, size_t num_subquantizers, - std::mt19937 &gen) { - std::uniform_int_distribution dist(0, 15); - size_t code_bytes = num_subquantizers / 2; - std::memset(codes, 0, code_bytes); - for (size_t m = 0; m < num_subquantizers; ++m) { - uint8_t idx = static_cast(dist(gen)); - if (m & 1u) { - codes[m >> 1] |= static_cast(idx << 4); - } else { - codes[m >> 1] |= static_cast(idx & 0x0Fu); - } - } -} - -void fill_random_int4_lut(float *lut, size_t num_subquantizers, - std::mt19937 &gen) { - constexpr size_t kNumCentroids = 16; - std::uniform_real_distribution dist(0.0f, 1.0f); - for (size_t m = 0; m < num_subquantizers; ++m) { - for (size_t c = 0; c < kNumCentroids; ++c) { - lut[m * kNumCentroids + c] = dist(gen); - } - } -} - -} // anonymous namespace - -// Test ADC int4 kernel directly. -TEST(PqInt4Kernel, AdcDistance) { - std::mt19937 gen(2024); - - for (size_t nsq : {4, 8, 12, 16}) { - constexpr size_t kNumCentroids = 16; - std::vector codes(nsq / 2); - std::vector lut(nsq * kNumCentroids); - - fill_random_int4_codes(codes.data(), nsq, gen); - fill_random_int4_lut(lut.data(), nsq, gen); - - // Reference: manual sum. - float ref = 0.0f; - for (size_t m = 0; m < nsq; ++m) { - uint8_t idx = test_decode_nibble(codes.data(), m); - ref += lut[m * kNumCentroids + idx]; - } - - float result = 0.0f; - zvec::turbo::scalar::pq_adc_int4_distance(codes.data(), lut.data(), nsq, - &result); - EXPECT_NEAR(ref, result, 1e-5f) << "ADC mismatch for nsq=" << nsq; - } -} - -// Test SDC int4 kernel directly. -TEST(PqInt4Kernel, SdcDistance) { - std::mt19937 gen(2025); - constexpr size_t kNumCentroids = 16; - constexpr size_t kTablePerSub = kNumCentroids * kNumCentroids; // 256 - - for (size_t nsq : {4, 8}) { - std::vector codes_a(nsq / 2); - std::vector codes_b(nsq / 2); - std::vector dist_table(nsq * kTablePerSub); - - fill_random_int4_codes(codes_a.data(), nsq, gen); - fill_random_int4_codes(codes_b.data(), nsq, gen); - - std::uniform_real_distribution dist(0.0f, 1.0f); - for (auto &v : dist_table) v = dist(gen); - - // Reference. - float ref = 0.0f; - for (size_t m = 0; m < nsq; ++m) { - uint8_t ai = test_decode_nibble(codes_a.data(), m); - uint8_t bi = test_decode_nibble(codes_b.data(), m); - size_t idx = m * kTablePerSub + static_cast(ai) * kNumCentroids + - static_cast(bi); - ref += dist_table[idx]; - } - - float result = 0.0f; - zvec::turbo::scalar::pq_sdc_int4_distance(codes_a.data(), codes_b.data(), - dist_table.data(), nsq, &result); - EXPECT_NEAR(ref, result, 1e-5f) << "SDC mismatch for nsq=" << nsq; - } -} - -// Test Batch ADC int4 kernel. -TEST(PqInt4Kernel, BatchAdcDistance) { - std::mt19937 gen(2026); - constexpr size_t kNumCentroids = 16; - const size_t nsq = 8; - const size_t num_candidates = 7; // not multiple of 4 to test leftover - - std::vector> codes(num_candidates, - std::vector(nsq / 2)); - std::vector lut(nsq * kNumCentroids); - - for (auto &c : codes) fill_random_int4_codes(c.data(), nsq, gen); - fill_random_int4_lut(lut.data(), nsq, gen); - - // Build pointer array. - std::vector ptrs(num_candidates); - for (size_t i = 0; i < num_candidates; ++i) { - ptrs[i] = codes[i].data(); - } - - // Reference per-candidate distances. - std::vector ref(num_candidates, 0.0f); - for (size_t i = 0; i < num_candidates; ++i) { - for (size_t m = 0; m < nsq; ++m) { - uint8_t idx = test_decode_nibble(codes[i].data(), m); - ref[i] += lut[m * kNumCentroids + idx]; - } - } - - std::vector result(num_candidates, 0.0f); - zvec::turbo::scalar::pq_adc_int4_batch_distance( - ptrs.data(), lut.data(), num_candidates, nsq, result.data()); - - for (size_t i = 0; i < num_candidates; ++i) { - EXPECT_NEAR(ref[i], result[i], 1e-5f) << "batch ADC mismatch at i=" << i; - } -} - -// --------------------------------------------------------------------------- -// SIMD Consistency Tests -// --------------------------------------------------------------------------- - -// Helper to generate random SDC table for int4 (kTablePerSub = 256) -namespace { -void fill_random_int4_sdc_table(float *table, size_t num_subquantizers, - std::mt19937 &gen) { - constexpr size_t kTablePerSub = 16 * 16; - std::uniform_real_distribution dist(0.0f, 1.0f); - for (size_t m = 0; m < num_subquantizers; ++m) { - for (size_t i = 0; i < kTablePerSub; ++i) { - table[m * kTablePerSub + i] = dist(gen); - } - } -} -} // anonymous namespace - -// Test ADC SIMD consistency across multiple M values (must be even for int4). -TEST(PqInt4SimdConsistency, AdcDistance) { - std::mt19937 gen(2024); - constexpr size_t kNumCentroids = 16; - // M=4: less than AVX2/AVX512 chunk (16) — scalar only - // M=8: less than chunk — scalar only - // M=16: exact chunk boundary - // M=20: chunk + scalar leftover - for (size_t num_sq : {4, 8, 16, 20}) { - std::vector codes(num_sq / 2); - std::vector lut(num_sq * kNumCentroids); - - fill_random_int4_codes(codes.data(), num_sq, gen); - fill_random_int4_lut(lut.data(), num_sq, gen); - - float scalar_result = 0.0f; - zvec::turbo::scalar::pq_adc_int4_distance(codes.data(), lut.data(), num_sq, - &scalar_result); - -#if defined(__AVX2__) - { - float avx2_result = 0.0f; - zvec::turbo::avx2::pq_adc_int4_distance_avx2(codes.data(), lut.data(), - num_sq, &avx2_result); - EXPECT_NEAR(scalar_result, avx2_result, 1e-5f) - << "AVX2 ADC mismatch for M=" << num_sq; - } -#endif - -#if defined(__AVX512F__) - { - float avx512_result = 0.0f; - zvec::turbo::avx512::pq_adc_int4_distance_avx512(codes.data(), lut.data(), - num_sq, &avx512_result); - EXPECT_NEAR(scalar_result, avx512_result, 1e-5f) - << "AVX512 ADC mismatch for M=" << num_sq; - } -#endif - } -} - -// Test SDC SIMD consistency across multiple M values. -TEST(PqInt4SimdConsistency, SdcDistance) { - std::mt19937 gen(2025); - constexpr size_t kTablePerSub = 16 * 16; - - for (size_t num_sq : {4, 8, 16, 20}) { - std::vector codes_a(num_sq / 2); - std::vector codes_b(num_sq / 2); - std::vector dist_table(num_sq * kTablePerSub); - - fill_random_int4_codes(codes_a.data(), num_sq, gen); - fill_random_int4_codes(codes_b.data(), num_sq, gen); - fill_random_int4_sdc_table(dist_table.data(), num_sq, gen); - - float scalar_result = 0.0f; - zvec::turbo::scalar::pq_sdc_int4_distance(codes_a.data(), codes_b.data(), - dist_table.data(), num_sq, - &scalar_result); - -#if defined(__AVX2__) - { - float avx2_result = 0.0f; - zvec::turbo::avx2::pq_sdc_int4_distance_avx2( - codes_a.data(), codes_b.data(), dist_table.data(), num_sq, - &avx2_result); - EXPECT_NEAR(scalar_result, avx2_result, 1e-5f) - << "AVX2 SDC mismatch for M=" << num_sq; - } -#endif - -#if defined(__AVX512F__) - { - float avx512_result = 0.0f; - zvec::turbo::avx512::pq_sdc_int4_distance_avx512( - codes_a.data(), codes_b.data(), dist_table.data(), num_sq, - &avx512_result); - EXPECT_NEAR(scalar_result, avx512_result, 1e-5f) - << "AVX512 SDC mismatch for M=" << num_sq; - } -#endif - } -} - -// Test Batch ADC SIMD consistency. -TEST(PqInt4SimdConsistency, BatchAdcDistance) { - std::mt19937 gen(2026); - constexpr size_t kNumCentroids = 16; - const size_t nsq = 16; // chunk-aligned - const size_t num_candidates = 7; // not multiple of 4 to test leftover - - std::vector> codes(num_candidates, - std::vector(nsq / 2)); - std::vector lut(nsq * kNumCentroids); - - for (auto &c : codes) fill_random_int4_codes(c.data(), nsq, gen); - fill_random_int4_lut(lut.data(), nsq, gen); - - // Reference via scalar. - std::vector ptrs(num_candidates); - std::vector scalar_result(num_candidates, 0.0f); - for (size_t i = 0; i < num_candidates; ++i) { - ptrs[i] = codes[i].data(); - } - zvec::turbo::scalar::pq_adc_int4_batch_distance( - ptrs.data(), lut.data(), num_candidates, nsq, scalar_result.data()); - -#if defined(__AVX2__) - { - std::vector avx2_result(num_candidates, 0.0f); - zvec::turbo::avx2::pq_adc_int4_batch_distance_avx2( - ptrs.data(), lut.data(), num_candidates, nsq, avx2_result.data()); - for (size_t i = 0; i < num_candidates; ++i) { - EXPECT_NEAR(scalar_result[i], avx2_result[i], 1e-5f) - << "AVX2 batch ADC mismatch at i=" << i; - } - } -#endif - -#if defined(__AVX512F__) - { - std::vector avx512_result(num_candidates, 0.0f); - zvec::turbo::avx512::pq_adc_int4_batch_distance_avx512( - ptrs.data(), lut.data(), num_candidates, nsq, avx512_result.data()); - for (size_t i = 0; i < num_candidates; ++i) { - EXPECT_NEAR(scalar_result[i], avx512_result[i], 1e-5f) - << "AVX512 batch ADC mismatch at i=" << i; - } - } -#endif -} From 128dcfd508ca31d949991d504649ca97a63660b6 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Tue, 14 Jul 2026 15:25:40 +0800 Subject: [PATCH 056/100] add spherical --- .../quantizer/pq_int8_quantizer/pq_int8_quantizer.cc | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 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 f2ea03348..d0a7e19e2 100644 --- a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc @@ -138,8 +138,14 @@ void PqInt8Quantizer::train_subquantizer(const float *data, size_t num, const size_t d = sub_dim_; float *centroids_m = centroids_.data() + static_cast(sub_idx) * k * d; - // Create KMeans algorithm (L2 distance via NumericalKmeansContext) - ailego::NumericalKmeans algorithm(k, d); + // Create KMeans algorithm (L2 distance via NumericalKmeansContext). + // Enable spherical mode for IP / Cosine so that centroids are updated + // on the unit sphere, improving convergence for angular objectives. + const bool spherical = + (meta_.metric_name() == "InnerProduct" || + meta_.metric_name() == "Cosine"); + ailego::NumericalKmeans algorithm( + k, d, spherical); // Append sub-vectors (NumericalKmeans handles transpose internally) for (size_t i = 0; i < num; ++i) { From 389d36f2da32241916ae4c2e6d1512dd5df028ed Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Tue, 14 Jul 2026 15:56:17 +0800 Subject: [PATCH 057/100] remove sym_func --- src/turbo/quantizer/distance.h | 22 ------------- .../pq_int8_quantizer/pq_int8_quantizer.cc | 33 ++++++++++++++----- .../pq_int8_quantizer/pq_int8_quantizer.h | 3 ++ src/turbo/quantizer/quantizer.h | 5 +++ tests/turbo/turbo_pq_int8_quantizer_test.cc | 13 +------- 5 files changed, 33 insertions(+), 43 deletions(-) diff --git a/src/turbo/quantizer/distance.h b/src/turbo/quantizer/distance.h index ecc43f21a..19658e641 100644 --- a/src/turbo/quantizer/distance.h +++ b/src/turbo/quantizer/distance.h @@ -34,26 +34,12 @@ 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: sym_func_ is the SDC kernel (distinct - //! from func_ which is the ADC kernel). Non-PQ quantizers should use - //! the constructors above where sym_func_ defaults to func_. - DistanceImpl(DistanceFunc func, DistanceFunc sym_func, - BatchDistanceFunc batch_func, std::string quantized_query, - size_t dim) - : func_(std::move(func)), - sym_func_(std::move(sym_func)), batch_func_(std::move(batch_func)), query_storage_(std::move(quantized_query)), dim_(dim) {} @@ -104,13 +90,6 @@ class DistanceImpl { return func_; } - //! Symmetric (SDC) distance function. For non-PQ quantizers this - //! defaults to func(); for PQ quantizers it points to the SDC kernel - //! that operates on two PQ codes via a centroid-to-centroid table. - const DistanceFunc &sym_func() const { - return sym_func_; - } - //! Raw batch distance function. const BatchDistanceFunc &batch_func() const { return batch_func_; @@ -118,7 +97,6 @@ class DistanceImpl { private: DistanceFunc func_{}; - 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 d0a7e19e2..0a77fd48f 100644 --- a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc @@ -526,13 +526,6 @@ DistanceImpl PqInt8Quantizer::distance(const void *query, // (both use void*), so we can assign without a lambda wrapper. DistanceFunc adc_func = adc_fn_; - // SDC: still needs a lambda because the kernel has 5 parameters - // (extra dist_table pointer), vs DistanceFunc's 4. - auto sdc = sdc_fn_; - const void *dt = dist_table_.data(); - DistanceFunc sdc_func = [sdc, dt](const void *a, const void *b, size_t dim, - float *out) { sdc(a, b, dt, dim, out); }; - // Batch ADC: ISA-dispatched batch4 kernel, no lambda needed. BatchDistanceFunc batch_func = batch_adc_fn_; @@ -541,8 +534,30 @@ DistanceImpl PqInt8Quantizer::distance(const void *query, 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), + return DistanceImpl(std::move(adc_func), std::move(batch_func), + std::move(lut_storage), + static_cast(num_subquantizers_)); +} + +DistanceImpl PqInt8Quantizer::sym_distance(const void *query, + const IndexQueryMeta &qmeta) const { + (void)qmeta; + + // SDC kernel needs a lambda: 5 parameters (extra dist_table pointer) + // vs DistanceFunc's 4. + auto sdc = sdc_fn_; + const void *dt = dist_table_.data(); + DistanceFunc sdc_func = [sdc, dt](const void *a, const void *b, size_t dim, + float *out) { sdc(a, b, dt, dim, out); }; + + // The query is a PQ code (uint8_t[num_subquantizers]), NOT a LUT. + // SDC computes distance between two PQ codes via the centroid-to-centroid + // dist_table_, so both sides must be PQ codes. + size_t code_bytes = quantized_datapoint_vector_length(); + std::string code_storage(static_cast(query), code_bytes); + + // SDC has no batch kernel — use the 3-arg constructor. + return DistanceImpl(std::move(sdc_func), std::move(code_storage), static_cast(num_subquantizers_)); } 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 83a25c4e6..e05e36429 100644 --- a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h @@ -105,6 +105,9 @@ class PqInt8Quantizer : public Quantizer { DistanceImpl distance(const void *query, const IndexQueryMeta &qmeta) const override; + DistanceImpl sym_distance(const void *query, + const IndexQueryMeta &qmeta) const override; + int serialize(std::string *out) const override; int deserialize(std::string &in) override; diff --git a/src/turbo/quantizer/quantizer.h b/src/turbo/quantizer/quantizer.h index 4ef239bd5..5d3d11237 100644 --- a/src/turbo/quantizer/quantizer.h +++ b/src/turbo/quantizer/quantizer.h @@ -159,6 +159,11 @@ class Quantizer { return DistanceImpl{}; } + virtual DistanceImpl sym_distance(const void *query, + const IndexQueryMeta &qmeta) const { + return distance(query, qmeta); + } + //! Serialize quantizer parameters virtual int serialize(std::string * /*out*/) const { return 0; diff --git a/tests/turbo/turbo_pq_int8_quantizer_test.cc b/tests/turbo/turbo_pq_int8_quantizer_test.cc index 438b19c96..43cb86c96 100644 --- a/tests/turbo/turbo_pq_int8_quantizer_test.cc +++ b/tests/turbo/turbo_pq_int8_quantizer_test.cc @@ -223,9 +223,8 @@ TEST(PqInt8Quantizer, DistanceImplAdcAndSdc) { auto dist_impl = quantizer->distance(lut_storage.data(), qmeta); ASSERT_TRUE(dist_impl.valid()); - // func() and sym_func() should both be set. + // func() should be set (ADC path). EXPECT_TRUE(static_cast(dist_impl.func())); - EXPECT_TRUE(static_cast(dist_impl.sym_func())); // Encode a candidate and compute distance via DistanceImpl (ADC path). iter->next(); @@ -235,16 +234,6 @@ TEST(PqInt8Quantizer, DistanceImplAdcAndSdc) { float d = dist_impl(code.data()); EXPECT_GE(d, 0.0f); - - // SDC (symmetric) via sym_func() directly. - std::vector code2(quantizer->quantized_datapoint_vector_length()); - iter->next(); - iter->is_valid(); - quantizer->quantize_data(iter->data(), code2.data()); - - float sdc_d = 0.0f; - dist_impl.sym_func()(code.data(), code2.data(), NSQ, &sdc_d); - EXPECT_GE(sdc_d, 0.0f); } TEST(PqInt8Quantizer, SerializeDeserialize) { From 428713154e48ba7326f0091c4f0221096e2eb2f9 Mon Sep 17 00:00:00 2001 From: ray Date: Wed, 15 Jul 2026 10:33:01 +0800 Subject: [PATCH 058/100] fix: fix alignment --- src/turbo/turbo.cc | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/turbo/turbo.cc b/src/turbo/turbo.cc index 80f2ee77d..4fe8309d0 100644 --- a/src/turbo/turbo.cc +++ b/src/turbo/turbo.cc @@ -56,7 +56,9 @@ DistanceFunc get_distance_func(MetricType metric_type, DataType data_type, } } if (quantize_type == QuantizeType::kUniform) { - 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::uniform_squared_euclidean_int8_distance; } @@ -99,7 +101,9 @@ BatchDistanceFunc get_batch_distance_func(MetricType metric_type, } } if (quantize_type == QuantizeType::kUniform) { - 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::uniform_squared_euclidean_int8_batch_distance; } From 394356679dd61357b0376e19f54132c758692e57 Mon Sep 17 00:00:00 2001 From: ray Date: Wed, 15 Jul 2026 11:12:55 +0800 Subject: [PATCH 059/100] refactor: move header file --- examples/c++/CMakeLists.txt | 6 ++---- src/core/framework/index_factory.cc | 1 + src/include/zvec/core/framework/index_factory.h | 7 ++++++- tests/turbo/turbo_fp32_quantizer_test.cc | 1 + 4 files changed, 10 insertions(+), 5 deletions(-) diff --git a/examples/c++/CMakeLists.txt b/examples/c++/CMakeLists.txt index 4de64b45c..1e111cd48 100644 --- a/examples/c++/CMakeLists.txt +++ b/examples/c++/CMakeLists.txt @@ -14,10 +14,8 @@ if(NOT DEFINED HOST_BUILD_DIR) endif() get_filename_component(ZVEC_ROOT_DIR "${CMAKE_CURRENT_LIST_DIR}/../.." ABSOLUTE) -# 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) +# Public headers are self-contained under src/include. +set(ZVEC_INCLUDE_DIR ${ZVEC_ROOT_DIR}/src/include) set(ZVEC_LIB_DIR ${ZVEC_ROOT_DIR}/${HOST_BUILD_DIR}/lib) # Add include and library search paths diff --git a/src/core/framework/index_factory.cc b/src/core/framework/index_factory.cc index e93f57bc7..d26192eba 100644 --- a/src/core/framework/index_factory.cc +++ b/src/core/framework/index_factory.cc @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include #include namespace zvec { diff --git a/src/include/zvec/core/framework/index_factory.h b/src/include/zvec/core/framework/index_factory.h index 00e77894c..1c9e287ae 100644 --- a/src/include/zvec/core/framework/index_factory.h +++ b/src/include/zvec/core/framework/index_factory.h @@ -14,7 +14,6 @@ #pragma once -#include #include #include #include @@ -30,6 +29,12 @@ #include #include +namespace zvec { +namespace turbo { +class Quantizer; +} // namespace turbo +} // namespace zvec + namespace zvec { namespace core { diff --git a/tests/turbo/turbo_fp32_quantizer_test.cc b/tests/turbo/turbo_fp32_quantizer_test.cc index c180de4a4..80610ca94 100644 --- a/tests/turbo/turbo_fp32_quantizer_test.cc +++ b/tests/turbo/turbo_fp32_quantizer_test.cc @@ -16,6 +16,7 @@ #include #include #include +#include #include #include #include "zvec/core/framework/index_factory.h" From 0d27ec3f506a18566314bde51f94373f00d3a734 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Wed, 15 Jul 2026 13:44:15 +0800 Subject: [PATCH 060/100] num_chunk --- src/include/zvec/core/interface/index_param.h | 4 +- src/include/zvec/turbo/turbo.h | 14 ++-- .../avx2/pq_quantizer_int8/pq_distance.cc | 20 ++--- .../avx2/pq_quantizer_int8/pq_distance.h | 6 +- .../avx512/pq_quantizer_int8/pq_distance.cc | 20 ++--- .../avx512/pq_quantizer_int8/pq_distance.h | 6 +- .../scalar/pq_quantizer_int8/pq_distance.cc | 14 ++-- .../scalar/pq_quantizer_int8/pq_distance.h | 12 +-- .../pq_int8_quantizer/pq_int8_quantizer.cc | 79 ++++++++++--------- .../pq_int8_quantizer/pq_int8_quantizer.h | 14 ++-- tests/turbo/turbo_pq_int8_quantizer_test.cc | 26 +++--- 11 files changed, 108 insertions(+), 107 deletions(-) diff --git a/src/include/zvec/core/interface/index_param.h b/src/include/zvec/core/interface/index_param.h index 78bebd570..70a2e5474 100644 --- a/src/include/zvec/core/interface/index_param.h +++ b/src/include/zvec/core/interface/index_param.h @@ -121,7 +121,7 @@ struct SerializableBase { // TODO: maybe a base class for quantizer? struct QuantizerParam : public SerializableBase { QuantizerType type = QuantizerType::kNone; - int num_subquantizers = 8; // M + int num_chunk = 8; // M int num_bits = 8; // bits per subquantizer bool enable_rotate = false; // rotate vectors before quantization to reduce error @@ -131,7 +131,7 @@ struct QuantizerParam : public SerializableBase { QuantizerParam(QuantizerType t = QuantizerType::kNone, int subquantizers = 8, int bits = 8, bool rotate = false) : type(t), - num_subquantizers(subquantizers), + num_chunk(subquantizers), num_bits(bits), enable_rotate(rotate) {} diff --git a/src/include/zvec/turbo/turbo.h b/src/include/zvec/turbo/turbo.h index 50732ed31..336125800 100644 --- a/src/include/zvec/turbo/turbo.h +++ b/src/include/zvec/turbo/turbo.h @@ -44,24 +44,24 @@ 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 +// pq_code: [num_chunk] uint8_t +// lut: [num_chunk * 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_chunk, 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 +// a, b: [num_chunk] uint8_t +// dist_table: [num_chunk * 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); + size_t num_chunk, 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, + size_t num, size_t num_chunk, float *out); // Aggregate of all FHT kernels needed by FhtRotator, dispatched by ISA. 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 18c177069..a6ca691f8 100644 --- a/src/turbo/distance/avx2/pq_quantizer_int8/pq_distance.cc +++ b/src/turbo/distance/avx2/pq_quantizer_int8/pq_distance.cc @@ -37,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_chunk, 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); @@ -54,7 +54,7 @@ void pq_adc_int8_distance_avx2(const void *pq_code_v, const void *lut_v, size_t m = 0; // Main loop: process 8 subquantizers per iteration - for (; m + kChunkSize <= num_subquantizers; m += kChunkSize) { + for (; m + kChunkSize <= num_chunk; m += kChunkSize) { // Load 8 uint8 codes and zero-extend to int32 // pq_code[m..m+7] -> 8 int32 indices __m128i codes_8x8 = @@ -74,7 +74,7 @@ void pq_adc_int8_distance_avx2(const void *pq_code_v, const void *lut_v, float sum = horizontal_sum_avx2(acc); // Scalar leftover: process remaining subquantizers - for (; m < num_subquantizers; ++m) { + for (; m < num_chunk; ++m) { sum += lut[m * kNumCentroids + pq_code[m]]; } @@ -83,7 +83,7 @@ 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) { + size_t num_chunk, float *out) { constexpr int kNumCentroids = 256; constexpr int kTablePerSub = kNumCentroids * kNumCentroids; // 65536 constexpr int kChunkSize = 8; @@ -104,7 +104,7 @@ void pq_sdc_int8_distance_avx2(const void *a_v, const void *b_v, size_t m = 0; // Main loop: process 8 subquantizers per iteration - for (; m + kChunkSize <= num_subquantizers; m += kChunkSize) { + for (; m + kChunkSize <= num_chunk; 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)); @@ -125,7 +125,7 @@ void pq_sdc_int8_distance_avx2(const void *a_v, const void *b_v, float sum = horizontal_sum_avx2(acc); // Scalar leftover - for (; m < num_subquantizers; ++m) { + for (; m < num_chunk; ++m) { size_t idx = m * kTablePerSub + static_cast(a[m]) * kNumCentroids + static_cast(b[m]); sum += dist_table[idx]; @@ -136,7 +136,7 @@ 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) { + size_t num_chunk, float *out) { constexpr int kNumCentroids = 256; constexpr int kChunkSize = 8; constexpr int kBatch = 4; @@ -161,7 +161,7 @@ void pq_adc_int8_batch_distance_avx2(const void **candidates_v, __m256 acc3 = _mm256_setzero_ps(); size_t m = 0; - for (; m + kChunkSize <= num_subquantizers; m += kChunkSize) { + for (; m + kChunkSize <= num_chunk; m += kChunkSize) { const float *lut_base = lut + m * kNumCentroids; __m128i codes0 = @@ -194,7 +194,7 @@ void pq_adc_int8_batch_distance_avx2(const void **candidates_v, float s3 = horizontal_sum_avx2(acc3); // Scalar leftover for remaining subquantizers. - for (; m < num_subquantizers; ++m) { + for (; m < num_chunk; ++m) { const float *tab = lut + m * kNumCentroids; s0 += tab[c0[m]]; s1 += tab[c1[m]]; @@ -208,7 +208,7 @@ void pq_adc_int8_batch_distance_avx2(const void **candidates_v, } // Remaining candidates: use single ADC kernel. for (; i < num; ++i) { - pq_adc_int8_distance_avx2(candidates[i], lut, num_subquantizers, out + i); + pq_adc_int8_distance_avx2(candidates[i], lut, num_chunk, out + i); } } 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 1d1b412cd..e4420fdeb 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_chunk, 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, + const void *dist_table, size_t num_chunk, 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, + size_t num, size_t num_chunk, float *out); } // namespace zvec::turbo::avx2 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 8ad40e1e4..8b69d881d 100644 --- a/src/turbo/distance/avx512/pq_quantizer_int8/pq_distance.cc +++ b/src/turbo/distance/avx512/pq_quantizer_int8/pq_distance.cc @@ -28,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_chunk, 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); @@ -48,7 +48,7 @@ void pq_adc_int8_distance_avx512(const void *pq_code_v, const void *lut_v, size_t m = 0; // Main loop: process 16 subquantizers per iteration - for (; m + kChunkSize <= num_subquantizers; m += kChunkSize) { + for (; m + kChunkSize <= num_chunk; m += kChunkSize) { // Load 16 uint8 codes and zero-extend to int32 // Use unaligned load of 16 bytes __m128i codes_16x8 = @@ -67,7 +67,7 @@ void pq_adc_int8_distance_avx512(const void *pq_code_v, const void *lut_v, float sum = horizontal_sum_avx512(acc); // Scalar leftover: process remaining subquantizers - for (; m < num_subquantizers; ++m) { + for (; m < num_chunk; ++m) { sum += lut[m * kNumCentroids + pq_code[m]]; } @@ -76,7 +76,7 @@ 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) { + size_t num_chunk, float *out) { constexpr int kNumCentroids = 256; constexpr int kTablePerSub = kNumCentroids * kNumCentroids; // 65536 constexpr int kChunkSize = 16; @@ -100,7 +100,7 @@ void pq_sdc_int8_distance_avx512(const void *a_v, const void *b_v, size_t m = 0; // Main loop: process 16 subquantizers per iteration - for (; m + kChunkSize <= num_subquantizers; m += kChunkSize) { + for (; m + kChunkSize <= num_chunk; 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)); @@ -121,7 +121,7 @@ void pq_sdc_int8_distance_avx512(const void *a_v, const void *b_v, float sum = horizontal_sum_avx512(acc); // Scalar leftover - for (; m < num_subquantizers; ++m) { + for (; m < num_chunk; ++m) { size_t idx = m * kTablePerSub + static_cast(a[m]) * kNumCentroids + static_cast(b[m]); sum += dist_table[idx]; @@ -132,7 +132,7 @@ 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) { + size_t num_chunk, float *out) { constexpr int kNumCentroids = 256; constexpr int kChunkSize = 16; constexpr int kBatch = 4; @@ -161,7 +161,7 @@ void pq_adc_int8_batch_distance_avx512(const void **candidates_v, __m512 acc3 = _mm512_setzero_ps(); size_t m = 0; - for (; m + kChunkSize <= num_subquantizers; m += kChunkSize) { + for (; m + kChunkSize <= num_chunk; m += kChunkSize) { const float *lut_base = lut + m * kNumCentroids; __m128i codes0 = @@ -194,7 +194,7 @@ void pq_adc_int8_batch_distance_avx512(const void **candidates_v, float s3 = _mm512_reduce_add_ps(acc3); // Scalar leftover for remaining subquantizers. - for (; m < num_subquantizers; ++m) { + for (; m < num_chunk; ++m) { const float *tab = lut + m * kNumCentroids; s0 += tab[c0[m]]; s1 += tab[c1[m]]; @@ -208,7 +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_chunk, 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 2c2ea92d5..0a3cc3987 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_chunk, 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); + size_t num_chunk, 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, + size_t num, size_t num_chunk, 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 3fb46f1a8..77ee86cb5 100644 --- a/src/turbo/distance/scalar/pq_quantizer_int8/pq_distance.cc +++ b/src/turbo/distance/scalar/pq_quantizer_int8/pq_distance.cc @@ -17,19 +17,19 @@ namespace zvec::turbo::scalar { void pq_adc_int8_distance(const void *pq_code_v, const void *lut_v, - size_t num_subquantizers, float *out) { + size_t num_chunk, 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) { + for (size_t m = 0; m < num_chunk; ++m) { sum += lut[m * kNumCentroids + pq_code[m]]; } *out = sum; } void pq_sdc_int8_distance(const void *a_v, const void *b_v, - const void *dist_table_v, size_t num_subquantizers, + const void *dist_table_v, size_t num_chunk, float *out) { constexpr size_t kNumCentroids = 256; constexpr size_t kTablePerSub = kNumCentroids * kNumCentroids; // 65536 @@ -37,7 +37,7 @@ void pq_sdc_int8_distance(const void *a_v, const void *b_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) { + for (size_t m = 0; m < num_chunk; ++m) { size_t idx = m * kTablePerSub + static_cast(a[m]) * kNumCentroids + static_cast(b[m]); sum += dist_table[idx]; @@ -46,7 +46,7 @@ void pq_sdc_int8_distance(const void *a_v, const void *b_v, } void pq_adc_int8_batch_distance(const void **candidates_v, const void *lut_v, - size_t num, size_t num_subquantizers, + size_t num, size_t num_chunk, float *out) { constexpr size_t kNumCentroids = 256; const auto *lut = reinterpret_cast(lut_v); @@ -64,7 +64,7 @@ void pq_adc_int8_batch_distance(const void **candidates_v, const void *lut_v, 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) { + for (size_t m = 0; m < num_chunk; ++m) { const float *tab = lut + m * kNumCentroids; d0 += tab[c0[m]]; d1 += tab[c1[m]]; @@ -80,7 +80,7 @@ void pq_adc_int8_batch_distance(const void **candidates_v, const void *lut_v, for (; i < num; ++i) { const uint8_t *code = candidates[i]; float d = 0.0f; - for (size_t m = 0; m < num_subquantizers; ++m) { + for (size_t m = 0; m < num_chunk; ++m) { d += lut[m * kNumCentroids + code[m]]; } out[i] = d; 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 d968aad7d..7ec8c1849 100644 --- a/src/turbo/distance/scalar/pq_quantizer_int8/pq_distance.h +++ b/src/turbo/distance/scalar/pq_quantizer_int8/pq_distance.h @@ -22,29 +22,29 @@ 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]] +// distance = sum_{m=0}^{num_chunk-1} lut[m * 256 + pq_code[m]] void pq_adc_int8_distance(const void *pq_code, const void *lut, - size_t num_subquantizers, float *out); + size_t num_chunk, 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 layout: [num_chunk * 256 * 256] // dist_table[m * 256 * 256 + i * 256 + j] = // ||centroid[m][i] - centroid[m][j]||^2 // -// distance = sum_{m=0}^{num_subquantizers-1} +// distance = sum_{m=0}^{num_chunk-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); + size_t num_chunk, 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, + size_t num, size_t num_chunk, 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 0a77fd48f..1cddd5d91 100644 --- a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc @@ -34,7 +34,7 @@ namespace turbo { // --------------------------------------------------------------------------- struct PqInt8SerPayload { uint32_t original_dim; - uint32_t num_subquantizers; + uint32_t num_chunk; uint32_t sub_dim; uint32_t num_centroids; // always 256 for int8 }; @@ -45,25 +45,25 @@ int PqInt8Quantizer::init(const IndexMeta &meta, const ailego::Params ¶ms) { uint32_t d = meta.dimension(); original_dim_ = d; - // Read num_subquantizers from params (required). + // Read num_chunk from params (required). uint32_t nsq = 0; - if (!params.get("num_subquantizers", &nsq) || nsq == 0) { - LOG_ERROR("PqInt8Quantizer: num_subquantizers not set or zero"); + if (!params.get("num_chunk", &nsq) || nsq == 0) { + LOG_ERROR("PqInt8Quantizer: num_chunk not set or zero"); return kErrUnsupported; } if (d % nsq != 0) { LOG_ERROR( - "PqInt8Quantizer: dim (%u) is not divisible by num_subquantizers (%u)", + "PqInt8Quantizer: dim (%u) is not divisible by num_chunk (%u)", d, nsq); return kErrUnsupported; } - num_subquantizers_ = nsq; + num_chunk_ = nsq; sub_dim_ = d / nsq; // Pre-allocate centroids (filled by train()). centroids_.resize( - static_cast(num_subquantizers_) * kNumCentroids * sub_dim_, 0.0f); + static_cast(num_chunk_) * kNumCentroids * sub_dim_, 0.0f); // Dispatch ISA kernels (scalar only for now). auto pq_k = get_pq_kernels(DataType::kInt8); @@ -120,8 +120,8 @@ int PqInt8Quantizer::init(const IndexMeta &meta, const ailego::Params ¶ms) { 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) { + centroid_ptrs_cache_.resize(num_chunk_); + for (uint32_t m = 0; m < num_chunk_; ++m) { const float *centroids_m = centroids_.data() + static_cast(m) * k * d; auto &ptrs = centroid_ptrs_cache_[m]; @@ -139,11 +139,12 @@ void PqInt8Quantizer::train_subquantizer(const float *data, size_t num, float *centroids_m = centroids_.data() + static_cast(sub_idx) * k * d; // Create KMeans algorithm (L2 distance via NumericalKmeansContext). - // Enable spherical mode for IP / Cosine so that centroids are updated - // on the unit sphere, improving convergence for angular objectives. - const bool spherical = - (meta_.metric_name() == "InnerProduct" || - meta_.metric_name() == "Cosine"); + // Enable spherical mode for Cosine only. Cosine training data is + // pre-normalized in train(), so spherical centroid constraint keeps + // L2 assignment equivalent to argmax dot(x,c). For IP, data is NOT + // normalized, so spherical would create a mismatch between assignment + // (norm-dependent L2) and update (norm-discarding normalization). + const bool spherical = (meta_.metric_name() == "Cosine"); ailego::NumericalKmeans algorithm( k, d, spherical); @@ -252,7 +253,7 @@ int PqInt8Quantizer::train(IndexHolder::Pointer holder, int 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, + num, original_dim_, num_chunk_, sub_dim_, kMaxKmeansIters, thread_count); // Create thread pool (aligned with multi_chunk_cluster L183-185) @@ -267,7 +268,7 @@ int PqInt8Quantizer::train(IndexHolder::Pointer holder, int thread_count) { task_group->submit(ailego::Closure::New( [this, &all_data, num, data_stride, i, pool_count, &finished]() { for (uint32_t m = static_cast(i); - m < num_subquantizers_; + m < num_chunk_; m += static_cast(pool_count)) { train_subquantizer(all_data.data(), num, data_stride, m); finished++; @@ -277,7 +278,7 @@ int PqInt8Quantizer::train(IndexHolder::Pointer holder, int thread_count) { task_group->wait_finish(); LOG_INFO(" all %u sub-quantizers trained (%zu threads)", - num_subquantizers_, pool_count); + num_chunk_, pool_count); // Pre-build centroid pointer cache (needed by compute_dist_table). build_centroid_ptrs_cache(); @@ -293,14 +294,14 @@ 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_chunk_) * k * k, 0.0f); // For each sub-quantizer, compute centroid-to-centroid distances via // the metric-aware fp32 batch distance function (fp32_batch_fn_). // L2: dist_table[m][i][j] = ||c_m[i] - c_m[j]||^2 // IP: dist_table[m][i][j] = -dot(c_m[i], c_m[j]) // Cosine: centroids trained on normalized data, same as IP. - for (uint32_t m = 0; m < num_subquantizers_; ++m) { + for (uint32_t m = 0; m < num_chunk_; ++m) { const float *centroids_m = centroids_.data() + m * k * d; float *table_m = dist_table_.data() + m * k * k; @@ -338,7 +339,7 @@ void PqInt8Quantizer::quantize_data(const void *input, void *output) const { // a second pass over the data. float dists[kNumCentroids]; - for (uint32_t m = 0; m < num_subquantizers_; ++m) { + for (uint32_t m = 0; m < num_chunk_; ++m) { const float *sub_vec = vec + m * sub_dim_; const auto ¢roid_ptrs = centroid_ptrs_cache_[m]; @@ -361,7 +362,7 @@ void PqInt8Quantizer::quantize_data(const void *input, void *output) const { // Store norm after PQ code for Cosine dequantize support. if (meta_.metric_name() == "Cosine") { - float *norm_out = reinterpret_cast(code + num_subquantizers_); + float *norm_out = reinterpret_cast(code + num_chunk_); *norm_out = vec_norm; } } @@ -389,7 +390,7 @@ void PqInt8Quantizer::quantize_query(const void *input, void *output) const { // 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) { + for (uint32_t m = 0; m < num_chunk_; ++m) { const auto ¢roid_ptrs = centroid_ptrs_cache_[m]; fp32_batch_fn_(const_cast(centroid_ptrs.data()), reinterpret_cast(query + m * sub_dim_), @@ -399,11 +400,11 @@ void PqInt8Quantizer::quantize_query(const void *input, void *output) const { 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]) + // dp = pq_code (uint8_t[num_chunk]) + // query = LUT (float[num_chunk * 256]) float d = 0.0f; adc_fn_(reinterpret_cast(dp), - reinterpret_cast(query), num_subquantizers_, &d); + reinterpret_cast(query), num_chunk_, &d); // For Cosine: ADC sum = -cos_sim; convert to cosine distance = 1 - cos_sim. if (meta_.metric_name() == "Cosine") { d = 1.0f + d; @@ -420,7 +421,7 @@ void PqInt8Quantizer::calc_distance_dp_query_batch(const void *const *dp_list, // 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_chunk_, dist_list); // For Cosine: convert -cos_sim to cosine distance. if (meta_.metric_name() == "Cosine") { for (int i = 0; i < dp_num; ++i) { @@ -432,11 +433,11 @@ void PqInt8Quantizer::calc_distance_dp_query_batch(const void *const *dp_list, 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_) * + std::vector lut(static_cast(num_chunk_) * kNumCentroids); quantize_query(query, lut.data()); float d = 0.0f; - adc_fn_(reinterpret_cast(dp), lut.data(), num_subquantizers_, + adc_fn_(reinterpret_cast(dp), lut.data(), num_chunk_, &d); if (meta_.metric_name() == "Cosine") { d = 1.0f + d; @@ -447,13 +448,13 @@ float PqInt8Quantizer::calc_distance_dp_query_unquantized( 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_) * + std::vector lut(static_cast(num_chunk_) * kNumCentroids); quantize_query(query, lut.data()); // Use ISA-dispatched batch4 ADC kernel (4-way ILP + SIMD gather). // const_cast: see calc_distance_dp_query_batch for rationale. batch_adc_fn_(const_cast(dp_list), lut.data(), - static_cast(dp_num), num_subquantizers_, dist_list); + static_cast(dp_num), num_chunk_, dist_list); if (meta_.metric_name() == "Cosine") { for (int i = 0; i < dp_num; ++i) { dist_list[i] = 1.0f + dist_list[i]; @@ -466,7 +467,7 @@ float PqInt8Quantizer::calc_distance_dp_dp(const void *dp1, float d = 0.0f; sdc_fn_(reinterpret_cast(dp1), reinterpret_cast(dp2), dist_table_.data(), - num_subquantizers_, &d); + num_chunk_, &d); return d; } @@ -499,7 +500,7 @@ int PqInt8Quantizer::dequantize(const void *in, const IndexQueryMeta &qmeta, // 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) { + for (uint32_t m = 0; m < num_chunk_; ++m) { const float *centroids_m = centroids_.data() + static_cast(m) * k * d; const float *centroid = centroids_m + static_cast(code[m]) * d; @@ -510,7 +511,7 @@ int PqInt8Quantizer::dequantize(const void *in, const IndexQueryMeta &qmeta, // 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)); + std::memcpy(&norm, code + num_chunk_, sizeof(float)); for (uint32_t j = 0; j < original_dim_; ++j) { result[j] *= norm; } @@ -536,7 +537,7 @@ DistanceImpl PqInt8Quantizer::distance(const void *query, return DistanceImpl(std::move(adc_func), std::move(batch_func), std::move(lut_storage), - static_cast(num_subquantizers_)); + static_cast(num_chunk_)); } DistanceImpl PqInt8Quantizer::sym_distance(const void *query, @@ -550,7 +551,7 @@ DistanceImpl PqInt8Quantizer::sym_distance(const void *query, DistanceFunc sdc_func = [sdc, dt](const void *a, const void *b, size_t dim, float *out) { sdc(a, b, dt, dim, out); }; - // The query is a PQ code (uint8_t[num_subquantizers]), NOT a LUT. + // The query is a PQ code (uint8_t[num_chunk]), NOT a LUT. // SDC computes distance between two PQ codes via the centroid-to-centroid // dist_table_, so both sides must be PQ codes. size_t code_bytes = quantized_datapoint_vector_length(); @@ -558,7 +559,7 @@ DistanceImpl PqInt8Quantizer::sym_distance(const void *query, // SDC has no batch kernel — use the 3-arg constructor. return DistanceImpl(std::move(sdc_func), std::move(code_storage), - static_cast(num_subquantizers_)); + static_cast(num_chunk_)); } // --------------------------------------------------------------------------- @@ -576,7 +577,7 @@ int PqInt8Quantizer::serialize(std::string *out) const { PqInt8SerPayload payload{}; payload.original_dim = original_dim_; - payload.num_subquantizers = num_subquantizers_; + payload.num_chunk = num_chunk_; payload.sub_dim = sub_dim_; payload.num_centroids = kNumCentroids; @@ -616,12 +617,12 @@ int PqInt8Quantizer::deserialize(const void *data, size_t len) { ptr += sizeof(payload); original_dim_ = payload.original_dim; - num_subquantizers_ = payload.num_subquantizers; + num_chunk_ = payload.num_chunk; sub_dim_ = payload.sub_dim; meta_.set_meta(IndexMeta::DataType::DT_FP32, original_dim_); - size_t centroids_bytes = static_cast(num_subquantizers_) * + size_t centroids_bytes = static_cast(num_chunk_) * kNumCentroids * sub_dim_ * sizeof(float); centroids_.resize(centroids_bytes / sizeof(float)); 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 e05e36429..1d4250bee 100644 --- a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h @@ -28,8 +28,8 @@ 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] +//! Datapoints are encoded as uint8_t[num_chunk] codes. +//! Queries are encoded as a float LUT of size [num_chunk * 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). @@ -68,11 +68,11 @@ class PqInt8Quantizer : public Quantizer { int train(IndexHolder::Pointer holder, int thread_count) override; size_t quantized_datapoint_vector_length() const override { - return num_subquantizers_ + extra_meta_size_; + return num_chunk_ + extra_meta_size_; } size_t quantized_query_vector_length() const override { - return static_cast(num_subquantizers_) * kNumCentroids * + return static_cast(num_chunk_) * kNumCentroids * sizeof(float); } @@ -144,14 +144,14 @@ class PqInt8Quantizer : public Quantizer { IndexMeta meta_{}; uint32_t original_dim_{0}; - uint32_t num_subquantizers_{0}; + uint32_t num_chunk_{0}; uint32_t sub_dim_{0}; - //! Centroids: [num_subquantizers * kNumCentroids * sub_dim] + //! Centroids: [num_chunk * kNumCentroids * sub_dim] std::vector centroids_; //! Centroid-to-centroid distance table for SDC: - //! [num_subquantizers * kNumCentroids * kNumCentroids] + //! [num_chunk * kNumCentroids * kNumCentroids] std::vector dist_table_; //! Pre-built centroid pointer arrays for each sub-quantizer. diff --git a/tests/turbo/turbo_pq_int8_quantizer_test.cc b/tests/turbo/turbo_pq_int8_quantizer_test.cc index 43cb86c96..c6569f179 100644 --- a/tests/turbo/turbo_pq_int8_quantizer_test.cc +++ b/tests/turbo/turbo_pq_int8_quantizer_test.cc @@ -48,7 +48,7 @@ static float reference_sq_euclidean(const float *a, const float *b, // Helper to create a PqInt8Quantizer via the factory. static std::shared_ptr make_pq_quantizer( - size_t dim, size_t num_subquantizers) { + size_t dim, size_t num_chunk) { auto q = IndexFactory::CreateQuantizer("PqInt8Quantizer"); if (!q) return nullptr; @@ -57,7 +57,7 @@ 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_chunk", static_cast(num_chunk)); if (q->init(meta, params) != 0) return nullptr; return q; } @@ -82,18 +82,18 @@ make_random_holder(size_t count, size_t dim, uint32_t seed = 42) { // --------------------------------------------------------------------------- TEST(PqInt8Quantizer, InitInvalidParams) { - // dim not divisible by num_subquantizers + // dim not divisible by num_chunk auto q = make_pq_quantizer(10, 3); EXPECT_EQ(q, nullptr); - // num_subquantizers = 0 + // num_chunk = 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)); + params.set("num_chunk", static_cast(0)); EXPECT_NE(0, q2->init(meta, params)); } @@ -300,10 +300,10 @@ 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_chunk, 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 m = 0; m < num_chunk; ++m) { for (size_t c = 0; c < kNumCentroids; ++c) { lut[m * kNumCentroids + c] = dist(gen); } @@ -311,11 +311,11 @@ void fill_random_lut(float *lut, size_t num_subquantizers, std::mt19937 &gen) { } // Helper to generate random dist_table (SDC) -void fill_random_sdc_table(float *table, size_t num_subquantizers, +void fill_random_sdc_table(float *table, size_t num_chunk, 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 m = 0; m < num_chunk; ++m) { for (size_t i = 0; i < kTablePerSub; ++i) { table[m * kTablePerSub + i] = dist(gen); } @@ -452,7 +452,7 @@ TEST(PqInt8SimdConsistency, AdcDistanceM1) { // Helper to create a PqInt8Quantizer with Cosine metric. static std::shared_ptr make_pq_cosine_quantizer( - size_t dim, size_t num_subquantizers) { + size_t dim, size_t num_chunk) { auto q = IndexFactory::CreateQuantizer("PqInt8Quantizer"); if (!q) return nullptr; @@ -461,7 +461,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_chunk", static_cast(num_chunk)); if (q->init(meta, params) != 0) return nullptr; return q; } @@ -694,7 +694,7 @@ TEST(PqInt8Quantizer, L2Dequantize) { // Helper to create a PqInt8Quantizer with InnerProduct metric. static std::shared_ptr make_pq_ip_quantizer( - size_t dim, size_t num_subquantizers) { + size_t dim, size_t num_chunk) { auto q = IndexFactory::CreateQuantizer("PqInt8Quantizer"); if (!q) return nullptr; @@ -703,7 +703,7 @@ 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_chunk", static_cast(num_chunk)); if (q->init(meta, params) != 0) return nullptr; return q; } From 084b475b4a81157bea84bb283f7e8890e80e2dd4 Mon Sep 17 00:00:00 2001 From: ray Date: Wed, 15 Jul 2026 16:44:01 +0800 Subject: [PATCH 061/100] fix: remove comment --- examples/c++/CMakeLists.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/c++/CMakeLists.txt b/examples/c++/CMakeLists.txt index 1e111cd48..0786bdb24 100644 --- a/examples/c++/CMakeLists.txt +++ b/examples/c++/CMakeLists.txt @@ -14,7 +14,7 @@ if(NOT DEFINED HOST_BUILD_DIR) endif() get_filename_component(ZVEC_ROOT_DIR "${CMAKE_CURRENT_LIST_DIR}/../.." ABSOLUTE) -# Public headers are self-contained under src/include. + set(ZVEC_INCLUDE_DIR ${ZVEC_ROOT_DIR}/src/include) set(ZVEC_LIB_DIR ${ZVEC_ROOT_DIR}/${HOST_BUILD_DIR}/lib) From 931c61da2fcd016c540f18934c10cfd38b787d1b Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Wed, 15 Jul 2026 17:24:32 +0800 Subject: [PATCH 062/100] tmp --- .../preprocessor/fht_rotator/fht_rotator.cc | 41 +++++++++---------- .../preprocessor/fht_rotator/fht_rotator.h | 4 +- 2 files changed, 20 insertions(+), 25 deletions(-) diff --git a/src/turbo/preprocessor/fht_rotator/fht_rotator.cc b/src/turbo/preprocessor/fht_rotator/fht_rotator.cc index 979112305..783564d9e 100644 --- a/src/turbo/preprocessor/fht_rotator/fht_rotator.cc +++ b/src/turbo/preprocessor/fht_rotator/fht_rotator.cc @@ -17,6 +17,7 @@ #include #include #include "quantizer/quantizer.h" + namespace zvec { namespace turbo { @@ -150,54 +151,50 @@ void FhtRotator::apply(const float *in, float *out) const { 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); + std::memcpy(out, in, sizeof(float) * 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); + inplace_fn_(out, trunc_dim_); + rescale_fn_(out, trunc_dim_, inv_fac); flip_sign_fn_(flip_.data() + static_cast(round) * flip_offset_, - data.data(), dim); + out, 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); + rescale_fn_(out, 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; + float *trunc_ptr = out + start; // Undo Round 4 (FHT on [start, start+trunc_dim)) - inv_kacs_walk_fn_(data.data(), dim); + inv_kacs_walk_fn_(out, 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); + flip_sign_fn_(flip_.data() + 3 * flip_offset_, out, 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); + inv_kacs_walk_fn_(out, dim); + inplace_fn_(out, trunc_dim_); + rescale_fn_(out, trunc_dim_, inv_fac); + flip_sign_fn_(flip_.data() + 2 * flip_offset_, out, dim); // Undo Round 2 (FHT on [start, start+trunc_dim)) - inv_kacs_walk_fn_(data.data(), dim); + inv_kacs_walk_fn_(out, 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); + flip_sign_fn_(flip_.data() + flip_offset_, out, 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)); + inv_kacs_walk_fn_(out, dim); + inplace_fn_(out, trunc_dim_); + rescale_fn_(out, trunc_dim_, inv_fac); + flip_sign_fn_(flip_.data(), out, dim); } // --------------------------------------------------------------------------- diff --git a/src/turbo/preprocessor/fht_rotator/fht_rotator.h b/src/turbo/preprocessor/fht_rotator/fht_rotator.h index 450420276..c914636a9 100644 --- a/src/turbo/preprocessor/fht_rotator/fht_rotator.h +++ b/src/turbo/preprocessor/fht_rotator/fht_rotator.h @@ -14,8 +14,6 @@ #pragma once -#include -#include #include #include #include @@ -64,7 +62,7 @@ class FhtRotator : public Preprocessor { int serialize(std::string *out) const override; int deserialize(const void *data, size_t len) override; - //! Rotator type tag (kFht = 1). + //! Rotator type tag (kFht = 0). RotatorType rotator_type() const { return RotatorType::kFht; } From 7552df5231c70187c28de824b3e7b272adbd8fd0 Mon Sep 17 00:00:00 2001 From: ray Date: Wed, 15 Jul 2026 17:26:57 +0800 Subject: [PATCH 063/100] fix: fix empty line --- examples/c++/CMakeLists.txt | 1 - 1 file changed, 1 deletion(-) diff --git a/examples/c++/CMakeLists.txt b/examples/c++/CMakeLists.txt index 0786bdb24..f5a19f056 100644 --- a/examples/c++/CMakeLists.txt +++ b/examples/c++/CMakeLists.txt @@ -12,7 +12,6 @@ set(CMAKE_EXPORT_COMPILE_COMMANDS ON) if(NOT DEFINED HOST_BUILD_DIR) set(HOST_BUILD_DIR "build") endif() - get_filename_component(ZVEC_ROOT_DIR "${CMAKE_CURRENT_LIST_DIR}/../.." ABSOLUTE) set(ZVEC_INCLUDE_DIR ${ZVEC_ROOT_DIR}/src/include) From 28fe49ac8e370c3dcc20dd1a45c4e1ce160e7a8c Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Wed, 15 Jul 2026 17:48:21 +0800 Subject: [PATCH 064/100] add RotateType --- src/include/zvec/turbo/turbo.h | 14 ++- .../preprocessor/fht_rotator/fht_rotator.cc | 10 +- .../preprocessor/fht_rotator/fht_rotator.h | 5 +- src/turbo/preprocessor/preprocessor.h | 10 +- src/turbo/turbo.cc | 114 ++++++++++-------- tests/turbo/turbo_fht_rotator_test.cc | 4 +- 6 files changed, 87 insertions(+), 70 deletions(-) diff --git a/src/include/zvec/turbo/turbo.h b/src/include/zvec/turbo/turbo.h index f594c3273..1b3fd11f0 100644 --- a/src/include/zvec/turbo/turbo.h +++ b/src/include/zvec/turbo/turbo.h @@ -16,6 +16,7 @@ #include #include #include +#include #include namespace zvec::turbo { @@ -50,6 +51,10 @@ struct FhtKernels { FhtVecRescaleFunc rescale; }; +// Variant over all rotator kernel structs. Add new types here as they are +// introduced (e.g. MatrixKernels). +using RotatorKernels = std::variant; + enum class MetricType { kSquaredEuclidean, kCosine, @@ -76,6 +81,10 @@ enum class QuantizeType { kRabit }; +enum class RotateType : uint16_t { + kFht = 0, //!< O(d log d) FHT-based Kac random rotation +}; + enum class CpuArchType { kAuto, kScalar, @@ -112,7 +121,8 @@ 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(CpuArchType cpu_arch_type = CpuArchType::kAuto); +// Returns rotator kernels dispatched for the current CPU. +RotatorKernels get_rotator_kernels( + RotateType rotate_type, CpuArchType cpu_arch_type = CpuArchType::kAuto); } // namespace zvec::turbo diff --git a/src/turbo/preprocessor/fht_rotator/fht_rotator.cc b/src/turbo/preprocessor/fht_rotator/fht_rotator.cc index 783564d9e..54b02cee2 100644 --- a/src/turbo/preprocessor/fht_rotator/fht_rotator.cc +++ b/src/turbo/preprocessor/fht_rotator/fht_rotator.cc @@ -41,7 +41,7 @@ FhtRotator::Pointer FhtRotator::create(int 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(); + auto k = std::get(get_rotator_kernels(RotateType::kFht)); r->flip_sign_fn_ = k.flip_sign; r->kacs_walk_fn_ = k.kacs_walk; r->inv_kacs_walk_fn_ = k.inv_kacs_walk; @@ -65,7 +65,7 @@ FhtRotator::Pointer FhtRotator::from_blob(const void *data, size_t len) { 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) { + if (static_cast(hdr->rotator_type) != RotateType::kFht) { return nullptr; } @@ -208,7 +208,7 @@ int FhtRotator::serialize(std::string *out) const { RotatorSerHeader hdr{}; hdr.magic = kRotatorMagic; hdr.version = kRotatorSerVersion; - hdr.rotator_type = static_cast(RotatorType::kFht); + hdr.rotator_type = static_cast(RotateType::kFht); hdr.in_dim = static_cast(in_dim_); hdr.out_dim = static_cast(out_dim_); hdr.payload_size = static_cast(flip_.size()); @@ -226,7 +226,7 @@ int FhtRotator::deserialize(const void *data, size_t len) { const auto *hdr = reinterpret_cast(data); if (hdr->magic != kRotatorMagic) return kErrUnsupported; if (hdr->version != kRotatorSerVersion) return kErrUnsupported; - if (static_cast(hdr->rotator_type) != RotatorType::kFht) { + if (static_cast(hdr->rotator_type) != RotateType::kFht) { return kErrUnsupported; } @@ -238,7 +238,7 @@ int FhtRotator::deserialize(const void *data, size_t len) { 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(); + auto k = std::get(get_rotator_kernels(RotateType::kFht)); flip_sign_fn_ = k.flip_sign; kacs_walk_fn_ = k.kacs_walk; inv_kacs_walk_fn_ = k.inv_kacs_walk; diff --git a/src/turbo/preprocessor/fht_rotator/fht_rotator.h b/src/turbo/preprocessor/fht_rotator/fht_rotator.h index c914636a9..40aadf42a 100644 --- a/src/turbo/preprocessor/fht_rotator/fht_rotator.h +++ b/src/turbo/preprocessor/fht_rotator/fht_rotator.h @@ -16,7 +16,6 @@ #include #include -#include #include "preprocessor/preprocessor.h" namespace zvec { @@ -63,8 +62,8 @@ class FhtRotator : public Preprocessor { int deserialize(const void *data, size_t len) override; //! Rotator type tag (kFht = 0). - RotatorType rotator_type() const { - return RotatorType::kFht; + RotateType rotate_type() const { + return RotateType::kFht; } private: diff --git a/src/turbo/preprocessor/preprocessor.h b/src/turbo/preprocessor/preprocessor.h index 07ed313ea..c51f19b14 100644 --- a/src/turbo/preprocessor/preprocessor.h +++ b/src/turbo/preprocessor/preprocessor.h @@ -14,10 +14,9 @@ #pragma once -#include -#include #include #include +#include namespace zvec { namespace turbo { @@ -27,18 +26,13 @@ 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 + uint16_t rotator_type; // RotateType uint32_t in_dim; // input dimensionality uint32_t out_dim; // output dimensionality uint32_t payload_size; // bytes following the header diff --git a/src/turbo/turbo.cc b/src/turbo/turbo.cc index 87826f12f..7d58cd6be 100644 --- a/src/turbo/turbo.cc +++ b/src/turbo/turbo.cc @@ -157,66 +157,80 @@ UniformQuantizeFunc get_uniform_quantize_func(DataType data_type) { return nullptr; } -FhtKernels get_fht_kernels(CpuArchType cpu_arch_type) { - // Suppress unused-parameter warning when no SIMD #if blocks are compiled in. +RotatorKernels get_rotator_kernels(RotateType rotate_type, + CpuArchType cpu_arch_type) { (void)cpu_arch_type; - 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; + + switch (rotate_type) { + case RotateType::kFht: { + 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 && - (cpu_arch_type == CpuArchType::kAuto || - cpu_arch_type == CpuArchType::kAVX512)) { - 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; - } + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F && + zvec::ailego::internal::CpuFeatures::static_flags_.AVX512DQ && + (cpu_arch_type == CpuArchType::kAuto || + cpu_arch_type == CpuArchType::kAVX512)) { + 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 && - (cpu_arch_type == CpuArchType::kAuto || - cpu_arch_type == CpuArchType::kAVX2)) { - 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; - } + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2 && + (cpu_arch_type == CpuArchType::kAuto || + cpu_arch_type == CpuArchType::kAVX2)) { + 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 && - (cpu_arch_type == CpuArchType::kAuto || - cpu_arch_type == CpuArchType::kSSE)) { - 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; - } + if (zvec::ailego::internal::CpuFeatures::static_flags_.SSE2 && + (cpu_arch_type == CpuArchType::kAuto || + cpu_arch_type == CpuArchType::kSSE)) { + 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 #if defined(__ARM_NEON) && defined(__aarch64__) - if (cpu_arch_type == CpuArchType::kAuto || - cpu_arch_type == CpuArchType::kNEON) { - k.flip_sign = neon::fht_flip_sign_neon; - k.kacs_walk = neon::fht_kacs_walk_neon; - k.inv_kacs_walk = neon::fht_inv_kacs_walk_neon; - k.rescale = neon::fht_vec_rescale_neon; - // inplace fallback to scalar (NEON has no fht_inplace) - return k; - } + if (cpu_arch_type == CpuArchType::kAuto || + cpu_arch_type == CpuArchType::kNEON) { + k.flip_sign = neon::fht_flip_sign_neon; + k.kacs_walk = neon::fht_kacs_walk_neon; + k.inv_kacs_walk = neon::fht_inv_kacs_walk_neon; + k.rescale = neon::fht_vec_rescale_neon; + // inplace fallback to scalar (NEON has no fht_inplace) + return k; + } #endif - return k; // scalar + return k; // scalar + } + } + + // Fallback (unreachable for valid RotateType values). + FhtKernels k; + 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; + return k; } } // namespace zvec::turbo diff --git a/tests/turbo/turbo_fht_rotator_test.cc b/tests/turbo/turbo_fht_rotator_test.cc index f1ad4877b..1e454ed9b 100644 --- a/tests/turbo/turbo_fht_rotator_test.cc +++ b/tests/turbo/turbo_fht_rotator_test.cc @@ -152,7 +152,7 @@ TEST(FhtRotator, CreateGeneratesFlip) { 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->rotator_type), RotateType::kFht); EXPECT_EQ(static_cast(hdr->in_dim), dim); EXPECT_EQ(static_cast(hdr->out_dim), dim); EXPECT_GT(hdr->payload_size, 0u); @@ -195,7 +195,7 @@ TEST(FhtRotator, FromBlobMalformed) { RotatorSerHeader hdr{}; hdr.magic = 0xDEADBEEF; hdr.version = kRotatorSerVersion; - hdr.rotator_type = static_cast(RotatorType::kFht); + hdr.rotator_type = static_cast(RotateType::kFht); hdr.payload_size = 0; EXPECT_EQ(FhtRotator::from_blob(&hdr, sizeof(hdr)), nullptr); } From bc9149d92da6960a1df7db9bd24597e7a4bd5bb3 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Wed, 15 Jul 2026 17:48:30 +0800 Subject: [PATCH 065/100] add RotateType --- src/turbo/turbo.cc | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/src/turbo/turbo.cc b/src/turbo/turbo.cc index 7d58cd6be..db21f99e7 100644 --- a/src/turbo/turbo.cc +++ b/src/turbo/turbo.cc @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include #include #include #include "avx512_vnni/record_quantized_int8/cosine.h" @@ -224,13 +225,8 @@ RotatorKernels get_rotator_kernels(RotateType rotate_type, } // Fallback (unreachable for valid RotateType values). - FhtKernels k; - 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; - return k; + assert(false && "unsupported RotateType"); + return FhtKernels{}; } } // namespace zvec::turbo From 8997541b55db74a94ae39cb8f37fdfececc537f7 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Thu, 16 Jul 2026 17:32:21 +0800 Subject: [PATCH 066/100] Unified Scheduling Mode --- src/include/zvec/turbo/turbo.h | 29 ++- src/turbo/distance/avx2/fht/fht.cc | 99 ++++++++++ src/turbo/distance/avx2/fht/fht.h | 8 + src/turbo/distance/avx512/fht/fht.cc | 99 ++++++++++ src/turbo/distance/avx512/fht/fht.h | 8 + src/turbo/distance/neon/fht/fht.cc | 100 ++++++++++ src/turbo/distance/neon/fht/fht.h | 8 + src/turbo/distance/scalar/fht/fht.cc | 101 ++++++++++ src/turbo/distance/scalar/fht/fht.h | 9 + src/turbo/distance/sse/fht/fht.cc | 100 ++++++++++ src/turbo/distance/sse/fht/fht.h | 8 + .../preprocessor/fht_rotator/fht_rotator.cc | 173 +++++------------- .../preprocessor/fht_rotator/fht_rotator.h | 45 +++-- src/turbo/turbo.cc | 41 ++--- 14 files changed, 639 insertions(+), 189 deletions(-) diff --git a/src/include/zvec/turbo/turbo.h b/src/include/zvec/turbo/turbo.h index 1b3fd11f0..3face81bd 100644 --- a/src/include/zvec/turbo/turbo.h +++ b/src/include/zvec/turbo/turbo.h @@ -16,7 +16,6 @@ #include #include #include -#include #include namespace zvec::turbo { @@ -36,25 +35,19 @@ 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; +// Generic rotate / unrotate function pointer types. +// ctx is an opaque context (e.g. FhtCtx*) managed by the caller. +using RotateFunc = void (*)(const float *in, float *out, size_t in_dim, + size_t out_dim, void *ctx); +using UnrotateFunc = void (*)(const float *in, float *out, size_t in_dim, + size_t out_dim, void *ctx); + +// ISA-dispatched rotate/unrotate kernels. +struct RotatorKernels { + RotateFunc rotate = nullptr; + UnrotateFunc unrotate = nullptr; }; -// Variant over all rotator kernel structs. Add new types here as they are -// introduced (e.g. MatrixKernels). -using RotatorKernels = std::variant; - enum class MetricType { kSquaredEuclidean, kCosine, diff --git a/src/turbo/distance/avx2/fht/fht.cc b/src/turbo/distance/avx2/fht/fht.cc index a6d3edc36..10c101b6f 100644 --- a/src/turbo/distance/avx2/fht/fht.cc +++ b/src/turbo/distance/avx2/fht/fht.cc @@ -132,6 +132,105 @@ void fht_vec_rescale_avx2(float *data, size_t n, float factor) { } } +void fht_rotate_avx2(const float * /*in*/, float *out, size_t in_dim, + size_t /*out_dim*/, void *ctx) { + auto *base = reinterpret_cast(ctx); + const size_t flip_offset = *reinterpret_cast(base); + const size_t trunc_dim = *reinterpret_cast(base + 8); + const float fac = *reinterpret_cast(base + 16); + const uint8_t *flip = base + 24; + const size_t dim = in_dim; + + if (trunc_dim == dim) { + fht_flip_sign_avx2(flip, out, dim); + fht_inplace_avx2(out, trunc_dim); + fht_vec_rescale_avx2(out, trunc_dim, fac); + + fht_flip_sign_avx2(flip + flip_offset, out, dim); + fht_inplace_avx2(out, trunc_dim); + fht_vec_rescale_avx2(out, trunc_dim, fac); + + fht_flip_sign_avx2(flip + 2 * flip_offset, out, dim); + fht_inplace_avx2(out, trunc_dim); + fht_vec_rescale_avx2(out, trunc_dim, fac); + + fht_flip_sign_avx2(flip + 3 * flip_offset, out, dim); + fht_inplace_avx2(out, trunc_dim); + fht_vec_rescale_avx2(out, trunc_dim, fac); + return; + } + + size_t start = dim - trunc_dim; + float *trunc_ptr = out + start; + + fht_flip_sign_avx2(flip, out, dim); + fht_inplace_avx2(out, trunc_dim); + fht_vec_rescale_avx2(out, trunc_dim, fac); + fht_kacs_walk_avx2(out, dim); + + fht_flip_sign_avx2(flip + flip_offset, out, dim); + fht_inplace_avx2(trunc_ptr, trunc_dim); + fht_vec_rescale_avx2(trunc_ptr, trunc_dim, fac); + fht_kacs_walk_avx2(out, dim); + + fht_flip_sign_avx2(flip + 2 * flip_offset, out, dim); + fht_inplace_avx2(out, trunc_dim); + fht_vec_rescale_avx2(out, trunc_dim, fac); + fht_kacs_walk_avx2(out, dim); + + fht_flip_sign_avx2(flip + 3 * flip_offset, out, dim); + fht_inplace_avx2(trunc_ptr, trunc_dim); + fht_vec_rescale_avx2(trunc_ptr, trunc_dim, fac); + fht_kacs_walk_avx2(out, dim); + + fht_vec_rescale_avx2(out, dim, 0.25f); +} + +void fht_unrotate_avx2(const float * /*in*/, float *out, size_t in_dim, + size_t /*out_dim*/, void *ctx) { + auto *base = reinterpret_cast(ctx); + const size_t flip_offset = *reinterpret_cast(base); + const size_t trunc_dim = *reinterpret_cast(base + 8); + const float fac = *reinterpret_cast(base + 16); + const uint8_t *flip = base + 24; + const size_t dim = in_dim; + + if (trunc_dim == dim) { + for (int round = 3; round >= 0; --round) { + fht_inplace_avx2(out, trunc_dim); + fht_vec_rescale_avx2(out, trunc_dim, fac); + fht_flip_sign_avx2(flip + static_cast(round) * flip_offset, out, + dim); + } + return; + } + + fht_vec_rescale_avx2(out, dim, 4.0f); + + size_t start = dim - trunc_dim; + float *trunc_ptr = out + start; + + fht_inv_kacs_walk_avx2(out, dim); + fht_inplace_avx2(trunc_ptr, trunc_dim); + fht_vec_rescale_avx2(trunc_ptr, trunc_dim, fac); + fht_flip_sign_avx2(flip + 3 * flip_offset, out, dim); + + fht_inv_kacs_walk_avx2(out, dim); + fht_inplace_avx2(out, trunc_dim); + fht_vec_rescale_avx2(out, trunc_dim, fac); + fht_flip_sign_avx2(flip + 2 * flip_offset, out, dim); + + fht_inv_kacs_walk_avx2(out, dim); + fht_inplace_avx2(trunc_ptr, trunc_dim); + fht_vec_rescale_avx2(trunc_ptr, trunc_dim, fac); + fht_flip_sign_avx2(flip + flip_offset, out, dim); + + fht_inv_kacs_walk_avx2(out, dim); + fht_inplace_avx2(out, trunc_dim); + fht_vec_rescale_avx2(out, trunc_dim, fac); + fht_flip_sign_avx2(flip, out, dim); +} + } // namespace zvec::turbo::avx2 #endif // __AVX2__ diff --git a/src/turbo/distance/avx2/fht/fht.h b/src/turbo/distance/avx2/fht/fht.h index 5e636a54f..600842108 100644 --- a/src/turbo/distance/avx2/fht/fht.h +++ b/src/turbo/distance/avx2/fht/fht.h @@ -34,4 +34,12 @@ 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); +//! Forward FHT rotation (AVX2). +void fht_rotate_avx2(const float *in, float *out, size_t in_dim, size_t out_dim, + void *ctx); + +//! Inverse FHT rotation (AVX2). +void fht_unrotate_avx2(const float *in, float *out, size_t in_dim, + size_t out_dim, void *ctx); + } // namespace zvec::turbo::avx2 diff --git a/src/turbo/distance/avx512/fht/fht.cc b/src/turbo/distance/avx512/fht/fht.cc index 2ba9e4b3b..a50ad8637 100644 --- a/src/turbo/distance/avx512/fht/fht.cc +++ b/src/turbo/distance/avx512/fht/fht.cc @@ -137,6 +137,105 @@ void fht_vec_rescale_avx512(float *data, size_t n, float factor) { } } +void fht_rotate_avx512(const float * /*in*/, float *out, size_t in_dim, + size_t /*out_dim*/, void *ctx) { + auto *base = reinterpret_cast(ctx); + const size_t flip_offset = *reinterpret_cast(base); + const size_t trunc_dim = *reinterpret_cast(base + 8); + const float fac = *reinterpret_cast(base + 16); + const uint8_t *flip = base + 24; + const size_t dim = in_dim; + + if (trunc_dim == dim) { + fht_flip_sign_avx512(flip, out, dim); + fht_inplace_avx512(out, trunc_dim); + fht_vec_rescale_avx512(out, trunc_dim, fac); + + fht_flip_sign_avx512(flip + flip_offset, out, dim); + fht_inplace_avx512(out, trunc_dim); + fht_vec_rescale_avx512(out, trunc_dim, fac); + + fht_flip_sign_avx512(flip + 2 * flip_offset, out, dim); + fht_inplace_avx512(out, trunc_dim); + fht_vec_rescale_avx512(out, trunc_dim, fac); + + fht_flip_sign_avx512(flip + 3 * flip_offset, out, dim); + fht_inplace_avx512(out, trunc_dim); + fht_vec_rescale_avx512(out, trunc_dim, fac); + return; + } + + size_t start = dim - trunc_dim; + float *trunc_ptr = out + start; + + fht_flip_sign_avx512(flip, out, dim); + fht_inplace_avx512(out, trunc_dim); + fht_vec_rescale_avx512(out, trunc_dim, fac); + fht_kacs_walk_avx512(out, dim); + + fht_flip_sign_avx512(flip + flip_offset, out, dim); + fht_inplace_avx512(trunc_ptr, trunc_dim); + fht_vec_rescale_avx512(trunc_ptr, trunc_dim, fac); + fht_kacs_walk_avx512(out, dim); + + fht_flip_sign_avx512(flip + 2 * flip_offset, out, dim); + fht_inplace_avx512(out, trunc_dim); + fht_vec_rescale_avx512(out, trunc_dim, fac); + fht_kacs_walk_avx512(out, dim); + + fht_flip_sign_avx512(flip + 3 * flip_offset, out, dim); + fht_inplace_avx512(trunc_ptr, trunc_dim); + fht_vec_rescale_avx512(trunc_ptr, trunc_dim, fac); + fht_kacs_walk_avx512(out, dim); + + fht_vec_rescale_avx512(out, dim, 0.25f); +} + +void fht_unrotate_avx512(const float * /*in*/, float *out, size_t in_dim, + size_t /*out_dim*/, void *ctx) { + auto *base = reinterpret_cast(ctx); + const size_t flip_offset = *reinterpret_cast(base); + const size_t trunc_dim = *reinterpret_cast(base + 8); + const float fac = *reinterpret_cast(base + 16); + const uint8_t *flip = base + 24; + const size_t dim = in_dim; + + if (trunc_dim == dim) { + for (int round = 3; round >= 0; --round) { + fht_inplace_avx512(out, trunc_dim); + fht_vec_rescale_avx512(out, trunc_dim, fac); + fht_flip_sign_avx512(flip + static_cast(round) * flip_offset, out, + dim); + } + return; + } + + fht_vec_rescale_avx512(out, dim, 4.0f); + + size_t start = dim - trunc_dim; + float *trunc_ptr = out + start; + + fht_inv_kacs_walk_avx512(out, dim); + fht_inplace_avx512(trunc_ptr, trunc_dim); + fht_vec_rescale_avx512(trunc_ptr, trunc_dim, fac); + fht_flip_sign_avx512(flip + 3 * flip_offset, out, dim); + + fht_inv_kacs_walk_avx512(out, dim); + fht_inplace_avx512(out, trunc_dim); + fht_vec_rescale_avx512(out, trunc_dim, fac); + fht_flip_sign_avx512(flip + 2 * flip_offset, out, dim); + + fht_inv_kacs_walk_avx512(out, dim); + fht_inplace_avx512(trunc_ptr, trunc_dim); + fht_vec_rescale_avx512(trunc_ptr, trunc_dim, fac); + fht_flip_sign_avx512(flip + flip_offset, out, dim); + + fht_inv_kacs_walk_avx512(out, dim); + fht_inplace_avx512(out, trunc_dim); + fht_vec_rescale_avx512(out, trunc_dim, fac); + fht_flip_sign_avx512(flip, out, dim); +} + } // namespace zvec::turbo::avx512 #endif // __AVX512F__ diff --git a/src/turbo/distance/avx512/fht/fht.h b/src/turbo/distance/avx512/fht/fht.h index 53241a0c5..6eecca959 100644 --- a/src/turbo/distance/avx512/fht/fht.h +++ b/src/turbo/distance/avx512/fht/fht.h @@ -34,4 +34,12 @@ 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); +//! Forward FHT rotation (AVX512). +void fht_rotate_avx512(const float *in, float *out, size_t in_dim, + size_t out_dim, void *ctx); + +//! Inverse FHT rotation (AVX512). +void fht_unrotate_avx512(const float *in, float *out, size_t in_dim, + size_t out_dim, void *ctx); + } // namespace zvec::turbo::avx512 diff --git a/src/turbo/distance/neon/fht/fht.cc b/src/turbo/distance/neon/fht/fht.cc index 8a94e7e2f..cadfd0a79 100644 --- a/src/turbo/distance/neon/fht/fht.cc +++ b/src/turbo/distance/neon/fht/fht.cc @@ -20,6 +20,7 @@ #include #include #include +#include "scalar/fht/fht.h" namespace zvec::turbo::neon { @@ -114,6 +115,105 @@ void fht_vec_rescale_neon(float *data, size_t n, float factor) { } } +void fht_rotate_neon(const float * /*in*/, float *out, size_t in_dim, + size_t /*out_dim*/, void *ctx) { + auto *base = reinterpret_cast(ctx); + const size_t flip_offset = *reinterpret_cast(base); + const size_t trunc_dim = *reinterpret_cast(base + 8); + const float fac = *reinterpret_cast(base + 16); + const uint8_t *flip = base + 24; + const size_t dim = in_dim; + + if (trunc_dim == dim) { + fht_flip_sign_neon(flip, out, dim); + scalar::fht_inplace(out, trunc_dim); + fht_vec_rescale_neon(out, trunc_dim, fac); + + fht_flip_sign_neon(flip + flip_offset, out, dim); + scalar::fht_inplace(out, trunc_dim); + fht_vec_rescale_neon(out, trunc_dim, fac); + + fht_flip_sign_neon(flip + 2 * flip_offset, out, dim); + scalar::fht_inplace(out, trunc_dim); + fht_vec_rescale_neon(out, trunc_dim, fac); + + fht_flip_sign_neon(flip + 3 * flip_offset, out, dim); + scalar::fht_inplace(out, trunc_dim); + fht_vec_rescale_neon(out, trunc_dim, fac); + return; + } + + size_t start = dim - trunc_dim; + float *trunc_ptr = out + start; + + fht_flip_sign_neon(flip, out, dim); + scalar::fht_inplace(out, trunc_dim); + fht_vec_rescale_neon(out, trunc_dim, fac); + fht_kacs_walk_neon(out, dim); + + fht_flip_sign_neon(flip + flip_offset, out, dim); + scalar::fht_inplace(trunc_ptr, trunc_dim); + fht_vec_rescale_neon(trunc_ptr, trunc_dim, fac); + fht_kacs_walk_neon(out, dim); + + fht_flip_sign_neon(flip + 2 * flip_offset, out, dim); + scalar::fht_inplace(out, trunc_dim); + fht_vec_rescale_neon(out, trunc_dim, fac); + fht_kacs_walk_neon(out, dim); + + fht_flip_sign_neon(flip + 3 * flip_offset, out, dim); + scalar::fht_inplace(trunc_ptr, trunc_dim); + fht_vec_rescale_neon(trunc_ptr, trunc_dim, fac); + fht_kacs_walk_neon(out, dim); + + fht_vec_rescale_neon(out, dim, 0.25f); +} + +void fht_unrotate_neon(const float * /*in*/, float *out, size_t in_dim, + size_t /*out_dim*/, void *ctx) { + auto *base = reinterpret_cast(ctx); + const size_t flip_offset = *reinterpret_cast(base); + const size_t trunc_dim = *reinterpret_cast(base + 8); + const float fac = *reinterpret_cast(base + 16); + const uint8_t *flip = base + 24; + const size_t dim = in_dim; + + if (trunc_dim == dim) { + for (int round = 3; round >= 0; --round) { + scalar::fht_inplace(out, trunc_dim); + fht_vec_rescale_neon(out, trunc_dim, fac); + fht_flip_sign_neon(flip + static_cast(round) * flip_offset, out, + dim); + } + return; + } + + fht_vec_rescale_neon(out, dim, 4.0f); + + size_t start = dim - trunc_dim; + float *trunc_ptr = out + start; + + fht_inv_kacs_walk_neon(out, dim); + scalar::fht_inplace(trunc_ptr, trunc_dim); + fht_vec_rescale_neon(trunc_ptr, trunc_dim, fac); + fht_flip_sign_neon(flip + 3 * flip_offset, out, dim); + + fht_inv_kacs_walk_neon(out, dim); + scalar::fht_inplace(out, trunc_dim); + fht_vec_rescale_neon(out, trunc_dim, fac); + fht_flip_sign_neon(flip + 2 * flip_offset, out, dim); + + fht_inv_kacs_walk_neon(out, dim); + scalar::fht_inplace(trunc_ptr, trunc_dim); + fht_vec_rescale_neon(trunc_ptr, trunc_dim, fac); + fht_flip_sign_neon(flip + flip_offset, out, dim); + + fht_inv_kacs_walk_neon(out, dim); + scalar::fht_inplace(out, trunc_dim); + fht_vec_rescale_neon(out, trunc_dim, fac); + fht_flip_sign_neon(flip, out, dim); +} + } // namespace zvec::turbo::neon #endif // __ARM_NEON && __aarch64__ diff --git a/src/turbo/distance/neon/fht/fht.h b/src/turbo/distance/neon/fht/fht.h index 409e8b875..e7f4792d1 100644 --- a/src/turbo/distance/neon/fht/fht.h +++ b/src/turbo/distance/neon/fht/fht.h @@ -31,4 +31,12 @@ void fht_inv_kacs_walk_neon(float *data, size_t len); //! Element-wise rescale: data[i] *= factor (NEON). void fht_vec_rescale_neon(float *data, size_t n, float factor); +//! Forward FHT rotation (NEON). Inplace falls back to scalar. +void fht_rotate_neon(const float *in, float *out, size_t in_dim, size_t out_dim, + void *ctx); + +//! Inverse FHT rotation (NEON). Inplace falls back to scalar. +void fht_unrotate_neon(const float *in, float *out, size_t in_dim, + size_t out_dim, void *ctx); + } // namespace zvec::turbo::neon diff --git a/src/turbo/distance/scalar/fht/fht.cc b/src/turbo/distance/scalar/fht/fht.cc index f1cdffd78..76aaf0896 100644 --- a/src/turbo/distance/scalar/fht/fht.cc +++ b/src/turbo/distance/scalar/fht/fht.cc @@ -76,4 +76,105 @@ void fht_vec_rescale(float *data, size_t n, float factor) { } } +void fht_rotate(const float * /*in*/, float *out, size_t in_dim, + size_t /*out_dim*/, void *ctx) { + // FhtCtx layout: [0:flip_offset, 8:trunc_dim, 16:fac, 24:flip_data] + auto *base = reinterpret_cast(ctx); + const size_t flip_offset = *reinterpret_cast(base); + const size_t trunc_dim = *reinterpret_cast(base + 8); + const float fac = *reinterpret_cast(base + 16); + const uint8_t *flip = base + 24; + const size_t dim = in_dim; + + if (trunc_dim == dim) { + fht_flip_sign(flip, out, dim); + fht_inplace(out, trunc_dim); + fht_vec_rescale(out, trunc_dim, fac); + + fht_flip_sign(flip + flip_offset, out, dim); + fht_inplace(out, trunc_dim); + fht_vec_rescale(out, trunc_dim, fac); + + fht_flip_sign(flip + 2 * flip_offset, out, dim); + fht_inplace(out, trunc_dim); + fht_vec_rescale(out, trunc_dim, fac); + + fht_flip_sign(flip + 3 * flip_offset, out, dim); + fht_inplace(out, trunc_dim); + fht_vec_rescale(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; + + fht_flip_sign(flip, out, dim); + fht_inplace(out, trunc_dim); + fht_vec_rescale(out, trunc_dim, fac); + fht_kacs_walk(out, dim); + + fht_flip_sign(flip + flip_offset, out, dim); + fht_inplace(trunc_ptr, trunc_dim); + fht_vec_rescale(trunc_ptr, trunc_dim, fac); + fht_kacs_walk(out, dim); + + fht_flip_sign(flip + 2 * flip_offset, out, dim); + fht_inplace(out, trunc_dim); + fht_vec_rescale(out, trunc_dim, fac); + fht_kacs_walk(out, dim); + + fht_flip_sign(flip + 3 * flip_offset, out, dim); + fht_inplace(trunc_ptr, trunc_dim); + fht_vec_rescale(trunc_ptr, trunc_dim, fac); + fht_kacs_walk(out, dim); + + fht_vec_rescale(out, dim, 0.25f); +} + +void fht_unrotate(const float * /*in*/, float *out, size_t in_dim, + size_t /*out_dim*/, void *ctx) { + auto *base = reinterpret_cast(ctx); + const size_t flip_offset = *reinterpret_cast(base); + const size_t trunc_dim = *reinterpret_cast(base + 8); + const float fac = *reinterpret_cast(base + 16); + const uint8_t *flip = base + 24; + const size_t dim = in_dim; + + if (trunc_dim == dim) { + for (int round = 3; round >= 0; --round) { + fht_inplace(out, trunc_dim); + fht_vec_rescale(out, trunc_dim, fac); + fht_flip_sign(flip + static_cast(round) * flip_offset, out, dim); + } + return; + } + + // Non-power-of-2: undo final rescale(0.25) first + fht_vec_rescale(out, dim, 4.0f); + + size_t start = dim - trunc_dim; + float *trunc_ptr = out + start; + + fht_inv_kacs_walk(out, dim); + fht_inplace(trunc_ptr, trunc_dim); + fht_vec_rescale(trunc_ptr, trunc_dim, fac); + fht_flip_sign(flip + 3 * flip_offset, out, dim); + + fht_inv_kacs_walk(out, dim); + fht_inplace(out, trunc_dim); + fht_vec_rescale(out, trunc_dim, fac); + fht_flip_sign(flip + 2 * flip_offset, out, dim); + + fht_inv_kacs_walk(out, dim); + fht_inplace(trunc_ptr, trunc_dim); + fht_vec_rescale(trunc_ptr, trunc_dim, fac); + fht_flip_sign(flip + flip_offset, out, dim); + + fht_inv_kacs_walk(out, dim); + fht_inplace(out, trunc_dim); + fht_vec_rescale(out, trunc_dim, fac); + fht_flip_sign(flip, out, dim); +} + } // namespace zvec::turbo::scalar diff --git a/src/turbo/distance/scalar/fht/fht.h b/src/turbo/distance/scalar/fht/fht.h index 7ba98be5d..c31b82543 100644 --- a/src/turbo/distance/scalar/fht/fht.h +++ b/src/turbo/distance/scalar/fht/fht.h @@ -35,4 +35,13 @@ 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); +//! Forward FHT rotation (compose flip -> FHT -> rescale, 4 rounds). +//! ctx is a FhtCtx* defined in preprocessor/fht_rotator/fht_rotator.h. +void fht_rotate(const float *in, float *out, size_t in_dim, size_t out_dim, + void *ctx); + +//! Inverse FHT rotation (undo 4 rounds in reverse order). +void fht_unrotate(const float *in, float *out, size_t in_dim, size_t out_dim, + void *ctx); + } // namespace zvec::turbo::scalar diff --git a/src/turbo/distance/sse/fht/fht.cc b/src/turbo/distance/sse/fht/fht.cc index 18b12c545..e732f1e83 100644 --- a/src/turbo/distance/sse/fht/fht.cc +++ b/src/turbo/distance/sse/fht/fht.cc @@ -20,6 +20,7 @@ #include #include #include +#include "scalar/fht/fht.h" namespace zvec::turbo::sse { @@ -113,6 +114,105 @@ void fht_vec_rescale_sse(float *data, size_t n, float factor) { } } +void fht_rotate_sse(const float * /*in*/, float *out, size_t in_dim, + size_t /*out_dim*/, void *ctx) { + auto *base = reinterpret_cast(ctx); + const size_t flip_offset = *reinterpret_cast(base); + const size_t trunc_dim = *reinterpret_cast(base + 8); + const float fac = *reinterpret_cast(base + 16); + const uint8_t *flip = base + 24; + const size_t dim = in_dim; + + if (trunc_dim == dim) { + fht_flip_sign_sse(flip, out, dim); + scalar::fht_inplace(out, trunc_dim); + fht_vec_rescale_sse(out, trunc_dim, fac); + + fht_flip_sign_sse(flip + flip_offset, out, dim); + scalar::fht_inplace(out, trunc_dim); + fht_vec_rescale_sse(out, trunc_dim, fac); + + fht_flip_sign_sse(flip + 2 * flip_offset, out, dim); + scalar::fht_inplace(out, trunc_dim); + fht_vec_rescale_sse(out, trunc_dim, fac); + + fht_flip_sign_sse(flip + 3 * flip_offset, out, dim); + scalar::fht_inplace(out, trunc_dim); + fht_vec_rescale_sse(out, trunc_dim, fac); + return; + } + + size_t start = dim - trunc_dim; + float *trunc_ptr = out + start; + + fht_flip_sign_sse(flip, out, dim); + scalar::fht_inplace(out, trunc_dim); + fht_vec_rescale_sse(out, trunc_dim, fac); + fht_kacs_walk_sse(out, dim); + + fht_flip_sign_sse(flip + flip_offset, out, dim); + scalar::fht_inplace(trunc_ptr, trunc_dim); + fht_vec_rescale_sse(trunc_ptr, trunc_dim, fac); + fht_kacs_walk_sse(out, dim); + + fht_flip_sign_sse(flip + 2 * flip_offset, out, dim); + scalar::fht_inplace(out, trunc_dim); + fht_vec_rescale_sse(out, trunc_dim, fac); + fht_kacs_walk_sse(out, dim); + + fht_flip_sign_sse(flip + 3 * flip_offset, out, dim); + scalar::fht_inplace(trunc_ptr, trunc_dim); + fht_vec_rescale_sse(trunc_ptr, trunc_dim, fac); + fht_kacs_walk_sse(out, dim); + + fht_vec_rescale_sse(out, dim, 0.25f); +} + +void fht_unrotate_sse(const float * /*in*/, float *out, size_t in_dim, + size_t /*out_dim*/, void *ctx) { + auto *base = reinterpret_cast(ctx); + const size_t flip_offset = *reinterpret_cast(base); + const size_t trunc_dim = *reinterpret_cast(base + 8); + const float fac = *reinterpret_cast(base + 16); + const uint8_t *flip = base + 24; + const size_t dim = in_dim; + + if (trunc_dim == dim) { + for (int round = 3; round >= 0; --round) { + scalar::fht_inplace(out, trunc_dim); + fht_vec_rescale_sse(out, trunc_dim, fac); + fht_flip_sign_sse(flip + static_cast(round) * flip_offset, out, + dim); + } + return; + } + + fht_vec_rescale_sse(out, dim, 4.0f); + + size_t start = dim - trunc_dim; + float *trunc_ptr = out + start; + + fht_inv_kacs_walk_sse(out, dim); + scalar::fht_inplace(trunc_ptr, trunc_dim); + fht_vec_rescale_sse(trunc_ptr, trunc_dim, fac); + fht_flip_sign_sse(flip + 3 * flip_offset, out, dim); + + fht_inv_kacs_walk_sse(out, dim); + scalar::fht_inplace(out, trunc_dim); + fht_vec_rescale_sse(out, trunc_dim, fac); + fht_flip_sign_sse(flip + 2 * flip_offset, out, dim); + + fht_inv_kacs_walk_sse(out, dim); + scalar::fht_inplace(trunc_ptr, trunc_dim); + fht_vec_rescale_sse(trunc_ptr, trunc_dim, fac); + fht_flip_sign_sse(flip + flip_offset, out, dim); + + fht_inv_kacs_walk_sse(out, dim); + scalar::fht_inplace(out, trunc_dim); + fht_vec_rescale_sse(out, trunc_dim, fac); + fht_flip_sign_sse(flip, out, dim); +} + } // namespace zvec::turbo::sse #endif // __SSE2__ diff --git a/src/turbo/distance/sse/fht/fht.h b/src/turbo/distance/sse/fht/fht.h index 9e1cd20a5..996c3008d 100644 --- a/src/turbo/distance/sse/fht/fht.h +++ b/src/turbo/distance/sse/fht/fht.h @@ -31,4 +31,12 @@ 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); +//! Forward FHT rotation (SSE). Inplace falls back to scalar. +void fht_rotate_sse(const float *in, float *out, size_t in_dim, size_t out_dim, + void *ctx); + +//! Inverse FHT rotation (SSE). Inplace falls back to scalar. +void fht_unrotate_sse(const float *in, float *out, size_t in_dim, + size_t out_dim, void *ctx); + } // namespace zvec::turbo::sse diff --git a/src/turbo/preprocessor/fht_rotator/fht_rotator.cc b/src/turbo/preprocessor/fht_rotator/fht_rotator.cc index 54b02cee2..119b4a7ac 100644 --- a/src/turbo/preprocessor/fht_rotator/fht_rotator.cc +++ b/src/turbo/preprocessor/fht_rotator/fht_rotator.cc @@ -14,6 +14,7 @@ #include "fht_rotator.h" #include +#include #include #include #include "quantizer/quantizer.h" @@ -32,29 +33,35 @@ size_t FhtRotator::floor_pow2(size_t n) { return p; } +FhtRotator::~FhtRotator() { + std::free(fht_ctx_); +} + 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 = std::get(get_rotator_kernels(RotateType::kFht)); - r->flip_sign_fn_ = k.flip_sign; - r->kacs_walk_fn_ = k.kacs_walk; - r->inv_kacs_walk_fn_ = k.inv_kacs_walk; - r->inplace_fn_ = k.inplace; - r->rescale_fn_ = k.rescale; - - // Generate 4 rounds of random flip-sign arrays so the rotator is - // immediately usable after create(). - r->flip_.resize(4 * r->flip_offset_); + r->kernels_ = get_rotator_kernels(RotateType::kFht); + + const size_t trunc_dim = floor_pow2(static_cast(dim)); + const float fac = 1.0f / std::sqrt(static_cast(trunc_dim)); + const size_t flip_size = 4 * r->flip_offset_; + + // Single allocation: FhtCtx header + trailing flip data. + r->fht_ctx_ = static_cast(std::malloc(sizeof(FhtCtx) + flip_size)); + r->fht_ctx_->flip_offset = r->flip_offset_; + r->fht_ctx_->trunc_dim = trunc_dim; + r->fht_ctx_->fac = fac; + + // Generate 4 rounds of random flip-sign arrays. 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)); + for (size_t i = 0; i < flip_size; ++i) + r->fht_ctx_->flip[i] = static_cast(dist(gen)); return r; } @@ -89,60 +96,9 @@ void FhtRotator::train(const void * /*data*/, size_t /*num*/, // --------------------------------------------------------------------------- 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); + std::memcpy(out, in, sizeof(float) * static_cast(in_dim_)); + kernels_.rotate(in, out, static_cast(in_dim_), + static_cast(out_dim_), static_cast(fht_ctx_)); } // --------------------------------------------------------------------------- @@ -150,51 +106,10 @@ void FhtRotator::apply(const float *in, float *out) const { // --------------------------------------------------------------------------- void FhtRotator::apply_inverse(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: 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_(out, trunc_dim_); - rescale_fn_(out, trunc_dim_, inv_fac); - flip_sign_fn_(flip_.data() + static_cast(round) * flip_offset_, - out, dim); - } - return; - } - - // Non-power-of-2: undo final rescale(0.25) first - rescale_fn_(out, dim, 4.0f); - - const float inv_fac = 1.0f / std::sqrt(static_cast(trunc_dim_)); - size_t start = dim - trunc_dim_; - float *trunc_ptr = out + start; - - // Undo Round 4 (FHT on [start, start+trunc_dim)) - inv_kacs_walk_fn_(out, dim); - inplace_fn_(trunc_ptr, trunc_dim_); - rescale_fn_(trunc_ptr, trunc_dim_, inv_fac); - flip_sign_fn_(flip_.data() + 3 * flip_offset_, out, dim); - - // Undo Round 3 (FHT on [0, trunc_dim)) - inv_kacs_walk_fn_(out, dim); - inplace_fn_(out, trunc_dim_); - rescale_fn_(out, trunc_dim_, inv_fac); - flip_sign_fn_(flip_.data() + 2 * flip_offset_, out, dim); - - // Undo Round 2 (FHT on [start, start+trunc_dim)) - inv_kacs_walk_fn_(out, dim); - inplace_fn_(trunc_ptr, trunc_dim_); - rescale_fn_(trunc_ptr, trunc_dim_, inv_fac); - flip_sign_fn_(flip_.data() + flip_offset_, out, dim); - - // Undo Round 1 (FHT on [0, trunc_dim)) - inv_kacs_walk_fn_(out, dim); - inplace_fn_(out, trunc_dim_); - rescale_fn_(out, trunc_dim_, inv_fac); - flip_sign_fn_(flip_.data(), out, dim); + std::memcpy(out, in, sizeof(float) * static_cast(in_dim_)); + kernels_.unrotate(in, out, static_cast(in_dim_), + static_cast(out_dim_), + static_cast(fht_ctx_)); } // --------------------------------------------------------------------------- @@ -203,7 +118,9 @@ void FhtRotator::apply_inverse(const float *in, float *out) const { int FhtRotator::serialize(std::string *out) const { if (!out) return kErrInvalidArgument; - if (flip_.empty()) return kErrRuntime; + if (!fht_ctx_) return kErrRuntime; + + const size_t flip_size = 4 * flip_offset_; RotatorSerHeader hdr{}; hdr.magic = kRotatorMagic; @@ -211,12 +128,12 @@ int FhtRotator::serialize(std::string *out) const { hdr.rotator_type = static_cast(RotateType::kFht); hdr.in_dim = static_cast(in_dim_); hdr.out_dim = static_cast(out_dim_); - hdr.payload_size = static_cast(flip_.size()); + hdr.payload_size = static_cast(flip_size); hdr.reserved = 0; - out->resize(sizeof(hdr) + flip_.size()); + out->resize(sizeof(hdr) + flip_size); std::memcpy(&(*out)[0], &hdr, sizeof(hdr)); - std::memcpy(&(*out)[sizeof(hdr)], flip_.data(), flip_.size()); + std::memcpy(&(*out)[sizeof(hdr)], fht_ctx_->flip, flip_size); return 0; } @@ -235,18 +152,20 @@ int FhtRotator::deserialize(const void *data, size_t len) { 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 = std::get(get_rotator_kernels(RotateType::kFht)); - 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(), + kernels_ = get_rotator_kernels(RotateType::kFht); + + const size_t trunc_dim = floor_pow2(static_cast(in_dim_)); + const float fac = 1.0f / std::sqrt(static_cast(trunc_dim)); + + // Free old ctx, allocate new one with trailing flip data. + std::free(fht_ctx_); + fht_ctx_ = + static_cast(std::malloc(sizeof(FhtCtx) + hdr->payload_size)); + fht_ctx_->flip_offset = flip_offset_; + fht_ctx_->trunc_dim = trunc_dim; + fht_ctx_->fac = fac; + std::memcpy(fht_ctx_->flip, reinterpret_cast(data) + sizeof(RotatorSerHeader), hdr->payload_size); diff --git a/src/turbo/preprocessor/fht_rotator/fht_rotator.h b/src/turbo/preprocessor/fht_rotator/fht_rotator.h index 40aadf42a..ac8768758 100644 --- a/src/turbo/preprocessor/fht_rotator/fht_rotator.h +++ b/src/turbo/preprocessor/fht_rotator/fht_rotator.h @@ -14,13 +14,34 @@ #pragma once +#include +#include #include -#include +#include #include "preprocessor/preprocessor.h" namespace zvec { namespace turbo { +// FHT context passed to ISA-level rotate/unrotate via void*. +// Layout (ISA code accesses by address, NOT by type): +// offset 0: size_t flip_offset (bytes per round) +// offset 8: size_t trunc_dim (largest power-of-2 <= in_dim) +// offset 16: float fac (1 / sqrt(trunc_dim)) +// offset 20: uint8_t pad[4] (explicit padding) +// offset 24: uint8_t flip[] (4 * flip_offset bytes, trailing data) +// +// Allocated as a single block: malloc(sizeof(FhtCtx) + 4 * flip_offset). +struct FhtCtx { + size_t flip_offset; + size_t trunc_dim; + float fac; + uint8_t pad_[4]; + uint8_t flip[]; // trailing flexible array (C++ extension) +}; +static_assert(offsetof(FhtCtx, flip) == 24, "FhtCtx flip offset must be 24"); +static_assert(sizeof(FhtCtx) == 24, "FhtCtx sizeof must be 24"); + // ============================================================================ // FhtRotator - O(d log d) FHT-based Kac random rotation // @@ -66,6 +87,8 @@ class FhtRotator : public Preprocessor { return RotateType::kFht; } + ~FhtRotator() override; + private: FhtRotator() = default; @@ -75,24 +98,14 @@ class FhtRotator : public Preprocessor { 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). + //! Bytes per round: ceil(in_dim / 8). Kept for serialization. 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 rotate/unrotate kernels. + RotatorKernels kernels_{}; - //! 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}; + //! FHT state (flip_offset, trunc_dim, fac, flip[]) -- single allocation. + FhtCtx *fht_ctx_{nullptr}; static constexpr size_t kByteLen = 8; }; diff --git a/src/turbo/turbo.cc b/src/turbo/turbo.cc index 387ee71e7..0218e27ed 100644 --- a/src/turbo/turbo.cc +++ b/src/turbo/turbo.cc @@ -168,24 +168,18 @@ RotatorKernels get_rotator_kernels(RotateType rotate_type, switch (rotate_type) { case RotateType::kFht: { - 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; + RotatorKernels k; + // Default: scalar + k.rotate = scalar::fht_rotate; + k.unrotate = scalar::fht_unrotate; #if defined(__AVX512F__) if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F && zvec::ailego::internal::CpuFeatures::static_flags_.AVX512DQ && (cpu_arch_type == CpuArchType::kAuto || cpu_arch_type == CpuArchType::kAVX512)) { - 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; + k.rotate = avx512::fht_rotate_avx512; + k.unrotate = avx512::fht_unrotate_avx512; return k; } #endif @@ -193,11 +187,8 @@ RotatorKernels get_rotator_kernels(RotateType rotate_type, if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2 && (cpu_arch_type == CpuArchType::kAuto || cpu_arch_type == CpuArchType::kAVX2)) { - 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; + k.rotate = avx2::fht_rotate_avx2; + k.unrotate = avx2::fht_unrotate_avx2; return k; } #endif @@ -205,22 +196,16 @@ RotatorKernels get_rotator_kernels(RotateType rotate_type, if (zvec::ailego::internal::CpuFeatures::static_flags_.SSE2 && (cpu_arch_type == CpuArchType::kAuto || cpu_arch_type == CpuArchType::kSSE)) { - 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) + k.rotate = sse::fht_rotate_sse; + k.unrotate = sse::fht_unrotate_sse; return k; } #endif #if defined(__ARM_NEON) && defined(__aarch64__) if (cpu_arch_type == CpuArchType::kAuto || cpu_arch_type == CpuArchType::kNEON) { - k.flip_sign = neon::fht_flip_sign_neon; - k.kacs_walk = neon::fht_kacs_walk_neon; - k.inv_kacs_walk = neon::fht_inv_kacs_walk_neon; - k.rescale = neon::fht_vec_rescale_neon; - // inplace fallback to scalar (NEON has no fht_inplace) + k.rotate = neon::fht_rotate_neon; + k.unrotate = neon::fht_unrotate_neon; return k; } #endif @@ -230,7 +215,7 @@ RotatorKernels get_rotator_kernels(RotateType rotate_type, // Fallback (unreachable for valid RotateType values). assert(false && "unsupported RotateType"); - return FhtKernels{}; + return {}; } } // namespace zvec::turbo From ec946b17a7ff50d87f04a8fc42d3513993bb4fdc Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Thu, 16 Jul 2026 18:03:11 +0800 Subject: [PATCH 067/100] Adjust Directory --- .../distance/avx2/{ => rotate}/fht/fht.cc | 100 +--------- .../distance/avx2/{ => rotate}/fht/fht.h | 0 .../distance/avx512/{ => rotate}/fht/fht.cc | 100 +--------- .../distance/avx512/{ => rotate}/fht/fht.h | 0 src/turbo/distance/common/fht_common.h | 126 ++++++++++++ .../distance/neon/{ => rotate}/fht/fht.cc | 102 +--------- .../distance/neon/{ => rotate}/fht/fht.h | 0 src/turbo/distance/scalar/fht/fht.cc | 180 ------------------ src/turbo/distance/scalar/rotate/fht/fht.cc | 96 ++++++++++ .../distance/scalar/{ => rotate}/fht/fht.h | 0 .../distance/sse/{ => rotate}/fht/fht.cc | 102 +--------- src/turbo/distance/sse/{ => rotate}/fht/fht.h | 0 src/turbo/turbo.cc | 10 +- 13 files changed, 265 insertions(+), 551 deletions(-) rename src/turbo/distance/avx2/{ => rotate}/fht/fht.cc (59%) rename src/turbo/distance/avx2/{ => rotate}/fht/fht.h (100%) rename src/turbo/distance/avx512/{ => rotate}/fht/fht.cc (60%) rename src/turbo/distance/avx512/{ => rotate}/fht/fht.h (100%) create mode 100644 src/turbo/distance/common/fht_common.h rename src/turbo/distance/neon/{ => rotate}/fht/fht.cc (54%) rename src/turbo/distance/neon/{ => rotate}/fht/fht.h (100%) delete mode 100644 src/turbo/distance/scalar/fht/fht.cc create mode 100644 src/turbo/distance/scalar/rotate/fht/fht.cc rename src/turbo/distance/scalar/{ => rotate}/fht/fht.h (100%) rename src/turbo/distance/sse/{ => rotate}/fht/fht.cc (54%) rename src/turbo/distance/sse/{ => rotate}/fht/fht.h (100%) diff --git a/src/turbo/distance/avx2/fht/fht.cc b/src/turbo/distance/avx2/rotate/fht/fht.cc similarity index 59% rename from src/turbo/distance/avx2/fht/fht.cc rename to src/turbo/distance/avx2/rotate/fht/fht.cc index 10c101b6f..a77dc087b 100644 --- a/src/turbo/distance/avx2/fht/fht.cc +++ b/src/turbo/distance/avx2/rotate/fht/fht.cc @@ -20,6 +20,7 @@ #include #include #include +#include "common/fht_common.h" namespace zvec::turbo::avx2 { @@ -134,101 +135,18 @@ void fht_vec_rescale_avx2(float *data, size_t n, float factor) { void fht_rotate_avx2(const float * /*in*/, float *out, size_t in_dim, size_t /*out_dim*/, void *ctx) { - auto *base = reinterpret_cast(ctx); - const size_t flip_offset = *reinterpret_cast(base); - const size_t trunc_dim = *reinterpret_cast(base + 8); - const float fac = *reinterpret_cast(base + 16); - const uint8_t *flip = base + 24; - const size_t dim = in_dim; - - if (trunc_dim == dim) { - fht_flip_sign_avx2(flip, out, dim); - fht_inplace_avx2(out, trunc_dim); - fht_vec_rescale_avx2(out, trunc_dim, fac); - - fht_flip_sign_avx2(flip + flip_offset, out, dim); - fht_inplace_avx2(out, trunc_dim); - fht_vec_rescale_avx2(out, trunc_dim, fac); - - fht_flip_sign_avx2(flip + 2 * flip_offset, out, dim); - fht_inplace_avx2(out, trunc_dim); - fht_vec_rescale_avx2(out, trunc_dim, fac); - - fht_flip_sign_avx2(flip + 3 * flip_offset, out, dim); - fht_inplace_avx2(out, trunc_dim); - fht_vec_rescale_avx2(out, trunc_dim, fac); - return; - } - - size_t start = dim - trunc_dim; - float *trunc_ptr = out + start; - - fht_flip_sign_avx2(flip, out, dim); - fht_inplace_avx2(out, trunc_dim); - fht_vec_rescale_avx2(out, trunc_dim, fac); - fht_kacs_walk_avx2(out, dim); - - fht_flip_sign_avx2(flip + flip_offset, out, dim); - fht_inplace_avx2(trunc_ptr, trunc_dim); - fht_vec_rescale_avx2(trunc_ptr, trunc_dim, fac); - fht_kacs_walk_avx2(out, dim); - - fht_flip_sign_avx2(flip + 2 * flip_offset, out, dim); - fht_inplace_avx2(out, trunc_dim); - fht_vec_rescale_avx2(out, trunc_dim, fac); - fht_kacs_walk_avx2(out, dim); - - fht_flip_sign_avx2(flip + 3 * flip_offset, out, dim); - fht_inplace_avx2(trunc_ptr, trunc_dim); - fht_vec_rescale_avx2(trunc_ptr, trunc_dim, fac); - fht_kacs_walk_avx2(out, dim); - - fht_vec_rescale_avx2(out, dim, 0.25f); + static constexpr FhtPrimitives kPrim = { + fht_flip_sign_avx2, fht_inplace_avx2, fht_kacs_walk_avx2, + fht_inv_kacs_walk_avx2, fht_vec_rescale_avx2}; + fht_rotate_impl(out, in_dim, ctx, kPrim); } void fht_unrotate_avx2(const float * /*in*/, float *out, size_t in_dim, size_t /*out_dim*/, void *ctx) { - auto *base = reinterpret_cast(ctx); - const size_t flip_offset = *reinterpret_cast(base); - const size_t trunc_dim = *reinterpret_cast(base + 8); - const float fac = *reinterpret_cast(base + 16); - const uint8_t *flip = base + 24; - const size_t dim = in_dim; - - if (trunc_dim == dim) { - for (int round = 3; round >= 0; --round) { - fht_inplace_avx2(out, trunc_dim); - fht_vec_rescale_avx2(out, trunc_dim, fac); - fht_flip_sign_avx2(flip + static_cast(round) * flip_offset, out, - dim); - } - return; - } - - fht_vec_rescale_avx2(out, dim, 4.0f); - - size_t start = dim - trunc_dim; - float *trunc_ptr = out + start; - - fht_inv_kacs_walk_avx2(out, dim); - fht_inplace_avx2(trunc_ptr, trunc_dim); - fht_vec_rescale_avx2(trunc_ptr, trunc_dim, fac); - fht_flip_sign_avx2(flip + 3 * flip_offset, out, dim); - - fht_inv_kacs_walk_avx2(out, dim); - fht_inplace_avx2(out, trunc_dim); - fht_vec_rescale_avx2(out, trunc_dim, fac); - fht_flip_sign_avx2(flip + 2 * flip_offset, out, dim); - - fht_inv_kacs_walk_avx2(out, dim); - fht_inplace_avx2(trunc_ptr, trunc_dim); - fht_vec_rescale_avx2(trunc_ptr, trunc_dim, fac); - fht_flip_sign_avx2(flip + flip_offset, out, dim); - - fht_inv_kacs_walk_avx2(out, dim); - fht_inplace_avx2(out, trunc_dim); - fht_vec_rescale_avx2(out, trunc_dim, fac); - fht_flip_sign_avx2(flip, out, dim); + static constexpr FhtPrimitives kPrim = { + fht_flip_sign_avx2, fht_inplace_avx2, fht_kacs_walk_avx2, + fht_inv_kacs_walk_avx2, fht_vec_rescale_avx2}; + fht_unrotate_impl(out, in_dim, ctx, kPrim); } } // namespace zvec::turbo::avx2 diff --git a/src/turbo/distance/avx2/fht/fht.h b/src/turbo/distance/avx2/rotate/fht/fht.h similarity index 100% rename from src/turbo/distance/avx2/fht/fht.h rename to src/turbo/distance/avx2/rotate/fht/fht.h diff --git a/src/turbo/distance/avx512/fht/fht.cc b/src/turbo/distance/avx512/rotate/fht/fht.cc similarity index 60% rename from src/turbo/distance/avx512/fht/fht.cc rename to src/turbo/distance/avx512/rotate/fht/fht.cc index a50ad8637..9a237214d 100644 --- a/src/turbo/distance/avx512/fht/fht.cc +++ b/src/turbo/distance/avx512/rotate/fht/fht.cc @@ -20,6 +20,7 @@ #include #include #include +#include "common/fht_common.h" namespace zvec::turbo::avx512 { @@ -139,101 +140,18 @@ void fht_vec_rescale_avx512(float *data, size_t n, float factor) { void fht_rotate_avx512(const float * /*in*/, float *out, size_t in_dim, size_t /*out_dim*/, void *ctx) { - auto *base = reinterpret_cast(ctx); - const size_t flip_offset = *reinterpret_cast(base); - const size_t trunc_dim = *reinterpret_cast(base + 8); - const float fac = *reinterpret_cast(base + 16); - const uint8_t *flip = base + 24; - const size_t dim = in_dim; - - if (trunc_dim == dim) { - fht_flip_sign_avx512(flip, out, dim); - fht_inplace_avx512(out, trunc_dim); - fht_vec_rescale_avx512(out, trunc_dim, fac); - - fht_flip_sign_avx512(flip + flip_offset, out, dim); - fht_inplace_avx512(out, trunc_dim); - fht_vec_rescale_avx512(out, trunc_dim, fac); - - fht_flip_sign_avx512(flip + 2 * flip_offset, out, dim); - fht_inplace_avx512(out, trunc_dim); - fht_vec_rescale_avx512(out, trunc_dim, fac); - - fht_flip_sign_avx512(flip + 3 * flip_offset, out, dim); - fht_inplace_avx512(out, trunc_dim); - fht_vec_rescale_avx512(out, trunc_dim, fac); - return; - } - - size_t start = dim - trunc_dim; - float *trunc_ptr = out + start; - - fht_flip_sign_avx512(flip, out, dim); - fht_inplace_avx512(out, trunc_dim); - fht_vec_rescale_avx512(out, trunc_dim, fac); - fht_kacs_walk_avx512(out, dim); - - fht_flip_sign_avx512(flip + flip_offset, out, dim); - fht_inplace_avx512(trunc_ptr, trunc_dim); - fht_vec_rescale_avx512(trunc_ptr, trunc_dim, fac); - fht_kacs_walk_avx512(out, dim); - - fht_flip_sign_avx512(flip + 2 * flip_offset, out, dim); - fht_inplace_avx512(out, trunc_dim); - fht_vec_rescale_avx512(out, trunc_dim, fac); - fht_kacs_walk_avx512(out, dim); - - fht_flip_sign_avx512(flip + 3 * flip_offset, out, dim); - fht_inplace_avx512(trunc_ptr, trunc_dim); - fht_vec_rescale_avx512(trunc_ptr, trunc_dim, fac); - fht_kacs_walk_avx512(out, dim); - - fht_vec_rescale_avx512(out, dim, 0.25f); + static constexpr FhtPrimitives kPrim = { + fht_flip_sign_avx512, fht_inplace_avx512, fht_kacs_walk_avx512, + fht_inv_kacs_walk_avx512, fht_vec_rescale_avx512}; + fht_rotate_impl(out, in_dim, ctx, kPrim); } void fht_unrotate_avx512(const float * /*in*/, float *out, size_t in_dim, size_t /*out_dim*/, void *ctx) { - auto *base = reinterpret_cast(ctx); - const size_t flip_offset = *reinterpret_cast(base); - const size_t trunc_dim = *reinterpret_cast(base + 8); - const float fac = *reinterpret_cast(base + 16); - const uint8_t *flip = base + 24; - const size_t dim = in_dim; - - if (trunc_dim == dim) { - for (int round = 3; round >= 0; --round) { - fht_inplace_avx512(out, trunc_dim); - fht_vec_rescale_avx512(out, trunc_dim, fac); - fht_flip_sign_avx512(flip + static_cast(round) * flip_offset, out, - dim); - } - return; - } - - fht_vec_rescale_avx512(out, dim, 4.0f); - - size_t start = dim - trunc_dim; - float *trunc_ptr = out + start; - - fht_inv_kacs_walk_avx512(out, dim); - fht_inplace_avx512(trunc_ptr, trunc_dim); - fht_vec_rescale_avx512(trunc_ptr, trunc_dim, fac); - fht_flip_sign_avx512(flip + 3 * flip_offset, out, dim); - - fht_inv_kacs_walk_avx512(out, dim); - fht_inplace_avx512(out, trunc_dim); - fht_vec_rescale_avx512(out, trunc_dim, fac); - fht_flip_sign_avx512(flip + 2 * flip_offset, out, dim); - - fht_inv_kacs_walk_avx512(out, dim); - fht_inplace_avx512(trunc_ptr, trunc_dim); - fht_vec_rescale_avx512(trunc_ptr, trunc_dim, fac); - fht_flip_sign_avx512(flip + flip_offset, out, dim); - - fht_inv_kacs_walk_avx512(out, dim); - fht_inplace_avx512(out, trunc_dim); - fht_vec_rescale_avx512(out, trunc_dim, fac); - fht_flip_sign_avx512(flip, out, dim); + static constexpr FhtPrimitives kPrim = { + fht_flip_sign_avx512, fht_inplace_avx512, fht_kacs_walk_avx512, + fht_inv_kacs_walk_avx512, fht_vec_rescale_avx512}; + fht_unrotate_impl(out, in_dim, ctx, kPrim); } } // namespace zvec::turbo::avx512 diff --git a/src/turbo/distance/avx512/fht/fht.h b/src/turbo/distance/avx512/rotate/fht/fht.h similarity index 100% rename from src/turbo/distance/avx512/fht/fht.h rename to src/turbo/distance/avx512/rotate/fht/fht.h diff --git a/src/turbo/distance/common/fht_common.h b/src/turbo/distance/common/fht_common.h new file mode 100644 index 000000000..66377e756 --- /dev/null +++ b/src/turbo/distance/common/fht_common.h @@ -0,0 +1,126 @@ +// 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 { + +/// ISA-level FHT primitive function pointers. +/// Each ISA fills in its own SIMD-optimized (or scalar-fallback) functions. +struct FhtPrimitives { + void (*flip_sign)(const uint8_t *flip, float *data, size_t dim); + void (*inplace)(float *data, size_t n); + void (*kacs_walk)(float *data, size_t len); + void (*inv_kacs_walk)(float *data, size_t len); + void (*rescale)(float *data, size_t n, float factor); +}; + +/// FhtCtx memory layout (accessed by address, NOT by type): +/// offset 0: size_t flip_offset +/// offset 8: size_t trunc_dim +/// offset 16: float fac +/// offset 20: uint8_t pad[4] +/// offset 24: uint8_t flip[] +inline void fht_rotate_impl(float *out, size_t in_dim, void *ctx, + const FhtPrimitives &p) { + auto *base = reinterpret_cast(ctx); + const size_t flip_offset = *reinterpret_cast(base); + const size_t trunc_dim = *reinterpret_cast(base + 8); + const float fac = *reinterpret_cast(base + 16); + const uint8_t *flip = base + 24; + const size_t dim = in_dim; + + if (trunc_dim == dim) { + for (size_t r = 0; r < 4; ++r) { + p.flip_sign(flip + r * flip_offset, out, dim); + p.inplace(out, trunc_dim); + p.rescale(out, trunc_dim, fac); + } + return; + } + + size_t start = dim - trunc_dim; + float *trunc_ptr = out + start; + + p.flip_sign(flip, out, dim); + p.inplace(out, trunc_dim); + p.rescale(out, trunc_dim, fac); + p.kacs_walk(out, dim); + + p.flip_sign(flip + flip_offset, out, dim); + p.inplace(trunc_ptr, trunc_dim); + p.rescale(trunc_ptr, trunc_dim, fac); + p.kacs_walk(out, dim); + + p.flip_sign(flip + 2 * flip_offset, out, dim); + p.inplace(out, trunc_dim); + p.rescale(out, trunc_dim, fac); + p.kacs_walk(out, dim); + + p.flip_sign(flip + 3 * flip_offset, out, dim); + p.inplace(trunc_ptr, trunc_dim); + p.rescale(trunc_ptr, trunc_dim, fac); + p.kacs_walk(out, dim); + + p.rescale(out, dim, 0.25f); +} + +inline void fht_unrotate_impl(float *out, size_t in_dim, void *ctx, + const FhtPrimitives &p) { + auto *base = reinterpret_cast(ctx); + const size_t flip_offset = *reinterpret_cast(base); + const size_t trunc_dim = *reinterpret_cast(base + 8); + const float fac = *reinterpret_cast(base + 16); + const uint8_t *flip = base + 24; + const size_t dim = in_dim; + + if (trunc_dim == dim) { + for (int round = 3; round >= 0; --round) { + p.inplace(out, trunc_dim); + p.rescale(out, trunc_dim, fac); + p.flip_sign(flip + static_cast(round) * flip_offset, out, dim); + } + return; + } + + p.rescale(out, dim, 4.0f); + + size_t start = dim - trunc_dim; + float *trunc_ptr = out + start; + + p.inv_kacs_walk(out, dim); + p.inplace(trunc_ptr, trunc_dim); + p.rescale(trunc_ptr, trunc_dim, fac); + p.flip_sign(flip + 3 * flip_offset, out, dim); + + p.inv_kacs_walk(out, dim); + p.inplace(out, trunc_dim); + p.rescale(out, trunc_dim, fac); + p.flip_sign(flip + 2 * flip_offset, out, dim); + + p.inv_kacs_walk(out, dim); + p.inplace(trunc_ptr, trunc_dim); + p.rescale(trunc_ptr, trunc_dim, fac); + p.flip_sign(flip + flip_offset, out, dim); + + p.inv_kacs_walk(out, dim); + p.inplace(out, trunc_dim); + p.rescale(out, trunc_dim, fac); + p.flip_sign(flip, out, dim); +} + +} // namespace zvec::turbo diff --git a/src/turbo/distance/neon/fht/fht.cc b/src/turbo/distance/neon/rotate/fht/fht.cc similarity index 54% rename from src/turbo/distance/neon/fht/fht.cc rename to src/turbo/distance/neon/rotate/fht/fht.cc index cadfd0a79..cc736a0a2 100644 --- a/src/turbo/distance/neon/fht/fht.cc +++ b/src/turbo/distance/neon/rotate/fht/fht.cc @@ -20,7 +20,8 @@ #include #include #include -#include "scalar/fht/fht.h" +#include "common/fht_common.h" +#include "scalar/rotate/fht/fht.h" namespace zvec::turbo::neon { @@ -117,101 +118,18 @@ void fht_vec_rescale_neon(float *data, size_t n, float factor) { void fht_rotate_neon(const float * /*in*/, float *out, size_t in_dim, size_t /*out_dim*/, void *ctx) { - auto *base = reinterpret_cast(ctx); - const size_t flip_offset = *reinterpret_cast(base); - const size_t trunc_dim = *reinterpret_cast(base + 8); - const float fac = *reinterpret_cast(base + 16); - const uint8_t *flip = base + 24; - const size_t dim = in_dim; - - if (trunc_dim == dim) { - fht_flip_sign_neon(flip, out, dim); - scalar::fht_inplace(out, trunc_dim); - fht_vec_rescale_neon(out, trunc_dim, fac); - - fht_flip_sign_neon(flip + flip_offset, out, dim); - scalar::fht_inplace(out, trunc_dim); - fht_vec_rescale_neon(out, trunc_dim, fac); - - fht_flip_sign_neon(flip + 2 * flip_offset, out, dim); - scalar::fht_inplace(out, trunc_dim); - fht_vec_rescale_neon(out, trunc_dim, fac); - - fht_flip_sign_neon(flip + 3 * flip_offset, out, dim); - scalar::fht_inplace(out, trunc_dim); - fht_vec_rescale_neon(out, trunc_dim, fac); - return; - } - - size_t start = dim - trunc_dim; - float *trunc_ptr = out + start; - - fht_flip_sign_neon(flip, out, dim); - scalar::fht_inplace(out, trunc_dim); - fht_vec_rescale_neon(out, trunc_dim, fac); - fht_kacs_walk_neon(out, dim); - - fht_flip_sign_neon(flip + flip_offset, out, dim); - scalar::fht_inplace(trunc_ptr, trunc_dim); - fht_vec_rescale_neon(trunc_ptr, trunc_dim, fac); - fht_kacs_walk_neon(out, dim); - - fht_flip_sign_neon(flip + 2 * flip_offset, out, dim); - scalar::fht_inplace(out, trunc_dim); - fht_vec_rescale_neon(out, trunc_dim, fac); - fht_kacs_walk_neon(out, dim); - - fht_flip_sign_neon(flip + 3 * flip_offset, out, dim); - scalar::fht_inplace(trunc_ptr, trunc_dim); - fht_vec_rescale_neon(trunc_ptr, trunc_dim, fac); - fht_kacs_walk_neon(out, dim); - - fht_vec_rescale_neon(out, dim, 0.25f); + static constexpr FhtPrimitives kPrim = { + fht_flip_sign_neon, scalar::fht_inplace, fht_kacs_walk_neon, + fht_inv_kacs_walk_neon, fht_vec_rescale_neon}; + fht_rotate_impl(out, in_dim, ctx, kPrim); } void fht_unrotate_neon(const float * /*in*/, float *out, size_t in_dim, size_t /*out_dim*/, void *ctx) { - auto *base = reinterpret_cast(ctx); - const size_t flip_offset = *reinterpret_cast(base); - const size_t trunc_dim = *reinterpret_cast(base + 8); - const float fac = *reinterpret_cast(base + 16); - const uint8_t *flip = base + 24; - const size_t dim = in_dim; - - if (trunc_dim == dim) { - for (int round = 3; round >= 0; --round) { - scalar::fht_inplace(out, trunc_dim); - fht_vec_rescale_neon(out, trunc_dim, fac); - fht_flip_sign_neon(flip + static_cast(round) * flip_offset, out, - dim); - } - return; - } - - fht_vec_rescale_neon(out, dim, 4.0f); - - size_t start = dim - trunc_dim; - float *trunc_ptr = out + start; - - fht_inv_kacs_walk_neon(out, dim); - scalar::fht_inplace(trunc_ptr, trunc_dim); - fht_vec_rescale_neon(trunc_ptr, trunc_dim, fac); - fht_flip_sign_neon(flip + 3 * flip_offset, out, dim); - - fht_inv_kacs_walk_neon(out, dim); - scalar::fht_inplace(out, trunc_dim); - fht_vec_rescale_neon(out, trunc_dim, fac); - fht_flip_sign_neon(flip + 2 * flip_offset, out, dim); - - fht_inv_kacs_walk_neon(out, dim); - scalar::fht_inplace(trunc_ptr, trunc_dim); - fht_vec_rescale_neon(trunc_ptr, trunc_dim, fac); - fht_flip_sign_neon(flip + flip_offset, out, dim); - - fht_inv_kacs_walk_neon(out, dim); - scalar::fht_inplace(out, trunc_dim); - fht_vec_rescale_neon(out, trunc_dim, fac); - fht_flip_sign_neon(flip, out, dim); + static constexpr FhtPrimitives kPrim = { + fht_flip_sign_neon, scalar::fht_inplace, fht_kacs_walk_neon, + fht_inv_kacs_walk_neon, fht_vec_rescale_neon}; + fht_unrotate_impl(out, in_dim, ctx, kPrim); } } // namespace zvec::turbo::neon diff --git a/src/turbo/distance/neon/fht/fht.h b/src/turbo/distance/neon/rotate/fht/fht.h similarity index 100% rename from src/turbo/distance/neon/fht/fht.h rename to src/turbo/distance/neon/rotate/fht/fht.h diff --git a/src/turbo/distance/scalar/fht/fht.cc b/src/turbo/distance/scalar/fht/fht.cc deleted file mode 100644 index 76aaf0896..000000000 --- a/src/turbo/distance/scalar/fht/fht.cc +++ /dev/null @@ -1,180 +0,0 @@ -// 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; - } -} - -void fht_rotate(const float * /*in*/, float *out, size_t in_dim, - size_t /*out_dim*/, void *ctx) { - // FhtCtx layout: [0:flip_offset, 8:trunc_dim, 16:fac, 24:flip_data] - auto *base = reinterpret_cast(ctx); - const size_t flip_offset = *reinterpret_cast(base); - const size_t trunc_dim = *reinterpret_cast(base + 8); - const float fac = *reinterpret_cast(base + 16); - const uint8_t *flip = base + 24; - const size_t dim = in_dim; - - if (trunc_dim == dim) { - fht_flip_sign(flip, out, dim); - fht_inplace(out, trunc_dim); - fht_vec_rescale(out, trunc_dim, fac); - - fht_flip_sign(flip + flip_offset, out, dim); - fht_inplace(out, trunc_dim); - fht_vec_rescale(out, trunc_dim, fac); - - fht_flip_sign(flip + 2 * flip_offset, out, dim); - fht_inplace(out, trunc_dim); - fht_vec_rescale(out, trunc_dim, fac); - - fht_flip_sign(flip + 3 * flip_offset, out, dim); - fht_inplace(out, trunc_dim); - fht_vec_rescale(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; - - fht_flip_sign(flip, out, dim); - fht_inplace(out, trunc_dim); - fht_vec_rescale(out, trunc_dim, fac); - fht_kacs_walk(out, dim); - - fht_flip_sign(flip + flip_offset, out, dim); - fht_inplace(trunc_ptr, trunc_dim); - fht_vec_rescale(trunc_ptr, trunc_dim, fac); - fht_kacs_walk(out, dim); - - fht_flip_sign(flip + 2 * flip_offset, out, dim); - fht_inplace(out, trunc_dim); - fht_vec_rescale(out, trunc_dim, fac); - fht_kacs_walk(out, dim); - - fht_flip_sign(flip + 3 * flip_offset, out, dim); - fht_inplace(trunc_ptr, trunc_dim); - fht_vec_rescale(trunc_ptr, trunc_dim, fac); - fht_kacs_walk(out, dim); - - fht_vec_rescale(out, dim, 0.25f); -} - -void fht_unrotate(const float * /*in*/, float *out, size_t in_dim, - size_t /*out_dim*/, void *ctx) { - auto *base = reinterpret_cast(ctx); - const size_t flip_offset = *reinterpret_cast(base); - const size_t trunc_dim = *reinterpret_cast(base + 8); - const float fac = *reinterpret_cast(base + 16); - const uint8_t *flip = base + 24; - const size_t dim = in_dim; - - if (trunc_dim == dim) { - for (int round = 3; round >= 0; --round) { - fht_inplace(out, trunc_dim); - fht_vec_rescale(out, trunc_dim, fac); - fht_flip_sign(flip + static_cast(round) * flip_offset, out, dim); - } - return; - } - - // Non-power-of-2: undo final rescale(0.25) first - fht_vec_rescale(out, dim, 4.0f); - - size_t start = dim - trunc_dim; - float *trunc_ptr = out + start; - - fht_inv_kacs_walk(out, dim); - fht_inplace(trunc_ptr, trunc_dim); - fht_vec_rescale(trunc_ptr, trunc_dim, fac); - fht_flip_sign(flip + 3 * flip_offset, out, dim); - - fht_inv_kacs_walk(out, dim); - fht_inplace(out, trunc_dim); - fht_vec_rescale(out, trunc_dim, fac); - fht_flip_sign(flip + 2 * flip_offset, out, dim); - - fht_inv_kacs_walk(out, dim); - fht_inplace(trunc_ptr, trunc_dim); - fht_vec_rescale(trunc_ptr, trunc_dim, fac); - fht_flip_sign(flip + flip_offset, out, dim); - - fht_inv_kacs_walk(out, dim); - fht_inplace(out, trunc_dim); - fht_vec_rescale(out, trunc_dim, fac); - fht_flip_sign(flip, out, dim); -} - -} // namespace zvec::turbo::scalar diff --git a/src/turbo/distance/scalar/rotate/fht/fht.cc b/src/turbo/distance/scalar/rotate/fht/fht.cc new file mode 100644 index 000000000..1d4e6e779 --- /dev/null +++ b/src/turbo/distance/scalar/rotate/fht/fht.cc @@ -0,0 +1,96 @@ +// 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 +#include "common/fht_common.h" + +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; + } +} + +void fht_rotate(const float * /*in*/, float *out, size_t in_dim, + size_t /*out_dim*/, void *ctx) { + static constexpr FhtPrimitives kPrim = {fht_flip_sign, fht_inplace, + fht_kacs_walk, fht_inv_kacs_walk, + fht_vec_rescale}; + fht_rotate_impl(out, in_dim, ctx, kPrim); +} + +void fht_unrotate(const float * /*in*/, float *out, size_t in_dim, + size_t /*out_dim*/, void *ctx) { + static constexpr FhtPrimitives kPrim = {fht_flip_sign, fht_inplace, + fht_kacs_walk, fht_inv_kacs_walk, + fht_vec_rescale}; + fht_unrotate_impl(out, in_dim, ctx, kPrim); +} + +} // namespace zvec::turbo::scalar diff --git a/src/turbo/distance/scalar/fht/fht.h b/src/turbo/distance/scalar/rotate/fht/fht.h similarity index 100% rename from src/turbo/distance/scalar/fht/fht.h rename to src/turbo/distance/scalar/rotate/fht/fht.h diff --git a/src/turbo/distance/sse/fht/fht.cc b/src/turbo/distance/sse/rotate/fht/fht.cc similarity index 54% rename from src/turbo/distance/sse/fht/fht.cc rename to src/turbo/distance/sse/rotate/fht/fht.cc index e732f1e83..279bb3e96 100644 --- a/src/turbo/distance/sse/fht/fht.cc +++ b/src/turbo/distance/sse/rotate/fht/fht.cc @@ -20,7 +20,8 @@ #include #include #include -#include "scalar/fht/fht.h" +#include "common/fht_common.h" +#include "scalar/rotate/fht/fht.h" namespace zvec::turbo::sse { @@ -116,101 +117,18 @@ void fht_vec_rescale_sse(float *data, size_t n, float factor) { void fht_rotate_sse(const float * /*in*/, float *out, size_t in_dim, size_t /*out_dim*/, void *ctx) { - auto *base = reinterpret_cast(ctx); - const size_t flip_offset = *reinterpret_cast(base); - const size_t trunc_dim = *reinterpret_cast(base + 8); - const float fac = *reinterpret_cast(base + 16); - const uint8_t *flip = base + 24; - const size_t dim = in_dim; - - if (trunc_dim == dim) { - fht_flip_sign_sse(flip, out, dim); - scalar::fht_inplace(out, trunc_dim); - fht_vec_rescale_sse(out, trunc_dim, fac); - - fht_flip_sign_sse(flip + flip_offset, out, dim); - scalar::fht_inplace(out, trunc_dim); - fht_vec_rescale_sse(out, trunc_dim, fac); - - fht_flip_sign_sse(flip + 2 * flip_offset, out, dim); - scalar::fht_inplace(out, trunc_dim); - fht_vec_rescale_sse(out, trunc_dim, fac); - - fht_flip_sign_sse(flip + 3 * flip_offset, out, dim); - scalar::fht_inplace(out, trunc_dim); - fht_vec_rescale_sse(out, trunc_dim, fac); - return; - } - - size_t start = dim - trunc_dim; - float *trunc_ptr = out + start; - - fht_flip_sign_sse(flip, out, dim); - scalar::fht_inplace(out, trunc_dim); - fht_vec_rescale_sse(out, trunc_dim, fac); - fht_kacs_walk_sse(out, dim); - - fht_flip_sign_sse(flip + flip_offset, out, dim); - scalar::fht_inplace(trunc_ptr, trunc_dim); - fht_vec_rescale_sse(trunc_ptr, trunc_dim, fac); - fht_kacs_walk_sse(out, dim); - - fht_flip_sign_sse(flip + 2 * flip_offset, out, dim); - scalar::fht_inplace(out, trunc_dim); - fht_vec_rescale_sse(out, trunc_dim, fac); - fht_kacs_walk_sse(out, dim); - - fht_flip_sign_sse(flip + 3 * flip_offset, out, dim); - scalar::fht_inplace(trunc_ptr, trunc_dim); - fht_vec_rescale_sse(trunc_ptr, trunc_dim, fac); - fht_kacs_walk_sse(out, dim); - - fht_vec_rescale_sse(out, dim, 0.25f); + static constexpr FhtPrimitives kPrim = { + fht_flip_sign_sse, scalar::fht_inplace, fht_kacs_walk_sse, + fht_inv_kacs_walk_sse, fht_vec_rescale_sse}; + fht_rotate_impl(out, in_dim, ctx, kPrim); } void fht_unrotate_sse(const float * /*in*/, float *out, size_t in_dim, size_t /*out_dim*/, void *ctx) { - auto *base = reinterpret_cast(ctx); - const size_t flip_offset = *reinterpret_cast(base); - const size_t trunc_dim = *reinterpret_cast(base + 8); - const float fac = *reinterpret_cast(base + 16); - const uint8_t *flip = base + 24; - const size_t dim = in_dim; - - if (trunc_dim == dim) { - for (int round = 3; round >= 0; --round) { - scalar::fht_inplace(out, trunc_dim); - fht_vec_rescale_sse(out, trunc_dim, fac); - fht_flip_sign_sse(flip + static_cast(round) * flip_offset, out, - dim); - } - return; - } - - fht_vec_rescale_sse(out, dim, 4.0f); - - size_t start = dim - trunc_dim; - float *trunc_ptr = out + start; - - fht_inv_kacs_walk_sse(out, dim); - scalar::fht_inplace(trunc_ptr, trunc_dim); - fht_vec_rescale_sse(trunc_ptr, trunc_dim, fac); - fht_flip_sign_sse(flip + 3 * flip_offset, out, dim); - - fht_inv_kacs_walk_sse(out, dim); - scalar::fht_inplace(out, trunc_dim); - fht_vec_rescale_sse(out, trunc_dim, fac); - fht_flip_sign_sse(flip + 2 * flip_offset, out, dim); - - fht_inv_kacs_walk_sse(out, dim); - scalar::fht_inplace(trunc_ptr, trunc_dim); - fht_vec_rescale_sse(trunc_ptr, trunc_dim, fac); - fht_flip_sign_sse(flip + flip_offset, out, dim); - - fht_inv_kacs_walk_sse(out, dim); - scalar::fht_inplace(out, trunc_dim); - fht_vec_rescale_sse(out, trunc_dim, fac); - fht_flip_sign_sse(flip, out, dim); + static constexpr FhtPrimitives kPrim = { + fht_flip_sign_sse, scalar::fht_inplace, fht_kacs_walk_sse, + fht_inv_kacs_walk_sse, fht_vec_rescale_sse}; + fht_unrotate_impl(out, in_dim, ctx, kPrim); } } // namespace zvec::turbo::sse diff --git a/src/turbo/distance/sse/fht/fht.h b/src/turbo/distance/sse/rotate/fht/fht.h similarity index 100% rename from src/turbo/distance/sse/fht/fht.h rename to src/turbo/distance/sse/rotate/fht/fht.h diff --git a/src/turbo/turbo.cc b/src/turbo/turbo.cc index 0218e27ed..d4112225f 100644 --- a/src/turbo/turbo.cc +++ b/src/turbo/turbo.cc @@ -19,22 +19,22 @@ #include "avx512_vnni/record_quantized_int8/squared_euclidean.h" #include "avx512_vnni/uniform_int8/quantize.h" #include "avx512_vnni/uniform_int8/squared_euclidean.h" -#include "scalar/fht/fht.h" #include "scalar/fp32/cosine.h" #include "scalar/fp32/inner_product.h" #include "scalar/fp32/squared_euclidean.h" +#include "scalar/rotate/fht/fht.h" #if defined(__SSE2__) -#include "sse/fht/fht.h" +#include "sse/rotate/fht/fht.h" #endif #if defined(__AVX2__) -#include "avx2/fht/fht.h" +#include "avx2/rotate/fht/fht.h" #endif #if defined(__AVX512F__) -#include "avx512/fht/fht.h" +#include "avx512/rotate/fht/fht.h" #endif #if defined(__ARM_NEON) && defined(__aarch64__) -#include "neon/fht/fht.h" +#include "neon/rotate/fht/fht.h" #endif namespace zvec::turbo { From 6b207473b543d945a282a06cd4156073d8a22278 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Fri, 17 Jul 2026 11:03:04 +0800 Subject: [PATCH 068/100] fix enum --- src/include/zvec/turbo/turbo.h | 2 +- src/turbo/preprocessor/fht_rotator/fht_rotator.h | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/include/zvec/turbo/turbo.h b/src/include/zvec/turbo/turbo.h index 3face81bd..8b48e36a6 100644 --- a/src/include/zvec/turbo/turbo.h +++ b/src/include/zvec/turbo/turbo.h @@ -75,7 +75,7 @@ enum class QuantizeType { }; enum class RotateType : uint16_t { - kFht = 0, //!< O(d log d) FHT-based Kac random rotation + kFht = 1, //!< O(d log d) FHT-based Kac random rotation }; enum class CpuArchType { diff --git a/src/turbo/preprocessor/fht_rotator/fht_rotator.h b/src/turbo/preprocessor/fht_rotator/fht_rotator.h index ac8768758..0ff6879d4 100644 --- a/src/turbo/preprocessor/fht_rotator/fht_rotator.h +++ b/src/turbo/preprocessor/fht_rotator/fht_rotator.h @@ -82,7 +82,7 @@ class FhtRotator : public Preprocessor { int serialize(std::string *out) const override; int deserialize(const void *data, size_t len) override; - //! Rotator type tag (kFht = 0). + //! Rotator type tag (kFht = 1). RotateType rotate_type() const { return RotateType::kFht; } From 406503533368451821e7f40a0478ffb2e1fc31d4 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Fri, 17 Jul 2026 12:01:45 +0800 Subject: [PATCH 069/100] add helper --- src/turbo/CMakeLists.txt | 3 +- src/turbo/distance/avx2/rotate/fht/fht.h | 4 ++ src/turbo/distance/avx512/rotate/fht/fht.h | 4 ++ src/turbo/distance/common/fht_common.h | 74 +++++++++++----------- src/turbo/distance/neon/rotate/fht/fht.h | 4 ++ src/turbo/distance/sse/rotate/fht/fht.h | 4 ++ src/turbo/turbo.cc | 72 +++++++-------------- 7 files changed, 76 insertions(+), 89 deletions(-) diff --git a/src/turbo/CMakeLists.txt b/src/turbo/CMakeLists.txt index 2ce56d40e..b5a116d1b 100644 --- a/src/turbo/CMakeLists.txt +++ b/src/turbo/CMakeLists.txt @@ -35,8 +35,9 @@ if(NOT ANDROID AND AUTO_DETECT_ARCH) COMPILE_FLAGS "${TURBO_MARCH_FLAG_AVX2}" ) - # AVX512 (avx512_vnni and avx512/fht both use AVX512 march) + # AVX512 (avx512_vnni, avx512/fht, and turbo.cc dispatch all use AVX512 march) file(GLOB_RECURSE AVX512_SRCS ${CMAKE_CURRENT_SOURCE_DIR}/distance/avx512_vnni/*.cc ${CMAKE_CURRENT_SOURCE_DIR}/distance/avx512/*.cc) + list(APPEND AVX512_SRCS ${CMAKE_CURRENT_SOURCE_DIR}/turbo.cc) set_source_files_properties( ${AVX512_SRCS} PROPERTIES diff --git a/src/turbo/distance/avx2/rotate/fht/fht.h b/src/turbo/distance/avx2/rotate/fht/fht.h index 600842108..2fb011236 100644 --- a/src/turbo/distance/avx2/rotate/fht/fht.h +++ b/src/turbo/distance/avx2/rotate/fht/fht.h @@ -14,6 +14,8 @@ #pragma once +#if defined(__AVX2__) + #include #include @@ -43,3 +45,5 @@ void fht_unrotate_avx2(const float *in, float *out, size_t in_dim, size_t out_dim, void *ctx); } // namespace zvec::turbo::avx2 + +#endif // __AVX2__ diff --git a/src/turbo/distance/avx512/rotate/fht/fht.h b/src/turbo/distance/avx512/rotate/fht/fht.h index 6eecca959..bcb3a9fe9 100644 --- a/src/turbo/distance/avx512/rotate/fht/fht.h +++ b/src/turbo/distance/avx512/rotate/fht/fht.h @@ -14,6 +14,8 @@ #pragma once +#if defined(__AVX512F__) + #include #include @@ -43,3 +45,5 @@ void fht_unrotate_avx512(const float *in, float *out, size_t in_dim, size_t out_dim, void *ctx); } // namespace zvec::turbo::avx512 + +#endif // __AVX512F__ diff --git a/src/turbo/distance/common/fht_common.h b/src/turbo/distance/common/fht_common.h index 66377e756..556338863 100644 --- a/src/turbo/distance/common/fht_common.h +++ b/src/turbo/distance/common/fht_common.h @@ -35,92 +35,90 @@ struct FhtPrimitives { /// offset 16: float fac /// offset 20: uint8_t pad[4] /// offset 24: uint8_t flip[] -inline void fht_rotate_impl(float *out, size_t in_dim, void *ctx, +inline void fht_rotate_impl(float *data, size_t dim, void *ctx, const FhtPrimitives &p) { auto *base = reinterpret_cast(ctx); const size_t flip_offset = *reinterpret_cast(base); const size_t trunc_dim = *reinterpret_cast(base + 8); const float fac = *reinterpret_cast(base + 16); const uint8_t *flip = base + 24; - const size_t dim = in_dim; if (trunc_dim == dim) { for (size_t r = 0; r < 4; ++r) { - p.flip_sign(flip + r * flip_offset, out, dim); - p.inplace(out, trunc_dim); - p.rescale(out, trunc_dim, fac); + p.flip_sign(flip + r * flip_offset, data, dim); + p.inplace(data, trunc_dim); + p.rescale(data, trunc_dim, fac); } return; } size_t start = dim - trunc_dim; - float *trunc_ptr = out + start; + float *trunc_ptr = data + start; - p.flip_sign(flip, out, dim); - p.inplace(out, trunc_dim); - p.rescale(out, trunc_dim, fac); - p.kacs_walk(out, dim); + p.flip_sign(flip, data, dim); + p.inplace(data, trunc_dim); + p.rescale(data, trunc_dim, fac); + p.kacs_walk(data, dim); - p.flip_sign(flip + flip_offset, out, dim); + p.flip_sign(flip + flip_offset, data, dim); p.inplace(trunc_ptr, trunc_dim); p.rescale(trunc_ptr, trunc_dim, fac); - p.kacs_walk(out, dim); + p.kacs_walk(data, dim); - p.flip_sign(flip + 2 * flip_offset, out, dim); - p.inplace(out, trunc_dim); - p.rescale(out, trunc_dim, fac); - p.kacs_walk(out, dim); + p.flip_sign(flip + 2 * flip_offset, data, dim); + p.inplace(data, trunc_dim); + p.rescale(data, trunc_dim, fac); + p.kacs_walk(data, dim); - p.flip_sign(flip + 3 * flip_offset, out, dim); + p.flip_sign(flip + 3 * flip_offset, data, dim); p.inplace(trunc_ptr, trunc_dim); p.rescale(trunc_ptr, trunc_dim, fac); - p.kacs_walk(out, dim); + p.kacs_walk(data, dim); - p.rescale(out, dim, 0.25f); + p.rescale(data, dim, 0.25f); } -inline void fht_unrotate_impl(float *out, size_t in_dim, void *ctx, +inline void fht_unrotate_impl(float *data, size_t dim, void *ctx, const FhtPrimitives &p) { auto *base = reinterpret_cast(ctx); const size_t flip_offset = *reinterpret_cast(base); const size_t trunc_dim = *reinterpret_cast(base + 8); const float fac = *reinterpret_cast(base + 16); const uint8_t *flip = base + 24; - const size_t dim = in_dim; if (trunc_dim == dim) { for (int round = 3; round >= 0; --round) { - p.inplace(out, trunc_dim); - p.rescale(out, trunc_dim, fac); - p.flip_sign(flip + static_cast(round) * flip_offset, out, dim); + p.inplace(data, trunc_dim); + p.rescale(data, trunc_dim, fac); + p.flip_sign(flip + static_cast(round) * flip_offset, data, dim); } return; } - p.rescale(out, dim, 4.0f); + p.rescale(data, dim, 4.0f); size_t start = dim - trunc_dim; - float *trunc_ptr = out + start; + float *trunc_ptr = data + start; - p.inv_kacs_walk(out, dim); + p.inv_kacs_walk(data, dim); p.inplace(trunc_ptr, trunc_dim); p.rescale(trunc_ptr, trunc_dim, fac); - p.flip_sign(flip + 3 * flip_offset, out, dim); + p.flip_sign(flip + 3 * flip_offset, data, dim); - p.inv_kacs_walk(out, dim); - p.inplace(out, trunc_dim); - p.rescale(out, trunc_dim, fac); - p.flip_sign(flip + 2 * flip_offset, out, dim); + p.inv_kacs_walk(data, dim); + p.inplace(data, trunc_dim); + p.rescale(data, trunc_dim, fac); + p.flip_sign(flip + 2 * flip_offset, data, dim); - p.inv_kacs_walk(out, dim); + p.inv_kacs_walk(data, dim); p.inplace(trunc_ptr, trunc_dim); p.rescale(trunc_ptr, trunc_dim, fac); - p.flip_sign(flip + flip_offset, out, dim); + p.flip_sign(flip + flip_offset, data, dim); - p.inv_kacs_walk(out, dim); - p.inplace(out, trunc_dim); - p.rescale(out, trunc_dim, fac); - p.flip_sign(flip, out, dim); + p.inv_kacs_walk(data, dim); + p.inplace(data, trunc_dim); + p.rescale(data, trunc_dim, fac); + p.flip_sign(flip, data, dim); } } // namespace zvec::turbo diff --git a/src/turbo/distance/neon/rotate/fht/fht.h b/src/turbo/distance/neon/rotate/fht/fht.h index e7f4792d1..ee6bcacdc 100644 --- a/src/turbo/distance/neon/rotate/fht/fht.h +++ b/src/turbo/distance/neon/rotate/fht/fht.h @@ -14,6 +14,8 @@ #pragma once +#if defined(__ARM_NEON) && defined(__aarch64__) + #include #include @@ -40,3 +42,5 @@ void fht_unrotate_neon(const float *in, float *out, size_t in_dim, size_t out_dim, void *ctx); } // namespace zvec::turbo::neon + +#endif // __ARM_NEON && __aarch64__ diff --git a/src/turbo/distance/sse/rotate/fht/fht.h b/src/turbo/distance/sse/rotate/fht/fht.h index 996c3008d..283ced493 100644 --- a/src/turbo/distance/sse/rotate/fht/fht.h +++ b/src/turbo/distance/sse/rotate/fht/fht.h @@ -14,6 +14,8 @@ #pragma once +#if defined(__SSE2__) + #include #include @@ -40,3 +42,5 @@ void fht_unrotate_sse(const float *in, float *out, size_t in_dim, size_t out_dim, void *ctx); } // namespace zvec::turbo::sse + +#endif // __SSE2__ diff --git a/src/turbo/turbo.cc b/src/turbo/turbo.cc index d4112225f..493b720d7 100644 --- a/src/turbo/turbo.cc +++ b/src/turbo/turbo.cc @@ -15,30 +15,26 @@ #include #include #include +#include "avx2/rotate/fht/fht.h" +#include "avx512/rotate/fht/fht.h" #include "avx512_vnni/record_quantized_int8/cosine.h" #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 "neon/rotate/fht/fht.h" #include "scalar/fp32/cosine.h" #include "scalar/fp32/inner_product.h" #include "scalar/fp32/squared_euclidean.h" #include "scalar/rotate/fht/fht.h" - -#if defined(__SSE2__) #include "sse/rotate/fht/fht.h" -#endif -#if defined(__AVX2__) -#include "avx2/rotate/fht/fht.h" -#endif -#if defined(__AVX512F__) -#include "avx512/rotate/fht/fht.h" -#endif -#if defined(__ARM_NEON) && defined(__aarch64__) -#include "neon/rotate/fht/fht.h" -#endif namespace zvec::turbo { +// Helper: check if the requested arch matches the target or is auto-detect. +static bool IsArchMatch(CpuArchType actual, CpuArchType target) { + return actual == CpuArchType::kAuto || actual == target; +} + DistanceFunc get_distance_func(MetricType metric_type, DataType data_type, QuantizeType quantize_type, CpuArchType cpu_arch_type) { @@ -60,8 +56,7 @@ DistanceFunc get_distance_func(MetricType metric_type, DataType data_type, if (data_type == DataType::kInt8) { if (quantize_type == QuantizeType::kDefault) { if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI && - (cpu_arch_type == CpuArchType::kAuto || - cpu_arch_type == CpuArchType::kAVX512VNNI)) { + IsArchMatch(cpu_arch_type, CpuArchType::kAVX512VNNI)) { if (metric_type == MetricType::kSquaredEuclidean) { return avx512_vnni::squared_euclidean_int8_distance; } @@ -72,8 +67,7 @@ DistanceFunc get_distance_func(MetricType metric_type, DataType data_type, } if (quantize_type == QuantizeType::kUniform) { if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI && - (cpu_arch_type == CpuArchType::kAuto || - cpu_arch_type == CpuArchType::kAVX512VNNI)) { + IsArchMatch(cpu_arch_type, CpuArchType::kAVX512VNNI)) { if (metric_type == MetricType::kSquaredEuclidean) { return avx512_vnni::uniform_squared_euclidean_int8_distance; } @@ -105,8 +99,7 @@ BatchDistanceFunc get_batch_distance_func(MetricType metric_type, if (data_type == DataType::kInt8) { if (quantize_type == QuantizeType::kDefault) { if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI && - (cpu_arch_type == CpuArchType::kAuto || - cpu_arch_type == CpuArchType::kAVX512VNNI)) { + IsArchMatch(cpu_arch_type, CpuArchType::kAVX512VNNI)) { if (metric_type == MetricType::kSquaredEuclidean) { return avx512_vnni::squared_euclidean_int8_batch_distance; } @@ -117,8 +110,7 @@ BatchDistanceFunc get_batch_distance_func(MetricType metric_type, } if (quantize_type == QuantizeType::kUniform) { if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI && - (cpu_arch_type == CpuArchType::kAuto || - cpu_arch_type == CpuArchType::kAVX512VNNI)) { + IsArchMatch(cpu_arch_type, CpuArchType::kAVX512VNNI)) { if (metric_type == MetricType::kSquaredEuclidean) { return avx512_vnni::uniform_squared_euclidean_int8_batch_distance; } @@ -136,8 +128,7 @@ QueryPreprocessFunc get_query_preprocess_func(MetricType metric_type, if (data_type == DataType::kInt8) { if (quantize_type == QuantizeType::kDefault) { if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI && - (cpu_arch_type == CpuArchType::kAuto || - cpu_arch_type == CpuArchType::kAVX512VNNI)) { + IsArchMatch(cpu_arch_type, CpuArchType::kAVX512VNNI)) { if (metric_type == MetricType::kSquaredEuclidean) { return avx512_vnni::squared_euclidean_int8_query_preprocess; } @@ -164,52 +155,33 @@ UniformQuantizeFunc get_uniform_quantize_func(DataType data_type) { RotatorKernels get_rotator_kernels(RotateType rotate_type, CpuArchType cpu_arch_type) { - (void)cpu_arch_type; - switch (rotate_type) { case RotateType::kFht: { - RotatorKernels k; - // Default: scalar - k.rotate = scalar::fht_rotate; - k.unrotate = scalar::fht_unrotate; - #if defined(__AVX512F__) if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F && zvec::ailego::internal::CpuFeatures::static_flags_.AVX512DQ && - (cpu_arch_type == CpuArchType::kAuto || - cpu_arch_type == CpuArchType::kAVX512)) { - k.rotate = avx512::fht_rotate_avx512; - k.unrotate = avx512::fht_unrotate_avx512; - return k; + IsArchMatch(cpu_arch_type, CpuArchType::kAVX512)) { + return {avx512::fht_rotate_avx512, avx512::fht_unrotate_avx512}; } #endif #if defined(__AVX2__) if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2 && - (cpu_arch_type == CpuArchType::kAuto || - cpu_arch_type == CpuArchType::kAVX2)) { - k.rotate = avx2::fht_rotate_avx2; - k.unrotate = avx2::fht_unrotate_avx2; - return k; + IsArchMatch(cpu_arch_type, CpuArchType::kAVX2)) { + return {avx2::fht_rotate_avx2, avx2::fht_unrotate_avx2}; } #endif #if defined(__SSE2__) if (zvec::ailego::internal::CpuFeatures::static_flags_.SSE2 && - (cpu_arch_type == CpuArchType::kAuto || - cpu_arch_type == CpuArchType::kSSE)) { - k.rotate = sse::fht_rotate_sse; - k.unrotate = sse::fht_unrotate_sse; - return k; + IsArchMatch(cpu_arch_type, CpuArchType::kSSE)) { + return {sse::fht_rotate_sse, sse::fht_unrotate_sse}; } #endif #if defined(__ARM_NEON) && defined(__aarch64__) - if (cpu_arch_type == CpuArchType::kAuto || - cpu_arch_type == CpuArchType::kNEON) { - k.rotate = neon::fht_rotate_neon; - k.unrotate = neon::fht_unrotate_neon; - return k; + if (IsArchMatch(cpu_arch_type, CpuArchType::kNEON)) { + return {neon::fht_rotate_neon, neon::fht_unrotate_neon}; } #endif - return k; // scalar + return {scalar::fht_rotate, scalar::fht_unrotate}; } } From e024bc8ebd0831e88e6f40079ea420162363d588 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Fri, 17 Jul 2026 16:16:01 +0800 Subject: [PATCH 070/100] debug --- src/turbo/CMakeLists.txt | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/turbo/CMakeLists.txt b/src/turbo/CMakeLists.txt index b5a116d1b..2ce56d40e 100644 --- a/src/turbo/CMakeLists.txt +++ b/src/turbo/CMakeLists.txt @@ -35,9 +35,8 @@ if(NOT ANDROID AND AUTO_DETECT_ARCH) COMPILE_FLAGS "${TURBO_MARCH_FLAG_AVX2}" ) - # AVX512 (avx512_vnni, avx512/fht, and turbo.cc dispatch all use AVX512 march) + # 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) - list(APPEND AVX512_SRCS ${CMAKE_CURRENT_SOURCE_DIR}/turbo.cc) set_source_files_properties( ${AVX512_SRCS} PROPERTIES From 52951dab64a2a03d3f52dfc61686dd8bed4fba91 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Fri, 17 Jul 2026 16:53:28 +0800 Subject: [PATCH 071/100] Runtime Distribution --- src/turbo/distance/avx2/rotate/fht/fht.cc | 48 +++++++++++++++++++-- src/turbo/distance/avx2/rotate/fht/fht.h | 4 -- src/turbo/distance/avx512/rotate/fht/fht.cc | 48 +++++++++++++++++++-- src/turbo/distance/avx512/rotate/fht/fht.h | 4 -- src/turbo/distance/neon/rotate/fht/fht.cc | 42 +++++++++++++++--- src/turbo/distance/neon/rotate/fht/fht.h | 4 -- src/turbo/distance/sse/rotate/fht/fht.cc | 43 ++++++++++++++++-- src/turbo/distance/sse/rotate/fht/fht.h | 4 -- src/turbo/turbo.cc | 8 ---- 9 files changed, 167 insertions(+), 38 deletions(-) diff --git a/src/turbo/distance/avx2/rotate/fht/fht.cc b/src/turbo/distance/avx2/rotate/fht/fht.cc index a77dc087b..a15ca6846 100644 --- a/src/turbo/distance/avx2/rotate/fht/fht.cc +++ b/src/turbo/distance/avx2/rotate/fht/fht.cc @@ -12,10 +12,15 @@ // See the License for the specific language governing permissions and // limitations under the License. -#if defined(__AVX2__) +// This file is compiled with per-file -march=core-avx2 (set in CMakeLists.txt) +// so that AVX2 intrinsics are available. When the build toolchain cannot emit +// AVX2 code, each function falls back to a no-op stub guarded by +// #if defined(__AVX2__). #include "fht.h" +#if defined(__AVX2__) #include +#endif #include #include #include @@ -25,6 +30,7 @@ namespace zvec::turbo::avx2 { void fht_flip_sign_avx2(const uint8_t *flip, float *data, size_t dim) { +#if defined(__AVX2__) size_t simd_end = dim & ~31u; constexpr size_t kChunk = 32; const __m256i bit_select = @@ -49,9 +55,15 @@ void fht_flip_sign_avx2(const uint8_t *flip, float *data, size_t dim) { data[i] = -data[i]; } } +#else + (void)flip; + (void)data; + (void)dim; +#endif } void fht_kacs_walk_avx2(float *data, size_t len) { +#if defined(__AVX2__) size_t half = len / 2; size_t base = len % 2; size_t offset = base + half; @@ -72,9 +84,14 @@ void fht_kacs_walk_avx2(float *data, size_t len) { if (base != 0) { data[half] *= std::sqrt(2.0f); } +#else + (void)data; + (void)len; +#endif } void fht_inv_kacs_walk_avx2(float *data, size_t len) { +#if defined(__AVX2__) size_t half = len / 2; size_t base = len % 2; size_t offset = base + half; @@ -97,9 +114,14 @@ void fht_inv_kacs_walk_avx2(float *data, size_t len) { data[i] = (a + b) * 0.5f; data[i + offset] = (a - b) * 0.5f; } +#else + (void)data; + (void)len; +#endif } void fht_inplace_avx2(float *data, size_t n) { +#if defined(__AVX2__) for (size_t len = 1; len < n; len <<= 1) { size_t step = len << 1; size_t simd_end = len & ~7u; @@ -118,9 +140,14 @@ void fht_inplace_avx2(float *data, size_t n) { } } } +#else + (void)data; + (void)n; +#endif } void fht_vec_rescale_avx2(float *data, size_t n, float factor) { +#if defined(__AVX2__) const __m256 fac = _mm256_set1_ps(factor); size_t simd_end = n & ~7u; for (size_t i = 0; i < simd_end; i += 8) { @@ -131,24 +158,39 @@ void fht_vec_rescale_avx2(float *data, size_t n, float factor) { for (size_t i = simd_end; i < n; ++i) { data[i] *= factor; } +#else + (void)data; + (void)n; + (void)factor; +#endif } void fht_rotate_avx2(const float * /*in*/, float *out, size_t in_dim, size_t /*out_dim*/, void *ctx) { +#if defined(__AVX2__) static constexpr FhtPrimitives kPrim = { fht_flip_sign_avx2, fht_inplace_avx2, fht_kacs_walk_avx2, fht_inv_kacs_walk_avx2, fht_vec_rescale_avx2}; fht_rotate_impl(out, in_dim, ctx, kPrim); +#else + (void)out; + (void)in_dim; + (void)ctx; +#endif } void fht_unrotate_avx2(const float * /*in*/, float *out, size_t in_dim, size_t /*out_dim*/, void *ctx) { +#if defined(__AVX2__) static constexpr FhtPrimitives kPrim = { fht_flip_sign_avx2, fht_inplace_avx2, fht_kacs_walk_avx2, fht_inv_kacs_walk_avx2, fht_vec_rescale_avx2}; fht_unrotate_impl(out, in_dim, ctx, kPrim); +#else + (void)out; + (void)in_dim; + (void)ctx; +#endif } } // namespace zvec::turbo::avx2 - -#endif // __AVX2__ diff --git a/src/turbo/distance/avx2/rotate/fht/fht.h b/src/turbo/distance/avx2/rotate/fht/fht.h index 2fb011236..600842108 100644 --- a/src/turbo/distance/avx2/rotate/fht/fht.h +++ b/src/turbo/distance/avx2/rotate/fht/fht.h @@ -14,8 +14,6 @@ #pragma once -#if defined(__AVX2__) - #include #include @@ -45,5 +43,3 @@ void fht_unrotate_avx2(const float *in, float *out, size_t in_dim, size_t out_dim, void *ctx); } // namespace zvec::turbo::avx2 - -#endif // __AVX2__ diff --git a/src/turbo/distance/avx512/rotate/fht/fht.cc b/src/turbo/distance/avx512/rotate/fht/fht.cc index 9a237214d..a892ac61f 100644 --- a/src/turbo/distance/avx512/rotate/fht/fht.cc +++ b/src/turbo/distance/avx512/rotate/fht/fht.cc @@ -12,10 +12,15 @@ // See the License for the specific language governing permissions and // limitations under the License. -#if defined(__AVX512F__) +// This file is compiled with per-file -march=icelake-server (set in +// CMakeLists.txt) so that AVX512 intrinsics are available. When the build +// toolchain cannot emit AVX-512 code, each function falls back to a no-op +// stub guarded by #if defined(__AVX512F__). #include "fht.h" +#if defined(__AVX512F__) #include +#endif #include #include #include @@ -25,6 +30,7 @@ namespace zvec::turbo::avx512 { void fht_flip_sign_avx512(const uint8_t *flip, float *data, size_t dim) { +#if defined(__AVX512F__) size_t simd_end = dim & ~63u; constexpr size_t kChunk = 64; const __m512 sign_flip = _mm512_castsi512_ps(_mm512_set1_epi32(0x80000000)); @@ -54,9 +60,15 @@ void fht_flip_sign_avx512(const uint8_t *flip, float *data, size_t dim) { data[i] = -data[i]; } } +#else + (void)flip; + (void)data; + (void)dim; +#endif } void fht_kacs_walk_avx512(float *data, size_t len) { +#if defined(__AVX512F__) size_t half = len / 2; size_t base = len % 2; size_t offset = base + half; @@ -77,9 +89,14 @@ void fht_kacs_walk_avx512(float *data, size_t len) { if (base != 0) { data[half] *= std::sqrt(2.0f); } +#else + (void)data; + (void)len; +#endif } void fht_inv_kacs_walk_avx512(float *data, size_t len) { +#if defined(__AVX512F__) size_t half = len / 2; size_t base = len % 2; size_t offset = base + half; @@ -102,9 +119,14 @@ void fht_inv_kacs_walk_avx512(float *data, size_t len) { data[i] = (a + b) * 0.5f; data[i + offset] = (a - b) * 0.5f; } +#else + (void)data; + (void)len; +#endif } void fht_inplace_avx512(float *data, size_t n) { +#if defined(__AVX512F__) for (size_t len = 1; len < n; len <<= 1) { size_t step = len << 1; size_t simd_end = len & ~15u; @@ -123,9 +145,14 @@ void fht_inplace_avx512(float *data, size_t n) { } } } +#else + (void)data; + (void)n; +#endif } void fht_vec_rescale_avx512(float *data, size_t n, float factor) { +#if defined(__AVX512F__) const __m512 fac = _mm512_set1_ps(factor); size_t simd_end = n & ~15u; for (size_t i = 0; i < simd_end; i += 16) { @@ -136,24 +163,39 @@ void fht_vec_rescale_avx512(float *data, size_t n, float factor) { for (size_t i = simd_end; i < n; ++i) { data[i] *= factor; } +#else + (void)data; + (void)n; + (void)factor; +#endif } void fht_rotate_avx512(const float * /*in*/, float *out, size_t in_dim, size_t /*out_dim*/, void *ctx) { +#if defined(__AVX512F__) static constexpr FhtPrimitives kPrim = { fht_flip_sign_avx512, fht_inplace_avx512, fht_kacs_walk_avx512, fht_inv_kacs_walk_avx512, fht_vec_rescale_avx512}; fht_rotate_impl(out, in_dim, ctx, kPrim); +#else + (void)out; + (void)in_dim; + (void)ctx; +#endif } void fht_unrotate_avx512(const float * /*in*/, float *out, size_t in_dim, size_t /*out_dim*/, void *ctx) { +#if defined(__AVX512F__) static constexpr FhtPrimitives kPrim = { fht_flip_sign_avx512, fht_inplace_avx512, fht_kacs_walk_avx512, fht_inv_kacs_walk_avx512, fht_vec_rescale_avx512}; fht_unrotate_impl(out, in_dim, ctx, kPrim); +#else + (void)out; + (void)in_dim; + (void)ctx; +#endif } } // namespace zvec::turbo::avx512 - -#endif // __AVX512F__ diff --git a/src/turbo/distance/avx512/rotate/fht/fht.h b/src/turbo/distance/avx512/rotate/fht/fht.h index bcb3a9fe9..6eecca959 100644 --- a/src/turbo/distance/avx512/rotate/fht/fht.h +++ b/src/turbo/distance/avx512/rotate/fht/fht.h @@ -14,8 +14,6 @@ #pragma once -#if defined(__AVX512F__) - #include #include @@ -45,5 +43,3 @@ void fht_unrotate_avx512(const float *in, float *out, size_t in_dim, size_t out_dim, void *ctx); } // namespace zvec::turbo::avx512 - -#endif // __AVX512F__ diff --git a/src/turbo/distance/neon/rotate/fht/fht.cc b/src/turbo/distance/neon/rotate/fht/fht.cc index cc736a0a2..07a15e46f 100644 --- a/src/turbo/distance/neon/rotate/fht/fht.cc +++ b/src/turbo/distance/neon/rotate/fht/fht.cc @@ -12,20 +12,21 @@ // See the License for the specific language governing permissions and // limitations under the License. -#if defined(__ARM_NEON) && defined(__aarch64__) - #include "fht.h" +#if defined(__ARM_NEON) && defined(__aarch64__) +#include +#endif #include #include #include #include -#include #include "common/fht_common.h" #include "scalar/rotate/fht/fht.h" namespace zvec::turbo::neon { void fht_flip_sign_neon(const uint8_t *flip, float *data, size_t dim) { +#if defined(__ARM_NEON) && defined(__aarch64__) const uint32x4_t sign_bit = vdupq_n_u32(0x80000000u); size_t simd_end = dim & ~3u; size_t flip_bytes = (dim + 7) / 8; @@ -54,9 +55,15 @@ void fht_flip_sign_neon(const uint8_t *flip, float *data, size_t dim) { data[i] = -data[i]; } } +#else + (void)flip; + (void)data; + (void)dim; +#endif } void fht_kacs_walk_neon(float *data, size_t len) { +#if defined(__ARM_NEON) && defined(__aarch64__) size_t half = len / 2; size_t base = len % 2; size_t offset = base + half; @@ -77,9 +84,14 @@ void fht_kacs_walk_neon(float *data, size_t len) { if (base != 0) { data[half] *= std::sqrt(2.0f); } +#else + (void)data; + (void)len; +#endif } void fht_inv_kacs_walk_neon(float *data, size_t len) { +#if defined(__ARM_NEON) && defined(__aarch64__) size_t half = len / 2; size_t base = len % 2; size_t offset = base + half; @@ -101,9 +113,14 @@ void fht_inv_kacs_walk_neon(float *data, size_t len) { data[i] = (a + b) * 0.5f; data[i + offset] = (a - b) * 0.5f; } +#else + (void)data; + (void)len; +#endif } void fht_vec_rescale_neon(float *data, size_t n, float factor) { +#if defined(__ARM_NEON) && defined(__aarch64__) const float32x4_t fac = vdupq_n_f32(factor); size_t simd_end = n & ~3u; for (size_t i = 0; i < simd_end; i += 4) { @@ -114,24 +131,39 @@ void fht_vec_rescale_neon(float *data, size_t n, float factor) { for (size_t i = simd_end; i < n; ++i) { data[i] *= factor; } +#else + (void)data; + (void)n; + (void)factor; +#endif } void fht_rotate_neon(const float * /*in*/, float *out, size_t in_dim, size_t /*out_dim*/, void *ctx) { +#if defined(__ARM_NEON) && defined(__aarch64__) static constexpr FhtPrimitives kPrim = { fht_flip_sign_neon, scalar::fht_inplace, fht_kacs_walk_neon, fht_inv_kacs_walk_neon, fht_vec_rescale_neon}; fht_rotate_impl(out, in_dim, ctx, kPrim); +#else + (void)out; + (void)in_dim; + (void)ctx; +#endif } void fht_unrotate_neon(const float * /*in*/, float *out, size_t in_dim, size_t /*out_dim*/, void *ctx) { +#if defined(__ARM_NEON) && defined(__aarch64__) static constexpr FhtPrimitives kPrim = { fht_flip_sign_neon, scalar::fht_inplace, fht_kacs_walk_neon, fht_inv_kacs_walk_neon, fht_vec_rescale_neon}; fht_unrotate_impl(out, in_dim, ctx, kPrim); +#else + (void)out; + (void)in_dim; + (void)ctx; +#endif } } // namespace zvec::turbo::neon - -#endif // __ARM_NEON && __aarch64__ diff --git a/src/turbo/distance/neon/rotate/fht/fht.h b/src/turbo/distance/neon/rotate/fht/fht.h index ee6bcacdc..e7f4792d1 100644 --- a/src/turbo/distance/neon/rotate/fht/fht.h +++ b/src/turbo/distance/neon/rotate/fht/fht.h @@ -14,8 +14,6 @@ #pragma once -#if defined(__ARM_NEON) && defined(__aarch64__) - #include #include @@ -42,5 +40,3 @@ void fht_unrotate_neon(const float *in, float *out, size_t in_dim, size_t out_dim, void *ctx); } // namespace zvec::turbo::neon - -#endif // __ARM_NEON && __aarch64__ diff --git a/src/turbo/distance/sse/rotate/fht/fht.cc b/src/turbo/distance/sse/rotate/fht/fht.cc index 279bb3e96..f020fcc1a 100644 --- a/src/turbo/distance/sse/rotate/fht/fht.cc +++ b/src/turbo/distance/sse/rotate/fht/fht.cc @@ -12,10 +12,15 @@ // See the License for the specific language governing permissions and // limitations under the License. -#if defined(__SSE2__) +// This file is compiled with per-file -march=corei7 (set in CMakeLists.txt) +// so that SSE2 intrinsics are available. When the build toolchain cannot emit +// SSE2 code, each function falls back to a no-op stub guarded by +// #if defined(__SSE2__). #include "fht.h" +#if defined(__SSE2__) #include +#endif #include #include #include @@ -26,6 +31,7 @@ namespace zvec::turbo::sse { void fht_flip_sign_sse(const uint8_t *flip, float *data, size_t dim) { +#if defined(__SSE2__) size_t simd_end = dim & ~3u; size_t flip_bytes = (dim + 7) / 8; for (size_t i = 0; i < simd_end; i += 4) { @@ -53,9 +59,15 @@ void fht_flip_sign_sse(const uint8_t *flip, float *data, size_t dim) { data[i] = -data[i]; } } +#else + (void)flip; + (void)data; + (void)dim; +#endif } void fht_kacs_walk_sse(float *data, size_t len) { +#if defined(__SSE2__) size_t half = len / 2; size_t base = len % 2; size_t offset = base + half; @@ -76,9 +88,14 @@ void fht_kacs_walk_sse(float *data, size_t len) { if (base != 0) { data[half] *= std::sqrt(2.0f); } +#else + (void)data; + (void)len; +#endif } void fht_inv_kacs_walk_sse(float *data, size_t len) { +#if defined(__SSE2__) size_t half = len / 2; size_t base = len % 2; size_t offset = base + half; @@ -100,9 +117,14 @@ void fht_inv_kacs_walk_sse(float *data, size_t len) { data[i] = (a + b) * 0.5f; data[i + offset] = (a - b) * 0.5f; } +#else + (void)data; + (void)len; +#endif } void fht_vec_rescale_sse(float *data, size_t n, float factor) { +#if defined(__SSE2__) const __m128 fac = _mm_set1_ps(factor); size_t simd_end = n & ~3u; for (size_t i = 0; i < simd_end; i += 4) { @@ -113,24 +135,39 @@ void fht_vec_rescale_sse(float *data, size_t n, float factor) { for (size_t i = simd_end; i < n; ++i) { data[i] *= factor; } +#else + (void)data; + (void)n; + (void)factor; +#endif } void fht_rotate_sse(const float * /*in*/, float *out, size_t in_dim, size_t /*out_dim*/, void *ctx) { +#if defined(__SSE2__) static constexpr FhtPrimitives kPrim = { fht_flip_sign_sse, scalar::fht_inplace, fht_kacs_walk_sse, fht_inv_kacs_walk_sse, fht_vec_rescale_sse}; fht_rotate_impl(out, in_dim, ctx, kPrim); +#else + (void)out; + (void)in_dim; + (void)ctx; +#endif } void fht_unrotate_sse(const float * /*in*/, float *out, size_t in_dim, size_t /*out_dim*/, void *ctx) { +#if defined(__SSE2__) static constexpr FhtPrimitives kPrim = { fht_flip_sign_sse, scalar::fht_inplace, fht_kacs_walk_sse, fht_inv_kacs_walk_sse, fht_vec_rescale_sse}; fht_unrotate_impl(out, in_dim, ctx, kPrim); +#else + (void)out; + (void)in_dim; + (void)ctx; +#endif } } // namespace zvec::turbo::sse - -#endif // __SSE2__ diff --git a/src/turbo/distance/sse/rotate/fht/fht.h b/src/turbo/distance/sse/rotate/fht/fht.h index 283ced493..996c3008d 100644 --- a/src/turbo/distance/sse/rotate/fht/fht.h +++ b/src/turbo/distance/sse/rotate/fht/fht.h @@ -14,8 +14,6 @@ #pragma once -#if defined(__SSE2__) - #include #include @@ -42,5 +40,3 @@ void fht_unrotate_sse(const float *in, float *out, size_t in_dim, size_t out_dim, void *ctx); } // namespace zvec::turbo::sse - -#endif // __SSE2__ diff --git a/src/turbo/turbo.cc b/src/turbo/turbo.cc index 493b720d7..b45a99617 100644 --- a/src/turbo/turbo.cc +++ b/src/turbo/turbo.cc @@ -157,30 +157,22 @@ RotatorKernels get_rotator_kernels(RotateType rotate_type, CpuArchType cpu_arch_type) { switch (rotate_type) { case RotateType::kFht: { -#if defined(__AVX512F__) if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F && zvec::ailego::internal::CpuFeatures::static_flags_.AVX512DQ && IsArchMatch(cpu_arch_type, CpuArchType::kAVX512)) { return {avx512::fht_rotate_avx512, avx512::fht_unrotate_avx512}; } -#endif -#if defined(__AVX2__) if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2 && IsArchMatch(cpu_arch_type, CpuArchType::kAVX2)) { return {avx2::fht_rotate_avx2, avx2::fht_unrotate_avx2}; } -#endif -#if defined(__SSE2__) if (zvec::ailego::internal::CpuFeatures::static_flags_.SSE2 && IsArchMatch(cpu_arch_type, CpuArchType::kSSE)) { return {sse::fht_rotate_sse, sse::fht_unrotate_sse}; } -#endif -#if defined(__ARM_NEON) && defined(__aarch64__) if (IsArchMatch(cpu_arch_type, CpuArchType::kNEON)) { return {neon::fht_rotate_neon, neon::fht_unrotate_neon}; } -#endif return {scalar::fht_rotate, scalar::fht_unrotate}; } } From 22134a01bdf39a1100b1efb52bec9d69508eeb49 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Fri, 17 Jul 2026 19:44:46 +0800 Subject: [PATCH 072/100] merge preprocessor --- .../avx2/pq_quantizer_int8/pq_distance.cc | 34 +++ .../avx512/pq_quantizer_int8/pq_distance.cc | 34 +++ .../neon/pq_quantizer_int8/pq_distance.cc | 207 ++++++++++++++++++ .../neon/pq_quantizer_int8/pq_distance.h | 41 ++++ src/turbo/turbo.cc | 46 ++-- 5 files changed, 338 insertions(+), 24 deletions(-) create mode 100644 src/turbo/distance/neon/pq_quantizer_int8/pq_distance.cc create mode 100644 src/turbo/distance/neon/pq_quantizer_int8/pq_distance.h 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 a6ca691f8..d9fbc26c0 100644 --- a/src/turbo/distance/avx2/pq_quantizer_int8/pq_distance.cc +++ b/src/turbo/distance/avx2/pq_quantizer_int8/pq_distance.cc @@ -12,11 +12,21 @@ // See the License for the specific language governing permissions and // limitations under the License. +// This file is compiled with per-file -march=core-avx2 (set in CMakeLists.txt) +// so that AVX2 intrinsics are available. When the build toolchain cannot emit +// AVX2 code, each function falls back to a no-op stub guarded by +// #if defined(__AVX2__). + #include "avx2/pq_quantizer_int8/pq_distance.h" +#if defined(__AVX2__) #include +#endif +#include +#include namespace zvec::turbo::avx2 { +#if defined(__AVX2__) namespace { // Horizontal sum of 8 floats in a __m256 register. @@ -35,9 +45,11 @@ inline float horizontal_sum_avx2(__m256 v) { } } // namespace +#endif void pq_adc_int8_distance_avx2(const void *pq_code_v, const void *lut_v, size_t num_chunk, float *out) { +#if defined(__AVX2__) constexpr int kNumCentroids = 256; constexpr int kChunkSize = 8; // AVX2 processes 8 floats at once const auto *pq_code = reinterpret_cast(pq_code_v); @@ -79,11 +91,18 @@ void pq_adc_int8_distance_avx2(const void *pq_code_v, const void *lut_v, } *out = sum; +#else + (void)pq_code_v; + (void)lut_v; + (void)num_chunk; + (void)out; +#endif } void pq_sdc_int8_distance_avx2(const void *a_v, const void *b_v, const void *dist_table_v, size_t num_chunk, float *out) { +#if defined(__AVX2__) constexpr int kNumCentroids = 256; constexpr int kTablePerSub = kNumCentroids * kNumCentroids; // 65536 constexpr int kChunkSize = 8; @@ -132,11 +151,19 @@ void pq_sdc_int8_distance_avx2(const void *a_v, const void *b_v, } *out = sum; +#else + (void)a_v; + (void)b_v; + (void)dist_table_v; + (void)num_chunk; + (void)out; +#endif } void pq_adc_int8_batch_distance_avx2(const void **candidates_v, const void *lut_v, size_t num, size_t num_chunk, float *out) { +#if defined(__AVX2__) constexpr int kNumCentroids = 256; constexpr int kChunkSize = 8; constexpr int kBatch = 4; @@ -210,6 +237,13 @@ void pq_adc_int8_batch_distance_avx2(const void **candidates_v, for (; i < num; ++i) { pq_adc_int8_distance_avx2(candidates[i], lut, num_chunk, out + i); } +#else + (void)candidates_v; + (void)lut_v; + (void)num; + (void)num_chunk; + (void)out; +#endif } } // namespace zvec::turbo::avx2 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 8b69d881d..5ece65e13 100644 --- a/src/turbo/distance/avx512/pq_quantizer_int8/pq_distance.cc +++ b/src/turbo/distance/avx512/pq_quantizer_int8/pq_distance.cc @@ -12,11 +12,21 @@ // See the License for the specific language governing permissions and // limitations under the License. +// This file is compiled with per-file -march=icelake-server (set in +// CMakeLists.txt) so that AVX512 intrinsics are available. When the build +// toolchain cannot emit AVX-512 code, each function falls back to a no-op +// stub guarded by #if defined(__AVX512F__). + #include "avx512/pq_quantizer_int8/pq_distance.h" +#if defined(__AVX512F__) #include +#endif +#include +#include namespace zvec::turbo::avx512 { +#if defined(__AVX512F__) namespace { // Horizontal sum of 16 floats in a __m512 register. @@ -26,9 +36,11 @@ inline float horizontal_sum_avx512(__m512 v) { } } // namespace +#endif void pq_adc_int8_distance_avx512(const void *pq_code_v, const void *lut_v, size_t num_chunk, float *out) { +#if defined(__AVX512F__) constexpr int kNumCentroids = 256; constexpr int kChunkSize = 16; // AVX512 processes 16 floats at once const auto *pq_code = reinterpret_cast(pq_code_v); @@ -72,11 +84,18 @@ void pq_adc_int8_distance_avx512(const void *pq_code_v, const void *lut_v, } *out = sum; +#else + (void)pq_code_v; + (void)lut_v; + (void)num_chunk; + (void)out; +#endif } void pq_sdc_int8_distance_avx512(const void *a_v, const void *b_v, const void *dist_table_v, size_t num_chunk, float *out) { +#if defined(__AVX512F__) constexpr int kNumCentroids = 256; constexpr int kTablePerSub = kNumCentroids * kNumCentroids; // 65536 constexpr int kChunkSize = 16; @@ -128,11 +147,19 @@ void pq_sdc_int8_distance_avx512(const void *a_v, const void *b_v, } *out = sum; +#else + (void)a_v; + (void)b_v; + (void)dist_table_v; + (void)num_chunk; + (void)out; +#endif } void pq_adc_int8_batch_distance_avx512(const void **candidates_v, const void *lut_v, size_t num, size_t num_chunk, float *out) { +#if defined(__AVX512F__) constexpr int kNumCentroids = 256; constexpr int kChunkSize = 16; constexpr int kBatch = 4; @@ -210,6 +237,13 @@ void pq_adc_int8_batch_distance_avx512(const void **candidates_v, for (; i < num; ++i) { pq_adc_int8_distance_avx512(candidates[i], lut, num_chunk, out + i); } +#else + (void)candidates_v; + (void)lut_v; + (void)num; + (void)num_chunk; + (void)out; +#endif } } // namespace zvec::turbo::avx512 diff --git a/src/turbo/distance/neon/pq_quantizer_int8/pq_distance.cc b/src/turbo/distance/neon/pq_quantizer_int8/pq_distance.cc new file mode 100644 index 000000000..abffbfe37 --- /dev/null +++ b/src/turbo/distance/neon/pq_quantizer_int8/pq_distance.cc @@ -0,0 +1,207 @@ +// 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. + +// NEON lacks hardware gather instructions (like x86 _mm256_i32gather_ps), +// so LUT lookups remain scalar. NEON accelerates the accumulation path: +// 4 floats are loaded into a float32x4_t register and pairwise-added into +// the running accumulator, halving the number of FP add operations. + +#include "neon/pq_quantizer_int8/pq_distance.h" +#if defined(__ARM_NEON) && defined(__aarch64__) +#include +#endif +#include +#include + +namespace zvec::turbo::neon { + +#if defined(__ARM_NEON) && defined(__aarch64__) +namespace { + +// Horizontal sum of 4 floats in a float32x4_t register via pairwise add. +// [a, b, c, d] → [a+b, c+d, a+b, c+d] → [a+b+c+d, ...] +inline float horizontal_sum_neon(float32x4_t v) { + float32x2_t lo = vget_low_f32(v); + float32x2_t hi = vget_high_f32(v); + float32x2_t sum2 = vadd_f32(lo, hi); // [a+b, c+d] + return vget_lane_f32(vpadd_f32(sum2, sum2), 0); // a+b+c+d +} + +} // namespace +#endif + +void pq_adc_int8_distance_neon(const void *pq_code_v, const void *lut_v, + size_t num_chunk, float *out) { +#if defined(__ARM_NEON) && defined(__aarch64__) + constexpr int kNumCentroids = 256; + constexpr int kChunkSize = 4; // NEON processes 4 floats at once + const auto *pq_code = reinterpret_cast(pq_code_v); + const auto *lut = reinterpret_cast(lut_v); + + float32x4_t acc = vdupq_n_f32(0.0f); + + size_t m = 0; + + // Main loop: process 4 subquantizers per iteration. + // Scalar LUT lookups (NEON has no gather), then NEON pairwise accumulation. + for (; m + kChunkSize <= num_chunk; m += kChunkSize) { + float d0 = lut[(m + 0) * kNumCentroids + pq_code[m + 0]]; + float d1 = lut[(m + 1) * kNumCentroids + pq_code[m + 1]]; + float d2 = lut[(m + 2) * kNumCentroids + pq_code[m + 2]]; + float d3 = lut[(m + 3) * kNumCentroids + pq_code[m + 3]]; + float32x4_t d = {d0, d1, d2, d3}; + acc = vaddq_f32(acc, d); + } + + float sum = horizontal_sum_neon(acc); + + // Scalar leftover: process remaining subquantizers + for (; m < num_chunk; ++m) { + sum += lut[m * kNumCentroids + pq_code[m]]; + } + + *out = sum; +#else + (void)pq_code_v; + (void)lut_v; + (void)num_chunk; + (void)out; +#endif +} + +void pq_sdc_int8_distance_neon(const void *a_v, const void *b_v, + const void *dist_table_v, + size_t num_chunk, float *out) { +#if defined(__ARM_NEON) && defined(__aarch64__) + constexpr int kNumCentroids = 256; + constexpr int kTablePerSub = kNumCentroids * kNumCentroids; // 65536 + constexpr int kChunkSize = 4; + const auto *a = reinterpret_cast(a_v); + const auto *b = reinterpret_cast(b_v); + const auto *dist_table = reinterpret_cast(dist_table_v); + + float32x4_t acc = vdupq_n_f32(0.0f); + + size_t m = 0; + + // Main loop: process 4 subquantizers per iteration. + for (; m + kChunkSize <= num_chunk; m += kChunkSize) { + float d0 = dist_table[(m + 0) * kTablePerSub + + static_cast(a[m + 0]) * kNumCentroids + + static_cast(b[m + 0])]; + float d1 = dist_table[(m + 1) * kTablePerSub + + static_cast(a[m + 1]) * kNumCentroids + + static_cast(b[m + 1])]; + float d2 = dist_table[(m + 2) * kTablePerSub + + static_cast(a[m + 2]) * kNumCentroids + + static_cast(b[m + 2])]; + float d3 = dist_table[(m + 3) * kTablePerSub + + static_cast(a[m + 3]) * kNumCentroids + + static_cast(b[m + 3])]; + float32x4_t d = {d0, d1, d2, d3}; + acc = vaddq_f32(acc, d); + } + + float sum = horizontal_sum_neon(acc); + + // Scalar leftover + for (; m < num_chunk; ++m) { + size_t idx = m * kTablePerSub + static_cast(a[m]) * kNumCentroids + + static_cast(b[m]); + sum += dist_table[idx]; + } + + *out = sum; +#else + (void)a_v; + (void)b_v; + (void)dist_table_v; + (void)num_chunk; + (void)out; +#endif +} + +void pq_adc_int8_batch_distance_neon(const void **candidates_v, + const void *lut_v, size_t num, + size_t num_chunk, float *out) { +#if defined(__ARM_NEON) && defined(__aarch64__) + constexpr int kNumCentroids = 256; + constexpr int kChunkSize = 4; + constexpr int kBatch = 4; + const auto *lut = reinterpret_cast(lut_v); + const auto *candidates = + reinterpret_cast(candidates_v); + + 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]; + float32x4_t acc0 = vdupq_n_f32(0.0f); + float32x4_t acc1 = vdupq_n_f32(0.0f); + float32x4_t acc2 = vdupq_n_f32(0.0f); + float32x4_t acc3 = vdupq_n_f32(0.0f); + + size_t m = 0; + for (; m + kChunkSize <= num_chunk; m += kChunkSize) { + const float *tab = lut + m * kNumCentroids; + + float32x4_t d0 = {tab[c0[m + 0]], tab[c0[m + 1]], + tab[c0[m + 2]], tab[c0[m + 3]]}; + float32x4_t d1 = {tab[c1[m + 0]], tab[c1[m + 1]], + tab[c1[m + 2]], tab[c1[m + 3]]}; + float32x4_t d2 = {tab[c2[m + 0]], tab[c2[m + 1]], + tab[c2[m + 2]], tab[c2[m + 3]]}; + float32x4_t d3 = {tab[c3[m + 0]], tab[c3[m + 1]], + tab[c3[m + 2]], tab[c3[m + 3]]}; + + acc0 = vaddq_f32(acc0, d0); + acc1 = vaddq_f32(acc1, d1); + acc2 = vaddq_f32(acc2, d2); + acc3 = vaddq_f32(acc3, d3); + } + + float s0 = horizontal_sum_neon(acc0); + float s1 = horizontal_sum_neon(acc1); + float s2 = horizontal_sum_neon(acc2); + float s3 = horizontal_sum_neon(acc3); + + // Scalar leftover for remaining subquantizers. + for (; m < num_chunk; ++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_neon(candidates[i], lut, num_chunk, out + i); + } +#else + (void)candidates_v; + (void)lut_v; + (void)num; + (void)num_chunk; + (void)out; +#endif +} + +} // namespace zvec::turbo::neon diff --git a/src/turbo/distance/neon/pq_quantizer_int8/pq_distance.h b/src/turbo/distance/neon/pq_quantizer_int8/pq_distance.h new file mode 100644 index 000000000..b1be2cf2e --- /dev/null +++ b/src/turbo/distance/neon/pq_quantizer_int8/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::neon { + +// ADC (Asymmetric Distance Computation) using NEON vector accumulation. +// LUT lookups are scalar (NEON lacks hardware gather); accumulation uses +// float32x4_t with pairwise horizontal reduction. +void pq_adc_int8_distance_neon(const void *pq_code, const void *lut, + size_t num_chunk, float *out); + +// SDC (Symmetric Distance Computation) via scalar lookup. +// The dist_table is too large (65536 floats per subquantizer) for NEON +// table-lookup instructions, so the implementation is scalar. +void pq_sdc_int8_distance_neon(const void *a, const void *b, + const void *dist_table, size_t num_chunk, + float *out); + +// Batch ADC: compute distances for multiple PQ codes against a shared LUT. +// Processes 4 candidates per iteration with NEON vector accumulation. +void pq_adc_int8_batch_distance_neon(const void **candidates, const void *lut, + size_t num, size_t num_chunk, + float *out); + +} // namespace zvec::turbo::neon diff --git a/src/turbo/turbo.cc b/src/turbo/turbo.cc index 7666fd17c..f27ad8716 100644 --- a/src/turbo/turbo.cc +++ b/src/turbo/turbo.cc @@ -15,16 +15,20 @@ #include #include #include +#include "avx2/pq_quantizer_int8/pq_distance.h" #include "avx2/rotate/fht/fht.h" +#include "avx512/pq_quantizer_int8/pq_distance.h" #include "avx512/rotate/fht/fht.h" #include "avx512_vnni/record_quantized_int8/cosine.h" #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 "neon/pq_quantizer_int8/pq_distance.h" #include "neon/rotate/fht/fht.h" #include "scalar/fp32/cosine.h" #include "scalar/fp32/inner_product.h" #include "scalar/fp32/squared_euclidean.h" +#include "scalar/pq_quantizer_int8/pq_distance.h" #include "scalar/rotate/fht/fht.h" #include "sse/rotate/fht/fht.h" @@ -184,36 +188,30 @@ RotatorKernels get_rotator_kernels(RotateType rotate_type, PqKernels get_pq_kernels(DataType data_type, QuantizeType quantize_type, CpuArchType cpu_arch_type) { - // Suppress unused-parameter warning when no SIMD #if blocks are compiled in. - (void)cpu_arch_type; - PqKernels k{}; + (void)data_type; if (quantize_type == QuantizeType::kPQ) { - // Default (kInt8): scalar int8 kernels. - k.adc_distance = scalar::pq_adc_int8_distance; - k.sdc_distance = scalar::pq_sdc_int8_distance; - k.batch_adc_distance = scalar::pq_adc_int8_batch_distance; - -#if defined(__AVX512F__) if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F && - (cpu_arch_type == CpuArchType::kAuto || - cpu_arch_type == CpuArchType::kAVX512)) { - 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; + IsArchMatch(cpu_arch_type, CpuArchType::kAVX512)) { + return {avx512::pq_adc_int8_distance_avx512, + avx512::pq_sdc_int8_distance_avx512, + avx512::pq_adc_int8_batch_distance_avx512}; } -#endif -#if defined(__AVX2__) if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2 && - (cpu_arch_type == CpuArchType::kAuto || - cpu_arch_type == CpuArchType::kAVX2)) { - 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; + IsArchMatch(cpu_arch_type, CpuArchType::kAVX2)) { + return {avx2::pq_adc_int8_distance_avx2, + avx2::pq_sdc_int8_distance_avx2, + avx2::pq_adc_int8_batch_distance_avx2}; + } + if (IsArchMatch(cpu_arch_type, CpuArchType::kNEON)) { + return {neon::pq_adc_int8_distance_neon, + neon::pq_sdc_int8_distance_neon, + neon::pq_adc_int8_batch_distance_neon}; } -#endif + return {scalar::pq_adc_int8_distance, + scalar::pq_sdc_int8_distance, + scalar::pq_adc_int8_batch_distance}; } - return k; + return {}; } } // namespace zvec::turbo \ No newline at end of file From d4cce4cd73f3c346befe94241191d837b5274021 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Mon, 20 Jul 2026 15:46:15 +0800 Subject: [PATCH 073/100] debug --- .../distance/avx2/pq_quantizer_int8/pq_distance.cc | 12 +++++++++--- .../distance/avx512/pq_quantizer_int8/pq_distance.cc | 12 +++++++++--- 2 files changed, 18 insertions(+), 6 deletions(-) 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 d9fbc26c0..e9e762c00 100644 --- a/src/turbo/distance/avx2/pq_quantizer_int8/pq_distance.cc +++ b/src/turbo/distance/avx2/pq_quantizer_int8/pq_distance.cc @@ -130,13 +130,19 @@ void pq_sdc_int8_distance_avx2(const void *a_v, const void *b_v, __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 + // Compute in-lane index: a[m] * 256 + b[m] + k * 65536 (k = lane, 0..7). + // The m * 65536 offset is applied via the gather base pointer below. __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); + // Gather 8 floats from dist_table. The gather base must include the + // per-iteration m * kTablePerSub offset; base_offsets only carries the + // in-lane k * kTablePerSub component (k = 0..7), so gathering from a + // fixed dist_table base would read the wrong subquantizer tables once + // num_chunk > 8 (m >= 8). + __m256 gathered = + _mm256_i32gather_ps(dist_table + m * kTablePerSub, indices, 4); acc = _mm256_add_ps(acc, gathered); } 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 5ece65e13..79ee4c9dc 100644 --- a/src/turbo/distance/avx512/pq_quantizer_int8/pq_distance.cc +++ b/src/turbo/distance/avx512/pq_quantizer_int8/pq_distance.cc @@ -126,13 +126,19 @@ void pq_sdc_int8_distance_avx512(const void *a_v, const void *b_v, __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 + // Compute in-lane index: a[m] * 256 + b[m] + k * 65536 (k = lane, 0..15). + // The m * 65536 offset is applied via the gather base pointer below. __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); + // Gather 16 floats from dist_table. The gather base must include the + // per-iteration m * kTablePerSub offset; base_offsets only carries the + // in-lane k * kTablePerSub component (k = 0..15), so gathering from a + // fixed dist_table base would read the wrong subquantizer tables once + // num_chunk > 16 (m >= 16). + __m512 gathered = + _mm512_i32gather_ps(indices, dist_table + m * kTablePerSub, 4); acc = _mm512_add_ps(acc, gathered); } From de5fa1431fad556167c4daefb922a481c0879340 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Mon, 20 Jul 2026 17:37:20 +0800 Subject: [PATCH 074/100] add use_zero_mean --- .../pq_int8_quantizer/pq_int8_quantizer.cc | 114 ++++++-- .../pq_int8_quantizer/pq_int8_quantizer.h | 12 +- tests/turbo/turbo_pq_int8_quantizer_test.cc | 275 +++++++++++++++++- 3 files changed, 373 insertions(+), 28 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 1cddd5d91..14165a95a 100644 --- a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc @@ -37,6 +37,8 @@ struct PqInt8SerPayload { uint32_t num_chunk; uint32_t sub_dim; uint32_t num_centroids; // always 256 for int8 + uint8_t use_zero_mean; + uint8_t reserved[3]; }; int PqInt8Quantizer::init(const IndexMeta &meta, const ailego::Params ¶ms) { @@ -52,9 +54,8 @@ int PqInt8Quantizer::init(const IndexMeta &meta, const ailego::Params ¶ms) { return kErrUnsupported; } if (d % nsq != 0) { - LOG_ERROR( - "PqInt8Quantizer: dim (%u) is not divisible by num_chunk (%u)", - d, nsq); + LOG_ERROR("PqInt8Quantizer: dim (%u) is not divisible by num_chunk (%u)", d, + nsq); return kErrUnsupported; } @@ -62,8 +63,8 @@ int PqInt8Quantizer::init(const IndexMeta &meta, const ailego::Params ¶ms) { sub_dim_ = d / nsq; // Pre-allocate centroids (filled by train()). - centroids_.resize( - static_cast(num_chunk_) * kNumCentroids * sub_dim_, 0.0f); + centroids_.resize(static_cast(num_chunk_) * kNumCentroids * sub_dim_, + 0.0f); // Dispatch ISA kernels (scalar only for now). auto pq_k = get_pq_kernels(DataType::kInt8); @@ -108,6 +109,7 @@ int PqInt8Quantizer::init(const IndexMeta &meta, const ailego::Params ¶ms) { params.get("thread_count", &thread_count_); params.get("markov_chain_length", &markov_chain_length_); params.get("epsilon", &epsilon_); + params.get("use_zero_mean", &use_zero_mean_); meta_.set_meta(IndexMeta::DataType::DT_FP32, d); return 0; @@ -145,8 +147,8 @@ void PqInt8Quantizer::train_subquantizer(const float *data, size_t num, // normalized, so spherical would create a mismatch between assignment // (norm-dependent L2) and update (norm-discarding normalization). const bool spherical = (meta_.metric_name() == "Cosine"); - ailego::NumericalKmeans algorithm( - k, d, spherical); + ailego::NumericalKmeans algorithm(k, d, + spherical); // Append sub-vectors (NumericalKmeans handles transpose internally) for (size_t i = 0; i < num; ++i) { @@ -250,11 +252,35 @@ int PqInt8Quantizer::train(IndexHolder::Pointer holder, int thread_count) { } } + // Zero-mean centering: compute per-dimension mean of training data and + // subtract it from all vectors. The centroid is saved for later use + // in quantize_data / quantize_query / dequantize. + // Reference: DiskANN PQ trainer prepare_pq_train_data. + if (use_zero_mean_) { + centroid_.assign(original_dim_, 0.0f); + for (size_t i = 0; i < num; ++i) { + const float *v = all_data.data() + i * original_dim_; + for (uint32_t d = 0; d < original_dim_; ++d) { + centroid_[d] += v[d]; + } + } + for (uint32_t d = 0; d < original_dim_; ++d) { + centroid_[d] /= static_cast(num); + } + // Subtract centroid from all training vectors. + for (size_t i = 0; i < num; ++i) { + float *v = all_data.data() + i * original_dim_; + for (uint32_t d = 0; d < original_dim_; ++d) { + v[d] -= centroid_[d]; + } + } + LOG_INFO("PQ centering: zero-mean centroid computed"); + } + LOG_INFO( "PQ training: %zu vectors, dim=%u, nsq=%u, sub_dim=%u, " "max_iters=%u, threads=%d", - num, original_dim_, num_chunk_, sub_dim_, kMaxKmeansIters, - thread_count); + num, original_dim_, num_chunk_, sub_dim_, kMaxKmeansIters, thread_count); // Create thread pool (aligned with multi_chunk_cluster L183-185) auto threads = std::make_shared( @@ -267,8 +293,7 @@ int PqInt8Quantizer::train(IndexHolder::Pointer holder, int thread_count) { for (size_t i = 0; i < pool_count; ++i) { task_group->submit(ailego::Closure::New( [this, &all_data, num, data_stride, i, pool_count, &finished]() { - for (uint32_t m = static_cast(i); - m < num_chunk_; + for (uint32_t m = static_cast(i); m < num_chunk_; m += static_cast(pool_count)) { train_subquantizer(all_data.data(), num, data_stride, m); finished++; @@ -277,8 +302,8 @@ int PqInt8Quantizer::train(IndexHolder::Pointer holder, int thread_count) { } task_group->wait_finish(); - LOG_INFO(" all %u sub-quantizers trained (%zu threads)", - num_chunk_, pool_count); + LOG_INFO(" all %u sub-quantizers trained (%zu threads)", num_chunk_, + pool_count); // Pre-build centroid pointer cache (needed by compute_dist_table). build_centroid_ptrs_cache(); @@ -321,6 +346,16 @@ void PqInt8Quantizer::quantize_data(const void *input, void *output) const { const float *vec = reinterpret_cast(input); uint8_t *code = reinterpret_cast(output); + // Zero-mean centering: subtract centroid before encoding. + std::vector centered_vec_storage; + if (use_zero_mean_) { + centered_vec_storage.assign(vec, vec + original_dim_); + for (uint32_t d = 0; d < original_dim_; ++d) { + centered_vec_storage[d] -= centroid_[d]; + } + vec = centered_vec_storage.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(). @@ -371,6 +406,16 @@ void PqInt8Quantizer::quantize_query(const void *input, void *output) const { const float *query = reinterpret_cast(input); float *lut = reinterpret_cast(output); + // Zero-mean centering: subtract centroid before LUT computation. + std::vector centered_query_storage; + if (use_zero_mean_) { + centered_query_storage.assign(query, query + original_dim_); + for (uint32_t d = 0; d < original_dim_; ++d) { + centered_query_storage[d] -= centroid_[d]; + } + query = centered_query_storage.data(); + } + // 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. @@ -433,12 +478,10 @@ void PqInt8Quantizer::calc_distance_dp_query_batch(const void *const *dp_list, 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_chunk_) * - kNumCentroids); + std::vector lut(static_cast(num_chunk_) * kNumCentroids); quantize_query(query, lut.data()); float d = 0.0f; - adc_fn_(reinterpret_cast(dp), lut.data(), num_chunk_, - &d); + adc_fn_(reinterpret_cast(dp), lut.data(), num_chunk_, &d); if (meta_.metric_name() == "Cosine") { d = 1.0f + d; } @@ -448,8 +491,7 @@ float PqInt8Quantizer::calc_distance_dp_query_unquantized( 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_chunk_) * - kNumCentroids); + std::vector lut(static_cast(num_chunk_) * kNumCentroids); quantize_query(query, lut.data()); // Use ISA-dispatched batch4 ADC kernel (4-way ILP + SIMD gather). // const_cast: see calc_distance_dp_query_batch for rationale. @@ -516,6 +558,13 @@ int PqInt8Quantizer::dequantize(const void *in, const IndexQueryMeta &qmeta, result[j] *= norm; } } + + // Zero-mean centering: add centroid back to reconstruct original space. + if (use_zero_mean_) { + for (uint32_t j = 0; j < original_dim_; ++j) { + result[j] += centroid_[j]; + } + } return 0; } @@ -536,8 +585,7 @@ DistanceImpl PqInt8Quantizer::distance(const void *query, std::string lut_storage(static_cast(query), lut_bytes); return DistanceImpl(std::move(adc_func), std::move(batch_func), - std::move(lut_storage), - static_cast(num_chunk_)); + std::move(lut_storage), static_cast(num_chunk_)); } DistanceImpl PqInt8Quantizer::sym_distance(const void *query, @@ -580,15 +628,23 @@ int PqInt8Quantizer::serialize(std::string *out) const { payload.num_chunk = num_chunk_; payload.sub_dim = sub_dim_; payload.num_centroids = kNumCentroids; + payload.use_zero_mean = use_zero_mean_ ? 1 : 0; size_t centroids_bytes = centroids_.size() * sizeof(float); - hdr.payload_size = static_cast(sizeof(payload) + centroids_bytes); + size_t centroid_bytes = use_zero_mean_ ? centroid_.size() * sizeof(float) : 0; + hdr.payload_size = + static_cast(sizeof(payload) + centroids_bytes + centroid_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); + // Append zero-mean centroid for centering support. + if (use_zero_mean_) { + out->append(reinterpret_cast(centroid_.data()), + centroid_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 @@ -622,11 +678,21 @@ 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_chunk_) * - kNumCentroids * sub_dim_ * sizeof(float); + size_t centroids_bytes = static_cast(num_chunk_) * kNumCentroids * + sub_dim_ * sizeof(float); centroids_.resize(centroids_bytes / sizeof(float)); std::memcpy(centroids_.data(), ptr, centroids_bytes); + ptr += centroids_bytes; + + // Restore zero-mean centroid if centering was enabled. + if (payload.use_zero_mean) { + use_zero_mean_ = true; + size_t centroid_bytes = static_cast(original_dim_) * sizeof(float); + centroid_.resize(original_dim_); + std::memcpy(centroid_.data(), ptr, centroid_bytes); + ptr += centroid_bytes; + } // dist_table_ is intentionally not restored: SDC is only needed during // offline build, not after deserialization (search uses ADC). 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 1d4250bee..e772fb6e3 100644 --- a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h @@ -72,8 +72,7 @@ class PqInt8Quantizer : public Quantizer { } size_t quantized_query_vector_length() const override { - return static_cast(num_chunk_) * kNumCentroids * - sizeof(float); + return static_cast(num_chunk_) * kNumCentroids * sizeof(float); } void quantize_data(const void *input, void *output) const override; @@ -142,6 +141,11 @@ class PqInt8Quantizer : public Quantizer { //! Cost-based convergence threshold (aligned with multi_chunk_cluster). double epsilon_{std::numeric_limits::epsilon()}; + //! Whether to apply zero-mean centering before training/encoding. + //! When enabled, the per-dimension mean of training data is subtracted + //! from all vectors (train, encode, query) and added back on dequantize. + bool use_zero_mean_{false}; + IndexMeta meta_{}; uint32_t original_dim_{0}; uint32_t num_chunk_{0}; @@ -150,6 +154,10 @@ class PqInt8Quantizer : public Quantizer { //! Centroids: [num_chunk * kNumCentroids * sub_dim] std::vector centroids_; + //! Global centroid (per-dimension mean) for zero-mean centering. + //! Size: original_dim_ floats. Only populated when use_centering_ = true. + std::vector centroid_; + //! Centroid-to-centroid distance table for SDC: //! [num_chunk * kNumCentroids * kNumCentroids] std::vector dist_table_; diff --git a/tests/turbo/turbo_pq_int8_quantizer_test.cc b/tests/turbo/turbo_pq_int8_quantizer_test.cc index c6569f179..754cf3fce 100644 --- a/tests/turbo/turbo_pq_int8_quantizer_test.cc +++ b/tests/turbo/turbo_pq_int8_quantizer_test.cc @@ -311,8 +311,7 @@ void fill_random_lut(float *lut, size_t num_chunk, std::mt19937 &gen) { } // Helper to generate random dist_table (SDC) -void fill_random_sdc_table(float *table, size_t num_chunk, - std::mt19937 &gen) { +void fill_random_sdc_table(float *table, size_t num_chunk, 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_chunk; ++m) { @@ -804,3 +803,275 @@ TEST(PqInt8Quantizer, InnerProductDequantize) { << "recon_err=" << recon_err << " orig_norm=" << orig_norm; } } + +// --------------------------------------------------------------------------- +// Zero-Mean Centering Tests +// --------------------------------------------------------------------------- + +// Helper to create a PqInt8Quantizer with zero-mean centering enabled. +static std::shared_ptr make_pq_zero_mean_quantizer( + size_t dim, size_t num_chunk) { + 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_chunk", static_cast(num_chunk)); + params.set("use_zero_mean", true); + if (q->init(meta, params) != 0) return nullptr; + return q; +} + +// Helper: build a holder with random fp32 vectors that have a large offset +// (non-zero mean). This simulates real-world data where centering helps. +static std::shared_ptr> +make_offset_holder(size_t count, size_t dim, float offset = 10.0f, + 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) + offset; + holder->emplace(i + 1, vec); + } + return holder; +} + +// Verify basic functionality: train, encode, ADC distance with centering. +TEST(PqInt8Quantizer, ZeroMeanTrainAndEncode) { + const size_t DIM = 16; + const size_t NSQ = 4; + const size_t COUNT = 1000; + + auto quantizer = make_pq_zero_mean_quantizer(DIM, NSQ); + ASSERT_TRUE(quantizer); + EXPECT_TRUE(quantizer->require_train()); + + // Use offset data to exercise the centering path meaningfully. + auto holder = make_offset_holder(COUNT, DIM, 10.0f); + 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()); + for (size_t m = 0; m < NSQ; ++m) { + EXPECT_LE(code[m], 255u); + } + } + EXPECT_EQ(10u, checked); +} + +// Verify ADC distance accuracy with centering on offset data. +TEST(PqInt8Quantizer, ZeroMeanAdcDistance) { + const size_t DIM = 32; + const size_t NSQ = 8; + const size_t COUNT = 2000; + + auto quantizer = make_pq_zero_mean_quantizer(DIM, NSQ); + ASSERT_TRUE(quantizer); + + // Offset data: centering should improve PQ accuracy for high-offset vectors. + auto holder = make_offset_holder(COUNT, DIM, 10.0f); + 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 approximate true L2 distances. + 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; + } + EXPECT_LT(max_rel_error, 1.0f) << "max_rel_error=" << max_rel_error; +} + +// Verify dequantize correctly adds centroid back. +TEST(PqInt8Quantizer, ZeroMeanDequantize) { + const size_t DIM = 16; + const size_t NSQ = 4; + const size_t COUNT = 1000; + const float OFFSET = 10.0f; + + auto quantizer = make_pq_zero_mean_quantizer(DIM, NSQ); + ASSERT_TRUE(quantizer); + + auto holder = make_offset_holder(COUNT, DIM, OFFSET); + 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()); + + // Reconstructed vector should be near the original (which has ~10 offset). + 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; + } + + // The reconstructed values should be around OFFSET (not near zero), + // confirming that the centroid was added back. + float recon_mean = 0.0f; + for (size_t j = 0; j < DIM; ++j) recon_mean += recon[j]; + recon_mean /= DIM; + EXPECT_GT(recon_mean, OFFSET * 0.5f) + << "Reconstructed mean too low; centroid may not be added back"; +} + +// Verify serialize/deserialize preserves the zero-mean centroid. +TEST(PqInt8Quantizer, ZeroMeanSerializeDeserialize) { + const size_t DIM = 16; + const size_t NSQ = 4; + const size_t COUNT = 500; + + auto quantizer = make_pq_zero_mean_quantizer(DIM, NSQ); + ASSERT_TRUE(quantizer); + + auto holder = make_offset_holder(COUNT, DIM, 5.0f); + ASSERT_EQ(0, quantizer->train(holder)); + + // Serialize. + std::string blob; + ASSERT_EQ(0, quantizer->serialize(&blob)); + EXPECT_GT(blob.size(), sizeof(zvec::turbo::QuantizerSerHeader)); + + // Deserialize into a fresh quantizer. + auto q2 = IndexFactory::CreateQuantizer("PqInt8Quantizer"); + ASSERT_TRUE(q2); + ASSERT_EQ(0, q2->deserialize(blob)); + + // Encode the same vector with both and compare codes. + auto iter = holder->create_iterator(); + iter->is_valid(); + std::vector code1(quantizer->quantized_datapoint_vector_length()); + std::vector code2(q2->quantized_datapoint_vector_length()); + quantizer->quantize_data(iter->data(), code1.data()); + q2->quantize_data(iter->data(), code2.data()); + + for (size_t m = 0; m < NSQ; ++m) { + EXPECT_EQ(code1[m], code2[m]) << "m=" << m; + } + + // ADC distances should also match. + 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); + + // Dequantize from q2 should also produce vectors in the offset range. + IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, DIM); + std::string decoded; + ASSERT_EQ(0, q2->dequantize(code2.data(), qmeta, &decoded)); + const float *recon = reinterpret_cast(decoded.data()); + float recon_mean = 0.0f; + for (size_t j = 0; j < DIM; ++j) recon_mean += recon[j]; + recon_mean /= DIM; + EXPECT_GT(recon_mean, 2.5f) + << "Deserialized quantizer centroid not restored properly"; +} + +// Verify that centering does not significantly degrade PQ accuracy. +// Centering benefits data with skewed distributions or high-mean non-uniform +// spread; for uniform+offset data, the error should remain comparable. +TEST(PqInt8Quantizer, ZeroMeanAccuracyComparable) { + const size_t DIM = 32; + const size_t NSQ = 8; + const size_t COUNT = 2000; + const float OFFSET = 50.0f; + + // Train without centering. + auto q_no_center = make_pq_quantizer(DIM, NSQ); + ASSERT_TRUE(q_no_center); + auto holder = make_offset_holder(COUNT, DIM, OFFSET); + ASSERT_EQ(0, q_no_center->train(holder)); + + // Train with centering. + auto q_center = make_pq_zero_mean_quantizer(DIM, NSQ); + ASSERT_TRUE(q_center); + ASSERT_EQ(0, q_center->train(holder)); + + // Compute average reconstruction error for both. + auto iter = holder->create_iterator(); + float err_no_center_sum = 0.0f; + float err_center_sum = 0.0f; + size_t checked = 0; + + size_t code_len_no = q_no_center->quantized_datapoint_vector_length(); + size_t code_len_yes = q_center->quantized_datapoint_vector_length(); + IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, DIM); + + for (; iter->is_valid() && checked < 100; iter->next(), ++checked) { + const float *orig = reinterpret_cast(iter->data()); + + // Without centering. + std::vector code1(code_len_no); + q_no_center->quantize_data(iter->data(), code1.data()); + std::string decoded1; + q_no_center->dequantize(code1.data(), qmeta, &decoded1); + err_no_center_sum += reference_sq_euclidean( + orig, reinterpret_cast(decoded1.data()), DIM); + + // With centering. + std::vector code2(code_len_yes); + q_center->quantize_data(iter->data(), code2.data()); + std::string decoded2; + q_center->dequantize(code2.data(), qmeta, &decoded2); + err_center_sum += reference_sq_euclidean( + orig, reinterpret_cast(decoded2.data()), DIM); + } + + float avg_err_no_center = err_no_center_sum / checked; + float avg_err_center = err_center_sum / checked; + + // Centering should not degrade accuracy by more than 2x. + EXPECT_LT(avg_err_center, avg_err_no_center * 2.0f) + << "Centering degraded accuracy too much: center_err=" << avg_err_center + << " no_center_err=" << avg_err_no_center; +} From fed115efccc58edf804fc8a782bc05591a28c817 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Mon, 20 Jul 2026 18:06:49 +0800 Subject: [PATCH 075/100] :tmp: --- .../pq_int8_quantizer/pq_int8_quantizer.cc | 14 +++----------- 1 file changed, 3 insertions(+), 11 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 14165a95a..04cd3d48f 100644 --- a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc @@ -83,12 +83,6 @@ int PqInt8Quantizer::init(const IndexMeta &meta, const ailego::Params ¶ms) { get_batch_distance_func(MetricType::kSquaredEuclidean, DataType::kFp32, QuantizeType::kDefault, CpuArchType::kAuto); - // Metric-aware batch distance: used for search-side LUT computation - // (quantize_query) and SDC dist_table. For IP/Cosine this computes - // inner-product-based distances instead of L2. - fp32_batch_fn_ = get_batch_distance_func( - mt, DataType::kFp32, QuantizeType::kDefault, CpuArchType::kAuto); - // For Cosine: fall back to IP for the metric-aware batch function, // because cosine = normalize + IP. The normalization is applied // explicitly in quantize_query. @@ -96,13 +90,11 @@ int PqInt8Quantizer::init(const IndexMeta &meta, const ailego::Params ¶ms) { 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_); + } else { + fp32_batch_fn_ = get_batch_distance_func( + mt, DataType::kFp32, QuantizeType::kDefault, CpuArchType::kAuto); } // Read optional training params (aligned with multi_chunk_cluster) From 8bd8068d7123b3624cb8e0fcec03d8cb51a5a776 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Mon, 20 Jul 2026 19:25:10 +0800 Subject: [PATCH 076/100] tmp --- .../quantizer/pq_int8_quantizer/pq_int8_quantizer.cc | 12 +++++++----- 1 file changed, 7 insertions(+), 5 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 04cd3d48f..c1a785dca 100644 --- a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc @@ -699,15 +699,17 @@ int PqInt8Quantizer::deserialize(const void *data, size_t len) { 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); + // Metric-aware batch distance for search LUT. if (meta_.metric_name() == "Cosine") { fp32_batch_fn_ = get_batch_distance_func(MetricType::kInnerProduct, DataType::kFp32, QuantizeType::kDefault, CpuArchType::kAuto); + extra_meta_size_ = kExtraMetaSizeCosine; + meta_.set_extra_meta_size(extra_meta_size_); + } else { + fp32_batch_fn_ = get_batch_distance_func( + metric_from_name(meta_.metric_name()), DataType::kFp32, + QuantizeType::kDefault, CpuArchType::kAuto); } // Pre-build centroid pointer cache for fast encode/search. From bed708cc738e132a68aaed780c419a7993bbc8b9 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Mon, 20 Jul 2026 20:30:36 +0800 Subject: [PATCH 077/100] tmp --- .../quantizer/pq_int8_quantizer/pq_int8_quantizer.cc | 9 +++++++++ 1 file changed, 9 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 c1a785dca..897706cc8 100644 --- a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc @@ -103,6 +103,15 @@ int PqInt8Quantizer::init(const IndexMeta &meta, const ailego::Params ¶ms) { params.get("epsilon", &epsilon_); params.get("use_zero_mean", &use_zero_mean_); + // Zero-mean centering is only mathematically valid for SquaredEuclidean + // (translation-invariant). IP, Cosine, MipsSquaredEuclidean all break. + if (use_zero_mean_ && meta_.metric_name() != "SquaredEuclidean") { + LOG_WARN("PqInt8Quantizer: use_zero_mean is incompatible with metric '%s', " + "disabling centering", + meta_.metric_name().c_str()); + use_zero_mean_ = false; + } + meta_.set_meta(IndexMeta::DataType::DT_FP32, d); return 0; } From 2ba6a597c86b53c99f11fe04f6fbd76bf58fbb9c Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Tue, 21 Jul 2026 23:15:49 +0800 Subject: [PATCH 078/100] remove spherical --- .../pq_int8_quantizer/pq_int8_quantizer.cc | 63 ++++++++----------- 1 file changed, 25 insertions(+), 38 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 897706cc8..d61f15706 100644 --- a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc @@ -83,12 +83,14 @@ int PqInt8Quantizer::init(const IndexMeta &meta, const ailego::Params ¶ms) { get_batch_distance_func(MetricType::kSquaredEuclidean, 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. + // For Cosine: cosine = normalize + L2. Vectors are L2-normalized in + // train() / quantize_data() / quantize_query(); after normalization + // cosine distance is monotonic with squared-Euclidean distance, so the + // search LUT uses the same SquaredEuclidean batch function as encoding + // and training — no IP anywhere. if (meta_.metric_name() == "Cosine") { fp32_batch_fn_ = - get_batch_distance_func(MetricType::kInnerProduct, DataType::kFp32, + get_batch_distance_func(MetricType::kSquaredEuclidean, DataType::kFp32, QuantizeType::kDefault, CpuArchType::kAuto); extra_meta_size_ = kExtraMetaSizeCosine; meta_.set_extra_meta_size(extra_meta_size_); @@ -142,14 +144,12 @@ void PqInt8Quantizer::train_subquantizer(const float *data, size_t num, float *centroids_m = centroids_.data() + static_cast(sub_idx) * k * d; // Create KMeans algorithm (L2 distance via NumericalKmeansContext). - // Enable spherical mode for Cosine only. Cosine training data is - // pre-normalized in train(), so spherical centroid constraint keeps - // L2 assignment equivalent to argmax dot(x,c). For IP, data is NOT - // normalized, so spherical would create a mismatch between assignment - // (norm-dependent L2) and update (norm-discarding normalization). - const bool spherical = (meta_.metric_name() == "Cosine"); - ailego::NumericalKmeans algorithm(k, d, - spherical); + // Always non-spherical: the PQ codebook must minimize L2 reconstruction + // error, so centroids are the true (magnitude-preserving) means. Cosine + // data is pre-normalized in train(), but a sub-vector of a unit vector is + // NOT itself unit norm, so forcing unit-norm centroids (spherical) would + // discard per-subspace magnitude and hurt reconstruction. + ailego::NumericalKmeans algorithm(k, d); // Append sub-vectors (NumericalKmeans handles transpose internally) for (size_t i = 0; i < num; ++i) { @@ -326,7 +326,7 @@ void PqInt8Quantizer::compute_dist_table() { // 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. + // Cosine: centroids trained on normalized data, uses L2 (same as L2). for (uint32_t m = 0; m < num_chunk_; ++m) { const float *centroids_m = centroids_.data() + m * k * d; float *table_m = dist_table_.data() + m * k * k; @@ -417,10 +417,10 @@ void PqInt8Quantizer::quantize_query(const void *input, void *output) const { query = centered_query_storage.data(); } - // 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. + // For Cosine: normalize the query to unit length first. Cosine now uses + // an L2 LUT (see init), so after normalization the per-subspace squared + // L2 distances are additive into the full squared L2 distance, which is + // monotonic with cosine distance. std::vector norm_query_storage; if (meta_.metric_name() == "Cosine") { norm_query_storage.assign(query, query + original_dim_); @@ -434,7 +434,7 @@ void PqInt8Quantizer::quantize_query(const void *input, void *output) const { // 256 centroids via the metric-aware fp32 batch distance function. // For L2: LUT[m][j] = ||q_m - c_m[j]||^2 // For IP: LUT[m][j] = -dot(q_m, c_m[j]) - // For Cosine: (query is normalized) LUT[m][j] = -dot(q_norm_m, c_m[j]) + // For Cosine: (query is normalized) LUT[m][j] = ||q_norm_m - c_m[j]||^2 // const_cast: see compute_dist_table for rationale. for (uint32_t m = 0; m < num_chunk_; ++m) { const auto ¢roid_ptrs = centroid_ptrs_cache_[m]; @@ -451,10 +451,8 @@ float PqInt8Quantizer::calc_distance_dp_query(const void *dp, float d = 0.0f; adc_fn_(reinterpret_cast(dp), reinterpret_cast(query), num_chunk_, &d); - // For Cosine: ADC sum = -cos_sim; convert to cosine distance = 1 - cos_sim. - if (meta_.metric_name() == "Cosine") { - d = 1.0f + d; - } + // Cosine uses an L2 LUT on normalized vectors, so the ADC sum is already + // a valid (squared-Euclidean) distance — no conversion needed. return d; } @@ -468,12 +466,8 @@ void PqInt8Quantizer::calc_distance_dp_query_batch(const void *const *dp_list, // on the pointer array. batch_adc_fn_(const_cast(dp_list), query, static_cast(dp_num), num_chunk_, 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]; - } - } + // Cosine uses an L2 LUT on normalized vectors — the batch ADC sums are + // already valid distances, so no conversion is applied. } float PqInt8Quantizer::calc_distance_dp_query_unquantized( @@ -483,9 +477,6 @@ float PqInt8Quantizer::calc_distance_dp_query_unquantized( quantize_query(query, lut.data()); float d = 0.0f; adc_fn_(reinterpret_cast(dp), lut.data(), num_chunk_, &d); - if (meta_.metric_name() == "Cosine") { - d = 1.0f + d; - } return d; } @@ -498,11 +489,6 @@ void PqInt8Quantizer::calc_distance_dp_query_batch_unquantized( // const_cast: see calc_distance_dp_query_batch for rationale. batch_adc_fn_(const_cast(dp_list), lut.data(), static_cast(dp_num), num_chunk_, 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, @@ -708,10 +694,11 @@ int PqInt8Quantizer::deserialize(const void *data, size_t len) { get_batch_distance_func(MetricType::kSquaredEuclidean, DataType::kFp32, QuantizeType::kDefault, CpuArchType::kAuto); - // Metric-aware batch distance for search LUT. + // Metric-aware batch distance for search LUT. Cosine = normalize + L2, + // so it uses SquaredEuclidean (same as encoding), not IP. if (meta_.metric_name() == "Cosine") { fp32_batch_fn_ = - get_batch_distance_func(MetricType::kInnerProduct, DataType::kFp32, + get_batch_distance_func(MetricType::kSquaredEuclidean, DataType::kFp32, QuantizeType::kDefault, CpuArchType::kAuto); extra_meta_size_ = kExtraMetaSizeCosine; meta_.set_extra_meta_size(extra_meta_size_); From ca529ddd52750a835ba50fa372cb6b1322569848 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Tue, 21 Jul 2026 23:38:34 +0800 Subject: [PATCH 079/100] allow cosine --- .../pq_int8_quantizer/pq_int8_quantizer.cc | 90 +++++++++++-------- 1 file changed, 51 insertions(+), 39 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 d61f15706..ccd3b6dc2 100644 --- a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc @@ -105,9 +105,14 @@ int PqInt8Quantizer::init(const IndexMeta &meta, const ailego::Params ¶ms) { params.get("epsilon", &epsilon_); params.get("use_zero_mean", &use_zero_mean_); - // Zero-mean centering is only mathematically valid for SquaredEuclidean - // (translation-invariant). IP, Cosine, MipsSquaredEuclidean all break. - if (use_zero_mean_ && meta_.metric_name() != "SquaredEuclidean") { + // Zero-mean centering is valid only where the search metric is + // translation-invariant. SquaredEuclidean qualifies directly. Cosine now + // uses an L2 LUT on L2-normalized vectors, so centering the *normalized* + // data (train/quantize_data/quantize_query all normalize BEFORE centering) + // is also safe and can lower reconstruction error. IP and + // MipsSquaredEuclidean are NOT translation-invariant and are rejected. + if (use_zero_mean_ && meta_.metric_name() != "SquaredEuclidean" && + meta_.metric_name() != "Cosine") { LOG_WARN("PqInt8Quantizer: use_zero_mean is incompatible with metric '%s', " "disabling centering", meta_.metric_name().c_str()); @@ -255,7 +260,10 @@ int PqInt8Quantizer::train(IndexHolder::Pointer holder, int thread_count) { // Zero-mean centering: compute per-dimension mean of training data and // subtract it from all vectors. The centroid is saved for later use - // in quantize_data / quantize_query / dequantize. + // in quantize_data / quantize_query / dequantize. NOTE: for Cosine this + // runs AFTER the normalization above, so centering happens in normalized + // space; quantize_data / quantize_query / dequantize must keep the same + // order (normalize -> center; dequantize: un-center -> rescale). // Reference: DiskANN PQ trainer prepare_pq_train_data. if (use_zero_mean_) { centroid_.assign(original_dim_, 0.0f); @@ -347,19 +355,10 @@ void PqInt8Quantizer::quantize_data(const void *input, void *output) const { const float *vec = reinterpret_cast(input); uint8_t *code = reinterpret_cast(output); - // Zero-mean centering: subtract centroid before encoding. - std::vector centered_vec_storage; - if (use_zero_mean_) { - centered_vec_storage.assign(vec, vec + original_dim_); - for (uint32_t d = 0; d < original_dim_; ++d) { - centered_vec_storage[d] -= centroid_[d]; - } - vec = centered_vec_storage.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(). + // For Cosine: normalize the vector to unit length FIRST, because the + // codebook is trained in normalized space. The original norm is stored + // after the PQ code for later dequantize(). Centering (below) then + // operates on the normalized vector, consistent with train(). std::vector norm_vec_storage; float vec_norm = 0.0f; if (meta_.metric_name() == "Cosine") { @@ -369,6 +368,17 @@ void PqInt8Quantizer::quantize_data(const void *input, void *output) const { vec = norm_vec_storage.data(); } + // Zero-mean centering: subtract centroid (computed in the same normalized + // space during train) before encoding. + std::vector centered_vec_storage; + if (use_zero_mean_) { + centered_vec_storage.assign(vec, vec + original_dim_); + for (uint32_t d = 0; d < original_dim_; ++d) { + centered_vec_storage[d] -= centroid_[d]; + } + vec = centered_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 @@ -407,6 +417,18 @@ 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 (see init: Cosine + // uses an L2 LUT on normalized data). Centering (below) then operates on + // the normalized query, consistent with train() / quantize_data(). + 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(); + } + // Zero-mean centering: subtract centroid before LUT computation. std::vector centered_query_storage; if (use_zero_mean_) { @@ -417,19 +439,6 @@ void PqInt8Quantizer::quantize_query(const void *input, void *output) const { query = centered_query_storage.data(); } - // For Cosine: normalize the query to unit length first. Cosine now uses - // an L2 LUT (see init), so after normalization the per-subspace squared - // L2 distances are additive into the full squared L2 distance, which is - // monotonic with cosine distance. - std::vector norm_query_storage; - if (meta_.metric_name() == "Cosine") { - norm_query_storage.assign(query, query + original_dim_); - float norm = 0.0f; - ailego::Normalizer::L2(norm_query_storage.data(), original_dim_, - &norm); - query = norm_query_storage.data(); - } - // For each sub-quantizer, compute distance from query sub-vector to all // 256 centroids via the metric-aware fp32 batch distance function. // For L2: LUT[m][j] = ||q_m - c_m[j]||^2 @@ -536,20 +545,23 @@ int PqInt8Quantizer::dequantize(const void *in, const IndexQueryMeta &qmeta, 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_chunk_, sizeof(float)); + // Zero-mean centering: add the centroid back FIRST to undo centering, + // which was applied in normalized space (after normalization) during + // encode. This yields the (approximate) normalized vector. + if (use_zero_mean_) { for (uint32_t j = 0; j < original_dim_; ++j) { - result[j] *= norm; + result[j] += centroid_[j]; } } - // Zero-mean centering: add centroid back to reconstruct original space. - if (use_zero_mean_) { + // For Cosine: the codebook was trained on normalized data; the stored norm + // rescales the (now un-centered) unit-space vector back to the original + // vector's magnitude. + if (meta_.metric_name() == "Cosine") { + float norm = 0.0f; + std::memcpy(&norm, code + num_chunk_, sizeof(float)); for (uint32_t j = 0; j < original_dim_; ++j) { - result[j] += centroid_[j]; + result[j] *= norm; } } return 0; From a3695b0c872a9d9bf1ee2c0aaefd407ddad74f43 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Wed, 22 Jul 2026 15:10:12 +0800 Subject: [PATCH 080/100] debug --- src/turbo/preprocessor/fht_rotator/fht_rotator.cc | 1 - 1 file changed, 1 deletion(-) diff --git a/src/turbo/preprocessor/fht_rotator/fht_rotator.cc b/src/turbo/preprocessor/fht_rotator/fht_rotator.cc index 119b4a7ac..6a60bbe2b 100644 --- a/src/turbo/preprocessor/fht_rotator/fht_rotator.cc +++ b/src/turbo/preprocessor/fht_rotator/fht_rotator.cc @@ -17,7 +17,6 @@ #include #include #include -#include "quantizer/quantizer.h" namespace zvec { namespace turbo { From ca6b526ab91f8f232905e403a98373bd33c6f549 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Wed, 22 Jul 2026 15:24:28 +0800 Subject: [PATCH 081/100] debug --- src/include/zvec/turbo/turbo.h | 22 ++++++++++++++++++++++ src/turbo/quantizer/quantizer.h | 22 ---------------------- 2 files changed, 22 insertions(+), 22 deletions(-) diff --git a/src/include/zvec/turbo/turbo.h b/src/include/zvec/turbo/turbo.h index 8b48e36a6..04d0e3aae 100644 --- a/src/include/zvec/turbo/turbo.h +++ b/src/include/zvec/turbo/turbo.h @@ -20,6 +20,28 @@ namespace zvec::turbo { +//! Error code literals mirroring core::IndexError::Code integer values. +//! +//! Turbo quantizer sources use these directly instead of the +//! `IndexError_NotImplemented` / `IndexError_Unsupported` const objects +//! because MSVC's WINDOWS_EXPORT_ALL_SYMBOLS does not export const data +//! with constructors from zvec_shared.dll. zvec_turbo is a static library +//! linked with /WHOLEARCHIVE, so referencing those unexported symbols across +//! the DLL boundary triggers LNK2019 on Windows. +//! +//! IndexError::Code stores -val in its constructor, so NotImplemented(11) +//! yields -11 and Unsupported(12) yields -12. +constexpr int kErrRuntime = -1; +constexpr int kErrNotImplemented = -11; +constexpr int kErrUnsupported = -12; +constexpr int kErrInvalidArgument = -31; + +//! Magic number ('QTZR') stamped at the start of a serialized quantizer blob. +constexpr uint32_t kQuantizerMagic = 0x52545A51u; + +//! Current quantizer serialization format version. +constexpr uint16_t kQuantizerSerVersion = 1; + using DistanceFunc = std::function; using BatchDistanceFunc = std::function Date: Wed, 22 Jul 2026 17:41:30 +0800 Subject: [PATCH 082/100] fix neon --- src/ailego/internal/cpu_features.cc | 18 +++++++++ src/ailego/internal/cpu_features.h | 6 +++ src/turbo/distance/avx512/rotate/fht/fht.cc | 29 +++++++------- .../preprocessor/fht_rotator/fht_rotator.cc | 38 ++++++++++++++++--- src/turbo/turbo.cc | 10 +++-- tests/turbo/turbo_fht_rotator_test.cc | 2 - 6 files changed, 79 insertions(+), 24 deletions(-) diff --git a/src/ailego/internal/cpu_features.cc b/src/ailego/internal/cpu_features.cc index 066b0a6a3..b3e791c95 100644 --- a/src/ailego/internal/cpu_features.cc +++ b/src/ailego/internal/cpu_features.cc @@ -23,6 +23,13 @@ #include #endif +#if defined(__aarch64__) && defined(__linux__) +#include +#ifndef HWCAP_ASIMD +#define HWCAP_ASIMD (1 << 1) +#endif +#endif + namespace zvec { namespace ailego { namespace internal { @@ -336,6 +343,17 @@ bool CpuFeatures::HYPERVISOR(void) { return !!(flags_.L1_ECX & (1u << 31)); } +//! ARM NEON (ASIMD) support +bool CpuFeatures::NEON(void) { +#if defined(__aarch64__) && defined(__linux__) + return !!(getauxval(AT_HWCAP) & HWCAP_ASIMD); +#elif defined(__ARM_NEON) + return true; +#else + return false; +#endif +} + const char *CpuFeatures::Intrinsics(void) { return "" #if defined(__ARM_NEON) diff --git a/src/ailego/internal/cpu_features.h b/src/ailego/internal/cpu_features.h index 3db1dee2a..02c8740de 100644 --- a/src/ailego/internal/cpu_features.h +++ b/src/ailego/internal/cpu_features.h @@ -182,6 +182,9 @@ class CpuFeatures { // ！Running on a hypervisor static bool HYPERVISOR(void); + //! ARM NEON (ASIMD) support + static bool NEON(void); + //! Intrinsics of compiling static const char *Intrinsics(void); @@ -361,6 +364,9 @@ class CpuFeatures { // ！Running on a hypervisor bool HYPERVISOR = CpuFeatures::HYPERVISOR(); + + //! ARM NEON (ASIMD) support + bool NEON = CpuFeatures::NEON(); }; static StaticFlags static_flags_; }; diff --git a/src/turbo/distance/avx512/rotate/fht/fht.cc b/src/turbo/distance/avx512/rotate/fht/fht.cc index a892ac61f..05fd2a5a0 100644 --- a/src/turbo/distance/avx512/rotate/fht/fht.cc +++ b/src/turbo/distance/avx512/rotate/fht/fht.cc @@ -33,7 +33,10 @@ void fht_flip_sign_avx512(const uint8_t *flip, float *data, size_t dim) { #if defined(__AVX512F__) size_t simd_end = dim & ~63u; constexpr size_t kChunk = 64; - const __m512 sign_flip = _mm512_castsi512_ps(_mm512_set1_epi32(0x80000000)); + // Sign-flip is a pure bitwise op (x ^ 0x80000000), done in the integer + // domain so only AVX512F is required. The float-domain VXORPS (zmm) would + // pull in an AVX512DQ dependency, which we deliberately avoid here. + const __m512i sign_bit = _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)); @@ -41,18 +44,18 @@ void fht_flip_sign_avx512(const uint8_t *flip, float *data, size_t dim) { 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); + __m512i v0 = _mm512_castps_si512(_mm512_loadu_ps(&data[i])); + v0 = _mm512_mask_xor_epi32(v0, m0, v0, sign_bit); + _mm512_storeu_ps(&data[i], _mm512_castsi512_ps(v0)); + __m512i v1 = _mm512_castps_si512(_mm512_loadu_ps(&data[i + 16])); + v1 = _mm512_mask_xor_epi32(v1, m1, v1, sign_bit); + _mm512_storeu_ps(&data[i + 16], _mm512_castsi512_ps(v1)); + __m512i v2 = _mm512_castps_si512(_mm512_loadu_ps(&data[i + 32])); + v2 = _mm512_mask_xor_epi32(v2, m2, v2, sign_bit); + _mm512_storeu_ps(&data[i + 32], _mm512_castsi512_ps(v2)); + __m512i v3 = _mm512_castps_si512(_mm512_loadu_ps(&data[i + 48])); + v3 = _mm512_mask_xor_epi32(v3, m3, v3, sign_bit); + _mm512_storeu_ps(&data[i + 48], _mm512_castsi512_ps(v3)); } // Scalar tail for (size_t i = simd_end; i < dim; ++i) { diff --git a/src/turbo/preprocessor/fht_rotator/fht_rotator.cc b/src/turbo/preprocessor/fht_rotator/fht_rotator.cc index 6a60bbe2b..bac1c5296 100644 --- a/src/turbo/preprocessor/fht_rotator/fht_rotator.cc +++ b/src/turbo/preprocessor/fht_rotator/fht_rotator.cc @@ -16,6 +16,7 @@ #include #include #include +#include #include namespace zvec { @@ -51,6 +52,12 @@ FhtRotator::Pointer FhtRotator::create(int dim) { // Single allocation: FhtCtx header + trailing flip data. r->fht_ctx_ = static_cast(std::malloc(sizeof(FhtCtx) + flip_size)); + if (r->fht_ctx_ == nullptr) { + // Allocation failed: bail out before dereferencing. Returning nullptr lets + // the smart Pointer clean up the partially-built object (~FhtRotator frees + // the null fht_ctx_ safely). + return nullptr; + } r->fht_ctx_->flip_offset = r->flip_offset_; r->fht_ctx_->trunc_dim = trunc_dim; r->fht_ctx_->fac = fac; @@ -146,12 +153,32 @@ int FhtRotator::deserialize(const void *data, size_t len) { return kErrUnsupported; } - const size_t total = sizeof(RotatorSerHeader) + hdr->payload_size; - if (len < total) return kErrInvalidArgument; + // Validate dimensions before any cast to int: must be strictly positive and + // representable as int. FHT keeps dimensionality unchanged (in == out). + if (hdr->in_dim == 0 || + hdr->in_dim > static_cast(std::numeric_limits::max())) { + return kErrInvalidArgument; + } + if (hdr->out_dim != hdr->in_dim) return kErrInvalidArgument; + + // Length check via subtraction to avoid size_t overflow on 32-bit + // (len >= sizeof(header) was guaranteed above, so the subtraction is safe). + if (hdr->payload_size > len - sizeof(RotatorSerHeader)) { + return kErrInvalidArgument; + } + + // Payload must hold exactly 4 rounds of ceil(in_dim/8) flip bytes. The + // rotation kernels read 4 * flip_offset bytes, so any smaller payload would + // read out of bounds; require an exact match. + const size_t new_flip_offset = + (static_cast(hdr->in_dim) + kByteLen - 1) / kByteLen; + const size_t expected_flip_size = 4 * new_flip_offset; + if (hdr->payload_size != expected_flip_size) return kErrInvalidArgument; + // All fields validated: commit to member state. in_dim_ = static_cast(hdr->in_dim); out_dim_ = static_cast(hdr->out_dim); - flip_offset_ = (static_cast(in_dim_) + kByteLen - 1) / kByteLen; + flip_offset_ = new_flip_offset; kernels_ = get_rotator_kernels(RotateType::kFht); const size_t trunc_dim = floor_pow2(static_cast(in_dim_)); @@ -160,13 +187,14 @@ int FhtRotator::deserialize(const void *data, size_t len) { // Free old ctx, allocate new one with trailing flip data. std::free(fht_ctx_); fht_ctx_ = - static_cast(std::malloc(sizeof(FhtCtx) + hdr->payload_size)); + static_cast(std::malloc(sizeof(FhtCtx) + expected_flip_size)); + if (fht_ctx_ == nullptr) return kErrRuntime; fht_ctx_->flip_offset = flip_offset_; fht_ctx_->trunc_dim = trunc_dim; fht_ctx_->fac = fac; std::memcpy(fht_ctx_->flip, reinterpret_cast(data) + sizeof(RotatorSerHeader), - hdr->payload_size); + expected_flip_size); return 0; } diff --git a/src/turbo/turbo.cc b/src/turbo/turbo.cc index b45a99617..5a6f33d49 100644 --- a/src/turbo/turbo.cc +++ b/src/turbo/turbo.cc @@ -158,7 +158,6 @@ RotatorKernels get_rotator_kernels(RotateType rotate_type, switch (rotate_type) { case RotateType::kFht: { if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F && - zvec::ailego::internal::CpuFeatures::static_flags_.AVX512DQ && IsArchMatch(cpu_arch_type, CpuArchType::kAVX512)) { return {avx512::fht_rotate_avx512, avx512::fht_unrotate_avx512}; } @@ -170,16 +169,19 @@ RotatorKernels get_rotator_kernels(RotateType rotate_type, IsArchMatch(cpu_arch_type, CpuArchType::kSSE)) { return {sse::fht_rotate_sse, sse::fht_unrotate_sse}; } - if (IsArchMatch(cpu_arch_type, CpuArchType::kNEON)) { + if (zvec::ailego::internal::CpuFeatures::static_flags_.NEON && + IsArchMatch(cpu_arch_type, CpuArchType::kNEON)) { return {neon::fht_rotate_neon, neon::fht_unrotate_neon}; } return {scalar::fht_rotate, scalar::fht_unrotate}; } } - // Fallback (unreachable for valid RotateType values). + // Unsupported RotateType: assert in debug for early detection, but always + // return the scalar kernels so release builds never hand back null function + // pointers (which would crash on the first call). assert(false && "unsupported RotateType"); - return {}; + return {scalar::fht_rotate, scalar::fht_unrotate}; } } // namespace zvec::turbo diff --git a/tests/turbo/turbo_fht_rotator_test.cc b/tests/turbo/turbo_fht_rotator_test.cc index 1e454ed9b..45c447942 100644 --- a/tests/turbo/turbo_fht_rotator_test.cc +++ b/tests/turbo/turbo_fht_rotator_test.cc @@ -206,7 +206,6 @@ TEST(FhtRotator, FromBlobMalformed) { 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); @@ -241,7 +240,6 @@ TEST(FhtRotator, L2DistancePreserved) { 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; From 17aa431f1bfed2e1d9ddcc7d830f836723bc764a Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Wed, 22 Jul 2026 19:42:37 +0800 Subject: [PATCH 083/100] copilot fix --- .../preprocessor/fht_rotator/fht_rotator.cc | 78 +++++++++++-------- src/turbo/preprocessor/preprocessor.h | 7 +- tests/turbo/turbo_fht_rotator_test.cc | 49 ++++++------ 3 files changed, 78 insertions(+), 56 deletions(-) diff --git a/src/turbo/preprocessor/fht_rotator/fht_rotator.cc b/src/turbo/preprocessor/fht_rotator/fht_rotator.cc index bac1c5296..d282b2b3d 100644 --- a/src/turbo/preprocessor/fht_rotator/fht_rotator.cc +++ b/src/turbo/preprocessor/fht_rotator/fht_rotator.cc @@ -75,16 +75,21 @@ FhtRotator::Pointer FhtRotator::create(int dim) { 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) != RotateType::kFht) { + // Copy the header into a properly aligned local before reading any field. + // `data` may point to an unaligned byte buffer (e.g. std::string::data()), + // so dereferencing a reinterpret_cast directly + // would be undefined behavior on architectures that require alignment. + RotatorSerHeader hdr; + std::memcpy(&hdr, data, sizeof(RotatorSerHeader)); + if (hdr.magic != kRotatorMagic) return nullptr; + if (hdr.version != kRotatorSerVersion) return nullptr; + if (static_cast(hdr.rotator_type) != RotateType::kFht) { return nullptr; } Pointer r(new FhtRotator()); const size_t expected_total = - sizeof(RotatorSerHeader) + static_cast(hdr->payload_size); + sizeof(RotatorSerHeader) + static_cast(hdr.payload_size); if (len < expected_total) return nullptr; int rc = r->deserialize(data, len); @@ -146,24 +151,29 @@ int FhtRotator::serialize(std::string *out) const { int FhtRotator::deserialize(const void *data, size_t len) { if (!data || len < sizeof(RotatorSerHeader)) return kErrInvalidArgument; - const auto *hdr = reinterpret_cast(data); - if (hdr->magic != kRotatorMagic) return kErrUnsupported; - if (hdr->version != kRotatorSerVersion) return kErrUnsupported; - if (static_cast(hdr->rotator_type) != RotateType::kFht) { + // Copy the header into a properly aligned local before reading any field. + // `data` may point to an unaligned byte buffer (e.g. std::string::data()), + // so dereferencing a reinterpret_cast directly + // would be undefined behavior on architectures that require alignment. + RotatorSerHeader hdr; + std::memcpy(&hdr, data, sizeof(RotatorSerHeader)); + if (hdr.magic != kRotatorMagic) return kErrUnsupported; + if (hdr.version != kRotatorSerVersion) return kErrUnsupported; + if (static_cast(hdr.rotator_type) != RotateType::kFht) { return kErrUnsupported; } // Validate dimensions before any cast to int: must be strictly positive and // representable as int. FHT keeps dimensionality unchanged (in == out). - if (hdr->in_dim == 0 || - hdr->in_dim > static_cast(std::numeric_limits::max())) { + if (hdr.in_dim == 0 || + hdr.in_dim > static_cast(std::numeric_limits::max())) { return kErrInvalidArgument; } - if (hdr->out_dim != hdr->in_dim) return kErrInvalidArgument; + if (hdr.out_dim != hdr.in_dim) return kErrInvalidArgument; // Length check via subtraction to avoid size_t overflow on 32-bit // (len >= sizeof(header) was guaranteed above, so the subtraction is safe). - if (hdr->payload_size > len - sizeof(RotatorSerHeader)) { + if (hdr.payload_size > len - sizeof(RotatorSerHeader)) { return kErrInvalidArgument; } @@ -171,31 +181,37 @@ int FhtRotator::deserialize(const void *data, size_t len) { // rotation kernels read 4 * flip_offset bytes, so any smaller payload would // read out of bounds; require an exact match. const size_t new_flip_offset = - (static_cast(hdr->in_dim) + kByteLen - 1) / kByteLen; + (static_cast(hdr.in_dim) + kByteLen - 1) / kByteLen; const size_t expected_flip_size = 4 * new_flip_offset; - if (hdr->payload_size != expected_flip_size) return kErrInvalidArgument; - - // All fields validated: commit to member state. - in_dim_ = static_cast(hdr->in_dim); - out_dim_ = static_cast(hdr->out_dim); - flip_offset_ = new_flip_offset; - kernels_ = get_rotator_kernels(RotateType::kFht); - - const size_t trunc_dim = floor_pow2(static_cast(in_dim_)); + if (hdr.payload_size != expected_flip_size) return kErrInvalidArgument; + + // Build the new context into a temporary first, so a failed allocation + // leaves the existing object completely untouched (deserialize may target + // an already-initialized rotator). Only commit member state once the + // allocation and payload copy have both succeeded. + const int new_in_dim = static_cast(hdr.in_dim); + const int new_out_dim = static_cast(hdr.out_dim); + const size_t trunc_dim = floor_pow2(static_cast(new_in_dim)); const float fac = 1.0f / std::sqrt(static_cast(trunc_dim)); - // Free old ctx, allocate new one with trailing flip data. - std::free(fht_ctx_); - fht_ctx_ = + FhtCtx *new_ctx = static_cast(std::malloc(sizeof(FhtCtx) + expected_flip_size)); - if (fht_ctx_ == nullptr) return kErrRuntime; - fht_ctx_->flip_offset = flip_offset_; - fht_ctx_->trunc_dim = trunc_dim; - fht_ctx_->fac = fac; - std::memcpy(fht_ctx_->flip, + if (new_ctx == nullptr) return kErrRuntime; + new_ctx->flip_offset = new_flip_offset; + new_ctx->trunc_dim = trunc_dim; + new_ctx->fac = fac; + std::memcpy(new_ctx->flip, reinterpret_cast(data) + sizeof(RotatorSerHeader), expected_flip_size); + // All allocations succeeded: commit new state and free the old context. + std::free(fht_ctx_); + fht_ctx_ = new_ctx; + in_dim_ = new_in_dim; + out_dim_ = new_out_dim; + flip_offset_ = new_flip_offset; + kernels_ = get_rotator_kernels(RotateType::kFht); + return 0; } diff --git a/src/turbo/preprocessor/preprocessor.h b/src/turbo/preprocessor/preprocessor.h index c51f19b14..c83235ce1 100644 --- a/src/turbo/preprocessor/preprocessor.h +++ b/src/turbo/preprocessor/preprocessor.h @@ -55,9 +55,10 @@ class Preprocessor { //! 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. + //! Output dimensionality produced by apply(). A future preprocessor may + //! change dimensionality (out_dim() != in_dim()), but FhtRotator keeps it + //! unchanged: it operates on floor_pow2(in_dim) via the Hadamard transform + //! plus a Kac's walk over the remainder, so out_dim() == in_dim(). virtual int out_dim() const = 0; //! Forward transform: map an input vector to the preprocessed space. diff --git a/tests/turbo/turbo_fht_rotator_test.cc b/tests/turbo/turbo_fht_rotator_test.cc index 45c447942..f8d60dcc7 100644 --- a/tests/turbo/turbo_fht_rotator_test.cc +++ b/tests/turbo/turbo_fht_rotator_test.cc @@ -13,6 +13,7 @@ // limitations under the License. #include +#include #include #include #include @@ -148,26 +149,26 @@ TEST(FhtRotator, CreateGeneratesFlip) { EXPECT_EQ(0, rot->serialize(&blob)) << "serialize failed after create for dim=" << dim; - // Verify the payload is not all zeros (extremely unlikely for random bits). - const auto *hdr = reinterpret_cast(blob.data()); - EXPECT_EQ(hdr->magic, kRotatorMagic); - EXPECT_EQ(hdr->version, kRotatorSerVersion); - EXPECT_EQ(static_cast(hdr->rotator_type), RotateType::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; + // Verify the serialized structure deterministically. Copy the header into + // an aligned local first: blob.data() may be unaligned and dereferencing a + // reinterpret_cast would be undefined behavior on + // architectures that require alignment. + RotatorSerHeader hdr; + std::memcpy(&hdr, blob.data(), sizeof(RotatorSerHeader)); + EXPECT_EQ(hdr.magic, kRotatorMagic); + EXPECT_EQ(hdr.version, kRotatorSerVersion); + EXPECT_EQ(static_cast(hdr.rotator_type), RotateType::kFht); + EXPECT_EQ(static_cast(hdr.in_dim), dim); + EXPECT_EQ(static_cast(hdr.out_dim), dim); + + // Payload holds 4 rounds of ceil(dim/8) flip bytes. Checking the exact size + // (and total blob length) keeps the test deterministic: we avoid asserting + // on random bit values, which could theoretically be all-zero on a platform + // where std::random_device has low entropy. + const uint32_t expected_flip_size = + 4u * ((static_cast(dim) + 7u) / 8u); + EXPECT_EQ(hdr.payload_size, expected_flip_size); + EXPECT_EQ(blob.size(), sizeof(RotatorSerHeader) + expected_flip_size); } } @@ -377,8 +378,12 @@ TEST(FhtRotator, DeserializeTruncatedPayload) { 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; + // Copy the header into an aligned local before reading fields (blob.data() + // may be unaligned; a direct reinterpret_cast dereference is UB on + // alignment-sensitive architectures). + RotatorSerHeader hdr; + std::memcpy(&hdr, blob.data(), sizeof(RotatorSerHeader)); + size_t truncated_len = sizeof(RotatorSerHeader) + hdr.payload_size / 2; auto rot2 = FhtRotator::create(64); ASSERT_TRUE(rot2); From 84726e8aecd6adb2c74147c2800e544b2fb45db3 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Thu, 23 Jul 2026 14:22:46 +0800 Subject: [PATCH 084/100] unify structure --- src/turbo/distance/avx2/rotate/fht/fht.cc | 10 ++++++---- src/turbo/distance/avx512/rotate/fht/fht.cc | 10 ++++++---- src/turbo/distance/common/fht_common.h | 17 ++++++++++++++--- src/turbo/distance/neon/rotate/fht/fht.cc | 10 ++++++---- src/turbo/distance/scalar/rotate/fht/fht.cc | 10 +++++----- src/turbo/distance/sse/rotate/fht/fht.cc | 10 ++++++---- .../preprocessor/fht_rotator/fht_rotator.cc | 2 -- 7 files changed, 43 insertions(+), 26 deletions(-) diff --git a/src/turbo/distance/avx2/rotate/fht/fht.cc b/src/turbo/distance/avx2/rotate/fht/fht.cc index a15ca6846..a613334d5 100644 --- a/src/turbo/distance/avx2/rotate/fht/fht.cc +++ b/src/turbo/distance/avx2/rotate/fht/fht.cc @@ -165,28 +165,30 @@ void fht_vec_rescale_avx2(float *data, size_t n, float factor) { #endif } -void fht_rotate_avx2(const float * /*in*/, float *out, size_t in_dim, +void fht_rotate_avx2(const float *in, float *out, size_t in_dim, size_t /*out_dim*/, void *ctx) { #if defined(__AVX2__) static constexpr FhtPrimitives kPrim = { fht_flip_sign_avx2, fht_inplace_avx2, fht_kacs_walk_avx2, fht_inv_kacs_walk_avx2, fht_vec_rescale_avx2}; - fht_rotate_impl(out, in_dim, ctx, kPrim); + fht_rotate_impl(in, out, in_dim, ctx, kPrim); #else + (void)in; (void)out; (void)in_dim; (void)ctx; #endif } -void fht_unrotate_avx2(const float * /*in*/, float *out, size_t in_dim, +void fht_unrotate_avx2(const float *in, float *out, size_t in_dim, size_t /*out_dim*/, void *ctx) { #if defined(__AVX2__) static constexpr FhtPrimitives kPrim = { fht_flip_sign_avx2, fht_inplace_avx2, fht_kacs_walk_avx2, fht_inv_kacs_walk_avx2, fht_vec_rescale_avx2}; - fht_unrotate_impl(out, in_dim, ctx, kPrim); + fht_unrotate_impl(in, out, in_dim, ctx, kPrim); #else + (void)in; (void)out; (void)in_dim; (void)ctx; diff --git a/src/turbo/distance/avx512/rotate/fht/fht.cc b/src/turbo/distance/avx512/rotate/fht/fht.cc index 05fd2a5a0..7c6d96b29 100644 --- a/src/turbo/distance/avx512/rotate/fht/fht.cc +++ b/src/turbo/distance/avx512/rotate/fht/fht.cc @@ -173,28 +173,30 @@ void fht_vec_rescale_avx512(float *data, size_t n, float factor) { #endif } -void fht_rotate_avx512(const float * /*in*/, float *out, size_t in_dim, +void fht_rotate_avx512(const float *in, float *out, size_t in_dim, size_t /*out_dim*/, void *ctx) { #if defined(__AVX512F__) static constexpr FhtPrimitives kPrim = { fht_flip_sign_avx512, fht_inplace_avx512, fht_kacs_walk_avx512, fht_inv_kacs_walk_avx512, fht_vec_rescale_avx512}; - fht_rotate_impl(out, in_dim, ctx, kPrim); + fht_rotate_impl(in, out, in_dim, ctx, kPrim); #else + (void)in; (void)out; (void)in_dim; (void)ctx; #endif } -void fht_unrotate_avx512(const float * /*in*/, float *out, size_t in_dim, +void fht_unrotate_avx512(const float *in, float *out, size_t in_dim, size_t /*out_dim*/, void *ctx) { #if defined(__AVX512F__) static constexpr FhtPrimitives kPrim = { fht_flip_sign_avx512, fht_inplace_avx512, fht_kacs_walk_avx512, fht_inv_kacs_walk_avx512, fht_vec_rescale_avx512}; - fht_unrotate_impl(out, in_dim, ctx, kPrim); + fht_unrotate_impl(in, out, in_dim, ctx, kPrim); #else + (void)in; (void)out; (void)in_dim; (void)ctx; diff --git a/src/turbo/distance/common/fht_common.h b/src/turbo/distance/common/fht_common.h index 556338863..7285c2c23 100644 --- a/src/turbo/distance/common/fht_common.h +++ b/src/turbo/distance/common/fht_common.h @@ -16,6 +16,7 @@ #include #include +#include namespace zvec::turbo { @@ -35,8 +36,13 @@ struct FhtPrimitives { /// offset 16: float fac /// offset 20: uint8_t pad[4] /// offset 24: uint8_t flip[] -inline void fht_rotate_impl(float *data, size_t dim, void *ctx, +inline void fht_rotate_impl(const float *in, float *out, size_t dim, void *ctx, const FhtPrimitives &p) { + // Seed the output buffer with the input, then rotate it in place. Keeping the + // copy here makes the kernel a self-contained out-of-place transform (read + // `in`, write `out`); callers no longer need to pre-copy `in` into `out`. + std::memcpy(out, in, sizeof(float) * dim); + float *data = out; auto *base = reinterpret_cast(ctx); const size_t flip_offset = *reinterpret_cast(base); const size_t trunc_dim = *reinterpret_cast(base + 8); @@ -78,8 +84,13 @@ inline void fht_rotate_impl(float *data, size_t dim, void *ctx, p.rescale(data, dim, 0.25f); } -inline void fht_unrotate_impl(float *data, size_t dim, void *ctx, - const FhtPrimitives &p) { +inline void fht_unrotate_impl(const float *in, float *out, size_t dim, + void *ctx, const FhtPrimitives &p) { + // Seed the output buffer with the input, then unrotate it in place. Keeping + // the copy here makes the kernel a self-contained out-of-place transform + // (read `in`, write `out`); callers no longer need to pre-copy `in`. + std::memcpy(out, in, sizeof(float) * dim); + float *data = out; auto *base = reinterpret_cast(ctx); const size_t flip_offset = *reinterpret_cast(base); const size_t trunc_dim = *reinterpret_cast(base + 8); diff --git a/src/turbo/distance/neon/rotate/fht/fht.cc b/src/turbo/distance/neon/rotate/fht/fht.cc index 07a15e46f..0ffbfaad4 100644 --- a/src/turbo/distance/neon/rotate/fht/fht.cc +++ b/src/turbo/distance/neon/rotate/fht/fht.cc @@ -138,28 +138,30 @@ void fht_vec_rescale_neon(float *data, size_t n, float factor) { #endif } -void fht_rotate_neon(const float * /*in*/, float *out, size_t in_dim, +void fht_rotate_neon(const float *in, float *out, size_t in_dim, size_t /*out_dim*/, void *ctx) { #if defined(__ARM_NEON) && defined(__aarch64__) static constexpr FhtPrimitives kPrim = { fht_flip_sign_neon, scalar::fht_inplace, fht_kacs_walk_neon, fht_inv_kacs_walk_neon, fht_vec_rescale_neon}; - fht_rotate_impl(out, in_dim, ctx, kPrim); + fht_rotate_impl(in, out, in_dim, ctx, kPrim); #else + (void)in; (void)out; (void)in_dim; (void)ctx; #endif } -void fht_unrotate_neon(const float * /*in*/, float *out, size_t in_dim, +void fht_unrotate_neon(const float *in, float *out, size_t in_dim, size_t /*out_dim*/, void *ctx) { #if defined(__ARM_NEON) && defined(__aarch64__) static constexpr FhtPrimitives kPrim = { fht_flip_sign_neon, scalar::fht_inplace, fht_kacs_walk_neon, fht_inv_kacs_walk_neon, fht_vec_rescale_neon}; - fht_unrotate_impl(out, in_dim, ctx, kPrim); + fht_unrotate_impl(in, out, in_dim, ctx, kPrim); #else + (void)in; (void)out; (void)in_dim; (void)ctx; diff --git a/src/turbo/distance/scalar/rotate/fht/fht.cc b/src/turbo/distance/scalar/rotate/fht/fht.cc index 1d4e6e779..5c9064be9 100644 --- a/src/turbo/distance/scalar/rotate/fht/fht.cc +++ b/src/turbo/distance/scalar/rotate/fht/fht.cc @@ -77,20 +77,20 @@ void fht_vec_rescale(float *data, size_t n, float factor) { } } -void fht_rotate(const float * /*in*/, float *out, size_t in_dim, - size_t /*out_dim*/, void *ctx) { +void fht_rotate(const float *in, float *out, size_t in_dim, size_t /*out_dim*/, + void *ctx) { static constexpr FhtPrimitives kPrim = {fht_flip_sign, fht_inplace, fht_kacs_walk, fht_inv_kacs_walk, fht_vec_rescale}; - fht_rotate_impl(out, in_dim, ctx, kPrim); + fht_rotate_impl(in, out, in_dim, ctx, kPrim); } -void fht_unrotate(const float * /*in*/, float *out, size_t in_dim, +void fht_unrotate(const float *in, float *out, size_t in_dim, size_t /*out_dim*/, void *ctx) { static constexpr FhtPrimitives kPrim = {fht_flip_sign, fht_inplace, fht_kacs_walk, fht_inv_kacs_walk, fht_vec_rescale}; - fht_unrotate_impl(out, in_dim, ctx, kPrim); + fht_unrotate_impl(in, out, in_dim, ctx, kPrim); } } // namespace zvec::turbo::scalar diff --git a/src/turbo/distance/sse/rotate/fht/fht.cc b/src/turbo/distance/sse/rotate/fht/fht.cc index f020fcc1a..797e1cd08 100644 --- a/src/turbo/distance/sse/rotate/fht/fht.cc +++ b/src/turbo/distance/sse/rotate/fht/fht.cc @@ -142,28 +142,30 @@ void fht_vec_rescale_sse(float *data, size_t n, float factor) { #endif } -void fht_rotate_sse(const float * /*in*/, float *out, size_t in_dim, +void fht_rotate_sse(const float *in, float *out, size_t in_dim, size_t /*out_dim*/, void *ctx) { #if defined(__SSE2__) static constexpr FhtPrimitives kPrim = { fht_flip_sign_sse, scalar::fht_inplace, fht_kacs_walk_sse, fht_inv_kacs_walk_sse, fht_vec_rescale_sse}; - fht_rotate_impl(out, in_dim, ctx, kPrim); + fht_rotate_impl(in, out, in_dim, ctx, kPrim); #else + (void)in; (void)out; (void)in_dim; (void)ctx; #endif } -void fht_unrotate_sse(const float * /*in*/, float *out, size_t in_dim, +void fht_unrotate_sse(const float *in, float *out, size_t in_dim, size_t /*out_dim*/, void *ctx) { #if defined(__SSE2__) static constexpr FhtPrimitives kPrim = { fht_flip_sign_sse, scalar::fht_inplace, fht_kacs_walk_sse, fht_inv_kacs_walk_sse, fht_vec_rescale_sse}; - fht_unrotate_impl(out, in_dim, ctx, kPrim); + fht_unrotate_impl(in, out, in_dim, ctx, kPrim); #else + (void)in; (void)out; (void)in_dim; (void)ctx; diff --git a/src/turbo/preprocessor/fht_rotator/fht_rotator.cc b/src/turbo/preprocessor/fht_rotator/fht_rotator.cc index d282b2b3d..a10d2f4c3 100644 --- a/src/turbo/preprocessor/fht_rotator/fht_rotator.cc +++ b/src/turbo/preprocessor/fht_rotator/fht_rotator.cc @@ -107,7 +107,6 @@ void FhtRotator::train(const void * /*data*/, size_t /*num*/, // --------------------------------------------------------------------------- void FhtRotator::apply(const float *in, float *out) const { - std::memcpy(out, in, sizeof(float) * static_cast(in_dim_)); kernels_.rotate(in, out, static_cast(in_dim_), static_cast(out_dim_), static_cast(fht_ctx_)); } @@ -117,7 +116,6 @@ void FhtRotator::apply(const float *in, float *out) const { // --------------------------------------------------------------------------- void FhtRotator::apply_inverse(const float *in, float *out) const { - std::memcpy(out, in, sizeof(float) * static_cast(in_dim_)); kernels_.unrotate(in, out, static_cast(in_dim_), static_cast(out_dim_), static_cast(fht_ctx_)); From 43dd41a5842ed20f95ffd9cd1843a73aa1e4f539 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Thu, 23 Jul 2026 19:16:09 +0800 Subject: [PATCH 085/100] neon --- src/turbo/turbo.cc | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/turbo/turbo.cc b/src/turbo/turbo.cc index 648630e16..c39182d2d 100644 --- a/src/turbo/turbo.cc +++ b/src/turbo/turbo.cc @@ -204,7 +204,8 @@ PqKernels get_pq_kernels(DataType data_type, QuantizeType quantize_type, avx2::pq_sdc_int8_distance_avx2, avx2::pq_adc_int8_batch_distance_avx2}; } - if (IsArchMatch(cpu_arch_type, CpuArchType::kNEON)) { + if (zvec::ailego::internal::CpuFeatures::static_flags_.NEON && + IsArchMatch(cpu_arch_type, CpuArchType::kNEON)) { return {neon::pq_adc_int8_distance_neon, neon::pq_sdc_int8_distance_neon, neon::pq_adc_int8_batch_distance_neon}; From c8b00de454f87aee081a29fedb4e94ed16fb41e2 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Thu, 23 Jul 2026 19:29:58 +0800 Subject: [PATCH 086/100] tmp --- .../quantizer/pq_int8_quantizer/pq_int8_quantizer.cc | 8 +------- 1 file changed, 1 insertion(+), 7 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 ccd3b6dc2..71af73e37 100644 --- a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc @@ -105,12 +105,6 @@ int PqInt8Quantizer::init(const IndexMeta &meta, const ailego::Params ¶ms) { params.get("epsilon", &epsilon_); params.get("use_zero_mean", &use_zero_mean_); - // Zero-mean centering is valid only where the search metric is - // translation-invariant. SquaredEuclidean qualifies directly. Cosine now - // uses an L2 LUT on L2-normalized vectors, so centering the *normalized* - // data (train/quantize_data/quantize_query all normalize BEFORE centering) - // is also safe and can lower reconstruction error. IP and - // MipsSquaredEuclidean are NOT translation-invariant and are rejected. if (use_zero_mean_ && meta_.metric_name() != "SquaredEuclidean" && meta_.metric_name() != "Cosine") { LOG_WARN("PqInt8Quantizer: use_zero_mean is incompatible with metric '%s', " @@ -178,7 +172,7 @@ void PqInt8Quantizer::train_subquantizer(const float *data, size_t num, gen.set_assumption_free(false); algorithm.init_centroids(*local_threads, gen); - // Lloyd iterations (aligned with multi_chunk_cluster.h L200-216) + // Lloyd iterations double cost = 0.0; for (uint32_t iter = 0; iter < kMaxKmeansIters; ++iter) { double old_cost = cost; From c48fcb3515efb04bb182f7fe09b32fb031e17928 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Thu, 23 Jul 2026 19:53:37 +0800 Subject: [PATCH 087/100] clang-format --- .../pq_int8_quantizer/pq_int8_quantizer.cc | 112 +++++++----------- src/turbo/turbo.cc | 58 ++++----- 2 files changed, 66 insertions(+), 104 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 71af73e37..42f7cbea1 100644 --- a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc @@ -75,19 +75,14 @@ int PqInt8Quantizer::init(const IndexMeta &meta, const ailego::Params ¶ms) { // 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. + // L2-only batch distance for encoding and KMeans training: the PQ codebook + // is trained in L2 space regardless of the search metric. fp32_l2_batch_fn_ = get_batch_distance_func(MetricType::kSquaredEuclidean, DataType::kFp32, QuantizeType::kDefault, CpuArchType::kAuto); - // For Cosine: cosine = normalize + L2. Vectors are L2-normalized in - // train() / quantize_data() / quantize_query(); after normalization - // cosine distance is monotonic with squared-Euclidean distance, so the - // search LUT uses the same SquaredEuclidean batch function as encoding - // and training — no IP anywhere. + // Cosine = normalize + L2: after normalization cosine distance is monotonic + // with squared-Euclidean, so the search LUT reuses SquaredEuclidean. if (meta_.metric_name() == "Cosine") { fp32_batch_fn_ = get_batch_distance_func(MetricType::kSquaredEuclidean, DataType::kFp32, @@ -142,12 +137,8 @@ void PqInt8Quantizer::train_subquantizer(const float *data, size_t num, const size_t d = sub_dim_; float *centroids_m = centroids_.data() + static_cast(sub_idx) * k * d; - // Create KMeans algorithm (L2 distance via NumericalKmeansContext). - // Always non-spherical: the PQ codebook must minimize L2 reconstruction - // error, so centroids are the true (magnitude-preserving) means. Cosine - // data is pre-normalized in train(), but a sub-vector of a unit vector is - // NOT itself unit norm, so forcing unit-norm centroids (spherical) would - // discard per-subspace magnitude and hurt reconstruction. + // Non-spherical L2 KMeans: the PQ codebook must minimize L2 reconstruction + // error, so centroids are the true (magnitude-preserving) means. ailego::NumericalKmeans algorithm(k, d); // Append sub-vectors (NumericalKmeans handles transpose internally) @@ -159,11 +150,10 @@ void PqInt8Quantizer::train_subquantizer(const float *data, size_t num, algorithm.append(sub_vec, d); } - // Single-threaded pool — parallelism is at sub-quantizer level - // (aligned with multi_chunk_cluster.h L184-185) + // Single-threaded pool — parallelism is at the sub-quantizer level. auto local_threads = std::make_shared(1, false); - // KMC2 centroid initialization (aligned with multi_chunk_cluster.h L191-196) + // KMC2 centroid initialization. ailego::Kmc2CentroidsGenerator< ailego::NumericalKmeans, SingleQueueIndexThreads> @@ -241,9 +231,8 @@ 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. + // For Cosine: normalize training data so centroids are learned in + // normalized space (L2 minimization == maximizing cosine similarity). if (meta_.metric_name() == "Cosine") { for (size_t i = 0; i < num; ++i) { float *v = all_data.data() + i * original_dim_; @@ -252,13 +241,10 @@ int PqInt8Quantizer::train(IndexHolder::Pointer holder, int thread_count) { } } - // Zero-mean centering: compute per-dimension mean of training data and - // subtract it from all vectors. The centroid is saved for later use - // in quantize_data / quantize_query / dequantize. NOTE: for Cosine this - // runs AFTER the normalization above, so centering happens in normalized - // space; quantize_data / quantize_query / dequantize must keep the same - // order (normalize -> center; dequantize: un-center -> rescale). - // Reference: DiskANN PQ trainer prepare_pq_train_data. + // Zero-mean centering: subtract the per-dimension mean; the centroid is + // saved for quantize_data / quantize_query / dequantize. For Cosine this + // runs AFTER normalization, so all paths keep the same order + // (normalize -> center; dequantize: un-center -> rescale). if (use_zero_mean_) { centroid_.assign(original_dim_, 0.0f); for (size_t i = 0; i < num; ++i) { @@ -285,12 +271,12 @@ int PqInt8Quantizer::train(IndexHolder::Pointer holder, int thread_count) { "max_iters=%u, threads=%d", num, original_dim_, num_chunk_, sub_dim_, kMaxKmeansIters, thread_count); - // Create thread pool (aligned with multi_chunk_cluster L183-185) + // Create thread pool. auto threads = std::make_shared( static_cast(thread_count), false); auto task_group = threads->make_group(); - // Distribute sub-quantizers across threads (aligned with L198-201) + // Distribute sub-quantizers across threads. std::atomic finished{0}; size_t pool_count = threads->count(); for (size_t i = 0; i < pool_count; ++i) { @@ -324,11 +310,10 @@ void PqInt8Quantizer::compute_dist_table() { const size_t d = sub_dim_; dist_table_.resize(static_cast(num_chunk_) * k * k, 0.0f); - // For each sub-quantizer, compute centroid-to-centroid distances via - // the metric-aware fp32 batch distance function (fp32_batch_fn_). + // Centroid-to-centroid distances via the metric-aware 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, uses L2 (same as L2). + // Cosine: centroids trained on normalized data, uses L2. for (uint32_t m = 0; m < num_chunk_; ++m) { const float *centroids_m = centroids_.data() + m * k * d; float *table_m = dist_table_.data() + m * k * k; @@ -349,10 +334,8 @@ void PqInt8Quantizer::quantize_data(const void *input, void *output) const { const float *vec = reinterpret_cast(input); uint8_t *code = reinterpret_cast(output); - // For Cosine: normalize the vector to unit length FIRST, because the - // codebook is trained in normalized space. The original norm is stored - // after the PQ code for later dequantize(). Centering (below) then - // operates on the normalized vector, consistent with train(). + // For Cosine: normalize FIRST (codebook is trained in normalized space); + // the original norm is stored after the PQ code for dequantize(). std::vector norm_vec_storage; float vec_norm = 0.0f; if (meta_.metric_name() == "Cosine") { @@ -373,10 +356,8 @@ void PqInt8Quantizer::quantize_data(const void *input, void *output) const { vec = centered_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. + // Encode with L2-only batch distance (search-metric independent), + // fusing argmin into the distance loop. float dists[kNumCentroids]; for (uint32_t m = 0; m < num_chunk_; ++m) { @@ -411,9 +392,8 @@ 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 (see init: Cosine - // uses an L2 LUT on normalized data). Centering (below) then operates on - // the normalized query, consistent with train() / quantize_data(). + // For Cosine: normalize FIRST (Cosine uses an L2 LUT on normalized data), + // consistent with train() / quantize_data(). std::vector norm_query_storage; if (meta_.metric_name() == "Cosine") { norm_query_storage.assign(query, query + original_dim_); @@ -433,11 +413,8 @@ void PqInt8Quantizer::quantize_query(const void *input, void *output) const { query = centered_query_storage.data(); } - // For each sub-quantizer, compute distance from query sub-vector to all - // 256 centroids via the metric-aware fp32 batch distance function. - // For L2: LUT[m][j] = ||q_m - c_m[j]||^2 - // For IP: LUT[m][j] = -dot(q_m, c_m[j]) - // For Cosine: (query is normalized) LUT[m][j] = ||q_norm_m - c_m[j]||^2 + // LUT[m][j] = distance(q_m, c_m[j]) via the metric-aware fp32_batch_fn_: + // L2/Cosine: ||q_m - c_m[j]||^2 IP: -dot(q_m, c_m[j]). // const_cast: see compute_dist_table for rationale. for (uint32_t m = 0; m < num_chunk_; ++m) { const auto ¢roid_ptrs = centroid_ptrs_cache_[m]; @@ -463,10 +440,8 @@ 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. + // ISA-dispatched batch4 ADC kernel (4-way ILP + SIMD gather). + // const_cast: batch_adc_fn_ expects const void**; kernel is read-only. batch_adc_fn_(const_cast(dp_list), query, static_cast(dp_num), num_chunk_, dist_list); // Cosine uses an L2 LUT on normalized vectors — the batch ADC sums are @@ -527,8 +502,7 @@ int PqInt8Quantizer::dequantize(const void *in, const IndexQueryMeta &qmeta, 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 + // Reconstruct by concatenating the selected centroids per sub-quantizer, // in the space the codebook was trained in (normalized for Cosine). const size_t k = kNumCentroids; const size_t d = sub_dim_; @@ -539,18 +513,16 @@ int PqInt8Quantizer::dequantize(const void *in, const IndexQueryMeta &qmeta, std::memcpy(result + m * d, centroid, d * sizeof(float)); } - // Zero-mean centering: add the centroid back FIRST to undo centering, - // which was applied in normalized space (after normalization) during - // encode. This yields the (approximate) normalized vector. + // Undo zero-mean centering: add the centroid back FIRST (centering was + // applied in normalized space during encode). if (use_zero_mean_) { for (uint32_t j = 0; j < original_dim_; ++j) { result[j] += centroid_[j]; } } - // For Cosine: the codebook was trained on normalized data; the stored norm - // rescales the (now un-centered) unit-space vector back to the original - // vector's magnitude. + // For Cosine: rescale the un-centered unit-space vector back to the + // original magnitude using the stored norm. if (meta_.metric_name() == "Cosine") { float norm = 0.0f; std::memcpy(&norm, code + num_chunk_, sizeof(float)); @@ -565,15 +537,13 @@ DistanceImpl PqInt8Quantizer::distance(const void *query, const IndexQueryMeta &qmeta) const { (void)qmeta; - // ADC: PqAdcDistanceFunc signature now matches DistanceFunc directly - // (both use void*), so we can assign without a lambda wrapper. + // ADC: PqAdcDistanceFunc matches DistanceFunc directly (no lambda needed). DistanceFunc adc_func = adc_fn_; // 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. + // The query is already quantized (LUT) by the caller; copy it directly. size_t lut_bytes = quantized_query_vector_length(); std::string lut_storage(static_cast(query), lut_bytes); @@ -592,9 +562,8 @@ DistanceImpl PqInt8Quantizer::sym_distance(const void *query, DistanceFunc sdc_func = [sdc, dt](const void *a, const void *b, size_t dim, float *out) { sdc(a, b, dt, dim, out); }; - // The query is a PQ code (uint8_t[num_chunk]), NOT a LUT. - // SDC computes distance between two PQ codes via the centroid-to-centroid - // dist_table_, so both sides must be PQ codes. + // The query is a PQ code (uint8_t[num_chunk]), NOT a LUT: SDC compares two + // PQ codes via the centroid-to-centroid dist_table_. size_t code_bytes = quantized_datapoint_vector_length(); std::string code_storage(static_cast(query), code_bytes); @@ -638,10 +607,9 @@ int PqInt8Quantizer::serialize(std::string *out) const { out->append(reinterpret_cast(centroid_.data()), centroid_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. + // dist_table_ is NOT serialized: it is a build-phase-only derivative of the + // codebook (used by SDC), recomputable on demand and unneeded after + // deserialization (search uses ADC). return 0; } diff --git a/src/turbo/turbo.cc b/src/turbo/turbo.cc index c39182d2d..20caa3ac0 100644 --- a/src/turbo/turbo.cc +++ b/src/turbo/turbo.cc @@ -157,6 +157,32 @@ UniformQuantizeFunc get_uniform_quantize_func(DataType data_type) { return nullptr; } +PqKernels get_pq_kernels(DataType data_type, QuantizeType quantize_type, + CpuArchType cpu_arch_type) { + (void)data_type; + if (quantize_type == QuantizeType::kPQ) { + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F && + IsArchMatch(cpu_arch_type, CpuArchType::kAVX512)) { + return {avx512::pq_adc_int8_distance_avx512, + avx512::pq_sdc_int8_distance_avx512, + avx512::pq_adc_int8_batch_distance_avx512}; + } + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2 && + IsArchMatch(cpu_arch_type, CpuArchType::kAVX2)) { + return {avx2::pq_adc_int8_distance_avx2, avx2::pq_sdc_int8_distance_avx2, + avx2::pq_adc_int8_batch_distance_avx2}; + } + if (zvec::ailego::internal::CpuFeatures::static_flags_.NEON && + IsArchMatch(cpu_arch_type, CpuArchType::kNEON)) { + return {neon::pq_adc_int8_distance_neon, neon::pq_sdc_int8_distance_neon, + neon::pq_adc_int8_batch_distance_neon}; + } + return {scalar::pq_adc_int8_distance, scalar::pq_sdc_int8_distance, + scalar::pq_adc_int8_batch_distance}; + } + return {}; +} + RotatorKernels get_rotator_kernels(RotateType rotate_type, CpuArchType cpu_arch_type) { switch (rotate_type) { @@ -181,40 +207,8 @@ RotatorKernels get_rotator_kernels(RotateType rotate_type, } } - // Unsupported RotateType: assert in debug for early detection, but always - // return the scalar kernels so release builds never hand back null function - // pointers (which would crash on the first call). assert(false && "unsupported RotateType"); return {scalar::fht_rotate, scalar::fht_unrotate}; } -PqKernels get_pq_kernels(DataType data_type, QuantizeType quantize_type, - CpuArchType cpu_arch_type) { - (void)data_type; - if (quantize_type == QuantizeType::kPQ) { - if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F && - IsArchMatch(cpu_arch_type, CpuArchType::kAVX512)) { - return {avx512::pq_adc_int8_distance_avx512, - avx512::pq_sdc_int8_distance_avx512, - avx512::pq_adc_int8_batch_distance_avx512}; - } - if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2 && - IsArchMatch(cpu_arch_type, CpuArchType::kAVX2)) { - return {avx2::pq_adc_int8_distance_avx2, - avx2::pq_sdc_int8_distance_avx2, - avx2::pq_adc_int8_batch_distance_avx2}; - } - if (zvec::ailego::internal::CpuFeatures::static_flags_.NEON && - IsArchMatch(cpu_arch_type, CpuArchType::kNEON)) { - return {neon::pq_adc_int8_distance_neon, - neon::pq_sdc_int8_distance_neon, - neon::pq_adc_int8_batch_distance_neon}; - } - return {scalar::pq_adc_int8_distance, - scalar::pq_sdc_int8_distance, - scalar::pq_adc_int8_batch_distance}; - } - return {}; -} - } // namespace zvec::turbo \ No newline at end of file From 9625382198e7ba0317f937ae6b54c59ca377e67b Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Thu, 23 Jul 2026 20:14:18 +0800 Subject: [PATCH 088/100] fix bug --- .../pq_int8_quantizer/pq_int8_quantizer.cc | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 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 42f7cbea1..62a79b5e0 100644 --- a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc @@ -422,6 +422,16 @@ void PqInt8Quantizer::quantize_query(const void *input, void *output) const { reinterpret_cast(query + m * sub_dim_), kNumCentroids, sub_dim_, lut + m * kNumCentroids); } + + // Cosine: the LUT holds ||q_m - c_m[j]||^2 on L2-normalized vectors, and + // ||q - c||^2 = 2 - 2*cos_sim = 2*(1 - cos_sim). Scale by 0.5 so the ADC + // sum yields cosine distance (1 - cos_sim) directly on every ADC path. + if (meta_.metric_name() == "Cosine") { + const size_t lut_size = static_cast(num_chunk_) * kNumCentroids; + for (size_t i = 0; i < lut_size; ++i) { + lut[i] *= 0.5f; + } + } } float PqInt8Quantizer::calc_distance_dp_query(const void *dp, @@ -431,8 +441,8 @@ float PqInt8Quantizer::calc_distance_dp_query(const void *dp, float d = 0.0f; adc_fn_(reinterpret_cast(dp), reinterpret_cast(query), num_chunk_, &d); - // Cosine uses an L2 LUT on normalized vectors, so the ADC sum is already - // a valid (squared-Euclidean) distance — no conversion needed. + // Cosine LUT is pre-scaled by 0.5 in quantize_query, so the ADC sum is + // already the cosine distance — no conversion needed here. return d; } @@ -444,8 +454,8 @@ void PqInt8Quantizer::calc_distance_dp_query_batch(const void *const *dp_list, // const_cast: batch_adc_fn_ expects const void**; kernel is read-only. batch_adc_fn_(const_cast(dp_list), query, static_cast(dp_num), num_chunk_, dist_list); - // Cosine uses an L2 LUT on normalized vectors — the batch ADC sums are - // already valid distances, so no conversion is applied. + // Cosine LUT is pre-scaled by 0.5 in quantize_query, so the batch ADC sums + // are already cosine distances — no conversion applied here. } float PqInt8Quantizer::calc_distance_dp_query_unquantized( From abf6c32602e271c3dd46e5f6d49b959d1ae2c372 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Thu, 23 Jul 2026 20:18:11 +0800 Subject: [PATCH 089/100] clang-format --- src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc | 7 ++++--- 1 file changed, 4 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 62a79b5e0..eb714eaac 100644 --- a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc @@ -102,9 +102,10 @@ int PqInt8Quantizer::init(const IndexMeta &meta, const ailego::Params ¶ms) { if (use_zero_mean_ && meta_.metric_name() != "SquaredEuclidean" && meta_.metric_name() != "Cosine") { - LOG_WARN("PqInt8Quantizer: use_zero_mean is incompatible with metric '%s', " - "disabling centering", - meta_.metric_name().c_str()); + LOG_WARN( + "PqInt8Quantizer: use_zero_mean is incompatible with metric '%s', " + "disabling centering", + meta_.metric_name().c_str()); use_zero_mean_ = false; } From e800e98545d01fc462bcc80bd6c025c0b7576f48 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Fri, 24 Jul 2026 10:40:05 +0800 Subject: [PATCH 090/100] clang-format --- src/include/zvec/core/interface/index_param.h | 2 +- src/include/zvec/turbo/turbo.h | 7 ++-- .../avx2/pq_quantizer_int8/pq_distance.cc | 4 +- .../avx2/pq_quantizer_int8/pq_distance.h | 3 +- .../avx512/pq_quantizer_int8/pq_distance.cc | 4 +- .../avx512/pq_quantizer_int8/pq_distance.h | 4 +- .../neon/pq_quantizer_int8/pq_distance.cc | 38 +++++++++---------- .../neon/pq_quantizer_int8/pq_distance.h | 3 +- .../scalar/pq_quantizer_int8/pq_distance.cc | 3 +- .../scalar/pq_quantizer_int8/pq_distance.h | 3 +- 10 files changed, 33 insertions(+), 38 deletions(-) diff --git a/src/include/zvec/core/interface/index_param.h b/src/include/zvec/core/interface/index_param.h index 70a2e5474..27fc52bdc 100644 --- a/src/include/zvec/core/interface/index_param.h +++ b/src/include/zvec/core/interface/index_param.h @@ -122,7 +122,7 @@ struct SerializableBase { struct QuantizerParam : public SerializableBase { QuantizerType type = QuantizerType::kNone; int num_chunk = 8; // M - int num_bits = 8; // bits per subquantizer + int num_bits = 8; // bits per subquantizer bool enable_rotate = false; // rotate vectors before quantization to reduce error diff --git a/src/include/zvec/turbo/turbo.h b/src/include/zvec/turbo/turbo.h index aecfb930b..1d20ce7ed 100644 --- a/src/include/zvec/turbo/turbo.h +++ b/src/include/zvec/turbo/turbo.h @@ -78,14 +78,13 @@ using PqAdcDistanceFunc = void (*)(const void *pq_code, const void *lut, // dist_table: [num_chunk * 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_chunk, float *out); + const void *dist_table, size_t num_chunk, + 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_chunk, - float *out); + size_t num, size_t num_chunk, float *out); // ISA-dispatched rotate/unrotate kernels. struct RotatorKernels { 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 e9e762c00..b63d8357d 100644 --- a/src/turbo/distance/avx2/pq_quantizer_int8/pq_distance.cc +++ b/src/turbo/distance/avx2/pq_quantizer_int8/pq_distance.cc @@ -100,8 +100,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_chunk, float *out) { + const void *dist_table_v, size_t num_chunk, + float *out) { #if defined(__AVX2__) constexpr int kNumCentroids = 256; constexpr int kTablePerSub = kNumCentroids * kNumCentroids; // 65536 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 e4420fdeb..f95aba404 100644 --- a/src/turbo/distance/avx2/pq_quantizer_int8/pq_distance.h +++ b/src/turbo/distance/avx2/pq_quantizer_int8/pq_distance.h @@ -36,7 +36,6 @@ void pq_sdc_int8_distance_avx2(const void *a, const void *b, // 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_chunk, - float *out); + size_t num, size_t num_chunk, float *out); } // namespace zvec::turbo::avx2 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 79ee4c9dc..28fbd1b30 100644 --- a/src/turbo/distance/avx512/pq_quantizer_int8/pq_distance.cc +++ b/src/turbo/distance/avx512/pq_quantizer_int8/pq_distance.cc @@ -93,8 +93,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_chunk, float *out) { + const void *dist_table_v, size_t num_chunk, + float *out) { #if defined(__AVX512F__) constexpr int kNumCentroids = 256; constexpr int kTablePerSub = kNumCentroids * kNumCentroids; // 65536 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 0a3cc3987..64b35c017 100644 --- a/src/turbo/distance/avx512/pq_quantizer_int8/pq_distance.h +++ b/src/turbo/distance/avx512/pq_quantizer_int8/pq_distance.h @@ -27,8 +27,8 @@ void pq_adc_int8_distance_avx512(const void *pq_code, const void *lut, // 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_chunk, float *out); + const void *dist_table, size_t num_chunk, + float *out); // Batch ADC via AVX512 gather: process 4 candidates per iteration, // each using 16-wide _mm512_i32gather_ps. 4 independent __m512 diff --git a/src/turbo/distance/neon/pq_quantizer_int8/pq_distance.cc b/src/turbo/distance/neon/pq_quantizer_int8/pq_distance.cc index abffbfe37..9cc313f8f 100644 --- a/src/turbo/distance/neon/pq_quantizer_int8/pq_distance.cc +++ b/src/turbo/distance/neon/pq_quantizer_int8/pq_distance.cc @@ -34,7 +34,7 @@ namespace { inline float horizontal_sum_neon(float32x4_t v) { float32x2_t lo = vget_low_f32(v); float32x2_t hi = vget_high_f32(v); - float32x2_t sum2 = vadd_f32(lo, hi); // [a+b, c+d] + float32x2_t sum2 = vadd_f32(lo, hi); // [a+b, c+d] return vget_lane_f32(vpadd_f32(sum2, sum2), 0); // a+b+c+d } @@ -81,8 +81,8 @@ void pq_adc_int8_distance_neon(const void *pq_code_v, const void *lut_v, } void pq_sdc_int8_distance_neon(const void *a_v, const void *b_v, - const void *dist_table_v, - size_t num_chunk, float *out) { + const void *dist_table_v, size_t num_chunk, + float *out) { #if defined(__ARM_NEON) && defined(__aarch64__) constexpr int kNumCentroids = 256; constexpr int kTablePerSub = kNumCentroids * kNumCentroids; // 65536 @@ -98,17 +98,17 @@ void pq_sdc_int8_distance_neon(const void *a_v, const void *b_v, // Main loop: process 4 subquantizers per iteration. for (; m + kChunkSize <= num_chunk; m += kChunkSize) { float d0 = dist_table[(m + 0) * kTablePerSub + - static_cast(a[m + 0]) * kNumCentroids + - static_cast(b[m + 0])]; + static_cast(a[m + 0]) * kNumCentroids + + static_cast(b[m + 0])]; float d1 = dist_table[(m + 1) * kTablePerSub + - static_cast(a[m + 1]) * kNumCentroids + - static_cast(b[m + 1])]; + static_cast(a[m + 1]) * kNumCentroids + + static_cast(b[m + 1])]; float d2 = dist_table[(m + 2) * kTablePerSub + - static_cast(a[m + 2]) * kNumCentroids + - static_cast(b[m + 2])]; + static_cast(a[m + 2]) * kNumCentroids + + static_cast(b[m + 2])]; float d3 = dist_table[(m + 3) * kTablePerSub + - static_cast(a[m + 3]) * kNumCentroids + - static_cast(b[m + 3])]; + static_cast(a[m + 3]) * kNumCentroids + + static_cast(b[m + 3])]; float32x4_t d = {d0, d1, d2, d3}; acc = vaddq_f32(acc, d); } @@ -158,14 +158,14 @@ void pq_adc_int8_batch_distance_neon(const void **candidates_v, for (; m + kChunkSize <= num_chunk; m += kChunkSize) { const float *tab = lut + m * kNumCentroids; - float32x4_t d0 = {tab[c0[m + 0]], tab[c0[m + 1]], - tab[c0[m + 2]], tab[c0[m + 3]]}; - float32x4_t d1 = {tab[c1[m + 0]], tab[c1[m + 1]], - tab[c1[m + 2]], tab[c1[m + 3]]}; - float32x4_t d2 = {tab[c2[m + 0]], tab[c2[m + 1]], - tab[c2[m + 2]], tab[c2[m + 3]]}; - float32x4_t d3 = {tab[c3[m + 0]], tab[c3[m + 1]], - tab[c3[m + 2]], tab[c3[m + 3]]}; + float32x4_t d0 = {tab[c0[m + 0]], tab[c0[m + 1]], tab[c0[m + 2]], + tab[c0[m + 3]]}; + float32x4_t d1 = {tab[c1[m + 0]], tab[c1[m + 1]], tab[c1[m + 2]], + tab[c1[m + 3]]}; + float32x4_t d2 = {tab[c2[m + 0]], tab[c2[m + 1]], tab[c2[m + 2]], + tab[c2[m + 3]]}; + float32x4_t d3 = {tab[c3[m + 0]], tab[c3[m + 1]], tab[c3[m + 2]], + tab[c3[m + 3]]}; acc0 = vaddq_f32(acc0, d0); acc1 = vaddq_f32(acc1, d1); diff --git a/src/turbo/distance/neon/pq_quantizer_int8/pq_distance.h b/src/turbo/distance/neon/pq_quantizer_int8/pq_distance.h index b1be2cf2e..d0c0e55d4 100644 --- a/src/turbo/distance/neon/pq_quantizer_int8/pq_distance.h +++ b/src/turbo/distance/neon/pq_quantizer_int8/pq_distance.h @@ -35,7 +35,6 @@ void pq_sdc_int8_distance_neon(const void *a, const void *b, // Batch ADC: compute distances for multiple PQ codes against a shared LUT. // Processes 4 candidates per iteration with NEON vector accumulation. void pq_adc_int8_batch_distance_neon(const void **candidates, const void *lut, - size_t num, size_t num_chunk, - float *out); + size_t num, size_t num_chunk, float *out); } // namespace zvec::turbo::neon 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 77ee86cb5..6d8ca0d4f 100644 --- a/src/turbo/distance/scalar/pq_quantizer_int8/pq_distance.cc +++ b/src/turbo/distance/scalar/pq_quantizer_int8/pq_distance.cc @@ -46,8 +46,7 @@ void pq_sdc_int8_distance(const void *a_v, const void *b_v, } void pq_adc_int8_batch_distance(const void **candidates_v, const void *lut_v, - size_t num, size_t num_chunk, - float *out) { + size_t num, size_t num_chunk, 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** 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 7ec8c1849..580361e76 100644 --- a/src/turbo/distance/scalar/pq_quantizer_int8/pq_distance.h +++ b/src/turbo/distance/scalar/pq_quantizer_int8/pq_distance.h @@ -44,7 +44,6 @@ void pq_sdc_int8_distance(const void *a, const void *b, const void *dist_table, // 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_chunk, - float *out); + size_t num, size_t num_chunk, float *out); } // namespace zvec::turbo::scalar From d8900618601740c4b10148b056f7b614c1ab56f8 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Fri, 24 Jul 2026 10:49:39 +0800 Subject: [PATCH 091/100] fix memcpy --- src/turbo/distance/common/fht_common.h | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/src/turbo/distance/common/fht_common.h b/src/turbo/distance/common/fht_common.h index 7285c2c23..308b1a8d3 100644 --- a/src/turbo/distance/common/fht_common.h +++ b/src/turbo/distance/common/fht_common.h @@ -38,10 +38,9 @@ struct FhtPrimitives { /// offset 24: uint8_t flip[] inline void fht_rotate_impl(const float *in, float *out, size_t dim, void *ctx, const FhtPrimitives &p) { - // Seed the output buffer with the input, then rotate it in place. Keeping the - // copy here makes the kernel a self-contained out-of-place transform (read - // `in`, write `out`); callers no longer need to pre-copy `in` into `out`. - std::memcpy(out, in, sizeof(float) * dim); + if (out != in) { + std::memcpy(out, in, sizeof(float) * dim); + } float *data = out; auto *base = reinterpret_cast(ctx); const size_t flip_offset = *reinterpret_cast(base); @@ -86,10 +85,9 @@ inline void fht_rotate_impl(const float *in, float *out, size_t dim, void *ctx, inline void fht_unrotate_impl(const float *in, float *out, size_t dim, void *ctx, const FhtPrimitives &p) { - // Seed the output buffer with the input, then unrotate it in place. Keeping - // the copy here makes the kernel a self-contained out-of-place transform - // (read `in`, write `out`); callers no longer need to pre-copy `in`. - std::memcpy(out, in, sizeof(float) * dim); + if (out != in) { + std::memcpy(out, in, sizeof(float) * dim); + } float *data = out; auto *base = reinterpret_cast(ctx); const size_t flip_offset = *reinterpret_cast(base); From a1ddd4d4cba40ad7ade116599795866e730cd408 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Fri, 24 Jul 2026 14:08:31 +0800 Subject: [PATCH 092/100] debug --- src/include/zvec/turbo/turbo.h | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/include/zvec/turbo/turbo.h b/src/include/zvec/turbo/turbo.h index 1d20ce7ed..6b7e83c6f 100644 --- a/src/include/zvec/turbo/turbo.h +++ b/src/include/zvec/turbo/turbo.h @@ -95,9 +95,9 @@ struct RotatorKernels { // data_type selects the code packing layout: // kInt8: one uint8 per sub-quantizer (256 centroids, stride=256) struct PqKernels { - PqAdcDistanceFunc adc_distance; - PqSdcKernelFunc sdc_distance; - PqBatchAdcFunc batch_adc_distance; + PqAdcDistanceFunc adc_distance = nullptr; + PqSdcKernelFunc sdc_distance = nullptr; + PqBatchAdcFunc batch_adc_distance = nullptr; }; enum class MetricType { From 7324b9bea04a38d6c349e2a9de56beb4154d7754 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Fri, 24 Jul 2026 14:08:40 +0800 Subject: [PATCH 093/100] debug --- src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 e772fb6e3..f930d94b5 100644 --- a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h @@ -30,7 +30,7 @@ using namespace zvec::core; //! //! Datapoints are encoded as uint8_t[num_chunk] codes. //! Queries are encoded as a float LUT of size [num_chunk * 256] -//! via compute_distance_table(). Distance between a PQ code and a +//! 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). class PqInt8Quantizer : public Quantizer { @@ -155,7 +155,7 @@ class PqInt8Quantizer : public Quantizer { std::vector centroids_; //! Global centroid (per-dimension mean) for zero-mean centering. - //! Size: original_dim_ floats. Only populated when use_centering_ = true. + //! Size: original_dim_ floats. Only populated when use_zero_mean_ = true. std::vector centroid_; //! Centroid-to-centroid distance table for SDC: From a5a36e0dadad32bf547be3c7288261743b635396 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Fri, 24 Jul 2026 14:48:18 +0800 Subject: [PATCH 094/100] Reduce unnecessary changes --- src/db/index/common/index_params.cc | 3 +- src/db/index/common/schema.cc | 6 +- src/db/index/common/stats.cc | 4 +- .../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 ++--- 17 files changed, 178 insertions(+), 180 deletions(-) diff --git a/src/db/index/common/index_params.cc b/src/db/index/common/index_params.cc index 8f08bcdbc..0b696956a 100644 --- a/src/db/index/common/index_params.cc +++ b/src/db/index/common/index_params.cc @@ -26,7 +26,8 @@ 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/schema.cc b/src/db/index/common/schema.cc index f38dccfe2..532696803 100644 --- a/src/db/index/common/schema.cc +++ b/src/db/index/common/schema.cc @@ -286,7 +286,8 @@ 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_; @@ -371,7 +372,8 @@ 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 ff1223ec1..fdb26f25a 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/tests/ailego/parallel/multi_thread_list_test.cc b/tests/ailego/parallel/multi_thread_list_test.cc index 3a94afebf..17a896f1b 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 439fdc3ae..9dc8fc7e7 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 1538da6b9..20cd45da7 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 81ee6e36e..372e2b463 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 46365cbaf..17c3c8704 100644 --- a/tests/core/algorithm/hnsw/hnsw_streamer_buffer_test.cc +++ b/tests/core/algorithm/hnsw/hnsw_streamer_buffer_test.cc @@ -334,8 +334,7 @@ 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 1079e01ff..894c68f1e 100644 --- a/tests/core/utility/buffer_storage_write_test.cc +++ b/tests/core/utility/buffer_storage_write_test.cc @@ -42,9 +42,7 @@ 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() { @@ -71,8 +69,7 @@ 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); @@ -278,7 +275,8 @@ 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()); } @@ -355,7 +353,8 @@ 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); @@ -882,8 +881,7 @@ 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(); @@ -894,8 +892,7 @@ 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())); @@ -911,8 +908,7 @@ 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())); @@ -1051,9 +1047,8 @@ 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); @@ -1117,8 +1112,7 @@ 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 ceb4b7364..826111ca2 100644 --- a/tests/db/crash_recovery/utility.h +++ b/tests/db/crash_recovery/utility.h @@ -22,6 +22,7 @@ #include #include #include + #include #include @@ -155,6 +156,7 @@ 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 e58f41793..620b4ef6f 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 c51e6f733..b26fe7123 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,8 +337,7 @@ 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 a64d662c3..07bc1e7da 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 ee634cff2..0d2bda43a 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,8 +698,7 @@ 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 9d1fbfd21..b1da86595 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 c5f4c7b1b..9582d2bf3 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,8 +1674,9 @@ 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}); @@ -1698,8 +1699,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(); @@ -1806,8 +1807,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(); @@ -1890,8 +1891,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(); @@ -1977,8 +1978,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(); @@ -2086,8 +2087,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( @@ -2202,8 +2203,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 7356438f5..ca8bed2c9 100644 --- a/tests/ios_test_sandbox.cc +++ b/tests/ios_test_sandbox.cc @@ -7,38 +7,39 @@ #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__ From 4e690470879f51aebd51a95a2e80a6f8918ece75 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Fri, 24 Jul 2026 14:51:50 +0800 Subject: [PATCH 095/100] Reduce unnecessary changes --- src/turbo/quantizer/distance.h | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/turbo/quantizer/distance.h b/src/turbo/quantizer/distance.h index 19658e641..bc8af6c1a 100644 --- a/src/turbo/quantizer/distance.h +++ b/src/turbo/quantizer/distance.h @@ -85,7 +85,8 @@ class DistanceImpl { } //! Raw scalar distance function (operates on already-quantized - //! candidates). This is the ADC path for PQ. + //! candidates). Useful for pairwise node-vs-node distance where no + //! stored query is involved. const DistanceFunc &func() const { return func_; } From bedd5c38361da23c3f0cff622602905a10223c3c Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Fri, 24 Jul 2026 17:57:06 +0800 Subject: [PATCH 096/100] fix LNK2005 --- .../pq_int8_quantizer/pq_int8_quantizer.cc | 22 ------------------- 1 file changed, 22 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 eb714eaac..dd4c4ce36 100644 --- a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc @@ -23,7 +23,6 @@ #include #include #include -#include #include namespace zvec { @@ -50,12 +49,9 @@ int PqInt8Quantizer::init(const IndexMeta &meta, const ailego::Params ¶ms) { // Read num_chunk from params (required). uint32_t nsq = 0; if (!params.get("num_chunk", &nsq) || nsq == 0) { - LOG_ERROR("PqInt8Quantizer: num_chunk not set or zero"); return kErrUnsupported; } if (d % nsq != 0) { - LOG_ERROR("PqInt8Quantizer: dim (%u) is not divisible by num_chunk (%u)", d, - nsq); return kErrUnsupported; } @@ -102,10 +98,6 @@ int PqInt8Quantizer::init(const IndexMeta &meta, const ailego::Params ¶ms) { if (use_zero_mean_ && meta_.metric_name() != "SquaredEuclidean" && meta_.metric_name() != "Cosine") { - LOG_WARN( - "PqInt8Quantizer: use_zero_mean is incompatible with metric '%s', " - "disabling centering", - meta_.metric_name().c_str()); use_zero_mean_ = false; } @@ -169,7 +161,6 @@ void PqInt8Quantizer::train_subquantizer(const float *data, size_t num, double old_cost = cost; bool result = algorithm.cluster_once(*local_threads, &cost); if (!result) { - LOG_ERROR("sub[%zu] cluster_once failed at iter %u", sub_idx, iter); break; } double new_epsilon = std::abs(cost - old_cost); @@ -209,8 +200,6 @@ int PqInt8Quantizer::train(IndexHolder::Pointer holder, int thread_count) { // 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. @@ -264,14 +253,8 @@ int PqInt8Quantizer::train(IndexHolder::Pointer holder, int thread_count) { v[d] -= centroid_[d]; } } - LOG_INFO("PQ centering: zero-mean centroid computed"); } - LOG_INFO( - "PQ training: %zu vectors, dim=%u, nsq=%u, sub_dim=%u, " - "max_iters=%u, threads=%d", - num, original_dim_, num_chunk_, sub_dim_, kMaxKmeansIters, thread_count); - // Create thread pool. auto threads = std::make_shared( static_cast(thread_count), false); @@ -292,17 +275,12 @@ int PqInt8Quantizer::train(IndexHolder::Pointer holder, int thread_count) { } task_group->wait_finish(); - LOG_INFO(" all %u sub-quantizers trained (%zu threads)", num_chunk_, - pool_count); - // Pre-build centroid pointer cache (needed by compute_dist_table). build_centroid_ptrs_cache(); // Pre-compute SDC dist_table. - LOG_INFO("Computing SDC dist_table ..."); compute_dist_table(); - LOG_INFO("PQ training complete."); return 0; } From 952521ee803e0f960c0a8d66754d50590a9b520e Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Tue, 28 Jul 2026 15:15:06 +0800 Subject: [PATCH 097/100] add fp16 in pq --- src/turbo/distance/scalar/fp16/cosine.cc | 40 ++ src/turbo/distance/scalar/fp16/cosine.h | 30 ++ .../distance/scalar/fp16/inner_product.cc | 41 ++ .../distance/scalar/fp16/inner_product.h | 30 ++ .../distance/scalar/fp16/squared_euclidean.cc | 42 ++ .../distance/scalar/fp16/squared_euclidean.h | 30 ++ .../pq_int8_quantizer/pq_int8_quantizer.cc | 414 ++++++++++++------ .../pq_int8_quantizer/pq_int8_quantizer.h | 48 +- src/turbo/turbo.cc | 33 ++ tests/turbo/turbo_pq_int8_quantizer_test.cc | 372 ++++++++++++++++ 10 files changed, 928 insertions(+), 152 deletions(-) create mode 100644 src/turbo/distance/scalar/fp16/cosine.cc create mode 100644 src/turbo/distance/scalar/fp16/cosine.h create mode 100644 src/turbo/distance/scalar/fp16/inner_product.cc create mode 100644 src/turbo/distance/scalar/fp16/inner_product.h create mode 100644 src/turbo/distance/scalar/fp16/squared_euclidean.cc create mode 100644 src/turbo/distance/scalar/fp16/squared_euclidean.h diff --git a/src/turbo/distance/scalar/fp16/cosine.cc b/src/turbo/distance/scalar/fp16/cosine.cc new file mode 100644 index 000000000..a6bd8d61c --- /dev/null +++ b/src/turbo/distance/scalar/fp16/cosine.cc @@ -0,0 +1,40 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "scalar/fp16/cosine.h" +#include "scalar/fp16/inner_product.h" + +namespace zvec::turbo::scalar { + +void cosine_fp16_distance(const void *a, const void *b, size_t dim, + float *distance) { + // inner_product_fp16_distance returns -real_IP; cosine = 1 - real_IP = 1 + + // ip. + float ip; + inner_product_fp16_distance(a, b, dim, &ip); + + *distance = 1 + ip; +} + +void cosine_fp16_batch_distance(const void *const *vectors, const void *query, + size_t n, size_t dim, float *distances) { + inner_product_fp16_batch_distance(vectors, query, n, dim, distances); + // inner_product batch returns -real_IP per element; cosine = 1 - real_IP = + // 1 + d. + for (size_t i = 0; i < n; i++) { + distances[i] = 1 + distances[i]; + } +} + +} // namespace zvec::turbo::scalar diff --git a/src/turbo/distance/scalar/fp16/cosine.h b/src/turbo/distance/scalar/fp16/cosine.h new file mode 100644 index 000000000..39be959b2 --- /dev/null +++ b/src/turbo/distance/scalar/fp16/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 between a single FP16 vector pair. +// Returns 1 - dot(a, b). +void cosine_fp16_distance(const void *a, const void *b, size_t dim, + float *distance); + +// Batch version of cosine_fp16_distance. +void cosine_fp16_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/fp16/inner_product.cc b/src/turbo/distance/scalar/fp16/inner_product.cc new file mode 100644 index 000000000..a8d944b0e --- /dev/null +++ b/src/turbo/distance/scalar/fp16/inner_product.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/fp16/inner_product.h" +#include + +namespace zvec::turbo::scalar { + +void inner_product_fp16_distance(const void *a, const void *b, size_t dim, + float *distance) { + const ailego::Float16 *m = reinterpret_cast(a); + const ailego::Float16 *q = reinterpret_cast(b); + + float sum = 0.0f; + for (size_t i = 0; i < dim; ++i) { + sum += static_cast(m[i]) * static_cast(q[i]); + } + + *distance = -sum; +} + +void inner_product_fp16_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_fp16_distance(vectors[i], query, dim, &distances[i]); + } +} + +} // namespace zvec::turbo::scalar diff --git a/src/turbo/distance/scalar/fp16/inner_product.h b/src/turbo/distance/scalar/fp16/inner_product.h new file mode 100644 index 000000000..08059ed3d --- /dev/null +++ b/src/turbo/distance/scalar/fp16/inner_product.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 negated inner product between a single FP16 vector pair. +void inner_product_fp16_distance(const void *a, const void *b, size_t dim, + float *distance); + +// Batch version of inner_product_fp16_distance. +void inner_product_fp16_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/fp16/squared_euclidean.cc b/src/turbo/distance/scalar/fp16/squared_euclidean.cc new file mode 100644 index 000000000..931028b40 --- /dev/null +++ b/src/turbo/distance/scalar/fp16/squared_euclidean.cc @@ -0,0 +1,42 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "scalar/fp16/squared_euclidean.h" +#include + +namespace zvec::turbo::scalar { + +void squared_euclidean_fp16_distance(const void *a, const void *b, size_t dim, + float *distance) { + const ailego::Float16 *m = reinterpret_cast(a); + const ailego::Float16 *q = reinterpret_cast(b); + + float sum = 0.0f; + for (size_t i = 0; i < dim; ++i) { + float diff = static_cast(m[i]) - static_cast(q[i]); + sum += diff * diff; + } + + *distance = sum; +} + +void squared_euclidean_fp16_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_fp16_distance(vectors[i], query, dim, &distances[i]); + } +} + +} // namespace zvec::turbo::scalar diff --git a/src/turbo/distance/scalar/fp16/squared_euclidean.h b/src/turbo/distance/scalar/fp16/squared_euclidean.h new file mode 100644 index 000000000..b395066e2 --- /dev/null +++ b/src/turbo/distance/scalar/fp16/squared_euclidean.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 squared euclidean distance between a single FP16 vector pair. +void squared_euclidean_fp16_distance(const void *a, const void *b, size_t dim, + float *distance); + +// Batch version of squared euclidean FP16. +void squared_euclidean_fp16_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/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc index dd4c4ce36..9b741f59a 100644 --- a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc @@ -37,12 +37,22 @@ struct PqInt8SerPayload { uint32_t sub_dim; uint32_t num_centroids; // always 256 for int8 uint8_t use_zero_mean; - uint8_t reserved[3]; + uint8_t input_data_type; // turbo DataType: kFp32=3, kFp16=2 + uint8_t reserved[2]; }; int PqInt8Quantizer::init(const IndexMeta &meta, const ailego::Params ¶ms) { meta_ = meta; + // Map core IndexMeta::DataType to turbo DataType. + if (meta.data_type() == IndexMeta::DataType::DT_FP16) { + input_data_type_ = DataType::kFp16; + } else if (meta.data_type() == IndexMeta::DataType::DT_FP32) { + input_data_type_ = DataType::kFp32; + } else { + return kErrUnsupported; + } + uint32_t d = meta.dimension(); original_dim_ = d; @@ -58,9 +68,9 @@ int PqInt8Quantizer::init(const IndexMeta &meta, const ailego::Params ¶ms) { num_chunk_ = nsq; sub_dim_ = d / nsq; - // Pre-allocate centroids (filled by train()). - centroids_.resize(static_cast(num_chunk_) * kNumCentroids * sub_dim_, - 0.0f); + // Pre-allocate centroids as raw bytes in the original data type. + centroids_.resize(static_cast(num_chunk_) * kNumCentroids * sub_dim_ * + element_size()); // Dispatch ISA kernels (scalar only for now). auto pq_k = get_pq_kernels(DataType::kInt8); @@ -68,26 +78,27 @@ int PqInt8Quantizer::init(const IndexMeta &meta, const ailego::Params ¶ms) { sdc_fn_ = pq_k.sdc_distance; batch_adc_fn_ = pq_k.batch_adc_distance; - // Resolve the configured metric type (local only — not stored). + // Resolve the configured metric type (local only -- not stored). auto mt = metric_from_name(meta_.metric_name()); - // L2-only batch distance for encoding and KMeans training: the PQ codebook - // is trained in L2 space regardless of the search metric. - fp32_l2_batch_fn_ = - get_batch_distance_func(MetricType::kSquaredEuclidean, DataType::kFp32, + // L2-only batch distance for encoding: the PQ codebook is trained in L2 + // space regardless of the search metric. Data type matches input. + l2_batch_fn_ = + get_batch_distance_func(MetricType::kSquaredEuclidean, input_data_type_, QuantizeType::kDefault, CpuArchType::kAuto); // Cosine = normalize + L2: after normalization cosine distance is monotonic // with squared-Euclidean, so the search LUT reuses SquaredEuclidean. if (meta_.metric_name() == "Cosine") { - fp32_batch_fn_ = - get_batch_distance_func(MetricType::kSquaredEuclidean, DataType::kFp32, - QuantizeType::kDefault, CpuArchType::kAuto); + batch_fn_ = + get_batch_distance_func(MetricType::kSquaredEuclidean, + input_data_type_, QuantizeType::kDefault, + CpuArchType::kAuto); extra_meta_size_ = kExtraMetaSizeCosine; meta_.set_extra_meta_size(extra_meta_size_); } else { - fp32_batch_fn_ = get_batch_distance_func( - mt, DataType::kFp32, QuantizeType::kDefault, CpuArchType::kAuto); + batch_fn_ = get_batch_distance_func( + mt, input_data_type_, QuantizeType::kDefault, CpuArchType::kAuto); } // Read optional training params (aligned with multi_chunk_cluster) @@ -101,54 +112,45 @@ int PqInt8Quantizer::init(const IndexMeta &meta, const ailego::Params ¶ms) { use_zero_mean_ = false; } - meta_.set_meta(IndexMeta::DataType::DT_FP32, d); + // Set output meta: the quantized representation is INT8 codes with + // num_chunk_ bytes (+ extra_meta_size_ for Cosine norm storage). + meta_.set_meta(IndexMeta::DataType::DT_INT8, num_chunk_); + meta_.set_extra_meta_size(extra_meta_size_); return 0; } // --------------------------------------------------------------------------- -// Simple Lloyd's KMeans for one sub-quantizer. +// Simple Lloyd's KMeans for one sub-quantizer (templated on data type T). // --------------------------------------------------------------------------- -void PqInt8Quantizer::build_centroid_ptrs_cache() { - const size_t k = kNumCentroids; - const size_t d = sub_dim_; - centroid_ptrs_cache_.resize(num_chunk_); - for (uint32_t m = 0; m < num_chunk_; ++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, +template +void PqInt8Quantizer::train_subquantizer(const T *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; + uint8_t *centroids_m = + centroids_.data() + + static_cast(sub_idx) * k * d * sizeof(T); // Non-spherical L2 KMeans: the PQ codebook must minimize L2 reconstruction // error, so centroids are the true (magnitude-preserving) means. - ailego::NumericalKmeans algorithm(k, d); + ailego::NumericalKmeans algorithm(k, d); // Append sub-vectors (NumericalKmeans handles transpose internally) for (size_t i = 0; i < num; ++i) { - const float *sub_vec = - reinterpret_cast( + const T *sub_vec = + reinterpret_cast( reinterpret_cast(data) + i * stride) + sub_idx * d; algorithm.append(sub_vec, d); } - // Single-threaded pool — parallelism is at the sub-quantizer level. + // Single-threaded pool -- parallelism is at the sub-quantizer level. auto local_threads = std::make_shared(1, false); // KMC2 centroid initialization. ailego::Kmc2CentroidsGenerator< - ailego::NumericalKmeans, + ailego::NumericalKmeans, SingleQueueIndexThreads> gen; gen.set_chain_length(markov_chain_length_); @@ -169,10 +171,10 @@ void PqInt8Quantizer::train_subquantizer(const float *data, size_t num, } } - // Extract centroids into the flat centroids_ buffer + // Extract centroids into the flat centroids_ byte buffer const auto ¢s = algorithm.centroids(); for (size_t c = 0; c < cents.count(); ++c) { - std::memcpy(centroids_m + c * d, cents[c], d * sizeof(float)); + std::memcpy(centroids_m + c * d * sizeof(T), cents[c], d * sizeof(T)); } } @@ -186,19 +188,19 @@ int PqInt8Quantizer::train(IndexHolder::Pointer holder, int thread_count) { } size_t num = holder->count(); + const uint32_t elem_sz = element_size(); - // Collect all data into a contiguous buffer. + // Collect all data into a contiguous byte buffer (original data type). auto iter = holder->create_iterator(); - std::vector all_data(num * original_dim_); + std::vector all_data(num * original_dim_ * elem_sz); 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)); + std::memcpy(all_data.data() + row * original_dim_ * elem_sz, iter->data(), + original_dim_ * elem_sz); } // Subsample if the dataset exceeds the training limit (aligned with - // faiss/vsag: 256 centroids * 256 max_points_per_centroid ≈ 65535). + // faiss/vsag: 256 centroids * 256 max_points_per_centroid ~= 65535). if (num > kMaxTrainVectors) { std::mt19937 rng(42); // Fisher-Yates partial shuffle: randomly place kMaxTrainVectors vectors @@ -207,27 +209,38 @@ int PqInt8Quantizer::train(IndexHolder::Pointer holder, int thread_count) { 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]); + // Swap full vectors (dim-sized chunks in bytes). + size_t vec_bytes = original_dim_ * elem_sz; + for (size_t b = 0; b < vec_bytes; ++b) { + std::swap(all_data[i * vec_bytes + b], all_data[j * vec_bytes + b]); } } } num = kMaxTrainVectors; - all_data.resize(num * original_dim_); + all_data.resize(num * original_dim_ * elem_sz); all_data.shrink_to_fit(); } - size_t data_stride = original_dim_ * sizeof(float); + size_t data_stride = original_dim_ * elem_sz; // For Cosine: normalize training data so centroids are learned in // normalized space (L2 minimization == maximizing cosine similarity). if (meta_.metric_name() == "Cosine") { - for (size_t i = 0; i < num; ++i) { - float *v = all_data.data() + i * original_dim_; - float norm = 0.0f; - ailego::Normalizer::L2(v, original_dim_, &norm); + if (input_data_type_ == DataType::kFp16) { + for (size_t i = 0; i < num; ++i) { + ailego::Float16 *v = + reinterpret_cast(all_data.data()) + + i * original_dim_; + float norm = 0.0f; + ailego::Normalizer::L2(v, original_dim_, &norm); + } + } else { + for (size_t i = 0; i < num; ++i) { + float *v = reinterpret_cast(all_data.data()) + + i * original_dim_; + float norm = 0.0f; + ailego::Normalizer::L2(v, original_dim_, &norm); + } } } @@ -237,20 +250,45 @@ int PqInt8Quantizer::train(IndexHolder::Pointer holder, int thread_count) { // (normalize -> center; dequantize: un-center -> rescale). if (use_zero_mean_) { centroid_.assign(original_dim_, 0.0f); - for (size_t i = 0; i < num; ++i) { - const float *v = all_data.data() + i * original_dim_; + if (input_data_type_ == DataType::kFp16) { + for (size_t i = 0; i < num; ++i) { + const ailego::Float16 *v = + reinterpret_cast(all_data.data()) + + i * original_dim_; + for (uint32_t d = 0; d < original_dim_; ++d) { + centroid_[d] += static_cast(v[d]); + } + } for (uint32_t d = 0; d < original_dim_; ++d) { - centroid_[d] += v[d]; + centroid_[d] /= static_cast(num); + } + // Subtract centroid from all training vectors. + for (size_t i = 0; i < num; ++i) { + ailego::Float16 *v = + reinterpret_cast(all_data.data()) + + i * original_dim_; + for (uint32_t d = 0; d < original_dim_; ++d) { + v[d] -= centroid_[d]; + } + } + } else { + for (size_t i = 0; i < num; ++i) { + const float *v = reinterpret_cast(all_data.data()) + + i * original_dim_; + for (uint32_t d = 0; d < original_dim_; ++d) { + centroid_[d] += v[d]; + } } - } - for (uint32_t d = 0; d < original_dim_; ++d) { - centroid_[d] /= static_cast(num); - } - // Subtract centroid from all training vectors. - for (size_t i = 0; i < num; ++i) { - float *v = all_data.data() + i * original_dim_; for (uint32_t d = 0; d < original_dim_; ++d) { - v[d] -= centroid_[d]; + centroid_[d] /= static_cast(num); + } + // Subtract centroid from all training vectors. + for (size_t i = 0; i < num; ++i) { + float *v = reinterpret_cast(all_data.data()) + + i * original_dim_; + for (uint32_t d = 0; d < original_dim_; ++d) { + v[d] -= centroid_[d]; + } } } } @@ -260,18 +298,36 @@ int PqInt8Quantizer::train(IndexHolder::Pointer holder, int thread_count) { static_cast(thread_count), false); auto task_group = threads->make_group(); - // Distribute sub-quantizers across threads. + // Distribute sub-quantizers across threads, dispatched by data type. std::atomic finished{0}; size_t pool_count = threads->count(); - for (size_t i = 0; i < pool_count; ++i) { - task_group->submit(ailego::Closure::New( - [this, &all_data, num, data_stride, i, pool_count, &finished]() { - for (uint32_t m = static_cast(i); m < num_chunk_; - m += static_cast(pool_count)) { - train_subquantizer(all_data.data(), num, data_stride, m); - finished++; - } - })); + if (input_data_type_ == DataType::kFp16) { + const auto *typed_data = + reinterpret_cast(all_data.data()); + for (size_t i = 0; i < pool_count; ++i) { + task_group->submit(ailego::Closure::New( + [this, typed_data, num, data_stride, i, pool_count, &finished]() { + for (uint32_t m = static_cast(i); m < num_chunk_; + m += static_cast(pool_count)) { + train_subquantizer(typed_data, num, data_stride, + m); + finished++; + } + })); + } + } else { + const auto *typed_data = + reinterpret_cast(all_data.data()); + for (size_t i = 0; i < pool_count; ++i) { + task_group->submit(ailego::Closure::New( + [this, typed_data, num, data_stride, i, pool_count, &finished]() { + for (uint32_t m = static_cast(i); m < num_chunk_; + m += static_cast(pool_count)) { + train_subquantizer(typed_data, num, data_stride, m); + finished++; + } + })); + } } task_group->wait_finish(); @@ -284,53 +340,89 @@ int PqInt8Quantizer::train(IndexHolder::Pointer holder, int thread_count) { return 0; } +void PqInt8Quantizer::build_centroid_ptrs_cache() { + const size_t k = kNumCentroids; + const size_t d = sub_dim_; + const size_t type_sz = element_size(); + const uint8_t *base = centroids_.data(); + + centroid_ptrs_cache_.resize(num_chunk_); + for (uint32_t m = 0; m < num_chunk_; ++m) { + auto &ptrs = centroid_ptrs_cache_[m]; + ptrs.resize(k); + for (size_t c = 0; c < k; ++c) { + ptrs[c] = base + (static_cast(m) * k * d + c * d) * type_sz; + } + } +} + void PqInt8Quantizer::compute_dist_table() { const size_t k = kNumCentroids; const size_t d = sub_dim_; dist_table_.resize(static_cast(num_chunk_) * k * k, 0.0f); - // Centroid-to-centroid distances via the metric-aware fp32_batch_fn_: + // Centroid-to-centroid distances via the metric-aware 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, uses L2. for (uint32_t m = 0; m < num_chunk_; ++m) { - const float *centroids_m = centroids_.data() + m * k * d; float *table_m = dist_table_.data() + m * k * k; // Use pre-built centroid pointer cache. - // const_cast: .data() returns const void* const* but fp32_batch_fn_ + // const_cast: .data() returns const void* const* but batch_fn_ // expects const void**. The kernel never modifies the pointer array. const auto ¢roid_ptrs = centroid_ptrs_cache_[m]; + const void *centroid_i = centroid_ptrs[0]; 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); + batch_fn_(const_cast(centroid_ptrs.data()), + reinterpret_cast(centroid_i) + + static_cast(i) * d * element_size(), + k, d, table_m + i * k); } } } void PqInt8Quantizer::quantize_data(const void *input, void *output) const { - const float *vec = reinterpret_cast(input); uint8_t *code = reinterpret_cast(output); + const uint32_t elem_sz = element_size(); // For Cosine: normalize FIRST (codebook is trained in normalized space); // the original norm is stored after the PQ code for dequantize(). - std::vector norm_vec_storage; + std::vector norm_vec_storage; float vec_norm = 0.0f; + const void *vec = input; + if (meta_.metric_name() == "Cosine") { - norm_vec_storage.assign(vec, vec + original_dim_); - ailego::Normalizer::L2(norm_vec_storage.data(), original_dim_, - &vec_norm); + norm_vec_storage.resize(original_dim_ * elem_sz); + std::memcpy(norm_vec_storage.data(), input, original_dim_ * elem_sz); + if (input_data_type_ == DataType::kFp16) { + ailego::Normalizer::L2( + reinterpret_cast(norm_vec_storage.data()), + original_dim_, &vec_norm); + } else { + ailego::Normalizer::L2( + reinterpret_cast(norm_vec_storage.data()), original_dim_, + &vec_norm); + } vec = norm_vec_storage.data(); } - // Zero-mean centering: subtract centroid (computed in the same normalized - // space during train) before encoding. - std::vector centered_vec_storage; + // Zero-mean centering: subtract centroid before encoding. + std::vector centered_vec_storage; if (use_zero_mean_) { - centered_vec_storage.assign(vec, vec + original_dim_); - for (uint32_t d = 0; d < original_dim_; ++d) { - centered_vec_storage[d] -= centroid_[d]; + centered_vec_storage.resize(original_dim_ * elem_sz); + std::memcpy(centered_vec_storage.data(), vec, original_dim_ * elem_sz); + if (input_data_type_ == DataType::kFp16) { + ailego::Float16 *cv = + reinterpret_cast(centered_vec_storage.data()); + for (uint32_t d = 0; d < original_dim_; ++d) { + cv[d] -= centroid_[d]; + } + } else { + float *cv = reinterpret_cast(centered_vec_storage.data()); + for (uint32_t d = 0; d < original_dim_; ++d) { + cv[d] -= centroid_[d]; + } } vec = centered_vec_storage.data(); } @@ -338,15 +430,15 @@ void PqInt8Quantizer::quantize_data(const void *input, void *output) const { // Encode with L2-only batch distance (search-metric independent), // fusing argmin into the distance loop. float dists[kNumCentroids]; + const uint8_t *vec_bytes = reinterpret_cast(vec); for (uint32_t m = 0; m < num_chunk_; ++m) { - const float *sub_vec = vec + m * sub_dim_; + const void *sub_vec = vec_bytes + static_cast(m) * sub_dim_ * elem_sz; 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); + l2_batch_fn_(const_cast(centroid_ptrs.data()), sub_vec, + kNumCentroids, sub_dim_, dists); // Argmin: find nearest centroid. float best_dist = dists[0]; @@ -368,38 +460,61 @@ void PqInt8Quantizer::quantize_data(const void *input, void *output) const { } void PqInt8Quantizer::quantize_query(const void *input, void *output) const { - const float *query = reinterpret_cast(input); float *lut = reinterpret_cast(output); + const uint32_t elem_sz = element_size(); // For Cosine: normalize FIRST (Cosine uses an L2 LUT on normalized data), // consistent with train() / quantize_data(). - std::vector norm_query_storage; + std::vector norm_query_storage; + const void *query = input; + 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); + norm_query_storage.resize(original_dim_ * elem_sz); + std::memcpy(norm_query_storage.data(), input, original_dim_ * elem_sz); + if (input_data_type_ == DataType::kFp16) { + float norm = 0.0f; + ailego::Normalizer::L2( + reinterpret_cast(norm_query_storage.data()), + original_dim_, &norm); + } else { + float norm = 0.0f; + ailego::Normalizer::L2( + reinterpret_cast(norm_query_storage.data()), original_dim_, + &norm); + } query = norm_query_storage.data(); } // Zero-mean centering: subtract centroid before LUT computation. - std::vector centered_query_storage; + std::vector centered_query_storage; if (use_zero_mean_) { - centered_query_storage.assign(query, query + original_dim_); - for (uint32_t d = 0; d < original_dim_; ++d) { - centered_query_storage[d] -= centroid_[d]; + centered_query_storage.resize(original_dim_ * elem_sz); + std::memcpy(centered_query_storage.data(), query, original_dim_ * elem_sz); + if (input_data_type_ == DataType::kFp16) { + ailego::Float16 *cq = + reinterpret_cast(centered_query_storage.data()); + for (uint32_t d = 0; d < original_dim_; ++d) { + cq[d] -= centroid_[d]; + } + } else { + float *cq = reinterpret_cast(centered_query_storage.data()); + for (uint32_t d = 0; d < original_dim_; ++d) { + cq[d] -= centroid_[d]; + } } query = centered_query_storage.data(); } - // LUT[m][j] = distance(q_m, c_m[j]) via the metric-aware fp32_batch_fn_: + // LUT[m][j] = distance(q_m, c_m[j]) via the metric-aware batch_fn_. // L2/Cosine: ||q_m - c_m[j]||^2 IP: -dot(q_m, c_m[j]). // const_cast: see compute_dist_table for rationale. + const uint8_t *query_bytes = reinterpret_cast(query); for (uint32_t m = 0; m < num_chunk_; ++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); + const void *sub_query = + query_bytes + static_cast(m) * sub_dim_ * elem_sz; + batch_fn_(const_cast(centroid_ptrs.data()), sub_query, + kNumCentroids, sub_dim_, lut + m * kNumCentroids); } // Cosine: the LUT holds ||q_m - c_m[j]||^2 on L2-normalized vectors, and @@ -421,7 +536,7 @@ float PqInt8Quantizer::calc_distance_dp_query(const void *dp, adc_fn_(reinterpret_cast(dp), reinterpret_cast(query), num_chunk_, &d); // Cosine LUT is pre-scaled by 0.5 in quantize_query, so the ADC sum is - // already the cosine distance — no conversion needed here. + // already the cosine distance -- no conversion needed here. return d; } @@ -434,7 +549,7 @@ void PqInt8Quantizer::calc_distance_dp_query_batch(const void *const *dp_list, batch_adc_fn_(const_cast(dp_list), query, static_cast(dp_num), num_chunk_, dist_list); // Cosine LUT is pre-scaled by 0.5 in quantize_query, so the batch ADC sums - // are already cosine distances — no conversion applied here. + // are already cosine distances -- no conversion applied here. } float PqInt8Quantizer::calc_distance_dp_query_unquantized( @@ -469,7 +584,11 @@ 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)) { + // Validate unit_size against the input data type. + size_t expected_unit = + (input_data_type_ == DataType::kFp16) ? sizeof(ailego::Float16) + : sizeof(float); + if (qmeta.unit_size() != expected_unit) { return kErrUnsupported; } @@ -492,18 +611,29 @@ int PqInt8Quantizer::dequantize(const void *in, const IndexQueryMeta &qmeta, float *result = reinterpret_cast(&(*out)[0]); // Reconstruct by concatenating the selected centroids per sub-quantizer, - // in the space the codebook was trained in (normalized for Cosine). + // converting from the original data type to float. const size_t k = kNumCentroids; const size_t d = sub_dim_; + const uint32_t elem_sz = element_size(); + for (uint32_t m = 0; m < num_chunk_; ++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)); + const uint8_t *centroids_m = + centroids_.data() + static_cast(m) * k * d * elem_sz; + const uint8_t *centroid = centroids_m + static_cast(code[m]) * d * elem_sz; + + if (input_data_type_ == DataType::kFp16) { + const ailego::Float16 *src = + reinterpret_cast(centroid); + for (size_t j = 0; j < d; ++j) { + result[m * d + j] = static_cast(src[j]); + } + } else { + std::memcpy(result + m * d, centroid, d * sizeof(float)); + } } // Undo zero-mean centering: add the centroid back FIRST (centering was - // applied in normalized space during encode). + // applied in the original-type space during encode). if (use_zero_mean_) { for (uint32_t j = 0; j < original_dim_; ++j) { result[j] += centroid_[j]; @@ -556,7 +686,7 @@ DistanceImpl PqInt8Quantizer::sym_distance(const void *query, size_t code_bytes = quantized_datapoint_vector_length(); std::string code_storage(static_cast(query), code_bytes); - // SDC has no batch kernel — use the 3-arg constructor. + // SDC has no batch kernel -- use the 3-arg constructor. return DistanceImpl(std::move(sdc_func), std::move(code_storage), static_cast(num_chunk_)); } @@ -580,8 +710,9 @@ int PqInt8Quantizer::serialize(std::string *out) const { payload.sub_dim = sub_dim_; payload.num_centroids = kNumCentroids; payload.use_zero_mean = use_zero_mean_ ? 1 : 0; + payload.input_data_type = static_cast(input_data_type_); - size_t centroids_bytes = centroids_.size() * sizeof(float); + size_t centroids_bytes = centroids_.size(); size_t centroid_bytes = use_zero_mean_ ? centroid_.size() * sizeof(float) : 0; hdr.payload_size = static_cast(sizeof(payload) + centroids_bytes + centroid_bytes); @@ -626,12 +757,20 @@ int PqInt8Quantizer::deserialize(const void *data, size_t len) { num_chunk_ = payload.num_chunk; sub_dim_ = payload.sub_dim; - meta_.set_meta(IndexMeta::DataType::DT_FP32, original_dim_); + // Restore input data type. Old payloads have input_data_type == 0 + // (was reserved), which maps to kInt4 -- treat as kFp32 for compat. + if (payload.input_data_type == 0 || + payload.input_data_type == static_cast(DataType::kInt4) || + payload.input_data_type == static_cast(DataType::kInt8)) { + input_data_type_ = DataType::kFp32; + } else { + input_data_type_ = static_cast(payload.input_data_type); + } + // Restore centroids (raw bytes in original data type). size_t centroids_bytes = static_cast(num_chunk_) * kNumCentroids * - sub_dim_ * sizeof(float); - - centroids_.resize(centroids_bytes / sizeof(float)); + sub_dim_ * element_size(); + centroids_.resize(centroids_bytes); std::memcpy(centroids_.data(), ptr, centroids_bytes); ptr += centroids_bytes; @@ -653,24 +792,29 @@ 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, + l2_batch_fn_ = + get_batch_distance_func(MetricType::kSquaredEuclidean, input_data_type_, QuantizeType::kDefault, CpuArchType::kAuto); // Metric-aware batch distance for search LUT. Cosine = normalize + L2, // so it uses SquaredEuclidean (same as encoding), not IP. if (meta_.metric_name() == "Cosine") { - fp32_batch_fn_ = - get_batch_distance_func(MetricType::kSquaredEuclidean, DataType::kFp32, - QuantizeType::kDefault, CpuArchType::kAuto); + batch_fn_ = + get_batch_distance_func(MetricType::kSquaredEuclidean, + input_data_type_, QuantizeType::kDefault, + CpuArchType::kAuto); extra_meta_size_ = kExtraMetaSizeCosine; - meta_.set_extra_meta_size(extra_meta_size_); } else { - fp32_batch_fn_ = get_batch_distance_func( - metric_from_name(meta_.metric_name()), DataType::kFp32, + batch_fn_ = get_batch_distance_func( + metric_from_name(meta_.metric_name()), input_data_type_, QuantizeType::kDefault, CpuArchType::kAuto); } + // Set output meta: the quantized representation is INT8 codes with + // num_chunk_ bytes (+ extra_meta_size_ for Cosine norm storage). + meta_.set_meta(IndexMeta::DataType::DT_INT8, num_chunk_); + meta_.set_extra_meta_size(extra_meta_size_); + // 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 f930d94b5..09d1d5d5f 100644 --- a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h @@ -17,6 +17,7 @@ #include #include #include +#include #include #include #include "quantizer/quantizer.h" @@ -48,7 +49,7 @@ class PqInt8Quantizer : public Quantizer { } DataType input_data_type() const override { - return DataType::kFp32; + return input_data_type_; } QuantizeType type() const override { @@ -114,9 +115,12 @@ class PqInt8Quantizer : public Quantizer { 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, + //! Train a single sub-quantizer (KMeans, k=256) on the sub-vectors. + //! Templated on the data type T (float or ailego::Float16) so that + //! NumericalKmeans operates natively in the input precision. + //! sub_idx selects which sub-quantizer to train. + template + void train_subquantizer(const T *data, size_t num, size_t stride, size_t sub_idx); //! Compute the centroid-to-centroid distance table for SDC. @@ -126,11 +130,21 @@ class PqInt8Quantizer : public Quantizer { //! Called after train() and deserialize() when centroids are available. void build_centroid_ptrs_cache(); + //! Byte size of one element in the original data type. + uint32_t element_size() const { + return (input_data_type_ == DataType::kFp16) + ? static_cast(sizeof(ailego::Float16)) + : static_cast(sizeof(float)); + } + static constexpr uint32_t kNumCentroids = 256; static constexpr uint32_t kMaxKmeansIters = 25; static constexpr size_t kMaxTrainVectors = 65536; static constexpr uint32_t kExtraMetaSizeCosine = sizeof(float); + //! Actual input data type (kFp32 or kFp16). + DataType input_data_type_{DataType::kFp32}; + //! Thread count for KMeans training (0 = hardware_concurrency). //! Read from params in init(), aligned with multi_chunk_cluster. uint32_t thread_count_{0}; @@ -151,8 +165,10 @@ class PqInt8Quantizer : public Quantizer { uint32_t num_chunk_{0}; uint32_t sub_dim_{0}; - //! Centroids: [num_chunk * kNumCentroids * sub_dim] - std::vector centroids_; + //! Centroids stored as raw bytes in the original data type: + //! [num_chunk * kNumCentroids * sub_dim * sizeof(T)] + //! T = float for kFp32, ailego::Float16 for kFp16. + std::vector centroids_; //! Global centroid (per-dimension mean) for zero-mean centering. //! Size: original_dim_ floats. Only populated when use_zero_mean_ = true. @@ -173,17 +189,15 @@ class PqInt8Quantizer : public Quantizer { PqSdcKernelFunc sdc_fn_{nullptr}; PqBatchAdcFunc batch_adc_fn_{nullptr}; - //! Metric-aware fp32 batch distance function for search-side LUT - //! computation and SDC dist_table. Obtained from - //! get_batch_distance_func() with the configured metric. - BatchDistanceFunc fp32_batch_fn_{}; - - //! L2-only fp32 batch distance function for encoding (quantize_data) - //! and KMeans training (train_subquantizer). PQ encoding must always - //! minimize L2 quantization error regardless of the search metric. - //! This separation ensures encoding and search LUT use the correct - //! distance semantics independently. - BatchDistanceFunc fp32_l2_batch_fn_{}; + //! Metric-aware batch distance function for search-side LUT + //! computation and SDC dist_table. Data type matches input_data_type_. + //! Obtained from get_batch_distance_func() with the configured metric. + BatchDistanceFunc batch_fn_{}; + + //! L2-only batch distance function for encoding (quantize_data). + //! Data type matches input_data_type_. PQ encoding always minimizes L2 + //! quantization error regardless of the search metric. + BatchDistanceFunc l2_batch_fn_{}; }; } // namespace turbo diff --git a/src/turbo/turbo.cc b/src/turbo/turbo.cc index 34cb9d794..5ddb55b13 100644 --- a/src/turbo/turbo.cc +++ b/src/turbo/turbo.cc @@ -25,6 +25,9 @@ #include "avx512_vnni/uniform_int8/squared_euclidean.h" #include "neon/pq_quantizer_int8/pq_distance.h" #include "neon/rotate/fht/fht.h" +#include "scalar/fp16/cosine.h" +#include "scalar/fp16/inner_product.h" +#include "scalar/fp16/squared_euclidean.h" #include "scalar/fp32/cosine.h" #include "scalar/fp32/inner_product.h" #include "scalar/fp32/squared_euclidean.h" @@ -57,6 +60,21 @@ DistanceFunc get_distance_func(MetricType metric_type, DataType data_type, } return nullptr; } + if (data_type == DataType::kFp16) { + if (quantize_type == QuantizeType::kDefault || + quantize_type == QuantizeType::kFp16) { + if (metric_type == MetricType::kSquaredEuclidean) { + return scalar::squared_euclidean_fp16_distance; + } + if (metric_type == MetricType::kInnerProduct) { + return scalar::inner_product_fp16_distance; + } + if (metric_type == MetricType::kCosine) { + return scalar::cosine_fp16_distance; + } + } + return nullptr; + } if (data_type == DataType::kInt8) { if (quantize_type == QuantizeType::kDefault) { if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI && @@ -100,6 +118,21 @@ BatchDistanceFunc get_batch_distance_func(MetricType metric_type, } return nullptr; } + if (data_type == DataType::kFp16) { + if (quantize_type == QuantizeType::kDefault || + quantize_type == QuantizeType::kFp16) { + if (metric_type == MetricType::kSquaredEuclidean) { + return scalar::squared_euclidean_fp16_batch_distance; + } + if (metric_type == MetricType::kInnerProduct) { + return scalar::inner_product_fp16_batch_distance; + } + if (metric_type == MetricType::kCosine) { + return scalar::cosine_fp16_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_pq_int8_quantizer_test.cc b/tests/turbo/turbo_pq_int8_quantizer_test.cc index 754cf3fce..627e1ec5f 100644 --- a/tests/turbo/turbo_pq_int8_quantizer_test.cc +++ b/tests/turbo/turbo_pq_int8_quantizer_test.cc @@ -34,6 +34,7 @@ using namespace zvec; using namespace zvec::core; using namespace zvec::ailego; +using zvec::turbo::DataType; // Reference squared Euclidean distance between two raw fp32 vectors. static float reference_sq_euclidean(const float *a, const float *b, @@ -1075,3 +1076,374 @@ TEST(PqInt8Quantizer, ZeroMeanAccuracyComparable) { << "Centering degraded accuracy too much: center_err=" << avg_err_center << " no_center_err=" << avg_err_no_center; } + +// --------------------------------------------------------------------------- +// FP16 Input Tests +// --------------------------------------------------------------------------- + +#include + +// Helper to create a PqInt8Quantizer with FP16 input type. +static std::shared_ptr make_pq_fp16_quantizer( + size_t dim, size_t num_chunk, const char *metric = "SquaredEuclidean") { + auto q = IndexFactory::CreateQuantizer("PqInt8Quantizer"); + if (!q) return nullptr; + + IndexMeta meta; + meta.set_meta(IndexMeta::DataType::DT_FP16, dim); + meta.set_metric(metric, 0, Params()); + + Params params; + params.set("num_chunk", static_cast(num_chunk)); + if (q->init(meta, params) != 0) return nullptr; + return q; +} + +// Helper: build a holder with random fp16 vectors (via Float16). +static std::shared_ptr> +make_random_fp16_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] = ailego::Float16(dist(gen)); + holder->emplace(i + 1, vec); + } + return holder; +} + +// Convert a Float16 vector to fp32 for reference distance computation. +static std::vector fp16_to_fp32(const ailego::Float16 *v, size_t dim) { + std::vector out(dim); + for (size_t i = 0; i < dim; ++i) out[i] = static_cast(v[i]); + return out; +} + +// Verify FP16 init succeeds and output meta is correct. +TEST(PqInt8Fp16, InitAndOutputMeta) { + auto q = make_pq_fp16_quantizer(16, 4); + ASSERT_TRUE(q); + + // input_data_type should be kFp16. + EXPECT_EQ(q->input_data_type(), DataType::kFp16); + + // Output meta: data_type = DT_INT8, dimension = num_chunk. + const auto &meta = q->meta(); + EXPECT_EQ(meta.data_type(), IndexMeta::DataType::DT_INT8); + EXPECT_EQ(meta.dimension(), 4u); + // L2 metric: no extra meta. + EXPECT_EQ(meta.extra_meta_size(), 0u); + EXPECT_EQ(meta.element_size(), 4u); + + // Cosine FP16: extra meta for norm storage. + auto q_cos = make_pq_fp16_quantizer(32, 8, "Cosine"); + ASSERT_TRUE(q_cos); + EXPECT_EQ(q_cos->meta().extra_meta_size(), sizeof(float)); + EXPECT_EQ(q_cos->meta().element_size(), 8u + sizeof(float)); +} + +// Verify basic train + encode with FP16 input. +TEST(PqInt8Fp16, TrainAndEncode) { + const size_t DIM = 16; + const size_t NSQ = 4; + const size_t COUNT = 1000; + + auto quantizer = make_pq_fp16_quantizer(DIM, NSQ); + ASSERT_TRUE(quantizer); + EXPECT_TRUE(quantizer->require_train()); + + auto holder = make_random_fp16_holder(COUNT, DIM); + ASSERT_EQ(0, quantizer->train(holder)); + + 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()); + for (size_t m = 0; m < NSQ; ++m) { + EXPECT_LE(code[m], 255u); + } + } + EXPECT_EQ(10u, checked); +} + +// Verify ADC distances with FP16 input are reasonable. +TEST(PqInt8Fp16, AdcDistance) { + const size_t DIM = 32; + const size_t NSQ = 8; + const size_t COUNT = 2000; + + auto quantizer = make_pq_fp16_quantizer(DIM, NSQ); + ASSERT_TRUE(quantizer); + + auto holder = make_random_fp16_holder(COUNT, DIM); + ASSERT_EQ(0, quantizer->train(holder)); + + // Collect raw vectors (as fp32 for reference) 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 ailego::Float16 *v = + reinterpret_cast(iter->data()); + raw_vecs[i] = fp16_to_fp32(v, DIM); + pq_codes[i].resize(code_len); + quantizer->quantize_data(iter->data(), pq_codes[i].data()); + } + + // Build LUT for query = first vector (must use FP16 data from holder). + auto iter2 = holder->create_iterator(); + iter2->is_valid(); + std::vector lut(lut_len / sizeof(float)); + quantizer->quantize_query(iter2->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; + } + // FP16 has lower precision than FP32, so allow slightly larger error. + EXPECT_LT(max_rel_error, 1.5f) << "max_rel_error=" << max_rel_error; +} + +// Verify dequantize from FP16 PQ produces reasonable fp32 reconstruction. +TEST(PqInt8Fp16, Dequantize) { + const size_t DIM = 16; + const size_t NSQ = 4; + const size_t COUNT = 1000; + + auto quantizer = make_pq_fp16_quantizer(DIM, NSQ); + ASSERT_TRUE(quantizer); + + auto holder = make_random_fp16_holder(COUNT, DIM); + ASSERT_EQ(0, quantizer->train(holder)); + + auto iter = holder->create_iterator(); + iter->is_valid(); + + const ailego::Float16 *orig_fp16 = + reinterpret_cast(iter->data()); + std::vector orig_fp32 = fp16_to_fp32(orig_fp16, DIM); + + std::vector code(quantizer->quantized_datapoint_vector_length()); + quantizer->quantize_data(iter->data(), code.data()); + + // Dequantize always outputs fp32. + IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP16, 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()); + float recon_err = reference_sq_euclidean(orig_fp32.data(), recon, DIM); + float orig_norm = + reference_sq_euclidean(orig_fp32.data(), orig_fp32.data(), DIM); + if (orig_norm > 1e-6f) { + EXPECT_LT(recon_err / orig_norm, 1.0f) + << "recon_err=" << recon_err << " orig_norm=" << orig_norm; + } +} + +// Verify serialize/deserialize round-trip preserves FP16 PQ codes. +TEST(PqInt8Fp16, SerializeDeserialize) { + const size_t DIM = 16; + const size_t NSQ = 4; + const size_t COUNT = 500; + + auto quantizer = make_pq_fp16_quantizer(DIM, NSQ); + ASSERT_TRUE(quantizer); + + auto holder = make_random_fp16_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)); + + // Verify deserialized quantizer reports FP16 input type. + EXPECT_EQ(q2->input_data_type(), DataType::kFp16); + + // Encode the same vector with both and compare codes. + auto iter = holder->create_iterator(); + iter->is_valid(); + std::vector code1(quantizer->quantized_datapoint_vector_length()); + std::vector code2(q2->quantized_datapoint_vector_length()); + quantizer->quantize_data(iter->data(), code1.data()); + q2->quantize_data(iter->data(), code2.data()); + + for (size_t m = 0; m < NSQ; ++m) { + EXPECT_EQ(code1[m], code2[m]) << "m=" << m; + } + + // ADC distances should also match (same codebook → same LUT → same ADC). + size_t lut_len = quantizer->quantized_query_vector_length(); + std::vector lut1(lut_len / sizeof(float)); + std::vector lut2(lut_len / sizeof(float)); + quantizer->quantize_query(iter->data(), lut1.data()); + q2->quantize_query(iter->data(), lut2.data()); + + float adc1 = quantizer->calc_distance_dp_query(code1.data(), lut1.data()); + float adc2 = q2->calc_distance_dp_query(code2.data(), lut2.data()); + EXPECT_NEAR(adc1, adc2, 1e-6f); +} + +// Helper: build a holder with random fp16 vectors with varying norms (Cosine). +static std::shared_ptr> +make_cosine_fp16_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] = ailego::Float16(dist(gen) * s); + holder->emplace(i + 1, vec); + } + return holder; +} + +// Verify FP16 with Cosine metric: train, encode, ADC distance range. +TEST(PqInt8Fp16, CosineAdcDistance) { + const size_t DIM = 32; + const size_t NSQ = 8; + const size_t COUNT = 2000; + + auto quantizer = make_pq_fp16_quantizer(DIM, NSQ, "Cosine"); + ASSERT_TRUE(quantizer); + + // Cosine FP16: extra meta = sizeof(float). + EXPECT_EQ(quantizer->quantized_datapoint_vector_length(), + NSQ + sizeof(float)); + + auto holder = make_cosine_fp16_holder(COUNT, DIM); + ASSERT_EQ(0, quantizer->train(holder)); + + // Collect raw vectors (as fp32) 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 ailego::Float16 *v = + reinterpret_cast(iter->data()); + raw_vecs[i] = fp16_to_fp32(v, DIM); + pq_codes[i].resize(code_len); + quantizer->quantize_data(iter->data(), pq_codes[i].data()); + } + + // Build LUT for query = first vector (must use FP16 data from holder). + auto iter2 = holder->create_iterator(); + iter2->is_valid(); + std::vector lut(lut_len / sizeof(float)); + quantizer->quantize_query(iter2->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 ADC distance should be in [0, 2] (with some tolerance). + EXPECT_GE(adc_dist, -0.05f) << "i=" << i; + EXPECT_LE(adc_dist, 2.05f) << "i=" << i; + + // PQ approximation: should be roughly correlated. + EXPECT_LT(std::fabs(adc_dist - true_dist), 0.5f) + << "i=" << i << " adc=" << adc_dist << " true=" << true_dist; + } +} + +// Verify FP16 PQ produces ADC distance rankings consistent with FP32 PQ. +// Both quantizers train independent codebooks, so raw codes differ, but the +// relative distance ordering should be largely preserved on the same data. +TEST(PqInt8Fp16, ConsistencyWithFp32) { + const size_t DIM = 32; + const size_t NSQ = 8; + const size_t COUNT = 2000; + + // Train FP32 quantizer. + auto q_fp32 = make_pq_quantizer(DIM, NSQ); + ASSERT_TRUE(q_fp32); + auto holder_fp32 = make_random_holder(COUNT, DIM, 123); + ASSERT_EQ(0, q_fp32->train(holder_fp32)); + + // Build FP16 quantizer with same seed data. + auto q_fp16 = make_pq_fp16_quantizer(DIM, NSQ); + ASSERT_TRUE(q_fp16); + auto holder_fp16 = make_random_fp16_holder(COUNT, DIM, 123); + ASSERT_EQ(0, q_fp16->train(holder_fp16)); + + // Build LUTs for the first vector in each holder. + size_t lut_len = q_fp32->quantized_query_vector_length(); + std::vector lut_fp32(lut_len / sizeof(float)); + std::vector lut_fp16(lut_len / sizeof(float)); + + auto iter32 = holder_fp32->create_iterator(); + auto iter16 = holder_fp16->create_iterator(); + iter32->is_valid(); + iter16->is_valid(); + q_fp32->quantize_query(iter32->data(), lut_fp32.data()); + q_fp16->quantize_query(iter16->data(), lut_fp16.data()); + + // Encode next 200 vectors and collect ADC distances. + size_t code_len = q_fp32->quantized_datapoint_vector_length(); + std::vector code32(code_len), code16(code_len); + + std::vector adc32_vec, adc16_vec; + iter32->next(); + iter16->next(); + for (size_t i = 1; i < 201 && iter32->is_valid() && iter16->is_valid(); + iter32->next(), iter16->next(), ++i) { + q_fp32->quantize_data(iter32->data(), code32.data()); + q_fp16->quantize_data(iter16->data(), code16.data()); + adc32_vec.push_back( + q_fp32->calc_distance_dp_query(code32.data(), lut_fp32.data())); + adc16_vec.push_back( + q_fp16->calc_distance_dp_query(code16.data(), lut_fp16.data())); + } + + // Compute Spearman rank correlation between FP32 and FP16 ADC distances. + size_t n = adc32_vec.size(); + ASSERT_GT(n, 10u); + + // Count concordant vs discordant pairs (Kendall tau). + size_t concordant = 0, discordant = 0; + for (size_t i = 0; i < n; ++i) { + for (size_t j = i + 1; j < n; ++j) { + bool same_order = + (adc32_vec[i] < adc32_vec[j]) == (adc16_vec[i] < adc16_vec[j]); + if (same_order) + ++concordant; + else + ++discordant; + } + } + double tau = static_cast(concordant - discordant) / + static_cast(concordant + discordant); + + // Kendall tau > 0.5 means strong rank correlation (generous threshold; + // FP16 precision loss and different codebooks cause some reordering). + EXPECT_GT(tau, 0.5) << "Kendall tau=" << tau; +} From 9ef03dce42f43a83163f8428079670680570278e Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Tue, 28 Jul 2026 16:00:46 +0800 Subject: [PATCH 098/100] fix format of fp16 --- .../pq_int8_quantizer/pq_int8_quantizer.cc | 284 +++++++++--------- .../pq_int8_quantizer/pq_int8_quantizer.h | 17 ++ 2 files changed, 164 insertions(+), 137 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 9b741f59a..ae4c495b7 100644 --- a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc @@ -91,9 +91,8 @@ int PqInt8Quantizer::init(const IndexMeta &meta, const ailego::Params ¶ms) { // with squared-Euclidean, so the search LUT reuses SquaredEuclidean. if (meta_.metric_name() == "Cosine") { batch_fn_ = - get_batch_distance_func(MetricType::kSquaredEuclidean, - input_data_type_, QuantizeType::kDefault, - CpuArchType::kAuto); + get_batch_distance_func(MetricType::kSquaredEuclidean, input_data_type_, + QuantizeType::kDefault, CpuArchType::kAuto); extra_meta_size_ = kExtraMetaSizeCosine; meta_.set_extra_meta_size(extra_meta_size_); } else { @@ -129,8 +128,7 @@ void PqInt8Quantizer::train_subquantizer(const T *data, size_t num, const size_t k = kNumCentroids; const size_t d = sub_dim_; uint8_t *centroids_m = - centroids_.data() + - static_cast(sub_idx) * k * d * sizeof(T); + centroids_.data() + static_cast(sub_idx) * k * d * sizeof(T); // Non-spherical L2 KMeans: the PQ codebook must minimize L2 reconstruction // error, so centroids are the true (magnitude-preserving) means. @@ -139,8 +137,8 @@ void PqInt8Quantizer::train_subquantizer(const T *data, size_t num, // Append sub-vectors (NumericalKmeans handles transpose internally) for (size_t i = 0; i < num; ++i) { const T *sub_vec = - reinterpret_cast( - reinterpret_cast(data) + i * stride) + + reinterpret_cast(reinterpret_cast(data) + + i * stride) + sub_idx * d; algorithm.append(sub_vec, d); } @@ -226,21 +224,14 @@ int PqInt8Quantizer::train(IndexHolder::Pointer holder, int thread_count) { // For Cosine: normalize training data so centroids are learned in // normalized space (L2 minimization == maximizing cosine similarity). if (meta_.metric_name() == "Cosine") { - if (input_data_type_ == DataType::kFp16) { - for (size_t i = 0; i < num; ++i) { - ailego::Float16 *v = - reinterpret_cast(all_data.data()) + - i * original_dim_; - float norm = 0.0f; - ailego::Normalizer::L2(v, original_dim_, &norm); - } - } else { - for (size_t i = 0; i < num; ++i) { - float *v = reinterpret_cast(all_data.data()) + - i * original_dim_; - float norm = 0.0f; - ailego::Normalizer::L2(v, original_dim_, &norm); - } + switch (input_data_type_) { + case DataType::kFp16: + normalize_batch(reinterpret_cast(all_data.data()), + num); + break; + case DataType::kFp32: + normalize_batch(reinterpret_cast(all_data.data()), num); + break; } } @@ -249,47 +240,15 @@ int PqInt8Quantizer::train(IndexHolder::Pointer holder, int thread_count) { // runs AFTER normalization, so all paths keep the same order // (normalize -> center; dequantize: un-center -> rescale). if (use_zero_mean_) { - centroid_.assign(original_dim_, 0.0f); - if (input_data_type_ == DataType::kFp16) { - for (size_t i = 0; i < num; ++i) { - const ailego::Float16 *v = - reinterpret_cast(all_data.data()) + - i * original_dim_; - for (uint32_t d = 0; d < original_dim_; ++d) { - centroid_[d] += static_cast(v[d]); - } - } - for (uint32_t d = 0; d < original_dim_; ++d) { - centroid_[d] /= static_cast(num); - } - // Subtract centroid from all training vectors. - for (size_t i = 0; i < num; ++i) { - ailego::Float16 *v = - reinterpret_cast(all_data.data()) + - i * original_dim_; - for (uint32_t d = 0; d < original_dim_; ++d) { - v[d] -= centroid_[d]; - } - } - } else { - for (size_t i = 0; i < num; ++i) { - const float *v = reinterpret_cast(all_data.data()) + - i * original_dim_; - for (uint32_t d = 0; d < original_dim_; ++d) { - centroid_[d] += v[d]; - } - } - for (uint32_t d = 0; d < original_dim_; ++d) { - centroid_[d] /= static_cast(num); - } - // Subtract centroid from all training vectors. - for (size_t i = 0; i < num; ++i) { - float *v = reinterpret_cast(all_data.data()) + - i * original_dim_; - for (uint32_t d = 0; d < original_dim_; ++d) { - v[d] -= centroid_[d]; - } - } + switch (input_data_type_) { + case DataType::kFp16: + compute_and_subtract_center( + reinterpret_cast(all_data.data()), num); + break; + case DataType::kFp32: + compute_and_subtract_center(reinterpret_cast(all_data.data()), + num); + break; } } @@ -301,34 +260,31 @@ int PqInt8Quantizer::train(IndexHolder::Pointer holder, int thread_count) { // Distribute sub-quantizers across threads, dispatched by data type. std::atomic finished{0}; size_t pool_count = threads->count(); - if (input_data_type_ == DataType::kFp16) { - const auto *typed_data = - reinterpret_cast(all_data.data()); - for (size_t i = 0; i < pool_count; ++i) { - task_group->submit(ailego::Closure::New( - [this, typed_data, num, data_stride, i, pool_count, &finished]() { - for (uint32_t m = static_cast(i); m < num_chunk_; - m += static_cast(pool_count)) { - train_subquantizer(typed_data, num, data_stride, - m); - finished++; - } - })); - } - } else { - const auto *typed_data = - reinterpret_cast(all_data.data()); + + auto submit_training = [&](const auto *typed_data) { + using T = std::remove_const_t>; for (size_t i = 0; i < pool_count; ++i) { task_group->submit(ailego::Closure::New( [this, typed_data, num, data_stride, i, pool_count, &finished]() { for (uint32_t m = static_cast(i); m < num_chunk_; m += static_cast(pool_count)) { - train_subquantizer(typed_data, num, data_stride, m); + train_subquantizer(typed_data, num, data_stride, m); finished++; } })); } + }; + + switch (input_data_type_) { + case DataType::kFp16: + submit_training( + reinterpret_cast(all_data.data())); + break; + case DataType::kFp32: + submit_training(reinterpret_cast(all_data.data())); + break; } + task_group->wait_finish(); // Pre-build centroid pointer cache (needed by compute_dist_table). @@ -395,14 +351,16 @@ void PqInt8Quantizer::quantize_data(const void *input, void *output) const { if (meta_.metric_name() == "Cosine") { norm_vec_storage.resize(original_dim_ * elem_sz); std::memcpy(norm_vec_storage.data(), input, original_dim_ * elem_sz); - if (input_data_type_ == DataType::kFp16) { - ailego::Normalizer::L2( - reinterpret_cast(norm_vec_storage.data()), - original_dim_, &vec_norm); - } else { - ailego::Normalizer::L2( - reinterpret_cast(norm_vec_storage.data()), original_dim_, - &vec_norm); + switch (input_data_type_) { + case DataType::kFp16: + normalize_single( + reinterpret_cast(norm_vec_storage.data()), + &vec_norm); + break; + case DataType::kFp32: + normalize_single(reinterpret_cast(norm_vec_storage.data()), + &vec_norm); + break; } vec = norm_vec_storage.data(); } @@ -412,17 +370,14 @@ void PqInt8Quantizer::quantize_data(const void *input, void *output) const { if (use_zero_mean_) { centered_vec_storage.resize(original_dim_ * elem_sz); std::memcpy(centered_vec_storage.data(), vec, original_dim_ * elem_sz); - if (input_data_type_ == DataType::kFp16) { - ailego::Float16 *cv = - reinterpret_cast(centered_vec_storage.data()); - for (uint32_t d = 0; d < original_dim_; ++d) { - cv[d] -= centroid_[d]; - } - } else { - float *cv = reinterpret_cast(centered_vec_storage.data()); - for (uint32_t d = 0; d < original_dim_; ++d) { - cv[d] -= centroid_[d]; - } + switch (input_data_type_) { + case DataType::kFp16: + subtract_center( + reinterpret_cast(centered_vec_storage.data())); + break; + case DataType::kFp32: + subtract_center(reinterpret_cast(centered_vec_storage.data())); + break; } vec = centered_vec_storage.data(); } @@ -433,7 +388,8 @@ void PqInt8Quantizer::quantize_data(const void *input, void *output) const { const uint8_t *vec_bytes = reinterpret_cast(vec); for (uint32_t m = 0; m < num_chunk_; ++m) { - const void *sub_vec = vec_bytes + static_cast(m) * sub_dim_ * elem_sz; + const void *sub_vec = + vec_bytes + static_cast(m) * sub_dim_ * elem_sz; const auto ¢roid_ptrs = centroid_ptrs_cache_[m]; // Compute L2 distances from this sub-vector to all 256 centroids. @@ -471,16 +427,14 @@ void PqInt8Quantizer::quantize_query(const void *input, void *output) const { if (meta_.metric_name() == "Cosine") { norm_query_storage.resize(original_dim_ * elem_sz); std::memcpy(norm_query_storage.data(), input, original_dim_ * elem_sz); - if (input_data_type_ == DataType::kFp16) { - float norm = 0.0f; - ailego::Normalizer::L2( - reinterpret_cast(norm_query_storage.data()), - original_dim_, &norm); - } else { - float norm = 0.0f; - ailego::Normalizer::L2( - reinterpret_cast(norm_query_storage.data()), original_dim_, - &norm); + switch (input_data_type_) { + case DataType::kFp16: + normalize_single( + reinterpret_cast(norm_query_storage.data())); + break; + case DataType::kFp32: + normalize_single(reinterpret_cast(norm_query_storage.data())); + break; } query = norm_query_storage.data(); } @@ -490,17 +444,15 @@ void PqInt8Quantizer::quantize_query(const void *input, void *output) const { if (use_zero_mean_) { centered_query_storage.resize(original_dim_ * elem_sz); std::memcpy(centered_query_storage.data(), query, original_dim_ * elem_sz); - if (input_data_type_ == DataType::kFp16) { - ailego::Float16 *cq = - reinterpret_cast(centered_query_storage.data()); - for (uint32_t d = 0; d < original_dim_; ++d) { - cq[d] -= centroid_[d]; - } - } else { - float *cq = reinterpret_cast(centered_query_storage.data()); - for (uint32_t d = 0; d < original_dim_; ++d) { - cq[d] -= centroid_[d]; - } + switch (input_data_type_) { + case DataType::kFp16: + subtract_center( + reinterpret_cast(centered_query_storage.data())); + break; + case DataType::kFp32: + subtract_center( + reinterpret_cast(centered_query_storage.data())); + break; } query = centered_query_storage.data(); } @@ -585,9 +537,15 @@ float PqInt8Quantizer::calc_distance_dp_dp(const void *dp1, int PqInt8Quantizer::quantize(const void *query, const IndexQueryMeta &qmeta, std::string *out, IndexQueryMeta *ometa) const { // Validate unit_size against the input data type. - size_t expected_unit = - (input_data_type_ == DataType::kFp16) ? sizeof(ailego::Float16) - : sizeof(float); + size_t expected_unit = 0; + switch (input_data_type_) { + case DataType::kFp16: + expected_unit = sizeof(ailego::Float16); + break; + case DataType::kFp32: + expected_unit = sizeof(float); + break; + } if (qmeta.unit_size() != expected_unit) { return kErrUnsupported; } @@ -619,16 +577,21 @@ int PqInt8Quantizer::dequantize(const void *in, const IndexQueryMeta &qmeta, for (uint32_t m = 0; m < num_chunk_; ++m) { const uint8_t *centroids_m = centroids_.data() + static_cast(m) * k * d * elem_sz; - const uint8_t *centroid = centroids_m + static_cast(code[m]) * d * elem_sz; - - if (input_data_type_ == DataType::kFp16) { - const ailego::Float16 *src = - reinterpret_cast(centroid); - for (size_t j = 0; j < d; ++j) { - result[m * d + j] = static_cast(src[j]); + const uint8_t *centroid = + centroids_m + static_cast(code[m]) * d * elem_sz; + + switch (input_data_type_) { + case DataType::kFp16: { + const ailego::Float16 *src = + reinterpret_cast(centroid); + for (size_t j = 0; j < d; ++j) { + result[m * d + j] = static_cast(src[j]); + } + break; } - } else { - std::memcpy(result + m * d, centroid, d * sizeof(float)); + case DataType::kFp32: + std::memcpy(result + m * d, centroid, d * sizeof(float)); + break; } } @@ -800,9 +763,8 @@ int PqInt8Quantizer::deserialize(const void *data, size_t len) { // so it uses SquaredEuclidean (same as encoding), not IP. if (meta_.metric_name() == "Cosine") { batch_fn_ = - get_batch_distance_func(MetricType::kSquaredEuclidean, - input_data_type_, QuantizeType::kDefault, - CpuArchType::kAuto); + get_batch_distance_func(MetricType::kSquaredEuclidean, input_data_type_, + QuantizeType::kDefault, CpuArchType::kAuto); extra_meta_size_ = kExtraMetaSizeCosine; } else { batch_fn_ = get_batch_distance_func( @@ -823,5 +785,53 @@ int PqInt8Quantizer::deserialize(const void *data, size_t len) { INDEX_FACTORY_REGISTER_QUANTIZER(PqInt8Quantizer); +// --------------------------------------------------------------------------- +// Template helper implementations (type-dispatched at call sites) +// --------------------------------------------------------------------------- + +template +void PqInt8Quantizer::normalize_batch(T *data, size_t num) const { + for (size_t i = 0; i < num; ++i) { + float norm = 0.0f; + ailego::Normalizer::L2(data + i * original_dim_, original_dim_, &norm); + } +} + +template +void PqInt8Quantizer::compute_and_subtract_center(T *data, size_t num) { + centroid_.assign(original_dim_, 0.0f); + for (size_t i = 0; i < num; ++i) { + const T *v = data + i * original_dim_; + for (uint32_t d = 0; d < original_dim_; ++d) { + centroid_[d] += static_cast(v[d]); + } + } + for (uint32_t d = 0; d < original_dim_; ++d) { + centroid_[d] /= static_cast(num); + } + for (size_t i = 0; i < num; ++i) { + T *v = data + i * original_dim_; + for (uint32_t d = 0; d < original_dim_; ++d) { + v[d] -= centroid_[d]; + } + } +} + +template +void PqInt8Quantizer::normalize_single(T *vec, float *norm_out) const { + float norm = 0.0f; + ailego::Normalizer::L2(vec, original_dim_, &norm); + if (norm_out) { + *norm_out = norm; + } +} + +template +void PqInt8Quantizer::subtract_center(T *vec) const { + for (uint32_t d = 0; d < original_dim_; ++d) { + vec[d] -= centroid_[d]; + } +} + } // 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 index 09d1d5d5f..0ad673c25 100644 --- a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h @@ -123,6 +123,23 @@ class PqInt8Quantizer : public Quantizer { void train_subquantizer(const T *data, size_t num, size_t stride, size_t sub_idx); + //! L2-normalize a batch of vectors (train-time use). + template + void normalize_batch(T *data, size_t num) const; + + //! Compute the per-dimension mean (accumulated in float to avoid FP16 + //! overflow) and subtract it from all training vectors in-place. + template + void compute_and_subtract_center(T *data, size_t num); + + //! L2-normalize a single vector; optionally writes the norm out. + template + void normalize_single(T *vec, float *norm_out = nullptr) const; + + //! Subtract the pre-computed centroid_ from a single vector. + template + void subtract_center(T *vec) const; + //! Compute the centroid-to-centroid distance table for SDC. void compute_dist_table(); From d52a50a6f48ab5e16ad2d3ae4d42c11a4ff1ae54 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Tue, 28 Jul 2026 18:10:34 +0800 Subject: [PATCH 099/100] remote warning --- .../pq_int8_quantizer/pq_int8_quantizer.cc | 18 ++++++++++++++++++ 1 file changed, 18 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 ae4c495b7..d0a5a191b 100644 --- a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc @@ -232,6 +232,8 @@ int PqInt8Quantizer::train(IndexHolder::Pointer holder, int thread_count) { case DataType::kFp32: normalize_batch(reinterpret_cast(all_data.data()), num); break; + default: + break; } } @@ -249,6 +251,8 @@ int PqInt8Quantizer::train(IndexHolder::Pointer holder, int thread_count) { compute_and_subtract_center(reinterpret_cast(all_data.data()), num); break; + default: + break; } } @@ -283,6 +287,8 @@ int PqInt8Quantizer::train(IndexHolder::Pointer holder, int thread_count) { case DataType::kFp32: submit_training(reinterpret_cast(all_data.data())); break; + default: + break; } task_group->wait_finish(); @@ -361,6 +367,8 @@ void PqInt8Quantizer::quantize_data(const void *input, void *output) const { normalize_single(reinterpret_cast(norm_vec_storage.data()), &vec_norm); break; + default: + break; } vec = norm_vec_storage.data(); } @@ -378,6 +386,8 @@ void PqInt8Quantizer::quantize_data(const void *input, void *output) const { case DataType::kFp32: subtract_center(reinterpret_cast(centered_vec_storage.data())); break; + default: + break; } vec = centered_vec_storage.data(); } @@ -435,6 +445,8 @@ void PqInt8Quantizer::quantize_query(const void *input, void *output) const { case DataType::kFp32: normalize_single(reinterpret_cast(norm_query_storage.data())); break; + default: + break; } query = norm_query_storage.data(); } @@ -453,6 +465,8 @@ void PqInt8Quantizer::quantize_query(const void *input, void *output) const { subtract_center( reinterpret_cast(centered_query_storage.data())); break; + default: + break; } query = centered_query_storage.data(); } @@ -545,6 +559,8 @@ int PqInt8Quantizer::quantize(const void *query, const IndexQueryMeta &qmeta, case DataType::kFp32: expected_unit = sizeof(float); break; + default: + break; } if (qmeta.unit_size() != expected_unit) { return kErrUnsupported; @@ -592,6 +608,8 @@ int PqInt8Quantizer::dequantize(const void *in, const IndexQueryMeta &qmeta, case DataType::kFp32: std::memcpy(result + m * d, centroid, d * sizeof(float)); break; + default: + break; } } From 8adc58e2e57a781cc3293be2f2318536db81235d Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Thu, 30 Jul 2026 11:42:06 +0800 Subject: [PATCH 100/100] restruct --- src/core/interface/index.cc | 6 +- src/core/interface/index_param.cc | 52 +++++++++++- .../column/vector_column/engine_helper.hpp | 2 +- src/include/zvec/core/interface/index_param.h | 57 ++++++++++--- .../core/interface/index_param_builders.h | 15 +++- tests/core/interface/index_interface_test.cc | 83 ++++++++++++++++++- 6 files changed, 196 insertions(+), 19 deletions(-) diff --git a/src/core/interface/index.cc b/src/core/interface/index.cc index 755f673f9..d0d9dd61a 100644 --- a/src/core/interface/index.cc +++ b/src/core/interface/index.cc @@ -258,7 +258,11 @@ int Index::Init(const BaseIndexParam ¶m) { return core::IndexError_Runtime; } - if (CreateAndInitConverterReformer(param.quantizer_param, param) != 0) { + // an absent quantizer param behaves the same as a kNone one + const auto quantizer_param = param.quantizer_param + ? param.quantizer_param + : std::make_shared(); + if (CreateAndInitConverterReformer(*quantizer_param, param) != 0) { LOG_ERROR("Failed to create and init converter"); return core::IndexError_Runtime; } diff --git a/src/core/interface/index_param.cc b/src/core/interface/index_param.cc index 5d75276fd..3560cd06d 100644 --- a/src/core/interface/index_param.cc +++ b/src/core/interface/index_param.cc @@ -56,9 +56,15 @@ ailego::JsonObject BaseIndexParam::SerializeToJsonObject( // if (preprocess_param) { // json.set("preprocess_param", preprocess_param->SerializeToJson()); // } - if (!omit_empty_value || quantizer_param.type != QuantizerType::kNone) { + if (quantizer_param) { + if (!omit_empty_value || quantizer_param->type != QuantizerType::kNone) { + json_obj.set("quantizer_param", + quantizer_param->SerializeToJsonObject(omit_empty_value)); + } + } else if (!omit_empty_value) { + // no quantizer configured, keep the default(kNone) object as before json_obj.set("quantizer_param", - quantizer_param.SerializeToJsonObject(omit_empty_value)); + QuantizerParam().SerializeToJsonObject(false)); } // if (refiner_param) { // json.set("refiner_param", refiner_param->SerializeToJson()); @@ -110,7 +116,20 @@ bool BaseIndexParam::DeserializeFromJsonObject( if (json_obj.has("quantizer_param")) { if (json_obj.get("quantizer_param", &tmp_json_value); tmp_json_value.is_object()) { - quantizer_param.DeserializeFromJsonObject(tmp_json_value.as_object()); + const auto &quantizer_json_obj = tmp_json_value.as_object(); + // the concrete param type is determined by the quantizer type + auto quantizer_type = QuantizerType::kNone; + ailego::JsonValue type_json_value; + if (!extract_enum_from_json( + quantizer_json_obj, "type", quantizer_type, type_json_value)) { + LOG_ERROR("Error when deserialize json - field:quantizer_param.type"); + return false; + } + auto quantizer = QuantizerParam::Create(quantizer_type); + if (!quantizer->DeserializeFromJsonObject(quantizer_json_obj)) { + LOG_ERROR("Error when deserialize json - field:quantizer_param"); + } + quantizer_param = std::move(quantizer); } } @@ -264,6 +283,33 @@ bool QuantizerParam::DeserializeFromJsonObject( return true; } +QuantizerParam::Pointer QuantizerParam::Create(QuantizerType t) { + switch (t) { + case QuantizerType::kPQ: + return std::make_shared(); + default: + return std::make_shared(t); + } +} + +ailego::JsonObject PqQuantizerParam::SerializeToJsonObject( + bool omit_empty_value) const { + auto json_obj = QuantizerParam::SerializeToJsonObject(omit_empty_value); + json_obj.set("num_chunk", ailego::JsonValue(num_chunk)); + json_obj.set("num_bits", ailego::JsonValue(num_bits)); + return json_obj; +} + +bool PqQuantizerParam::DeserializeFromJsonObject( + const ailego::JsonObject &json_obj) { + if (!QuantizerParam::DeserializeFromJsonObject(json_obj)) { + return false; + } + DESERIALIZE_VALUE_FIELD(json_obj, num_chunk); + DESERIALIZE_VALUE_FIELD(json_obj, num_bits); + return true; +} + // bool BaseIndexQueryParam::DeserializeFromJsonObject( // const ailego::JsonObject &json_obj) { // DESERIALIZE_ENUM_FIELD(json_obj, index_type, IndexType); diff --git a/src/db/index/column/vector_column/engine_helper.hpp b/src/db/index/column/vector_column/engine_helper.hpp index dec1177cc..1685fb443 100644 --- a/src/db/index/column/vector_column/engine_helper.hpp +++ b/src/db/index/column/vector_column/engine_helper.hpp @@ -357,7 +357,7 @@ class ProximaEngineHelper { convert_to_engine_quantize_type(db_index_params->quantize_type()); quantize_type.has_value()) { index_param_builder->WithQuantizerParam( - core_interface::QuantizerParam(quantize_type.value())); + core_interface::QuantizerParam::Create(quantize_type.value())); } else { return tl::make_unexpected( Status::InvalidArgument("unsupported quantize type")); diff --git a/src/include/zvec/core/interface/index_param.h b/src/include/zvec/core/interface/index_param.h index 27fc52bdc..83ef7a985 100644 --- a/src/include/zvec/core/interface/index_param.h +++ b/src/include/zvec/core/interface/index_param.h @@ -118,23 +118,49 @@ struct SerializableBase { const ailego::JsonObject &json_obj) = 0; }; -// TODO: maybe a base class for quantizer? +//! Common quantizer params shared by all quantizer types struct QuantizerParam : public SerializableBase { + using Pointer = std::shared_ptr; + QuantizerType type = QuantizerType::kNone; - int num_chunk = 8; // M - int num_bits = 8; // bits per subquantizer bool enable_rotate = false; // rotate vectors before quantization to reduce error // Constructors - // QuantizerParam() = default; - QuantizerParam(QuantizerType t = QuantizerType::kNone, int subquantizers = 8, - int bits = 8, bool rotate = false) - : type(t), - num_chunk(subquantizers), - num_bits(bits), - enable_rotate(rotate) {} + QuantizerParam(QuantizerType t = QuantizerType::kNone, bool rotate = false) + : type(t), enable_rotate(rotate) {} + virtual ~QuantizerParam() = default; + + //! Duplicate the param object, keeping the concrete type + virtual Pointer Clone() const { + return std::make_shared(*this); + } + + //! Create the param object matching the quantizer type + static Pointer Create(QuantizerType t); + + protected: + friend class BaseIndexParam; + ailego::JsonObject SerializeToJsonObject( + bool omit_empty_value = false) const override; + + bool DeserializeFromJsonObject(const ailego::JsonObject &json_obj) override; +}; + +//! Product-Quantization specific params +struct PqQuantizerParam : public QuantizerParam { + int num_chunk = 8; // M: number of sub-quantizers + int num_bits = 8; // bits per sub-quantizer + + // Constructors + PqQuantizerParam(int chunks = 8, int bits = 8, bool rotate = false) + : QuantizerParam(QuantizerType::kPQ, rotate), + num_chunk(chunks), + num_bits(bits) {} + QuantizerParam::Pointer Clone() const override { + return std::make_shared(*this); + } protected: friend class BaseIndexParam; @@ -278,7 +304,16 @@ class BaseIndexParam : public SerializableBase { // pipeline PreprocessorParam preprocess_param; - QuantizerParam quantizer_param; + //! nullptr means no quantizer is configured (equivalent to kNone) + QuantizerParam::Pointer quantizer_param{nullptr}; + + QuantizerType quantizer_type() const { + return quantizer_param ? quantizer_param->type : QuantizerType::kNone; + } + + bool enable_rotate() const { + return quantizer_param && quantizer_param->enable_rotate; + } BaseIndexQueryParam::Pointer default_query_param = nullptr; // virtual std::shared_ptr GetDefaultQueryParam() const diff --git a/src/include/zvec/core/interface/index_param_builders.h b/src/include/zvec/core/interface/index_param_builders.h index d88057a93..37d77c93a 100644 --- a/src/include/zvec/core/interface/index_param_builders.h +++ b/src/include/zvec/core/interface/index_param_builders.h @@ -59,7 +59,13 @@ class BaseIndexParamBuilder { // : public } ActualIndexParamBuilderType &WithQuantizerParam( const QuantizerParam &quantizer_param) { - param->quantizer_param = quantizer_param; + param->quantizer_param = quantizer_param.Clone(); + return static_cast(*this); + } + ActualIndexParamBuilderType &WithQuantizerParam( + const QuantizerParam::Pointer &quantizer_param) { + param->quantizer_param = + quantizer_param ? quantizer_param->Clone() : nullptr; return static_cast(*this); } // ActualIndexParamBuilderType &WithRefinerParam( @@ -88,7 +94,12 @@ class BaseIndexParamBuilder { // : public } ActualIndexParamBuilderType &WithEnableRotate(bool enable_rotate) { - param->quantizer_param.enable_rotate = enable_rotate; + // copy-on-write: never mutate a param object shared with others + auto quantizer_param = param->quantizer_param + ? param->quantizer_param->Clone() + : std::make_shared(); + quantizer_param->enable_rotate = enable_rotate; + param->quantizer_param = std::move(quantizer_param); return static_cast(*this); } diff --git a/tests/core/interface/index_interface_test.cc b/tests/core/interface/index_interface_test.cc index 4453756da..0c6a35ba2 100644 --- a/tests/core/interface/index_interface_test.cc +++ b/tests/core/interface/index_interface_test.cc @@ -2514,7 +2514,88 @@ TEST(IndexInterface, BuilderSetsAllBaseFields) { EXPECT_FALSE(param->use_id_map); EXPECT_TRUE(param->use_external_vector); EXPECT_EQ(PreprocessorType::kPCA, param->preprocess_param.type); - EXPECT_EQ(QuantizerType::kFP16, param->quantizer_param.type); + EXPECT_EQ(QuantizerType::kFP16, param->quantizer_param->type); +} + +TEST(IndexInterface, QuantizerParamDefaultIsNull) { + auto param = FlatIndexParamBuilder() + .WithMetricType(MetricType::kL2sq) + .WithDimension(64) + .Build(); + + ASSERT_NE(nullptr, param); + EXPECT_EQ(nullptr, param->quantizer_param); + EXPECT_EQ(QuantizerType::kNone, param->quantizer_type()); + EXPECT_FALSE(param->enable_rotate()); +} + +TEST(IndexInterface, QuantizerParamPqFields) { + auto param = FlatIndexParamBuilder() + .WithMetricType(MetricType::kL2sq) + .WithDimension(128) + .WithQuantizerParam(PqQuantizerParam(16, 8)) + .Build(); + + ASSERT_NE(nullptr, param); + ASSERT_NE(nullptr, param->quantizer_param); + EXPECT_EQ(QuantizerType::kPQ, param->quantizer_type()); + + // the builder must keep the concrete type instead of slicing it + auto pq_param = + std::dynamic_pointer_cast(param->quantizer_param); + ASSERT_NE(nullptr, pq_param); + EXPECT_EQ(16, pq_param->num_chunk); + EXPECT_EQ(8, pq_param->num_bits); + + // enable_rotate is a common field, setting it keeps the concrete type + auto rotated_param = FlatIndexParamBuilder() + .WithQuantizerParam(PqQuantizerParam(16, 8)) + .WithEnableRotate(true) + .Build(); + ASSERT_NE(nullptr, rotated_param->quantizer_param); + EXPECT_TRUE(rotated_param->enable_rotate()); + EXPECT_NE(nullptr, std::dynamic_pointer_cast( + rotated_param->quantizer_param)); +} + +TEST(IndexInterface, QuantizerParamPqJsonRoundTrip) { + auto param = FlatIndexParamBuilder() + .WithIndexType(IndexType::kFlat) + .WithMetricType(MetricType::kL2sq) + .WithDimension(128) + .WithDataType(DataType::DT_FP32) + .WithQuantizerParam(PqQuantizerParam(32, 4)) + .Build(); + + auto deserialized_param = + IndexFactory::DeserializeIndexParamFromJson(param->SerializeToJson()); + ASSERT_NE(nullptr, deserialized_param); + EXPECT_EQ(param->SerializeToJson(), deserialized_param->SerializeToJson()); + EXPECT_EQ(param->SerializeToJson(true), + deserialized_param->SerializeToJson(true)); + + auto pq_param = std::dynamic_pointer_cast( + deserialized_param->quantizer_param); + ASSERT_NE(nullptr, pq_param); + EXPECT_EQ(QuantizerType::kPQ, pq_param->type); + EXPECT_EQ(32, pq_param->num_chunk); + EXPECT_EQ(4, pq_param->num_bits); +} + +TEST(IndexInterface, QuantizerParamLegacyJsonCompat) { + // legacy json only carries the common fields + const std::string json_str = + R"({"index_type":"kFlat","metric_type":"kL2sq","dimension":128,)" + R"("data_type":"DT_FP32",)" + R"("quantizer_param":{"type":"kInt8","enable_rotate":true}})"; + + auto param = IndexFactory::DeserializeIndexParamFromJson(json_str); + ASSERT_NE(nullptr, param); + ASSERT_NE(nullptr, param->quantizer_param); + EXPECT_EQ(QuantizerType::kInt8, param->quantizer_type()); + EXPECT_TRUE(param->enable_rotate()); + EXPECT_EQ(nullptr, std::dynamic_pointer_cast( + param->quantizer_param)); } TEST(IndexInterface, BuilderChainingReturnsCorrectType) {